静态变量模板实例化膨胀问题

问题来源

工作中遇到这么一个问题，把部分类型的变量安排在特定地址空间内，暂且称为该类型为handle。这个地址空间不算小，但是也有一定的限制。handle类型的变量之前的实现为静态生命周期，当初没有实现为运行时是因为实现比较复杂，而对于常规C++的语法来说静态生命周期的变量会随着模板的不同实例化而产生不同的实例。而恰巧用户的代码用了非常大量的模板实例化，导致该类型的静态变量的数量急剧膨胀，远远超出了地址范围所能接受的大小。

使用动态实现会复杂一些，并且还是会多少有些性能开销，所以想要尝试在编译期间处理掉这个问题。最后想到的方案是在生成代码的时候，所有模板的普通实例化都指向同一个实例。

一般情况下，这个方案是会有问题的，不同类型的实例化互相调用的情况下（算是一种特殊的递归），每个类型的handle都是独自的地址空间，如果把它们强行关联起来，会导致对一个地址空间额外做了很多操作，导致结果错误。

但是我们是针对特定用途的数据类型的方案，还是有可行性的，觉得还是很有意思的一个想法也因此写下本文记录，不过距离考虑方案的时候有些久远，有些细节没法很详细，暂且作为一个思路记录。

可行性

handle的使用场景确保了最终的可行性。当某个地址的handle要初始化，或者要被占用的时候，该地址的handle必须是一个不被占用的变量，handle是一个绑定到硬件的类型，如果当前地址的handle还没有被用完就被再次初始化，很大概率会造成结果错误。因此该类型在递归的场景下是会出现问题的，如果该类型没被用完就递归到这里再次初始化，那么会导致值错误的被写入，导致结果全部错误。因此我们可以扩大约束，认为handle类型的变量应该是一个递归后无法保证正确值的语义。

由于handle特殊的性质，导致没有递归的使用场景，并且handle类型也和模板参数没有关系，因此并没有减少使用上的灵活性。

代码实现部分

在clang的codegen部分，emit静态变量的时候添加处理。

1. 如果变量所在函数不是模板实例化 ⇒ 走常规静态变量的处理流程
2. 如果是模板实例化
   1. 根据type name + var name，构建一个新的完整变量名，用于作为唯一id
   2. 通过变量在代码中的location来区分是否为偏特化。如果该location找不到这个变量名，那么代表这个变量是偏特化，那么我们无需特殊处理。偏特化的情况下，执行的逻辑与常规的模板实例化不同，大部分情况都是代码的特殊处理，这里不修改偏特化的变量，避免引发奇怪的bug。
   3. 如果找到相关信息，那么直接返回现有的静态变量
   4. 如果没找相关信息，那么走常规创建静态变量的流程

int countDigits(int number) {
  return std::to_string(std::abs(number)).length();
}
// TypeName + VarName
std::string CodeGenModule::getLocalNameWithType(const VarDecl &D) {
  std::string localVarName;
  if (D.hasAttr<AsmLabelAttr>())
    localVarName = std::string(getMangledName(GlobalDecl(&D)));
  else
    localVarName = getStaticDeclName(*this, D);
  std::string typeName = D.getType().getAsString();
  const std::string varName = D.getName().str();
  // mangle rule: E N Name
  // E: prefix
  // N: sizeof(Name)
  // Name: var name in local
  // e.g. TestA => E5TestA
  int varNameMangleSize = varName.size() + 1 + countDigits(varName.size());
  localVarName = typeName + localVarName.substr(
      localVarName.size() - varNameMangleSize, varNameMangleSize);
  return localVarName;
}
llvm::Constant *CodeGenModule::getOrCreateScopedLocalVarDecl(
  const VarDecl &D, const Decl *CurFuncDecl,
  llvm::GlobalValue::LinkageTypes Linkage) {
    const FunctionDecl* FuncDecl = dyn_cast<FunctionDecl>(CurFuncDecl);
    if (FuncDecl == nullptr || !FuncDecl->isTemplateInstantiation()) {
      return getOrCreateStaticVarDecl(D, Linkage);
    }
    std::string curVarNameInLocal = getLocalNameWithType(D);
    auto Loc = CurFuncDecl->getLocation();
    const std::string varName = D.getName().str();
    // should create var when first find var in Loc
    if (ScopedLocalPool[Loc].find(varName) == ScopedLocalPool[Loc].end()) {
      ScopedLocalPool[Loc].insert(varName);
      return getOrCreateStaticVarDecl(D, Linkage);
    }
    // check var name in local
    // Each scoped local variable has a unique global name associated with
    // its enclosing function, but its name is also
    // unique within each instantiated function.
    for (const auto &[decl, addr] : StaticLocalDeclMap) {
      if(D.getLocation() != decl->getLocation()) {
        continue;
      }
      const VarDecl* varDecl = dyn_cast<VarDecl>(decl);
      std::string globalNameInLocal = getLocalNameWithType(*varDecl);
      for(auto &s : ScopedLocalPool[Loc]) {
        // If a variable has different types in different instantiations,
        // then it is dependent on the template. If the type remains the same,
        // then it is not template-dependent.
        // Types that depend on templates will not undergo this optimization,
        // because the instantiated types will differ.
        if(globalNameInLocal == curVarNameInLocal) {
          return addr;
        }
      }
    }
    return getOrCreateStaticVarDecl(D, Linkage);
}
void CodeGenFunction::EmitStaticVarDecl(const VarDecl &D,
                                      llvm::GlobalValue::LinkageTypes Linkage) {
                                      llvm::GlobalValue::LinkageTypes Linkage) {
  // Check to see if we already have a global variable for this
  // declaration.  This can happen when double-emitting function
  // bodies, e.g. with complete and base constructors.
	llvm::Constant *addr = CGM.getOrCreateStaticVarDecl(D, Linkage);
  CharUnits alignment = getContext().getDeclAlign(&D);
	...
}

测试代码

对应的测试代码，针对不同情况的处理。

namespace testnm {

template<typename T>
__device__ void general_template(void **addr) {
    __attribute__((scoped_local)) int originA, originB, originC;
    addr[0] = &originA;
    addr[1] = &originA;
    addr[2] = &originA;
    printf("originA:%d %p\\n", originA, &originA);
    printf("originB:%d %p\\n", originB, &originB);
    printf("originC:%d %p\\n", originC, &originC);
}

// overload with different func args
template<typename T>
__device__ void general_template(T n, void **addr) {
    __attribute__((scoped_local)) int originA;
    printf("templateArgA:%d %p\\n", originA, &originA);
}

// specialization
template<>
__device__ void general_template<double>(void **addr) {
    __attribute__((scoped_local)) int originA;
    addr[0] = &originA;
    printf("doubleX:%d %p\\n", originA, &originA);
}

// overload with different template args
template<typename T1, typename T2>
__device__ void general_template(void **addr) {
    __attribute__((scoped_local)) int originA;
    addr[0] = &originA;
    printf("doubleTemplateA:%d %p\\n", originA, &originA);
}

template<typename T>
__device__ void run_test_another(void **addr) {
    __attribute__((scoped_local)) int originA;
    addr[0] = &originA;
    printf("another A:%d %p\\n", originA, &originA);
}
}

namespace another_nm {
    template<typename T>
    __device__ void general_template(void **addr) {
        __attribute__((scoped_local)) int originA, originB, originC, originD;
        addr[0] = &originA;
        addr[1] = &originB;
        addr[2] = &originC;
        addr[3] = &originD;
        printf("another nm:%p\\n", originA, &originA);
        printf("another nm:%d %p\\n", originB, &originB);
        printf("another nm:%d %p\\n", originC, &originC);
        printf("another nm:%d %p\\n", originD, &originD);
    }
}

__device__ int check_addr_eq(void **a, void **b, int size) {
    int count = 0;
    for (size_t i = 0; i < size; i++) {
        if (a[i] != b[i]) {
            count++;
        }
    }
    return count;
}

__device__ int check_addr_ne(void **a, void **b, int size) {
    int count = 0;
    for (size_t i = 0; i < size; i++) {
        if (a[i] == b[i]) {
            count++;
        }
    }
    return count;
}

__device__ void check_optimized(int *errnum, void **baseP) {
    void **updateP = (void**)malloc(4 * sizeof(void *));
    // instantiation
    testnm::general_template<int>(baseP);
    testnm::general_template<float>(updateP);
    errnum += check_addr_eq(baseP, updateP, 3);
    testnm::general_template<unsigned int>(updateP);
    errnum += check_addr_eq(baseP, updateP, 3);
    testnm::general_template<char>(updateP);
    errnum += check_addr_eq(baseP, updateP, 3);
    testnm::general_template<int>(updateP);
    errnum += check_addr_eq(baseP, updateP, 3);
    // specialization
    testnm::general_template<double>(updateP);
    errnum += check_addr_ne(baseP, updateP, 1);
    // overload with different  func args
    testnm::general_template<int>(4, updateP);
    errnum += check_addr_ne(baseP, updateP, 1);
    // overload with different template args
    testnm::general_template<int, float>(updateP);
    errnum += check_addr_ne(baseP, updateP, 1);
    // same var name in different name func
    testnm::run_test_another<char>(updateP);
    errnum += check_addr_ne(baseP, updateP, 1);
    // same var name in different namespace func
    another_nm::general_template<int>(updateP);
    errnum += check_addr_ne(baseP, updateP, 4);
}

template<typename T>
__device__ void test_template_local(void **addr) {
    __scoped_local__ T originA;
    printf("test_template_local %p\\n", &originA);
}

template<typename T>
struct RS {
    T a;
};

template<typename T>
inline __attribute__((always_inline))__device__ void test_template_structure(void **addr) {
    __scoped_local__ RS<T> originA;
    printf("test_template_structure %p\\n", &originA);
}

__device__ void check_should_not_opt(int *errnum, void **baseP) {
    test_template_local<int>(baseP);
    test_template_local<float>(baseP);
    test_template_structure<int>(baseP);
    test_template_structure<float>(baseP);
}
__global__ void test_scoped_local() {
    if(threadIdx != 0) return;
    void **baseP = (void**)malloc(4 * sizeof(void *));
    int errnum = 0;
    check_optimized(&errnum, baseP);
    check_should_not_opt(&errnum, baseP);
}