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

1 應用背景

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

2 束內規約

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

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

這個設備函數可以求出當前線程所在線程束的指定變量的規約和，原理涉及洗牌指令的計算邏輯，不再贅述。
當矩陣寬度 hidden_units 較小時，通常可以使用一個 warp 處理一行數據，一個 block 內可以處理多行數據，筆者給出具體的核函數如下：

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

3 塊內規約

塊內規約，就是在一個線程塊內求規約值，通常塊內規約會通過束內規約來實現，以下是 Nvidia 提供的一個塊內規約設備函數。

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

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

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

除了必要的待規約變量、block_size 以外，還需要傳入一個計算函數，筆者給出了示例 SumOp。
當矩陣寬度 hidden_units 較大時，通常可以使用一個 block 處理一行數據，筆者給出具體的核函數如下：

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 處理一行數據，總共需要 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 的四個元素，以此類推，這個計算過程是沒有并行的。循環計算一輪后，對 block 內每個線程的 val 進行塊內規約就可以得到一行元素的規約和。

4 向量化數據提升訪存帶寬

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

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

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

4.1 pack 后的束內規約示例代碼

將 matrix2DWarpReduceSum 改寫為 pack 版的核函數也很簡單，計算思路都是一致的，只不過原來一次訪問一個元素，現在一次訪問一個 pack 的元素，在執行核函數之前筆者加了一個斷言，保證 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 數量
  uint32_t num_packs = hidden_units / pack_size;
  assert(hidden_units % pack_size == 0);
  matrix2DWarpReduceSumPack<<>>(d_x, d_y, hidden_units, num_packs);
}

核函數內部就一句核心代碼，將 const T* 指針轉換成 const Packed*。

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

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

4.2 pack 后的塊內規約示例代碼

matrix2DBlockReduceSumPack 核函數的實現就更簡單了，直接上代碼。

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 小結

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