具有動態共享內存的模板化 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() ）。