Distributed Shared Memory

計算能力9.0中引入的線程塊集群為線程塊集群中的線程提供了訪問集群中所有參與線程塊的共享內存的能力。這種分區共享內存稱為 Distributed Shared Memory，對應的地址空間稱為分布式共享內存地址空間。屬于線程塊集群的線程可以在分布式地址空間中讀、寫或執行原子操作，而不管該地址屬于本地線程塊還是遠程線程塊。無論內核是否使用分布式共享內存，共享內存大小規格(靜態的或動態的)仍然是每個線程塊。分布式共享內存的大小就是每個集群的線程塊數量乘以每個線程塊的共享內存大小。

訪問分布式共享內存中的數據需要所有線程塊存在 。用戶可以使用cluster .sync()從Cluster Group API中保證所有線程塊已經開始執行。用戶還需要確保在線程塊退出之前完成所有分布式共享內存操作。

CUDA提供了一種訪問分布式共享內存的機制，應用程序可以從利用它的功能中獲益。讓我們看看一個簡單的直方圖計算，以及如何使用線程塊集群在GPU上優化它。 計算直方圖的標準方法是在每個線程塊的共享內存中進行計算，然后執行全局內存原子 。

這種方法的一個限制是共享內存容量。一旦直方圖容器不再適合共享內存，用戶就需要直接計算直方圖，從而計算全局內存中的原子。對于分布式共享內存，CUDA提供了一個中間步驟，根據直方圖桶的大小，直方圖可以直接在共享內存、分布式共享內存或全局內存中計算。

下面的CUDA內核示例展示了如何在共享內存或分布式共享內存中計算直方圖，具體取決于直方圖箱的數量。

#include 

// Distributed Shared memory histogram kernel
__global__ void clusterHist_kernel(int *bins, const int nbins, const int bins_per_block, const int *__restrict__ input,
                                   size_t array_size)
{
  extern __shared__ int smem[];
  namespace cg = cooperative_groups;
  int tid = cg::this_grid().thread_rank();

  // Cluster initialization, size and calculating local bin offsets.
  cg::cluster_group cluster = cg::this_cluster();
  unsigned int clusterBlockRank = cluster.block_rank();
  int cluster_size = cluster.dim_blocks().x;
 
  for (int i = threadIdx.x; i < bins_per_block; i += blockDim.x)
  {
    smem[i] = 0; //Initialize shared memory histogram to zeros
  }

  // cluster synchronization ensures that shared memory is initialized to zero in
  // all thread blocks in the cluster. It also ensures that all thread blocks
  // have started executing and they exist concurrently.
  cluster.sync();

  for (int i = tid; i < array_size; i += blockDim.x * gridDim.x)
  {
    int ldata = input[i];

    //Find the right histogram bin.
    int binid = ldata;
    if (ldata < 0)
      binid = 0;
    else if (ldata >= nbins)
      binid = nbins - 1;

    //Find destination block rank and offset for computing
    //distributed shared memory histogram
    int dst_block_rank = (int)(binid / bins_per_block);
    int dst_offset = binid % bins_per_block;

    //Pointer to target block shared memory
    int *dst_smem = cluster.map_shared_rank(smem, dst_block_rank);

    //Perform atomic update of the histogram bin
    atomicAdd(dst_smem + dst_offset, 1);
  }

  // cluster synchronization is required to ensure all distributed shared
  // memory operations are completed and no thread block exits while
  // other thread blocks are still accessing distributed shared memory
  cluster.sync();

  // Perform global memory histogram, using the local distributed memory histogram 
  int *lbins = bins + cluster.block_rank() * bins_per_block;
  for (int i = threadIdx.x; i < bins_per_block; i += blockDim.x)
  {
    atomicAdd(&lbins[i], smem[i]);
  }
}

上面的內核可以在運行時啟動，集群大小取決于所需的分布式共享內存的數量。如果直方圖足夠小，可以容納一個塊的共享內存，用戶可以啟動集群大小為1的內核。下面的代碼片段展示了如何根據共享內存需求動態啟動集群內核。

// Launch via extensible launch
{
  cudaLaunchConfig_t config = {0};
  config.gridDim = array_size / threads_per_block;
  config.blockDim = threads_per_block;

  // cluster_size depends on the histogram size.
  // ( cluster_size == 1 ) implies no distributed shared memory, just thread block local shared memory
  int cluster_size = 2; // size 2 is an example here
  int nbins_per_block = nbins / cluster_size;

  //dynamic shared memory size is per block. 
  //Distributed shared memory size =  cluster_size * nbins_per_block * sizeof(int)
  config.dynamicSmemBytes = nbins_per_block * sizeof(int);

  CUDA_CHECK(::cudaFuncSetAttribute((void *)clusterHist_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, config.dynamicSmemBytes));

  cudaLaunchAttribute attribute[1];
  attribute[0].id = cudaLaunchAttributeClusterDimension;
  attribute[0].val.clusterDim.x = cluster_size; 
  attribute[0].val.clusterDim.y = 1;
  attribute[0].val.clusterDim.z = 1;

  config.numAttrs = 1;
  config.attrs = attribute;

  cudaLaunchKernelEx(&config, clusterHist_kernel, bins, nbins, nbins_per_block, input, array_size);
}