CUDA編程分布式共享內(nèi)存

Distributed Shared Memory

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

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

CUDA提供了一種訪問分布式共享內(nèi)存的機(jī)制，應(yīng)用程序可以從利用它的功能中獲益。讓我們看看一個(gè)簡(jiǎn)單的直方圖計(jì)算，以及如何使用線程塊集群在GPU上優(yōu)化它。 計(jì)算直方圖的標(biāo)準(zhǔn)方法是在每個(gè)線程塊的共享內(nèi)存中進(jìn)行計(jì)算，然后執(zhí)行全局內(nèi)存原子 。

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

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

#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]);
  }
}

上面的內(nèi)核可以在運(yùn)行時(shí)啟動(dòng)，集群大小取決于所需的分布式共享內(nèi)存的數(shù)量。如果直方圖足夠小，可以容納一個(gè)塊的共享內(nèi)存，用戶可以啟動(dòng)集群大小為1的內(nèi)核。下面的代碼片段展示了如何根據(jù)共享內(nèi)存需求動(dòng)態(tài)啟動(dòng)集群內(nèi)核。

// 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);
}