带有动态共享内存的模板化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()）。