具有动态共享内存的模板化 CUDA 内核

Question

我想在一个程序中使用动态分配的共享内存调用模板化CUDA 内核的不同实例。 我第一个天真的方法是这样写：

template<typename T>
__global__ void kernel(T* ptr)
{
  extern __shared__ T smem[];
  // calculations here ...                                                                                                                                          
}

template<typename T>
void call_kernel( T* ptr, const int n )
{
  dim3 dimBlock(n), dimGrid;
  kernel<<<dimGrid, dimBlock, n*sizeof(T)>>>(ptr);
}

int main(int argc, char *argv[])
{
  const int n = 32;
  float *float_ptr;
  double *double_ptr;
  cudaMalloc( (void**)&float_ptr, n*sizeof(float) );
  cudaMalloc( (void**)&double_ptr, n*sizeof(double) );

  call_kernel( float_ptr, n );
  call_kernel( double_ptr, n ); // problem, 2nd instantiation

  cudaFree( (void*)float_ptr );
  cudaFree( (void*)double_ptr );
  return 0;
}

但是，无法编译此代码。 nvcc 给了我以下错误信息：

main.cu(4): error: declaration is incompatible with previous "smem"
(4): here
          detected during:
            instantiation of "void kernel(T *) [with T=double]"
(12): here
            instantiation of "void call_kernel(T *, int) [with T=double]"
(24): here

我知道我遇到了名称冲突，因为共享内存被声明为 extern。 然而，据我所知，如果我想在运行时定义它的大小，就没有办法解决这个问题。

所以，我的问题是：是否有任何优雅的方式来获得所需的行为？ 优雅我的意思是没有代码重复等。

Answer 1

动态分配的共享内存实际上只是一个大小（以字节为单位）和一个为内核设置的指针。 所以这样的事情应该有效：

替换这个：

extern __shared__ T smem[];

有了这个：

extern __shared__ __align__(sizeof(T)) unsigned char my_smem[];
T *smem = reinterpret_cast<T *>(my_smem);

您可以在编程指南中看到重新转换动态分配的共享内存指针的其他示例，这些示例可以满足其他需求。

编辑：更新我的答案以反映@njuffa 的评论。

Answer 2

_{（@RobertCrovella答案的变体）}

NVCC 不愿意接受两个同名但类型不同的extern __shared__数组——即使它们从来不在彼此的范围内。 我们需要通过让我们的模板实例都使用相同类型的共享内存来满足 NVCC，同时让使用它们的内核代码看到它喜欢的类型。

所以我们替换这个指令：

extern __shared__ T smem[];

有了这个：

auto smem = shared_memory_proxy<T>();

在哪里：

template <typename T>
__device__ T* shared_memory_proxy()
{
    // do we need an __align__() here? I don't think so...
    extern __shared__ unsigned char memory[];
    return reinterpret_cast<T*>(memory);
}

在一些设备端代码包含文件中。

好处：

• 使用现场的单衬。
• 更简单的语法要记住。
• 关注点分离 - 阅读内核的人不必考虑为什么他/她会看到extern ，或对齐说明符，或重新解释强制转换等。

编辑：这是作为我的CUDA 内核作者的工具头文件库的一部分实现的： shared_memory.cuh （它被命名为shared_memory::dynamic::proxy() ）。