本文為CUDA并行規約系列文章的下篇，本文介紹了5種并行規約算法的實現，并從硬件的角度對它們進行分析和優化，最終給出一個開箱即用的模板代碼及其使用示例。

本文首發于本人微信公眾號，鏈接：
https://mp.weixin.qq.com/s/XSXjwPZBwZCE1uy8gRZqWg

摘要

本文為CUDA并行規約系列文章的下篇，本文介紹了5種并行規約算法的實現，并從硬件的角度對它們進行分析和優化，最終給出一個開箱即用的模板代碼及其使用示例。

勘誤

首先是一個勘誤，在上篇中存在一個拼寫錯誤，線程束的正確單詞是Warp而不是Wrap，非常抱歉給讀者朋友們造成了誤解。

寫在前面

這是本系列文章的下篇。上篇介紹了一些CUDA并行規約優化涉及到的GPU硬件知識，并給出了兩種并不完美的實現。

本文將接著介紹剩下的五種實現，并給出一個開箱即用的CUDA并行規約模板。

算法具體實現（下篇）

V2: Sequential Addressing

先簡單回顧一下，在上篇的最后，我們發現V1版本的實現存在Bank Conflict的問題，具體表現為，當\(s=32\)時，\(T_0\)會訪問\(A_0\)，\(T_1\)會訪問\(A_{32}\)，\(T_2\)會訪問\(A_{64}\).....，造成一個Warp里所有線程都同時訪問Bank 0，導致這些訪問被串行化，嚴重影響性能。

造成這一問題的根本原因是：同一個Warp里的線程，它們訪問的地址存在可變的Stride

Bank機制設計之初的預期就是一個Warp里32個線程訪問連續的32個地址，而不是訪問分散在各處的地址。

明確了這一點之后，優化方法就很明確了：我們只需要讓每個線程都負責其threadIdx對應的內存地址的規約即可，比如\(T_0\)就只管\(A_0\)，\(T_1\)就只管\(A_1\)，這樣就不會出現Bank Conflict了。

由于要防止Warp Divergence，所以第一輪循環只有前一半線程在工作，以\(n=512\)為例，第一輪循環只有\(T_0\)到\(T_{255}\)在工作，那么很顯然的，\(T_0\)就需要規約\(A_0\)和\(A_{256}\)，\(T_1\)需要規約\(A_1\)和\(A_{257}\)，以此類推。這個過程如下圖所示：

在代碼實現上也不難，只需要做如下改變即可：

一個容易混淆的點

可能會有讀者感到疑惑：總會有一次循環，\(T_0\)會訪問\(A_0\)和\(A_{32}\)，這怎么能算解決了Bank Conflict呢？

這里需要明確的一點是：Bank Conflict是發生在線程間的，而不是一個線程內的指令間的。

從指令執行的角度來分析，就以\(A_0 = A_0 + A_{32}\)為例，這條語句會分4步完成，分別是：讀取\(A_0\)，讀取\(A_{32}\)，計算\(A_0 + A_{32}\)，寫入\(A_0\)。

對\(T_0\)而言，讀取\(A_0\)和讀取\(A_{32}\)一定是先后發生的，不會并行，所以也就不存在Bank Conflict。

這里我們再從線程間執行的角度來看看，在\(T_0\)讀取\(A_0\)時，\(T_1\)在讀取\(A_1\)，這不會發生沖突。但是如果是V1的實現，在\(T_0\)讀取\(A_0\)時，\(T_1\)可能會同時在讀取\(A_{32}\)，這就導致了Bank Conflict。

優化效果

根據NVIDIA的數據，這一優化將性能提升了接近1倍，也是相當可觀的。

V3: First Add During Load

V2版本已經把硬件上踩過的坑基本都填完了，接下來就是一些細節上的優化。其中有一個思路就是在把數據加載到共享內存的階段做一些預處理。

比如，我們可以先把\(A_0\)和\(A_1\)規約了，然后存儲到\(A_0\)中，把\(A_2\)和\(A_3\)規約了，存儲到\(A_1\)中。這樣就只需要啟動一半的線程，以幾乎一半的時間執行完所有的任務。

這一想法還可以進一步推廣，如果在加載階段就提前規約4個元素，那就可以把時間壓縮到\(\frac{1}{4}\)。

這一思想就是所謂的算法級聯，即結合并行執行和串行執行的策略。這部分更詳細的會在后面V6版本里分析。

由于這里NVIDIA的實現和我們的問題定義有沖突，這里就不詳細展開了，僅貼出NVIDIA的實現供參考。

實測的數據也驗證了這一理論的正確性，時間確實縮短了接近一半。

V4: Unroll the Last Warp

V3版本計算出的有效帶寬為17GB/s，遠沒有達到硬件的上限，所以有理由懷疑這里還存在指令執行上的瓶頸。

觀察之后發現，這里的部分循環指令也許可以優化掉。在上篇中，我們提到過：在同一個Warp里，指令執行可以認為是同步的，所以我們可以在最后32個線程工作時去掉__syncthreads。

此外，既然都已經知道這里是最后32個線程在工作了，那循環也不必需要了，可以直接硬編碼寫\(T_i = A_i + A_{i+32}\)，\(T_i = A_i + A_{i+16}\)，......，最終的修改如下圖所示

為什么不需要線程同步了？

因為這里可以保證進入warpReduce的線程是\(T_0\)到\(T_{31}\)，即一個Warp內的線程。

根據之前的內容，\(T_0\)在執行函數第二行\(A_0 += A_{16}\)時，\(T_{16}\)是一定也在執行第二行，換言之，\(T_{16}\)是一定執行完成了第一行的，這就保證了在執行取\(A_{16}\)這個指令時，\(A_{16}\)存儲的一定是規約了\(A_{16}\)和\(A_{48}\)之后的值。

那么還有一個疑問：這里\(T_{16}\)寫入了新的\(A_{16}\)的值是\(A_{16} + A_{32}\)，這不就破壞了整個規約過程了嗎？

事實上，這里\(T_{16}\)寫入了什么并不重要，這是基于兩點事實：

首先，\(T_{16}\)寫入\(A_{16}\)和\(T_0\)寫入\(A_0\)一定是在同一個時鐘周期內發生的，所以在一開始執行第二行時，\(T_0\)讀取到的\(A_{16}\)一定是沒被破壞的\(A_{16}\)，這確保了\(T_0\)執行的正確性；

其次，在執行完第一行之后，實際上只有\(T_0\)到\(T_{15}\)真正起作用了，后面16個線程只要不亂寫數據妨礙到前16個線程的工作，無論做什么都對結果沒影響。

所以理論上講，這里加個判斷，只讓\(tid < 16\)的線程執行，最終的結果都是正確的。這里沒有這么做主要是為了防止Warp Divergence。

優化效果

實驗數據表明，僅僅是對最后一個Warp做了循環展開，最終性能就優化了接近一倍。

V5: Completely Unrolled

既然只對最后一個Warp進行循環展開優化就這么明顯，那能不能再激進一些，把所有循環都丟掉呢？

答案是Yes。

因為一個線程塊內的線程數量上限是1024，并且我們要求線程數量是2的整數次冪，所以我們可以枚舉所有可能的線程數量，然后仿照之前展開最后一個循環的方法，針對每種情況直接寫出對應的展開后的代碼。

如果只是在參數列表里加一個blockSize，那么在運行時還是會經過多個無意義的if-else，影響執行效率。這里可以借助C++的模板機制，進一步地提升執行性能。具體的修改如下圖所示：

這里的blockSize是在調用時使用模板傳遞的，這里面的if會在編譯時就進行優化。

編譯器會使用Dead Code Elimination，根據blockSize刪掉不可達的代碼，所以最終編譯出的二進制里面是不會有這些紅色的if的，只有一串順序執行的指令。

但是模板要求我們在編譯階段就確定blockSize，在實際實現時是比較困難的，這個該如何解決呢？答案是在調用時使用switch來枚舉，簡單粗暴，如下圖所示：

優化效果

V6: Multiple Adds

理論分析

這里我們要從成本的角度來看待一下現有的方案。

和我們租用算力服務器一樣，成本可以用使用時間 * GPU核心數量來衡量，這里GPU核心數量可以認為是線程數量，兩者之間只有常數倍的差別。
我們假設這樣一個場景來計算成本：

• 只有1個Batch，共有\(N\)個元素
• 啟動\(P\)個線程來處理，在V5實現中，\(P=N\)

這里使用Brent定理：

其中\(t(n)\)為算法的估計實際執行時間；\(W(n)\)表示算法的總工作量，即一共執行多少次運算；\(p\)表示線程數量；\(T(n)\)表示算法在最理想的情況下，算法執行的最短用時

很顯然，這里\(p=N\)；由于最少也需要執行\(\log N\)步規約，所以\(T(n)=\log N\)；一通計算可以得到，\(W(n)=O(N)\)，這里具體計算過程見NVIDIA的PPT，這里不再贅述。

帶入公式可得，\(t(n)=O(\log N)\)

那么可以計算得到，V5版本算法的成本為\(O(N \log N)\)。

但是如果用一個線程串行處理所有數據，成本卻只有\(O(N)\)，也就是說，我們的并行算法的成本比純串行處理還要高。那么怎么把這個成本降下來呢？

這里成本變高的本質是：線程數量過多，使得每個線程的工作量太少，導致了整體成本的增大。

那么相對應的，我們可以通過減少線程數量來實現降本增效。具體來說，我們可以使用V3中提到的方法，即在加載數據的階段就提前做幾次規約。

那么具體應該做多少次呢？這里NVIDIA的PPT里給出的數據是應該做\(\log N\)次，如此優化之后，最終的成本能降到\(O(N)\)，可以和串行執行持平。

這里把串行和并行搭配使用的策略就是算法級聯。

實現方式

這里NVIDIA的實現和我們的問題定義相沖突，所以這里不再贅述了，后面在開箱即用的部分會解釋我們是如何處理這一沖突的。
下圖是V5到V6版本的變化以及V6的完整代碼實現。

優化效果

這里貼出7個版本的優化效果數據表格：

還能再優化嗎？

NVIDIA官方的PPT到這里就結束了。我們可以思考一下：V6版本還能進一步優化嗎？

也許還可以，比如用循環展開的技巧和最開始的while循環展開一下，這個就作為open issue供大家探討了。

開箱即用的實現

數據要求

• 輸入：\(m \times n\)矩陣按行優先展開成的一個向量\(x\)
• 輸出：\(m\)維向量\(out\)
• 要求線程塊內線程的數量（block_size） \(\leq 1024\)，并且為2的整數次冪
• 要求\(n\)一定要能整除block_size

NVIDIA官方的實現里似乎沒有batch的概念，感覺是想要把輸入的所有數據都規約到一個值，因此在V3和V6里面會有跨Block規約的情況。

我們這里就不采用這個方案，而是把\(n\)分為了若干個fold，最終的線程數量就是 \(\frac{n}{\text{fold_num}}\)，在一開始加載數據時就提前規約fold_num次。這個fold_num由調用方傳遞。

對于一個batch數量大于1024的情況，和數據數量不是2的整數次冪的情況，則需要調用CUDA的上層框架做處理了，這里暫時不考慮這些case。

代碼實現

這一實現使用了C++的模板特性，支持調用者自行選擇數據類型和規約函數。

template<typename T>
using ReduceFunc = T(*)(T, T);

template<typename T, ReduceFunc<T> reduce_func, size_t block_size>
__device__ void reduceWarp(volatile T *shared_memory, size_t tid) {
  if (block_size >= 64) shared_memory[tid] = reduce_func(shared_memory[tid], shared_memory[tid + 32]);
  if (block_size >= 32) shared_memory[tid] = reduce_func(shared_memory[tid], shared_memory[tid + 16]);
  if (block_size >= 16) shared_memory[tid] = reduce_func(shared_memory[tid], shared_memory[tid + 8]);
  if (block_size >= 8) shared_memory[tid] = reduce_func(shared_memory[tid], shared_memory[tid + 4]);
  if (block_size >= 4) shared_memory[tid] = reduce_func(shared_memory[tid], shared_memory[tid + 2]);
  if (block_size >= 2) shared_memory[tid] = reduce_func(shared_memory[tid], shared_memory[tid + 1]);
}

/**
 * Grid(m, 1, 1)
 * Block(block_size, 1, 1)
 * @tparam T 需要Reduce的數據類型
 * @tparam reduce_func Reduce函數
 * @tparam block_size 一個block里的線程數量，要求一定是2的整數次冪，最多支持1024
 * @param x 輸入矩陣，按行優先展開
 * @param out 輸出向量
 * @param n 一個batch里的元素數量，要求一定能整除block_size
 */
template<typename T, ReduceFunc<T> reduce_func, size_t block_size>
__global__ void ReduceKernelTemplateFinal(const T *x, T *out, const size_t n) {
  assert(block_size <= 1024);
  assert((block_size & (block_size - 1)) == 0); // 確保n是2的整數次冪
  assert((n & (block_size - 1)) == 0); // 確保n能夠整除batch_size
  extern __shared__ T shared_memory[];

  const size_t tid = threadIdx.x;
  const size_t fold_num = n >> __builtin_ctz(block_size); // 這里會直接在編譯階段就展開，更高效
  size_t i = 1;
  size_t last_x_index = blockIdx.x * blockDim.x + tid;
  shared_memory[tid] = x[last_x_index];
  while (i < fold_num) {
    const size_t current_x_index = last_x_index + block_size;
    shared_memory[tid] = reduce_func(shared_memory[tid], x[current_x_index]);
    last_x_index = current_x_index;
    i++;
  }

  __syncthreads();

  // 完成加載階段，開始規約
  if (block_size >= 1024) {
    if (tid < 512) {
      shared_memory[tid] = reduce_func(shared_memory[tid], shared_memory[tid + 512]);
      __syncthreads();
    }
  }

  if (block_size >= 512) {
    if (tid < 256) {
      shared_memory[tid] = reduce_func(shared_memory[tid], shared_memory[tid + 256]);
      __syncthreads();
    }
  }

  if (block_size >= 256) {
    if (tid < 128) {
      shared_memory[tid] = reduce_func(shared_memory[tid], shared_memory[tid + 128]);
      __syncthreads();
    }
  }

  if (block_size >= 128) {
    if (tid < 64) {
      shared_memory[tid] = reduce_func(shared_memory[tid], shared_memory[tid + 64]);
      __syncthreads();
    }
  }

  if (tid < 32) {
    reduceWarp<T, reduce_func, block_size>(shared_memory, tid);
  }

  // 規約完成，寫入結果
  if (tid == 0) {
    out[blockIdx.x] = shared_memory[0];
  }
}

template<typename T, ReduceFunc<T> reduce_func>
void InvokeReduceFunc(const T *x, T *out, const size_t m, const size_t n, size_t fold) {
  if (n < fold) {
    fold = 1;
  }
  assert(n % fold == 0);
  size_t block_size = n / fold;
  dim3 grid(m, 1, 1);
  dim3 block(block_size, 1, 1);
  switch (block_size) {
    case 1024:
      ReduceKernelTemplateFinal<scalar_t, reduce_func, 1024><<<grid, block, sizeof(scalar_t) * block_size>>>(x, out, n);
      break;
    case 512:
      ReduceKernelTemplateFinal<scalar_t, reduce_func, 512><<<grid, block, sizeof(scalar_t) * block_size>>>(x, out, n);
      break;
    case 256:
      ReduceKernelTemplateFinal<scalar_t, reduce_func, 256><<<grid, block, sizeof(scalar_t) * block_size>>>(x, out, n);
      break;
    case 128:
      ReduceKernelTemplateFinal<scalar_t, reduce_func, 128><<<grid, block, sizeof(scalar_t) * block_size>>>(x, out, n);
      break;
    case 64:
      ReduceKernelTemplateFinal<scalar_t, reduce_func, 64><<<grid, block, sizeof(scalar_t) * block_size>>>(x, out, n);
      break;
    case 32:
      ReduceKernelTemplateFinal<scalar_t, reduce_func, 32><<<grid, block, sizeof(scalar_t) * block_size>>>(x, out, n);
      break;
    case 16:
      ReduceKernelTemplateFinal<scalar_t, reduce_func, 16><<<grid, block, sizeof(scalar_t) * block_size>>>(x, out, n);
      break;
    case 8:
      ReduceKernelTemplateFinal<scalar_t, reduce_func, 8><<<grid, block, sizeof(scalar_t) * block_size>>>(x, out, n);
      break;
    case 4:
      ReduceKernelTemplateFinal<scalar_t, reduce_func, 4><<<grid, block, sizeof(scalar_t) * block_size>>>(x, out, n);
      break;
    case 2:
      ReduceKernelTemplateFinal<scalar_t, reduce_func, 2><<<grid, block, sizeof(scalar_t) * block_size>>>(x, out, n);
      break;
    case 1:
      ReduceKernelTemplateFinal<scalar_t, reduce_func, 1><<<grid, block, sizeof(scalar_t) * block_size>>>(x, out, n);
      break;
  }
}

使用示例

使用起來還是比較簡單的，如下所示：

__device__ inline scalar_t sum_reduce_func(scalar_t x, scalar_t y) {
  return x + y;
}

void ReduceMax(const CudaArray& a, CudaArray* out, size_t reduce_size) {
  /**
   * Reduce by taking maximum over `reduce_size` contiguous blocks.  Even though it is inefficient,
   * for simplicity you can perform each reduction in a single CUDA thread.
   *
   * Args:
   *   a: compact array of size a.size = out.size * reduce_size to reduce over
   *   out: compact array to write into
   *   redice_size: size of the dimension to reduce over
   */
  /// BEGIN SOLUTION
  const size_t fold = 4;
  const size_t m = a.size / reduce_size;
  const size_t n = reduce_size;

  InvokeReduceFunc<scalar_t, max_reduce_func>(a.ptr, out->ptr, m, n, fold);
  /// END SOLUTION
}

結語

CUDA并行計算入門比較容易，但是要做好很難。NVIDIA官方給了一些如何做性能優化的建議，這里貼在下面供大家參考

非常感謝NVIDIA的知識開源，這38頁PPT可以算得上一個比較好的并行計算優化的教材了。

以上就是本系列文章的全部內容，很感謝屏幕前的讀者朋友能耐心看到現在，希望這些內容能幫到大家，如果大家對文章內容有疑問也歡迎在評論區探討。