OneFlow Softmax算子源碼解讀之BlockSoftmax

寫在前面：筆者這段時間工作太忙，身心俱疲，博客停更了一段時間，現(xiàn)在重新?lián)炱饋怼１疚闹饕庾x OneFlow 框架的第二種 Softmax 源碼實現(xiàn)細(xì)節(jié)，即 block 級別的 Softmax。

1 整體邏輯

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

2 源碼解析

2.1 數(shù)據(jù) Pack 提升訪問帶寬

數(shù)據(jù) pack 方面筆者在上一篇文章已經(jīng)詳細(xì)介紹，主要目的是提升內(nèi)存訪問帶寬，這里不再贅述。

2.2 調(diào)用鏈

針對 BlockSMemSoftmax 這個分支，筆者對源碼中函數(shù)的調(diào)用關(guān)系梳理后整理如下：

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

接下來筆者將從上到下逐個解讀其實現(xiàn)細(xì)節(jié)。

2.3 TryDispatchSoftmaxBlockSMemImpl

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

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 參數(shù)是在這個結(jié)構(gòu)體內(nèi)部確定的。該結(jié)構(gòu)體內(nèi)部重載了一個小括號運(yùn)算符，其函數(shù)內(nèi)部只做了一件事，對矩陣的列數(shù)進(jìn)行判斷從而確定 pack_size，如果矩陣列數(shù)是偶數(shù)，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 參數(shù)是在該函數(shù)內(nèi)部確定的。關(guān)于 block_size 參數(shù)的確定方法筆者在另一篇文章（【CUDA編程】OneFlow Element-Wise 算子源碼解讀）中有詳細(xì)介紹，因為 blockSMemSoftmax 方案主要使用共享內(nèi)存而不是寄存器，所以這里我們?nèi)∠拇嫫飨拗?，可以得?block_size 參數(shù)在理想情況下應(yīng)在 128 和 1024 之間。因此函數(shù) TryDispatchSoftmaxBlockSMemImplBlockSize 中分別定義了 4 個變量，對應(yīng) 4 種情況。

// 設(shè)置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 是運(yùn)行在 SM 上的，筆者在前面的文章說過，單個 SM 所能承載的最大 block 數(shù)和最大 thread 數(shù)是有上限的，為了保證單個 SM 被 thread 填滿（即 SM 占用率 100%），我們要求 block_size 最小取 128。

2.5.2 線程塊同步機(jī)制限制

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

2.5.3 cudaOccupancyMaxActiveBlocksPerMultiprocessor 函數(shù)

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

2.5.4 代碼邏輯

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

template
inline cudaError_t TryDispatchSoftmaxBlockSMemImplBlockSize(cudaStream_t stream, LOAD load,
                                                            STORE store, const int64_t rows,
                                                            const int64_t cols, bool* success) {
  // 設(shè)置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方案需要的共享內(nèi)存大小
  const size_t smem = cols * sizeof(ComputeType);
  int max_active_blocks_conf_1;
  {
    // 占用計算器API cudaOccupancyMaxActiveBlocksPerMultiprocessor可以根據(jù) kernel 的 block 大小和共享內(nèi)存使用情況提供占用率預(yù)測。
    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 同時調(diào)度的 block 數(shù)量 max_active_blocks_conf_1，并以此作為 SM 同時調(diào)度的最大 block 數(shù)量，然后分別計算其他三種 block_size 的 max_active_blocks_conf 如果等于最大的 block 數(shù)量，則取較大的 block_size。

2.6 核函數(shù) SoftmaxBlockSMemImpl

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

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) {
    // 當(dāng)前線程的最大值初始化為 -inf
    ComputeType thread_max = -Inf();
    // 以向量化的方式加載一行數(shù)據(jù)，然后執(zhí)行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]);
      }
    }
    // 執(zhí)行block reduce獲取當(dāng)前行（由一個 Block 進(jìn)行處理）的最大值
    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);
      }
    }
    // 同理，獲得當(dāng)前行的sum
    const ComputeType row_sum = BlockAllReduce(thread_sum);
    // 計算結(jié)果并寫回到全局內(nèi)存中
    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 定義共享內(nèi)存變量

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

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

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

2.6.2 計算每個線程的 thread_max

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

2.6.3 Bank Conflicts 優(yōu)化

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

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

2.6.4 block reduce 獲取 row_max

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

const ComputeType row_max = BlockAllReduce(thread_max);

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

2.6.5 thread_sum 和 row_sum

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

2.6.6 計算 Softmax

首先對當(dāng)前線程以 block_size 為步長做一次循環(huán)，然后在 pack 內(nèi)利用上一步計算的 row_max 和 buf 計算出 Softmax 值，最后將 pack 內(nèi)的計算結(jié)果寫入全局內(nèi)存中。

3 小結(jié)

總結(jié)一下 BlockSoftmax 源碼中的一些值得學(xué)習(xí)的內(nèi)容：

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

BlockSoftmax 的核心是塊內(nèi)規(guī)約，利用共享內(nèi)存讀寫速度遠(yuǎn)大于全局內(nèi)存的特性，提高 kernel 性能。

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

審核編輯：湯梓紅