束內(nèi)規(guī)約與塊內(nèi)規(guī)約問題

寫在前面：規(guī)約問題在 CUDA 編程中應用非常廣泛，筆者最近在研究 Faster Transformer 源碼，趁此機會結(jié)合 Nivida 官方的代碼對規(guī)約手段進行總結(jié)。

1 應用背景

關于規(guī)約的定義，相信能讀到這篇文章的讀者都不陌生，筆者在早期的文章中也介紹過一些規(guī)約方法，基本思想都是折半規(guī)約，主要應用于較大元素規(guī)模的向量規(guī)約，有興趣的讀者可以移步【CUDA編程】CUDA編程中的并行規(guī)約問題。
本文要介紹的規(guī)約場景與之前有所不同，主要應用于矩陣規(guī)約，也就是說本文假設的輸入變量的維度是 2 維的，形狀為 [batch_size, hidden_units]，規(guī)約之后的輸出變量形狀為 [batch_size, ]。
接下來，本文將以規(guī)約求和為例介紹兩種規(guī)約方式：束內(nèi)規(guī)約、塊內(nèi)規(guī)約。

2 束內(nèi)規(guī)約

束內(nèi)規(guī)約，也就是在一個線程束內(nèi)對某個變量進行規(guī)約。我們知道 CUDA 架構(gòu)下指令是以線程束（相鄰的 32 個線程）為基本單元執(zhí)行的，線程束內(nèi)也可以通過束內(nèi)洗牌指令進行通信，所以這提供了一個很好的束內(nèi)規(guī)約思路。下面是 Nvidia 提供的基礎的一個規(guī)約設備函數(shù)。

template <typename T>
__inline__ __device__
T warpReduceSum(T val)
{
  for(int mask = 16; mask > 0; mask >>= 1)
    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
  return val;
}

這個設備函數(shù)可以求出當前線程所在線程束的指定變量的規(guī)約和，原理涉及洗牌指令的計算邏輯，不再贅述。
當矩陣寬度 hidden_units 較小時，通?？梢允褂靡粋€ warp 處理一行數(shù)據(jù)，一個 block 內(nèi)可以處理多行數(shù)據(jù)，筆者給出具體的核函數(shù)如下：

// 一個 warp 處理一行數(shù)據(jù)
template<typename T>
__global__ void matrix2DWarpReduceSum(const T* inp, T*out, const uint32_t hidden_units) {
    uint32_t tid = threadIdx.x;
    uint32_t lane_id = tid % 32;
    uint32_t warp_id = tid / 32;
    uint32_t warp_num = blockDim.x / 32;
    uint32_t offset = blockIdx.x * warp_num * hidden_units + warp_id * hidden_units;
    T val = 0.0f;
    for (uint32_t i=lane_id; i32) {
        val += inp[offset + i];
    }
    __syncwarp();
    T warpSum;
    warpSum = warpReduceSum(val);
    if (lane_id == 0) {
      out[blockIdx.x * warp_num + warp_id] = warpSum;
    }
}

template<typename T>
void launchMatrix2DWarpReduceSum(const T* d_x, T* d_y, const uint32_t batch_size, const uint32_t hidden_units) {
  constexpr uint32_t warp_num = BLOCK_SIZE / 32;
  uint32_t gird_size = (batch_size - 1) / (warp_num) + 1;
  matrix2DWarpReduceSum<<>>(d_x, d_y, hidden_units);
}

先確定 block_size，這里筆者直接取 128，由于是一個 warp 處理一行數(shù)據(jù)，所以一個 block 可以處理 warp_num 行數(shù)據(jù)，總共需要 grid_size 個 block。
核函數(shù)內(nèi)部首先計算當前線程所在的 warp 編號 warp_id 用來定位當前處理元素在哪一行，然后確定線程在 warp 內(nèi)的編號 lane_id 用來定位該線程具體處理那些元素。由于矩陣寬度 hidden_units 實際肯定還是比 32 大的，所以不可能說一個線程只處理一個元素，因此每個線程會處理多個元素，步長為 32，例如當 hidden_units 為 128 時，lane_id = 0 的線程將處理位置為 0、32、64、96 的四個元素，lane_id = 1 的線程將處理位置為 1、33、65、97 的四個元素，以此類推，這個計算過程是沒有并行的。循環(huán)計算一輪后，對線程束內(nèi)每個線程的 val 進行束內(nèi)規(guī)約就可以得到一行元素的規(guī)約和。

3 塊內(nèi)規(guī)約

塊內(nèi)規(guī)約，就是在一個線程塊內(nèi)求規(guī)約值，通常塊內(nèi)規(guī)約會通過束內(nèi)規(guī)約來實現(xiàn)，以下是 Nvidia 提供的一個塊內(nèi)規(guī)約設備函數(shù)。

template <typename T>
__inline__ __device__
T blockReduceSum(T val)
{
  static __shared__ T shared[32]; 
  int lane = threadIdx.x & 0x1f; 
  int wid = threadIdx.x >> 5;  

  val = warpReduceSum(val);

  if(lane == 0)
    shared[wid] = val;
  __syncthreads();
  
  val = (threadIdx.x < (blockDim.x >> 5 )) ? shared[lane] : (T)0.0f;
  val = warpReduceSum(val);
  return val;
}

規(guī)約思路分為兩步，首先通過束內(nèi)規(guī)約求出當前線程所在 warp 的規(guī)約值，存入 shared 中，然后把 warpSum 賦值給 threadIdx.x 小于 32 的線程內(nèi)的變量 val，這 32 個線程正好也在一個線程束內(nèi)，然后再執(zhí)行一次束內(nèi)規(guī)約就得到塊內(nèi)規(guī)約值，計算思路非常巧妙。
另外針對塊內(nèi)規(guī)約的問題，官方 cub 庫其實提供了 API，開發(fā)者可以導入頭文件 cub/cub.cuh 后直接使用，注意低版本的 cuda 不支持此 API。我們來看下 API 的調(diào)用方式。

#include 

template<typename T>
struct SumOp {
  __device__ __forceinline__ T operator()(const T& a, const T& b) const { return a + b; }
};

template<template<typename> class ReductionOp, typename T, int block_size>
__inline__ __device__ T BlockAllReduce(T val) {
  typedef cub::BlockReduce BlockReduce;
  __shared__ typename BlockReduce::TempStorage temp_storage;
  __shared__ T result_broadcast;
  T result = BlockReduce(temp_storage).Reduce(val, ReductionOp());
  if (threadIdx.x == 0) { result_broadcast = result; }
  __syncthreads();
  return result_broadcast;
}

除了必要的待規(guī)約變量、block_size 以外，還需要傳入一個計算函數(shù)，筆者給出了示例 SumOp。
當矩陣寬度 hidden_units 較大時，通?？梢允褂靡粋€ block 處理一行數(shù)據(jù)，筆者給出具體的核函數(shù)如下：

template<typename T>
__global__ void matrix2DBlockReduceSum(const T* inp, T*out, const uint32_t hidden_units) {
  T val = 0.0f;
  uint32_t offset = blockIdx.x * hidden_units;
  for (uint32_t i=threadIdx.x; i(val);
  if (threadIdx.x == 0) {
    out[blockIdx.x] = blockSum;
  }
}

template<typename T>
void launchMatrix2DBlockReduceSum(const T* d_x, T* d_y, const uint32_t batch_size, const uint32_t hidden_units) {
  uint32_t gird_size = batch_size;
  matrix2DBlockReduceSum<<>>(d_x, d_y, hidden_units);
}

同樣，block_size 這里筆者直接取 128，由于是一個 block 處理一行數(shù)據(jù)，總共需要 batch_size 個 block。
由于矩陣寬度 hidden_units 實際肯定還是比 block_size 大的，所以不可能說一個線程只處理一個元素，因此每個線程會處理多個元素，步長為 block_size，例如當 hidden_units 為 512 時，lane_id = 0 的線程將處理位置為 0、128、256、384 的四個元素，lane_id = 1 的線程將處理位置為 1、129、257、385 的四個元素，以此類推，這個計算過程是沒有并行的。循環(huán)計算一輪后，對 block 內(nèi)每個線程的 val 進行塊內(nèi)規(guī)約就可以得到一行元素的規(guī)約和。

4 向量化數(shù)據(jù)提升訪存帶寬

使用向量化操作能夠提升內(nèi)存讀寫的帶寬，而 CUDA 里也提供了一系列數(shù)據(jù)類型來支持向量化操作，如 float2、float4，就是將 2 個或 4 個 float 數(shù)據(jù)作為一個整體。為了增加代碼的復用性，筆者這里封裝了一個 Packed 數(shù)據(jù)結(jié)構(gòu)，用于對不同的數(shù)據(jù)類型進行打包。

template <typename T, int pack_size>
struct alignas(sizeof(T) * pack_size) Packed
{
    __device__ Packed()
    {
        // do nothing
    }
    union
    {
        T elem[pack_size]; // 這里聯(lián)合體只有一個成員，為了方便后期擴展
    };
};

結(jié)構(gòu)體內(nèi)有一個 elem 數(shù)組變量，整個結(jié)構(gòu)的內(nèi)存對齊設置為 sizeof(T) * pack_size，說白了其實就是把 pack_size 個 T 類型的數(shù)據(jù)“捆綁”在一起組成一個新的數(shù)據(jù)結(jié)構(gòu)，讀寫內(nèi)存的時候只需要一次讀寫就可以讀 pack_size 個數(shù)據(jù)，目的是減小內(nèi)存讀寫次數(shù)。
那么這個 pack_size 能不能無限大呢？顯然不能，CUDA 里最大支持 128 bit 的訪問粒度，也就是說對于 float 類型（占 4 個字節(jié)，32 bit），一次最多讀寫 4 個，也就是說 float 的 pack_size 最多取到 4，本文筆者的示例代碼中數(shù)據(jù)類型都以 float 為例，pack_size 取 4。

4.1 pack 后的束內(nèi)規(guī)約示例代碼

將 matrix2DWarpReduceSum 改寫為 pack 版的核函數(shù)也很簡單，計算思路都是一致的，只不過原來一次訪問一個元素，現(xiàn)在一次訪問一個 pack 的元素，在執(zhí)行核函數(shù)之前筆者加了一個斷言，保證 hidden_units 能夠被 pack_size 整除，具體代碼如下。

template <int pack_size, typename T>
__global__ void matrix2DWarpReduceSumPack(const T* d_x, T* d_y, const uint32_t hidden_units, const uint32_t num_packs) {
  const uint32_t warp_id = threadIdx.x / 32;
  const uint32_t lane_id = threadIdx.x & 0x1f;
  const uint32_t warp_num = blockDim.x / 32;
  const uint32_t offset = blockIdx.x * warp_num * hidden_units + warp_id * hidden_units;
  const Packed* buf = reinterpret_cast<const Packed*>(d_x + offset);
  Packed pack;
  T val = 0.0f;
  for (uint32_t pack_id=lane_id; pack_id32) {
    pack = buf[pack_id];
    for (uint32_t i=0; i(val);
  if (lane_id == 0) {
    d_y[blockIdx.x * warp_num + warp_id] = warpSum;
  }
}

template<typename T>
void launchMatrix2DWarpReduceSumPack(const T* d_x, T* d_y, const uint32_t batch_size, const uint32_t hidden_units) {
  constexpr uint32_t warp_num = BLOCK_SIZE / 32;
  uint32_t gird_size = (batch_size - 1) / (warp_num) + 1;
  constexpr uint32_t pack_size = 4;
  // 一行元素的 pack 數(shù)量
  uint32_t num_packs = hidden_units / pack_size;
  assert(hidden_units % pack_size == 0);
  matrix2DWarpReduceSumPack<<>>(d_x, d_y, hidden_units, num_packs);
}

核函數(shù)內(nèi)部就一句核心代碼，將 const T* 指針轉(zhuǎn)換成 const Packed*。

const Packed* buf = reinterpret_cast<const Packed*>(d_x + offset);

然后用 pack_id 索引一次取一個 pack 的數(shù)據(jù)，注意這里對 pack 索引的時候不要寫錯了。跟前面一樣，相鄰的線程處理相鄰的 pack 數(shù)據(jù)，這是為了全局內(nèi)存的合并訪問。加法計算次數(shù)還是那么多次，因為 Packed 結(jié)構(gòu)體并不能直接參與計算，還是要用 elem 里面的元素計算，這個核函數(shù)也就節(jié)省了訪存次數(shù)而已。

4.2 pack 后的塊內(nèi)規(guī)約示例代碼

matrix2DBlockReduceSumPack 核函數(shù)的實現(xiàn)就更簡單了，直接上代碼。

template <int pack_size, typename T>
__global__ void matrix2DBlockReduceSumPack(const T* d_x, T* d_y, const uint32_t hidden_units, const uint32_t num_packs) {
  T val = 0.0f;
  uint32_t offset = blockIdx.x * hidden_units;
  const Packed* buf = reinterpret_cast<const Packed*>(d_x + offset);
  Packed pack;
  for (uint32_t pack_id=threadIdx.x; pack_idfor (uint32_t i=0; i(val);
  if (threadIdx.x == 0) {
    d_y[blockIdx.x] = blockSum;
  }
}

template<typename T>
void launchMatrix2DBlockReduceSumPack(const T* d_x, T* d_y, const uint32_t batch_size, const uint32_t hidden_units) {
  uint32_t gird_size = batch_size;
  constexpr uint32_t pack_size = 4;
  assert(hidden_units % pack_size == 0);
  uint32_t num_packs = hidden_units / pack_size;
  matrix2DBlockReduceSumPack<<>>(d_x, d_y, hidden_units, num_packs);
}

5 小結(jié)

在深度學習算子的開發(fā)過程中，規(guī)約是一個非常常見的場景，以 Softmax 為例就有 reduceMax 和 reduceSum 的應用，本文給出了兩種規(guī)約實現(xiàn)方式，可供讀者參考使用。實際開發(fā)過程中，規(guī)約計算一般是隱藏在其他 kernel 中的，并不會奢侈到單獨寫個規(guī)約 kernel，所以要求開發(fā)人員領會思路活學活用。