寫在前面：筆者這段時間工作太忙，身心俱疲，博客停更了一段時間，現在重新撿起來。本文主要解讀 OneFlow 框架的第二種 Softmax 源碼實現細節，即 block 級別的 Softmax。

1 整體邏輯

我們知道，對于形狀為 (num_rows, num_cols) 的矩陣來說，其 Softmax 計算結果只與當前元素所在行的元素相關，所以實現 cuda kernel 的關鍵就是采用多大維度的線程組來處理一行元素。BlockSoftmax 的核心思想是使用一個 block 處理一行元素的計算，借助共享內存保存中間結果數據以及進行線程間的通信。有興趣的讀者可以去如下地址閱讀源碼：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/softmax.cuh
線程網絡結構如下：

2 源碼解析

2.1 數據 Pack 提升訪問帶寬

數據 pack 方面筆者在上一篇文章已經詳細介紹，主要目的是提升內存訪問帶寬，這里不再贅述。

2.2 調用鏈

針對 BlockSMemSoftmax 這個分支，筆者對源碼中函數的調用關系梳理后整理如下：

TryDispatchSoftmaxBlockSMemImpl
  ->TryDispatchSoftmaxBlockSMemImplPackSize
    ->TryDispatchSoftmaxBlockSMemImplBlockSize
      ->LaunchSoftmaxBlockSMemImpl
        ->SoftmaxBlockSMemImpl(kernel)

接下來筆者將從上到下逐個解讀其實現細節。

2.3 TryDispatchSoftmaxBlockSMemImpl

該函數被 DispatchSoftmax 函數調用，其內部邏輯非常簡單，實例化了一個 TryDispatchSoftmaxBlockSMemImplPackSize 類并調用了其重載的()運算符函數，所有的參數都是透傳，沒有其他邏輯。

template
inline cudaError_t TryDispatchSoftmaxBlockSMemImpl(cudaStream_t stream, LOAD load, STORE store,
                                                   const int64_t rows, const int64_t cols,
                                                   bool* success) {
  return TryDispatchSoftmaxBlockSMemImplPackSize()(
      stream, load, store, rows, cols, success);
}

2.4 TryDispatchSoftmaxBlockSMemImplPackSize

顧名思義，pack_size 參數是在這個結構體內部確定的。該結構體內部重載了一個小括號運算符，其函數內部只做了一件事，對矩陣的列數進行判斷從而確定 pack_size，如果矩陣列數是偶數，pack_size 取 2，否則取 1。

template
struct TryDispatchSoftmaxBlockSMemImplPackSize {
  cudaError_t operator()(cudaStream_t stream, LOAD load, STORE store, const int64_t rows,
                         const int64_t cols, bool* success) {
    if (cols % 2 == 0) {
      return TryDispatchSoftmaxBlockSMemImplBlockSize(
          stream, load, store, rows, cols, success);
    } else {
      return TryDispatchSoftmaxBlockSMemImplBlockSize(
          stream, load, store, rows, cols, success);
    }
  }
};

2.5 TryDispatchSoftmaxBlockSMemImplBlockSize

顧名思義，block_size 參數是在該函數內部確定的。關于 block_size 參數的確定方法筆者在另一篇文章（【CUDA編程】OneFlow Element-Wise 算子源碼解讀）中有詳細介紹，因為 blockSMemSoftmax 方案主要使用共享內存而不是寄存器，所以這里我們取消寄存器限制，可以得到 block_size 參數在理想情況下應在 128 和 1024 之間。因此函數 TryDispatchSoftmaxBlockSMemImplBlockSize 中分別定義了 4 個變量，對應 4 種情況。

// 設置4個不同的block_size
constexpr int block_size_conf_1 = 128;
constexpr int block_size_conf_2 = 256;
constexpr int block_size_conf_3 = 512;
constexpr int block_size_conf_4 = 1024;

2.5.1 SM 占有率限制

我們知道，block 是運行在 SM 上的，筆者在前面的文章說過，單個 SM 所能承載的最大 block 數和最大 thread 數是有上限的，為了保證單個 SM 被 thread 填滿（即 SM 占用率 100%），我們要求 block_size 最小取 128。

2.5.2 線程塊同步機制限制

我們知道單個 SM 上的線程總量等于單個 SM 上的 block 數量乘以 block_size，所以在確定 block_size 前我們不妨思考一個問題：單個 SM 上 block 的數量越多越好還是越少越好？
當一個 block 內的線程共同完成一項計算任務時，通常 block 內線程要做同步防止出現讀寫競爭問題。極端情況下，我們假設單個 SM 上只有一個 block，當 SM 中正在調度執行的一個 block 到達同步點時，SM 內可執行 warp 將逐漸減少至 0，會導致計算資源空閑，相當于整個 SM 在等待剩余的 warp 逐步執行，造成資源浪費。若此時 SM 上同時有其他 block 在執行，則在一個 block 到達同步點時仍然有其他 block 可以執行。所以從這個層面上來說，單個 SM 可同時調度的 block 越多越好，單個 SM 上 block 數量越多的同時，block_size 越小。

2.5.3 cudaOccupancyMaxActiveBlocksPerMultiprocessor 函數

前面說過單個 SM 上同時調度的 block 數量越多越好，那么我們如何取到這個最大的 block 數量？首先直接取官方給出的單個 SM 可承載的 block 數上限值肯定是不行的，這只是一個理論上限，實際未必能保證這些 block 被同時調度，同時調度的影響因素還有：共享內存、CUDA 函數、block_size 等等。這里 Nvidia 官方提供了一個預估函數 cudaOccupancyMaxActiveBlocksPerMultiprocessor，該函數會根據設備的硬件資源限制，例如 SM 上的計算核心數量、共享內存的大小等，來確定可調度塊的數量。隨后，它會計算出特定內核函數中每個線程塊的負載平衡，并選出在給定硬件約束下可調度線程塊數量的最大值。最后，該函數將返回一個值，表示在當前硬件資源限制下每個 SM 可允許的最大線程塊數量，這個值是一個整數。該函數的返回結果可以用來估算程序并行效率，幫助開發人員優化程序以使在 GPU 上運行更有效率。

2.5.4 代碼邏輯

終上所述，如何選取最合適的 block_size 參數呢？首先 SM 能同時調度的 block 數量越大越好；當 SM 能同時調度的 Block 數不變的情況下，block_size 越大越好，越大就有越高的并行度。因此代碼中在選擇 block_size 時，對不同 block_size 都計算了 cudaOccupancyMaxActiveBlocksPerMultiprocessor，若結果相同，使用較大的 block_size。
簡單來說，優先讓 SM 同時調度的 block 數量達到最大，其次讓 block_size 達到最大。

template
inline cudaError_t TryDispatchSoftmaxBlockSMemImplBlockSize(cudaStream_t stream, LOAD load,
                                                            STORE store, const int64_t rows,
                                                            const int64_t cols, bool* success) {
  // 設置4個不同的block_size
  constexpr int block_size_conf_1 = 128;
  constexpr int block_size_conf_2 = 256;
  constexpr int block_size_conf_3 = 512;
  constexpr int block_size_conf_4 = 1024;
  // 計算blockSoftmax方案需要的共享內存大小
  const size_t smem = cols * sizeof(ComputeType);
  int max_active_blocks_conf_1;
  {
    // 占用計算器API cudaOccupancyMaxActiveBlocksPerMultiprocessor可以根據 kernel 的 block 大小和共享內存使用情況提供占用率預測。
    cudaError_t err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &max_active_blocks_conf_1,
        SoftmaxBlockSMemImpl,
        block_size_conf_1, smem);
    if (err != cudaSuccess) { return err; }
  }
  if (max_active_blocks_conf_1 <= 0) {
    *success = false;
    return cudaSuccess;
  }
  int max_active_blocks_conf_4;
  {
    cudaError_t err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &max_active_blocks_conf_4,
        SoftmaxBlockSMemImpl,
        block_size_conf_4, smem);
    if (err != cudaSuccess) { return err; }
  }
  if (max_active_blocks_conf_4 == max_active_blocks_conf_1) {
    *success = true;
    return LaunchSoftmaxBlockSMemImpl(stream, load, store, smem, rows, cols);
  }
  int max_active_blocks_conf_3;
  {
    cudaError_t err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &max_active_blocks_conf_3,
        SoftmaxBlockSMemImpl,
        block_size_conf_3, smem);
    if (err != cudaSuccess) { return err; }
  }
  if (max_active_blocks_conf_3 == max_active_blocks_conf_1) {
    *success = true;
    return LaunchSoftmaxBlockSMemImpl(stream, load, store, smem, rows, cols);
  }
  int max_active_blocks_conf_2;
  {
    cudaError_t err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &max_active_blocks_conf_2,
        SoftmaxBlockSMemImpl,
        block_size_conf_2, smem);
    if (err != cudaSuccess) { return err; }
  }
  if (max_active_blocks_conf_2 == max_active_blocks_conf_1) {
    *success = true;
    return LaunchSoftmaxBlockSMemImpl(stream, load, store, smem, rows, cols);
  }
  *success = true;
  return LaunchSoftmaxBlockSMemImpl(stream, load, store, smem, rows, cols);
}

源碼中首先計算了 block_size = 128 時的 SM 同時調度的 block 數量 max_active_blocks_conf_1，并以此作為 SM 同時調度的最大 block 數量，然后分別計算其他三種 block_size 的 max_active_blocks_conf 如果等于最大的 block 數量，則取較大的 block_size。

2.6 核函數 SoftmaxBlockSMemImpl

接下來就是 BlockSoftmax 的核函數 SoftmaxBlockSMemImpl，先來看一下代碼，接下來筆者將逐步解讀源碼作者的實現意圖。

template
__global__ void SoftmaxBlockSMemImpl(LOAD load, STORE store, const int64_t rows,
                                     const int64_t cols) {
  extern __shared__ __align__(sizeof(double)) unsigned char shared_buf[];
  auto* buf = reinterpret_cast(shared_buf);
  const int tid = threadIdx.x;
  assert(cols % pack_size == 0);
  const int num_packs = cols / pack_size;
  // 一個 Block 處理一行元素
  for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
    // 當前線程的最大值初始化為 -inf
    ComputeType thread_max = -Inf();
    // 以向量化的方式加載一行數據，然后執行pack reduce操作
    for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
      ComputeType pack[pack_size];
      load.template load(pack, row, pack_id * pack_size);
#pragma unroll
      for (int i = 0; i < pack_size; ++i) {
        buf[i * num_packs + pack_id] = pack[i];
        thread_max = max(thread_max, pack[i]);
      }
    }
    // 執行block reduce獲取當前行（由一個 Block 進行處理）的最大值
    const ComputeType row_max = BlockAllReduce(thread_max);
    ComputeType thread_sum = 0;
    for (int col = tid; col < cols; col += block_size) {
      if (algorithm == Algorithm::kSoftmax) {
        const ComputeType exp_x = Exp(buf[col] - row_max);
        buf[col] = exp_x;
        thread_sum += exp_x;
      } else {
        const ComputeType x = buf[col] - row_max;
        buf[col] = x;
        thread_sum += Exp(x);
      }
    }
    // 同理，獲得當前行的sum
    const ComputeType row_sum = BlockAllReduce(thread_sum);
    // 計算結果并寫回到全局內存中
    for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
      ComputeType pack[pack_size];
#pragma unroll
      for (int i = 0; i < pack_size; ++i) {
        if (algorithm == Algorithm::kSoftmax) {
          pack[i] = Div(buf[i * num_packs + pack_id], row_sum);
        } else if (algorithm == Algorithm::kLogSoftmax) {
          pack[i] = buf[i * num_packs + pack_id] - Log(row_sum);
        } else {
          __trap();
        }
      }
      store.template store(pack, row, pack_id * pack_size);
    }
  }
}

2.6.1 定義共享內存變量

核函數內部首先定義了一個共享內存數組變量 shared_buf，內存對齊大小為 sizeof(double)，隨后在核函數內部做了一個斷言，校驗 pack_size 是否能夠整除 cols。在 shared_buf 變量定義語句中，extern 是使用動態共享內存的一種聲明方式，表示內存大小將在調用核函數時通過 <<<>>> 的第三個參數指定，這里為矩陣的一行元素對應的內存大小，即 cols * sizeof(ComputeType)。然后使用 reinterpret_cast 將 shared_buf 的首地址轉換為 ComputeType* 類型的指針并賦給指針 buf。這段代碼的意義是將共享內存 shared_buf 的首地址強制轉換為ComputeType 類型的指針，使得在后續代碼中可以通過 buf 來訪問共享內存，進行后續的 GPU 并行計算操作。這里筆者沒想明白為什么要繞個圈子，而不是直接用如下代碼直接定義。

extern __shared__ __align__(sizeof(double)) ComputeType buf[];

咨詢源碼作者后收到反饋，“如果采用上述方式定義會編譯報錯”，然后筆者親測沒有報錯，可能是兩邊開發環境不同導致，有興趣地讀者可以嘗試下，筆者的環境為 win10+cuda11.7+rtx2070super。

2.6.2 計算每個線程的 thread_max

主體部分是一個 Grip-loop 的循環，循環步長設置為網格大小。Grip-loop 內部定義了一個寄存器變量 thread_max，這個變量用來存儲當前線程處理的元素中的最大值。
接下來是兩層循環求出 thread_max，這里和 WarpSoftmax 一樣，一個線程處理的 pack 是不相連的，這里主要是因為 GPU 在從內存取數的時候，為了優化內存訪問性能，一次取數時會將臨近空間的數據也取出存入緩存，而 GPU 指令是按照線程束為基本單位進行執行的，這樣的話，一次取數可以滿足相鄰的多個線程的使用需求，直接去緩存取即可，無需再次訪問內存。因此為了最大程度提高訪問效率，相鄰線程訪問的數據是緊鄰的。在循環體內部定義了一個寄存器數組變量 pack[pack_size]，將內存中一個 pack 的數據加載到數組 pack 中，然后在 pack 內求 thread_max，順便將數據也加載到共享內存變量 buf 中。

2.6.3 Bank Conflicts 優化

從往 buf 里加載數據的代碼中可以發現，在 buf 中內存排列有些反常，一個 pack 內的元素不是相鄰排列的，而是隔了 num_packs 個元素，這里可以把 buf 想象成一個 pack_size 行 num_packs 列的矩陣，一個 pack 內的元素是按列存儲的。為什么要按列往共享內存里寫入數據？
這里涉及一個Bank Conflicts 概念，這也是在使用共享內存時需要重點關注的問題。為了獲得更高的內存帶寬，共享內存在物理上被分為了 32 個寬度相等（開普勒架構為8個字節，其他架構4個字節）的 bank，這些 bank 可以被同時訪問。為什么恰好是 32 個？因為前面說過 GPU 中執行指令是以線程束（32個線程）為基本單位的，這樣可以保證每個線程同時訪問一個 bank 的數據，帶寬達到最大。那么 Bank Conflicts 的定義來了，在一個 warp 內，有 2 個或以上的線程訪問了同一個 bank 上的不同地址的內存。
現在我們假設 buf 里使用如下方式加載數據，即一個 pack 里的元素相鄰排列，則代碼如下：

#pragma unroll
      for (int i = 0; i < pack_size; ++i) {
        buf[pack_id * pack_size + i] = pack[i];
        thread_max = max(thread_max, pack[i]);
      }

當 pack_size = 1 時，每個線程連續寫 4 個字節時，每個 warp 剛好完整訪問 shared memory 的一行，這個時候并不會出現 bank conflict。而當 pack_size = 2 時，每個線程寫連續 2 個 4 字節時（可以看成8個字節），此時 0 號線程訪問的地址在第 0 和第 1 個 bank，1 號線程訪問的地址在第 2 和第 3 個 bank，以此類推，16 號線程訪問地址又在第 0 和第 1 個 bank 內，此時 16 號線程和 0 號線程訪問了同一個 bank 的不同地址，此時即產生了 Bank Conflicts。見圖a。
為了避免 Bank Conflicts，可以將一個 pack 的數據按列存儲，這樣當矩陣元素的字節大于 4 或 pack_size >= 2 時，可以保證每個線程訪問的數據都在自己的 bank 內，如圖b。

2.6.4 block reduce 獲取 row_max

得到每一個線程處理的元素的最大值后，還需要計算每一行的最大值。在 WarpSoftmax 中，由于每一行是由一個 warp 處理的，所以我們使用束內洗牌指令即可得到矩陣單行最大值，而在 BlockSoftmax 中，每一行是由一個 block 處理的，這時候不能再使用束內指令來規約了，為了保證較高的性能，應使用共享內存進行線程間的數據通信。這里源碼封裝了一個 BlockAllReduce(thread_max) 函數，在函數內直接使用 Nvidia 官方 cub 庫進行計算，官方文檔見：https://nvlabs.github.io/cub/classcub_1_1_block_reduce.html#a089953b3bdfe7c48208632d0cc2ac1fb

const ComputeType row_max = BlockAllReduce(thread_max);

函數體中聲明了兩個共享內存變量 temp_storage 和 result_broadcast，可見該庫函數底層也是通過共享內存實現的，這里如果不用官方庫，也可以自定義規約函數，有興趣的讀者可以參考筆者的另一篇文章【CUDA編程】CUDA編程中的并行規約問題

2.6.5 thread_sum 和 row_sum

定義一個寄存器變量 thread_sum 存儲當前線程處理的元素的指數和。由于矩陣元素已經加載進共享內存 buf 中，所以這一次遍歷求和不需要再訪問全局內存，直接在 buf 中以 block_size 為步長求和即可，遍歷的同時也將 buf 的內存替換成 exp_x 這是為了方便后續的計算。這里為什么要以 block_size 為步長？也是因為前面為了避免 bank conflicts 將 buf 的內存排列設定為相鄰線程的元素相鄰存儲，所以我們只要以 block_size 為步長遍歷即可完成該線程處理的所有 pack 元素的遍歷。
獲取到 thread_sum 后，同樣使用前面 BlockAllReduce 函數進行塊內規約計算出 row_sum。

2.6.6 計算 Softmax

首先對當前線程以 block_size 為步長做一次循環，然后在 pack 內利用上一步計算的 row_max 和 buf 計算出 Softmax 值，最后將 pack 內的計算結果寫入全局內存中。

3 小結

總結一下 BlockSoftmax 源碼中的一些值得學習的內容：

在選取 block_size 時，要結合 kernel 的計算邏輯，靈活選取。比如作者在實現 BlockSoftmax 時考慮到這種實現方式不會大量使用寄存器，所以去除了寄存器限制；考慮到塊內同步和共享內存，所以使用了官方庫函數預估 SM 占用率。這些都是很值得讀者學習的，兵無常勢，水無常形。

BlockSoftmax 的核心是塊內規約，利用共享內存讀寫速度遠大于全局內存的特性，提高 kernel 性能。

在使用共享內存時，一定要避免 Bank Conflicts，可以提升至少 20% 的訪存效率。

審核編輯：湯梓紅