寫在前面：近來筆者偶然間接觸了一個(gè)深度學(xué)習(xí)框架 OneFlow，所以這段時(shí)間主要在閱讀 OneFlow 框架的 cuda 源碼。官方源碼基于不同場景分三種方式實(shí)現(xiàn) Softmax，本文主要介紹其中一種的實(shí)現(xiàn)過程，即 Warp 級別 Softmax，適用于矩陣寬度不超過 1024 的情況。

1 Softmax

Softmax 操作是深度學(xué)習(xí)模型中最常用的操作之一。在深度學(xué)習(xí)的多分類任務(wù)中，最后一層通常是一個(gè) Softmax 操作將 logits 映射成概率，然后結(jié)合交叉熵求損失。另外還有一些場景會用到 Softmax 做一個(gè)歸一化操作，比如 Transformer 結(jié)構(gòu)中 query 和 key 矩陣相乘并縮放后會執(zhí)行一個(gè) Softmax 操作，這一步的意義是求出 query 和 key 中每一項(xiàng)的兩兩相似度，具體筆者在另一篇文章有詳述——【ASR】基于DFCNN-CTC模型的語音識別系統(tǒng)（二）

圖1 Scaled Dot-Product Attention 結(jié)構(gòu)示意圖

深度學(xué)習(xí)框架中的所有算子底層都對應(yīng)著 GPU上的 CUDA kernel function，Softmax 操作也不例外。Softmax 作為一個(gè)被廣泛使用的算子，其 CUDA Kernel 的實(shí)現(xiàn)會影響很多網(wǎng)絡(luò)最終的訓(xùn)練速度。那么如何實(shí)現(xiàn)一個(gè)高效的 Softmax CUDA Kernel？本文將會介紹 OneFlow 中優(yōu)化的 Softmax CUDA Kernel 的技巧，在這之前我們先來看一下 Softmax 的計(jì)算公式。
定義 x 是一個(gè) n 維向量，其 Softmax 輸出 y 也是一個(gè) n 維向量，那么有如下計(jì)算公式：

從上面的公式可以發(fā)現(xiàn)一個(gè)問題，當(dāng) 為一個(gè)較大的正數(shù)時(shí)，取指數(shù)后 將會非常大，從而導(dǎo)致數(shù)值溢出，如何解決這個(gè)問題呢？
一般的處理方法是，讓每個(gè)分量去減掉向量的最大值，這樣可以保證取指數(shù)后的結(jié)果必然在 0~1 之間，可以有效避免數(shù)值溢出。處理后的公式如下：

根據(jù)公式可以看出，要執(zhí)行 Softmax 計(jì)算，需要實(shí)現(xiàn) 5 個(gè)業(yè)務(wù)邏輯：reduceMax、broadcastSub、exp、reduceSum、broadcastDiv。下面筆者將對源碼中的計(jì)算技巧進(jìn)行解讀，有興趣的讀者可以下載源碼來閱讀(https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/softmax/oneflow_softmax.cu)。

2 三種實(shí)現(xiàn)方式

Softmax 函數(shù)的輸入形狀為:(num_rows, num_cols)，num_cols 的變化會對有效帶寬產(chǎn)生影響。因?yàn)椋瑳]有一種通用的優(yōu)化方法可以實(shí)現(xiàn)在所有 num_cols 的情況下都是傳輸最優(yōu)的。所以，在 OneFlow 中采用分段函數(shù)優(yōu)化 SoftmaxKernel：對于不同 num_cols 范圍，選擇不同的實(shí)現(xiàn)，以期在所有情況下都能達(dá)到較高的有效帶寬。
針對不同的 Softmax 場景，OneFlow 提供了三種實(shí)現(xiàn)，分段對 Softmax kernel 進(jìn)行優(yōu)化：

(1) 一個(gè) Warp 處理一行的計(jì)算，適用于 num_cols <= 1024 情況

(2) 一個(gè) Block 處理一行的計(jì)算，借助 Shared Memory 保存中間結(jié)果數(shù)據(jù)，適用于需要的 Shared Memory 資源滿足 Kernel Launch 的可啟動條件的情況，在本測試環(huán)境中是 1024 < num_cols <= 4096。

(3) 一個(gè) Block 處理一行的計(jì)算，不使用 Shared Memory，重復(fù)讀輸入 x，適用于不支持(1)、(2)的情況。

分段處理邏輯在 DispatchSoftmax 函數(shù)中體現(xiàn)，主體代碼如下：

if (cols < 1024) {
  return DispatchSoftmaxWarpImpl(
      stream, load, store, rows, cols);
} else {
  bool dispatch_smem_impl_success;
  {
    cudaError_t err =
        TryDispatchSoftmaxBlockSMemImpl(
            stream, load, store, rows, cols, &dispatch_smem_impl_success);
    if (err != cudaSuccess) { return err; }
  }
  if (!dispatch_smem_impl_success) {
    return DispatchSoftmaxBlockUncachedImpl(
        stream, load, store, rows, cols);
  }
  return cudaSuccess;
}
,>,>,>

3 WarpSoftmax

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

在筆者上一篇文章【CUDA編程】OneFlow Element-Wise 算子源碼解讀中詳細(xì)地介紹了如何進(jìn)行向量化讀寫，有興趣的讀者可以移步，這里我們先看源碼。

template
struct GetPackType {
  using type = typename std::aligned_storage::type;
};

template
using PackType = typename GetPackType::type;

template
union Pack {
  static_assert(sizeof(PackType) == sizeof(T) * N, "");
  __device__  Pack() {
    // do nothing
  }
  PackType storage;
  T elem[N];
};
,>,>,>

oneflow 利用 union 共享空間的特性實(shí)現(xiàn)了一個(gè) Pack 類型，細(xì)心的讀者可能會發(fā)現(xiàn)，跟 elementwise.cuh 源碼相比，這里少了一個(gè) Packed 類，這是因?yàn)?elementwise.cuh 實(shí)現(xiàn)的時(shí)間晚于 softmax.cuh。可能考慮到 Pack 后類型的內(nèi)存對齊特性，重新定義了 Packed 類，并聲明了內(nèi)存對齊值為 pack_size * sizeof(T)。
接下來定義了兩個(gè)代表輸入和輸出的數(shù)據(jù)結(jié)構(gòu) DirectLoad 和 DirectStore，分別實(shí)現(xiàn)了 load 和 store 兩個(gè)函數(shù)用來把讀取和寫入一個(gè) pack 的數(shù)據(jù)。使用 DirectLoad 和 DirectStore 有兩個(gè)好處：

可以在CUDA Kernel中只關(guān)心計(jì)算類型ComputeType，而不用關(guān)心具體的數(shù)據(jù)類型T。

只需要加幾行代碼就可以快速支持 Softmax 和其他 Kernel Fuse，減少帶寬需求，提升整體性能。

/**
 * @brief 定義了輸入的數(shù)據(jù)結(jié)構(gòu)
 * 
 * @tparam SRC 輸入數(shù)據(jù)的類型
 * @tparam DST 計(jì)算數(shù)據(jù)的類型，ComputeType
 */
template
struct DirectLoad {
  /**
   * @brief Construct a new Direct Load object
   * 
   * @param src 輸入的數(shù)據(jù)源
   * @param row_size num of elements per row
   */
  DirectLoad(const SRC* src, int64_t row_size) : src(src), row_size(row_size) {}
  /**
   * @brief 從數(shù)據(jù)源 load 一個(gè) pack 數(shù)據(jù)到 dst 
   * 
   * @tparam N pack_size
   * @param dst 
   * @param row 數(shù)據(jù)源的第 row 行
   * @param col 數(shù)據(jù)源的第 col 列
   * @return __device__ 
   */
  template
  __device__ void load(DST* dst, int64_t row, int64_t col) const {
    Pack pack;
    const int64_t offset = (row * row_size + col) / N;  // pack 偏移量
    pack.storage = *(reinterpret_cast*>(src) + offset);
    #pragma unroll
    for (int i = 0; i < N; ++i) { dst[i] = static_cast(pack.elem[i]); }
  }
  const SRC* src;
  int64_t row_size;
};

template
struct DirectStore {
  DirectStore(DST* dst, int64_t row_size) : dst(dst), row_size(row_size) {}
  template
  __device__ void store(const SRC* src, int64_t row, int64_t col) {
    Pack pack;
    const int64_t offset = (row * row_size + col) / N;
#pragma unroll
    for (int i = 0; i < N; ++i) { pack.elem[i] = static_cast(src[i]); }
    *(reinterpret_cast*>(dst) + offset) = pack.storage;
  }
  DST* dst;
  int64_t row_size;
};
,>,>

3.2 調(diào)用鏈

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

DispatchSoftmaxWarpImpl
  ->DispatchSoftmaxWarpImplPackSize
    ->DispatchSoftmaxWarpImplCols
      ->DispatchSoftmaxWarpImplPadding
        ->LaunchSoftmaxWarpImpl
          ->SoftmaxWarpImpl(kernel)

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

3.3 DispatchSoftmaxWarpImpl

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

template
inline cudaError_t DispatchSoftmaxWarpImpl(cudaStream_t stream, LOAD load, STORE store, const int64_t rows, const int64_t cols) {
  return DispatchSoftmaxWarpImplPackSize()(stream, load, store, rows, cols);
}
,>

3.4 DispatchSoftmaxWarpImplPackSize

顧名思義，pack_size 參數(shù)是在這個(gè)結(jié)構(gòu)體內(nèi)部確定的。該結(jié)構(gòu)體內(nèi)部重載了一個(gè)小括號運(yùn)算符，其函數(shù)內(nèi)部只做了一件事，對矩陣的列數(shù)進(jìn)行判斷，如果是偶數(shù)，pack_size 取 2，否則取 1。

template
struct DispatchSoftmaxWarpImplPackSize {
  cudaError_t operator()(cudaStream_t stream, LOAD load, STORE store, const int64_t rows, const int64_t cols) {
    if (cols % 2 == 0) {
      return DispatchSoftmaxWarpImplCols(stream, load, store, rows, cols);
    } else {
      return DispatchSoftmaxWarpImplCols(stream, load, store, rows, cols);
    }
  }
};
,>,>

筆者讀到這里不禁產(chǎn)生了疑問，前面說過數(shù)據(jù) Pack 后可以提升 GPU 訪問帶寬，但是在該函數(shù)中 pack_size 最大也只能取到 2，在前面的文章中筆者提到過在 cuda 中最大支持一次 128 bit的讀寫，意味著針對 float 類型 pack_size 最大可以取 4，對 half 類型甚至可以取 8。所以帶著這個(gè)疑問筆者咨詢了官方源碼的作者俊丞大佬，答曰可以取更大的 pack_size，這里是考慮到更多的特化會導(dǎo)致編譯時(shí)間過長所以只實(shí)現(xiàn)了 2 個(gè)模板。獲得解答后，筆者自行實(shí)現(xiàn)了一個(gè) pack_size = 4 的模板，然后經(jīng)過實(shí)測（矩陣大小為 1024*1024, 32*16）發(fā)現(xiàn)， pack_size 取 4 和取 2 相比幾乎沒有提升。。。倒是取 2 相比取 1 有 6% 的提升。猜測可能是 pack_size 影響了 DispatchSoftmaxWarpImplCols 這個(gè) kernel 的啟動參數(shù)，所以間接影響了性能，這里官方肯定做過一系列測試。。。

3.5 DispatchSoftmaxWarpImplCols

DispatchSoftmaxWarpImplCols 函數(shù)代碼比較長，讀起來稍顯晦澀，要理解它的實(shí)現(xiàn)邏輯，我們可以換個(gè)思路，看它的目的是什么，然后倒推它的實(shí)現(xiàn)過程。很顯然，該函數(shù)在最后調(diào)用了 DispatchSoftmaxWarpImplPadding 函數(shù)，那么我們就來看被調(diào)用的函數(shù)需要哪些參數(shù)，DispatchSoftmaxWarpImplCols 的作用就是確定這些參數(shù)。讀了 DispatchSoftmaxWarpImplPadding 的參數(shù)列表我們可以發(fā)現(xiàn)，有三個(gè)重要參數(shù)需要傳入：cols_per_thread, thread_group_width, rows_per_access，這里先對這三個(gè)參數(shù)做一個(gè)解釋：

cols_per_thread：每個(gè)線程處理的元素列數(shù)

thread_group_width：線程組的大小，一個(gè)線程組要處理整行的數(shù)據(jù)

rows_per_access：每個(gè)線程組一次處理的行數(shù)

函數(shù)體內(nèi)主要是針對 cols 的大小做了分支，前后代碼有一個(gè)分水嶺，即 cols <= 32 * pack_size，可以分開來看。
當(dāng) cols <= 32 * pack_size 時(shí)，thread_group_width 取 2 的 n 次冪，從 1 到 32 一直判斷，如果 cols <= (thread_group_width)*pack_size 那么 thread_group_width 就取當(dāng)前的值。cols_per_thread 取 pack_size，就是說當(dāng)前一個(gè)線程只處理一個(gè) Pack 寬度的數(shù)據(jù)，這時(shí)候數(shù)據(jù)量也比較小，所以對 rows 也做了一層判斷，如果 rows 是偶數(shù)，那么 rows_per_access 取 2，每個(gè)線程一次處理 2 行數(shù)據(jù)，否則一次只處理 1 行。
當(dāng) cols > 32 * pack_size 時(shí)，這種屬于數(shù)據(jù)量比較大的情況，所以 thread_group_width 直接取能取到的最大值 32，即 Warp 的大小。每個(gè)線程也要處理多個(gè) Pack，cols_per_thread 取 pack_size 的整數(shù)倍，直到 32 * cols_per_thread = 1024，一直判斷 cols <= 32 * cols_per_thread，如果滿足條件，cols_per_thread 就取當(dāng)前值。對于 rows_per_access 參數(shù)，直接取 1，即每個(gè)線程一次只處理 1 行數(shù)據(jù)。
至此函數(shù)邏輯就介紹完了，這個(gè)函數(shù)里有兩個(gè)宏，不熟悉 C++ 的讀者讀起來可能沒那么順暢，這里推薦一個(gè)網(wǎng)站（https://cppinsights.io/），從編譯器的角度將 C++ 源碼展開顯示，對閱讀泛型編程和宏這類代碼很有幫助。

template
typename std::enable_if::type DispatchSoftmaxWarpImplCols(
    cudaStream_t stream, LOAD load, STORE store, const int64_t rows, const int64_t cols) {
  if (cols <= 0) { return cudaErrorInvalidValue; }
#define DEFINE_ONE_ELIF(thread_group_width)                                                        
  else if (cols <= (thread_group_width)*pack_size) {                                               
    if (rows % 2 == 0) {                                                                           
      return DispatchSoftmaxWarpImplPadding(stream, load, store, 
                                                                              rows, cols);         
    } else {                                                                                       
      return DispatchSoftmaxWarpImplPadding(stream, load, store, 
                                                                              rows, cols);         
    }                                                                                              
  }
  DEFINE_ONE_ELIF(1)
  DEFINE_ONE_ELIF(2)
  DEFINE_ONE_ELIF(4)
  DEFINE_ONE_ELIF(8)
  DEFINE_ONE_ELIF(16)
  DEFINE_ONE_ELIF(32)
#undef DEFINE_ONE_ELIF
#define DEFINE_ONE_ELIF(col)                                                                      
  else if (cols <= (col)*kWarpSize) {                                                             
    return DispatchSoftmaxWarpImplPadding(stream, load, store, rows, cols);            
  }
  DEFINE_ONE_ELIF(4)
  DEFINE_ONE_ELIF(6)
  DEFINE_ONE_ELIF(8)
  DEFINE_ONE_ELIF(10)
  DEFINE_ONE_ELIF(12)
  DEFINE_ONE_ELIF(14)
  DEFINE_ONE_ELIF(16)
  DEFINE_ONE_ELIF(18)
  DEFINE_ONE_ELIF(20)
  DEFINE_ONE_ELIF(22)
  DEFINE_ONE_ELIF(24)
  DEFINE_ONE_ELIF(26)
  DEFINE_ONE_ELIF(28)
  DEFINE_ONE_ELIF(30)
  DEFINE_ONE_ELIF(32)
#undef DEFINE_ONE_ELIF
  else {
    return cudaErrorInvalidValue;
  }
}
,>,>,>

3.6 DispatchSoftmaxWarpImplPadding

顧名思義，這個(gè)函數(shù)內(nèi)部的邏輯跟 padding 相關(guān)，實(shí)際上這個(gè)函數(shù)只做了一件事，當(dāng) cols == cols_per_thread * thread_group_width 時(shí)說明矩陣列數(shù)能被線程組均分，這時(shí)候不需要 padding，否則需要 padding。

template
inline cudaError_t DispatchSoftmaxWarpImplPadding(cudaStream_t stream, LOAD load, STORE store,
                                                  const int64_t rows, const int64_t cols) {
  if (cols == cols_per_thread * thread_group_width) {
    return LaunchSoftmaxWarpImpl(
        stream, load, store, rows, cols);
  } else {
    return LaunchSoftmaxWarpImpl(
        stream, load, store, rows, cols);
  }
}
,>,>

3.7 LaunchSoftmaxWarpImpl

該函數(shù)是核函數(shù)的啟動函數(shù)，函數(shù)內(nèi)主要是確定 block_size、num_blocks 這兩個(gè)參數(shù)。這兩個(gè)參數(shù)的確定筆者在上一篇文章【CUDA編程】OneFlow Element-Wise 算子源碼解讀中有詳細(xì)介紹，有興趣的讀者可以移步，這里不再贅述。
函數(shù)中定義了一個(gè) block_dim 對象，從初始化參數(shù)可以看出這是一個(gè)二維的 block，寬是 thread_group_width，高取 thread_groups_per_block。從核函數(shù)啟動參數(shù) grid_dim_x 可以看出網(wǎng)格是一維的，由此我們可以確定 cuda 線程網(wǎng)格的形狀。這里筆者給出示意圖如下。

圖2 線程網(wǎng)格示意圖

template
inline cudaError_t LaunchSoftmaxWarpImpl(cudaStream_t stream, LOAD load, STORE store,
                                         const int64_t rows, const int64_t cols) {
  constexpr int block_size = 128;
  constexpr int waves = 32;
  static_assert(block_size % thread_group_width == 0, "");
  constexpr int thread_groups_per_block = block_size / thread_group_width;
  dim3 block_dim(thread_group_width, thread_groups_per_block);
  const int64_t num_blocks =
      (rows / rows_per_access + thread_groups_per_block - 1) / thread_groups_per_block;
  int grid_dim_x;
  {
    cudaError_t err = GetNumBlocks(block_size, num_blocks, waves, &grid_dim_x);
    if (err != cudaSuccess) { return err; }
  }
  SoftmaxWarpImpl
      <<>>(load, store, rows, cols);
  return cudaPeekAtLastError();
}
,>,>

3.8 核函數(shù) SoftmaxWarpImpl

接下來就是 WarpSoftmax 的核函數(shù) SoftmaxWarpImpl，該函數(shù)體內(nèi)部實(shí)現(xiàn)了 Softmax 的核心計(jì)算邏輯。

template
__global__ void SoftmaxWarpImpl(LOAD load, STORE store, const int64_t rows, const int64_t cols) {
  static_assert(cols_per_thread % pack_size == 0, "");  // 確保每個(gè)thread處理的元素個(gè)數(shù)正好被完全pack
  static_assert(thread_group_width <= kWarpSize, "");   // 處理元素的線程組的寬度需要小于等于kWarpSize，并且需要被kWarpSize整除
  static_assert(kWarpSize % thread_group_width == 0, "");
  constexpr int num_packs = cols_per_thread / pack_size;  // 每個(gè)線程處理的 pack 的數(shù)目，即每個(gè)線程需要處理的元素個(gè)數(shù) / pack_size
  assert(cols <= cols_per_thread * thread_group_width);   // 確保一個(gè)thread group 能處理的列數(shù)大于等于一行
  ComputeType buf[rows_per_access][cols_per_thread];  // 聲明寄存器大小，這是一個(gè)二維數(shù)組
  const int global_thread_group_id = blockIdx.x * blockDim.y + threadIdx.y;   // 當(dāng)前warp的全局index
  const int num_global_thread_group = gridDim.x * blockDim.y;   // warp的總數(shù)量
  const int lane_id = threadIdx.x;    // warp內(nèi)的線程id
  const int64_t step = num_global_thread_group * rows_per_access;   // 處理的行數(shù)步長
  // for 循環(huán)的開始為 row = 全局的線程組id * 每個(gè)線程組一次處理的行數(shù)，結(jié)束為總行數(shù)
  for (int64_t row = global_thread_group_id * rows_per_access; row < rows; row += step) {
    // 寄存器中開辟一塊內(nèi)存記錄當(dāng)前線程組處理的每一行的最大值
    ComputeType thread_max[rows_per_access];
    // 對每一行的循環(huán)
    #pragma unroll
    for (int row_id = 0; row_id < rows_per_access; ++row_id) {
      // 把當(dāng)前行最小值初始化為 -inf
      thread_max[row_id] = -Inf();
      // 獲取第 row_id 行的指針
      ComputeType* row_buf = buf[row_id];
      #pragma unroll
      for (int pack_id = 0; pack_id < num_packs; ++pack_id) {
        const int pack_offset = pack_id * pack_size;
        // 相鄰的線程讀取相鄰的pack，也就是說同一個(gè)線程處理的相鄰pack間間隔是thread_group_width*pack_size
        const int col = (pack_id * thread_group_width + lane_id) * pack_size;
        if (!padding || col < cols) {
          // 使用 obj.template 調(diào)用函數(shù)模板防止歧義，load 一個(gè) pack 的數(shù)據(jù)到寄存器
          load.template load(row_buf + pack_offset, row + row_id, col);
          #pragma unroll
          for (int i = 0; i < pack_size; ++i) {
            thread_max[row_id] = max(thread_max[row_id], row_buf[pack_offset + i]);
          }
        } else {  // 需要 padding 且 col > cols，這種情況對于第 col 列的數(shù)據(jù)直接將 row_buf 賦最新小值，不影響 thread_max 計(jì)算即可
          #pragma unroll
          for (int i = 0; i < pack_size; ++i) { row_buf[pack_offset + i] = -Inf(); }
        }
      }
    }
    // 記錄屬于同一個(gè)warp的線程組的每一行的最大值，也就是需要進(jìn)行一次warpReduce max
    ComputeType warp_max[rows_per_access];
    #pragma unroll
    for (int row_id = 0; row_id < rows_per_access; ++row_id) {
      // 通過線程束洗牌函數(shù)對一個(gè)線程組內(nèi)的所有線程的 thread_max 求規(guī)約得到一個(gè)線程組處理的每一行的最大值
      warp_max[row_id] = WarpAllReduce(thread_max[row_id]);
    }
    // 記錄當(dāng)前線程組處理的每一行的sum
    ComputeType thread_sum[rows_per_access];
    #pragma unroll
    for (int row_id = 0; row_id < rows_per_access; ++row_id) {
      thread_sum[row_id] = 0;
      ComputeType* row_buf = buf[row_id];
      #pragma unroll
      for (int i = 0; i < cols_per_thread; ++i) {
        if (algorithm == Algorithm::kSoftmax) {
          row_buf[i] = Exp(row_buf[i] - warp_max[row_id]);
          thread_sum[row_id] += row_buf[i];
        } else if (algorithm == Algorithm::kLogSoftmax) {
          row_buf[i] -= warp_max[row_id];
          thread_sum[row_id] += Exp(row_buf[i]);
        } else {
          __trap();   // 內(nèi)核的執(zhí)行被中止并在主機(jī)程序中引發(fā)中斷。
        }
      }
    }
    ComputeType warp_sum[rows_per_access];
    #pragma unroll
    for (int row_id = 0; row_id < rows_per_access; ++row_id) {
      warp_sum[row_id] = WarpAllReduce(thread_sum[row_id]);
    }
    #pragma unroll
    for (int row_id = 0; row_id < rows_per_access; ++row_id) {
      ComputeType* row_buf = buf[row_id];
      #pragma unroll
      for (int i = 0; i < cols_per_thread; ++i) {
        if (algorithm == Algorithm::kSoftmax) {
          row_buf[i] = Div(row_buf[i], warp_sum[row_id]);
        } else if (algorithm == Algorithm::kLogSoftmax) {
          row_buf[i] -= Log(warp_sum[row_id]);
        } else {
          __trap();
        }
      }
      #pragma unroll
      for (int i = 0; i < num_packs; ++i) {
        const int col = (i * thread_group_width + lane_id) * pack_size;
        if (!padding || col < cols) {
          store.template store(row_buf + i * pack_size, row + row_id, col);
        }
      }
    }
  }
}
,>,>

具體代碼如上，在解讀之前，需要先介紹一下幾個(gè)重要參數(shù)的意義。

algorithm：代表所使用的的算法，有 Algorithm::kSoftmax 和 Algorithm::kLogSoftmax。

global_thread_group_id：當(dāng)前線程組的全局索引

lane_id：當(dāng)前線程在線程組內(nèi)的索引

首先在核函數(shù)內(nèi)部做了幾個(gè)編譯期斷言操作，確保核函數(shù)能夠正常啟動。然后在寄存器中定義了一個(gè)二維數(shù)組 buf[rows_per_access][cols_per_thread] 用來存儲矩陣中的數(shù)據(jù)，我們知道，寄存器中的變量只能對當(dāng)前線程可見，每個(gè)線程中都有一個(gè)變量 buf，但是存儲的值可以不同，這里是為了減少對全局內(nèi)存的讀取，所以給每個(gè)線程都定義一個(gè)寄存器變量用于存儲該線程處理的矩陣元素。
接著是一個(gè) Grip-loop 的循環(huán)，因?yàn)橛锌赡芫仃囆袛?shù)過大導(dǎo)致前面求 num_blocks 的時(shí)候是根據(jù)硬件參數(shù)選取的，這時(shí)候每個(gè)線程不止處理一次，所以循環(huán)步長設(shè)置為網(wǎng)格大小。Grip-loop 內(nèi)部定義了一個(gè)寄存器變量 thread_max[rows_per_access]，這個(gè)數(shù)組用來存儲當(dāng)前線程處理的元素中的每一行的最大值。接下來就是一個(gè) reduceMax 操作。
（1）reduceMax
如圖 2，每個(gè)線程處理了多個(gè) Pack 的數(shù)據(jù)，求最大值需要兩層循環(huán)。第一層循環(huán)中把一個(gè) Pack 的矩陣元素 load 到 buf 數(shù)組中，這里主要是要理解 col 變量的含義，結(jié)合圖 2 的示意圖不難理解，相鄰的線程讀取相鄰的 Pack 的目的是讓一個(gè)線程束中各線程單次訪問的數(shù)據(jù)在內(nèi)存中相鄰，這是一個(gè)合并訪問的概念，目的是提升訪問效率。第二層循環(huán)中對單個(gè) Pack 中的元素求最大值存到 thread_max 中。
注意，這時(shí)候 thread_max 中存的只是每個(gè)線程內(nèi)部處理的元素的最大值，但是 reduceMax 操作要獲取的是矩陣每一行的最大值，由于 WarpSoftmax 的應(yīng)用范圍就是一個(gè)線程組處理一行數(shù)據(jù)，所以再對線程組內(nèi)所有的 thread_max 求最大值即可。前面說過，每個(gè)線程內(nèi)部都有一個(gè) thread_max 變量，對這些變量求最大值，必然要在線程間進(jìn)行通信，源碼中使用了 WarpAllReduce() 函數(shù)完成了這一操作得到了矩陣每一行的最大值 warp_max，核心就是利用了線程束洗牌指令 __shfl_xor_sync 完成了一個(gè)束內(nèi)折半規(guī)約操作，筆者之前在另一篇文章也有介紹：【CUDA編程】CUDA編程中的并行規(guī)約問題。有興趣的讀者可以去 cuda 官網(wǎng)詳細(xì)了解一下束內(nèi)洗牌指令的用法，當(dāng)然了這里也可以直接使用共享內(nèi)存存儲數(shù)據(jù)，我們知道共享內(nèi)存在整個(gè) block 都是可見的，也就不需要使用束內(nèi)通信，但是從訪問性能出發(fā)，共享內(nèi)存是不如寄存器快的，所以 oneflow 選擇了寄存器。,>

template class ReductionOp, typename T, int thread_group_width = kWarpSize>
__inline__ __device__ T WarpAllReduce(T val) {
  for (int mask = thread_group_width / 2; mask > 0; mask /= 2) {
    val = ReductionOp()(val, __shfl_xor_sync(0xffffffff, val, mask));
  }
  return val;
}

（2）reduceSum
接下來就是 reduceSum 操作，這里源碼提供了兩種算法： Algorithm::kSoftmax 和 Algorithm::kLogSoftmax。kSoftmax 就是公式（2）中的計(jì)算公式，kLogSoftmax 計(jì)算的是 計(jì)算公式如下：

reduceSum 的計(jì)算思路和 reduceMax 相同，先在寄存器定義一個(gè)變量 thread_sum 然后求出各個(gè)線程內(nèi)的指數(shù)和，最后束內(nèi)規(guī)約求每一行的指數(shù)和 warp_sum。
broadcastSub、exp、broadcastDiv 這三個(gè)操作比較簡單，其邏輯就直接包含在兩個(gè)規(guī)約操作的實(shí)現(xiàn)代碼里，這里不再贅述，至此 WarpSoftmax 源碼解讀完畢，有興趣的讀者可以自行嘗試。調(diào)用時(shí)可以將矩陣 cols 限制在 1024 以內(nèi)調(diào)用 DispatchSoftmax 函數(shù)，也可以直接調(diào)用 DispatchSoftmaxWarpImpl 函數(shù)。

4 小結(jié)

總結(jié)一下 WarpSoftmax 源碼中的一些值得注意的內(nèi)容。

數(shù)據(jù) Pack 可以有效地提升訪問帶寬，pack_size 可以根據(jù) cuda 中最大支持一次 128 bit 的讀寫來確定最大值。

WarpSoftmax 的核心就是束內(nèi)規(guī)約，利用了束內(nèi)線程可互相訪問寄存器的特性提高效率，但受制于單個(gè)線程可使用的寄存器大小，所以 WarpSoftmax 不適用于矩陣列數(shù)比較大的場景。

源碼中對于 pack_size 和 row_per_access 的確定都比較簡單粗暴，可以進(jìn)行更細(xì)致的處理。

審核編輯：湯梓紅