實(shí)操教程｜CUDA WarpReduce 學(xué)習(xí)筆記

來(lái)源丨GiantPandaCV編輯丨極市平臺(tái)前言

之前看我司的 如何實(shí)現(xiàn)一個(gè)高效的Softmax CUDA kernel？多少還是有些細(xì)節(jié)沒(méi)有理解，恰好最近要做一個(gè)類(lèi)似的 Reduce+Scale Kernel，原理機(jī)制還是比較相似的，所以翻出來(lái)重新理解一下。

我們定義這么一個(gè)ReduceScale操作：假設(shè)Tensor是(N, C)，首先在C這個(gè)維度計(jì)算出 absMax 值，我們記作scale，然后將每一行除以各自 行的scale，并最終輸出。一段樸素的numpy代碼是這樣：

import numpy as np  
  
  
N = 1000  
C = 128  
x = np.random.randn(N, C)  
scale = np.expand_dims(np.max(np.abs(x), axis=1), 1)  
out = x / scale  
print(out.shape)  

這里我們BaseLine是直接調(diào)用cub庫(kù)中的 BlockReduce，一個(gè) threadBlock 處理一行數(shù)據(jù)，計(jì)算出AbsMaxVal，然后再縮放，代碼如下：

#include "cuda.h"  
#include "cub/cub.cuh"  
  
constexpr int kReduceBlockSize = 128;  
  
template<typename T>  
__device__ T abs_func(const T& a) {  
  return abs(a);  
}  
  
  
template<typename T>  
__device__ T max_func(const T a, const T b) {  
  return a > b ? a : b;  
}  
  
template<typename T>  
struct AbsMaxOp {  
  __device__ __forceinline__ T operator()(const T& a, const T& b) const {  
    return max_func(abs_func(a), abs_func(b));  
  }  
};  
  
template<typename T>  
__inline__ __device__ T BlockAllReduceAbsMax(T val) {  
  typedef cub::BlockReduce<T, kReduceBlockSize> BlockReduce;  
  __shared__ typename BlockReduce::TempStorage temp_storage;  
  __shared__ T final_result;  
  T result = BlockReduce(temp_storage).Reduce(val, AbsMaxOp<T>());  
  if (threadIdx.x == 0) { final_result = result; }  
  __syncthreads();  
  return final_result;  
}  
  
template<typename T, typename IDX>  
__global__ void ReduceScaleBlockKernel(T* x, IDX row_size, IDX col_size) {  
  for(int32_t row = blockIdx.x, step=gridDim.x; row < row_size; row+= step){  
    T thread_scale_factor = 0.0;   
    for(int32_t col=threadIdx.x; col < col_size; col+= blockDim.x){  
      IDX idx = row * col_size + col;   
      T x_val = x[idx];  
      thread_scale_factor = max_func(thread_scale_factor, abs_func(x_val));   
    }  
    T row_scale_factor = BlockAllReduceAbsMax<T>(thread_scale_factor);   
    for(int32_t col=threadIdx.x; col < col_size; col+=blockDim.x){  
      IDX idx = row * col_size + col;   
      x[idx] /= row_scale_factor;  
    }  
  }  
}  

參數(shù)中 x 是輸入數(shù)據(jù)，row_size是行的數(shù)量，col_size是列的大小測(cè)試機(jī)器是在 A100 40GB，為了讓結(jié)果區(qū)別比較明顯，我們將行數(shù)設(shè)置的比較大，輸入形狀為(55296*8, 128)，啟動(dòng)的線程塊數(shù)目根據(jù) 如何設(shè)置CUDA Kernel中的grid_size和block_size？這篇文章來(lái)指定，這里比較粗暴的設(shè)置為(55296, 128)，數(shù)據(jù)類(lèi)型為 Float，然后我們看下ncu的結(jié)果：

主要有這幾個(gè)指標(biāo)，耗時(shí)為577.95us，吞吐量為 748.78Gb/s下面我們就根據(jù) Softmax 優(yōu)化那篇文章所提及的點(diǎn)來(lái)逐步分析：

在之前的 高效、易用、可拓展我全都要：OneFlow CUDA Elementwise 模板庫(kù)的設(shè)計(jì)優(yōu)化思路 里很詳細(xì)的描述了如何做向量化讀寫(xiě)，cuda里最大支持 128bit的讀寫(xiě)，那么在數(shù)據(jù)類(lèi)型為 Float 時(shí)，我們即可以將連續(xù)的4個(gè) Float 打包到一起，一次性讀寫(xiě)，提升吞吐。有了解過(guò)這方面的讀者應(yīng)該就反應(yīng)過(guò)來(lái)，誒 CUDA 里 不是剛好有一個(gè)類(lèi)型叫 float4 就是干這件事的么，沒(méi)錯(cuò)，但是為了更靈活的支持其他數(shù)據(jù)類(lèi)型的向量化，我們利用union共享空間的特性實(shí)現(xiàn)了一個(gè) Pack 類(lèi)：

template<typename T, int N>  
struct GetPackType {  
  using type = typename std::aligned_storage<N * sizeof(T), N * sizeof(T)>::type;  
};  
  
template<typename T, int N>  
using PackType = typename GetPackType<T, N>::type;  
  
template<typename T, int N>  
union Pack {  
  static_assert(sizeof(PackType<T, N>) == sizeof(T) * N, "");  
  __device__ Pack() {  
    // do nothing  
  }  
  PackType<T, N> storage;  
  T elem[N];  
};  

整個(gè)算子邏輯是需要讀取一遍數(shù)據(jù)，計(jì)算scale，然后再讀取一遍數(shù)據(jù)，用scale進(jìn)行縮放。很顯然這里我們讀取了兩遍數(shù)據(jù)，而數(shù)據(jù)是放在 Global Memory，帶寬比較低，會(huì)帶來(lái)讀取耗時(shí)。

一個(gè)很自然的想法是緩存到寄存器/Shared Memory中。由于這里我們只實(shí)現(xiàn) WarpReduce 版本，所以我們是緩存到寄存器（其他版本可以參考開(kāi)頭的優(yōu)化 Softmax 文章）中，減少一次對(duì) Global Memory 的讀取。

template<typename T, typename IDX, int pack_size, int cols_per_thread>  
__global__ void ReduceScaleWarpKernel(T* x, IDX row_size, IDX col_size) {  
    // ...  
    T buf[cols_per_thread];  
    // ...  

相較 BaseLine，我們這里使用 warp 作為 Reduce 的單位進(jìn)行操作，首先我們簡(jiǎn)單看下 WarpReduce 的實(shí)現(xiàn)。

template<typename T>  
struct AbsMaxOp {  
  __device__ __forceinline__ T operator()(const T& a, const T& b) const {  
    return max_func(abs_func(a), abs_func(b));  
  }  
};  
  
template<typename T>  
__inline__ __device__ T WarpAbsMaxAllReduce(T val){  
    for(int lane_mask = kWarpSize/2; lane_mask > 0; lane_mask /= 2){  
        val = AbsMaxOp<T>()(val, __shfl_xor_sync(0xffffffff, val, lane_mask));   
    }  
    return val;   
}  

這段代碼在別的 BlockReduce 也經(jīng)常看到，他是借助 __shfl_xor_sync 來(lái)實(shí)現(xiàn)比較，shuffle 指令允許同一線程束的兩個(gè)線程直接讀取對(duì)方的寄存器。

T __shfl_xor_sync(unsigned mask, T var, int laneMask, int width=warpSize);  

其中 mask 是對(duì)線程的一個(gè)掩碼，我們一般所有線程都要參與計(jì)算，所以 mask 是 0xffffffffvar 則是寄存器值，laneMask 則是用來(lái)做按位異或的掩碼

這里引入一個(gè)概念叫 Lane，它表示線程束中的第幾號(hào)線程

示意圖如下：

當(dāng) laneMask = 16 時(shí)，其二進(jìn)制為 0001 0000，然后線程束每個(gè)線程與 laneMask 做異或操作如：

• 0000 0000 xor 0001 0000 = 0001 0000 = 16
• 0000 0001 xor 0001 0000 = 0001 0001 = 17
• 0000 0010 xor 0001 0000 = 0001 0010 = 18

以此類(lèi)推，最終得到一個(gè) Warp 中的 absmax 值。接下來(lái)我們開(kāi)始寫(xiě)Kernel，模板參數(shù)分別為：

• T 數(shù)據(jù)類(lèi)型
• IDX 索引類(lèi)型
• pack_size pack數(shù)，比如float可以pack成4個(gè)，那對(duì)應(yīng)pack_size=4
• cols_per_thread 每個(gè)線程需要處理的元素個(gè)數(shù)，比如一行大小是128，而我們一個(gè)warp有32個(gè)線程，那么這里就是128/32 = 4

template<typename T, typename IDX, int pack_size, int cols_per_thread>  
__global__ void ReduceScaleWarpKernel(T* x, IDX row_size, IDX col_size) {  
    // ...      
}  

跟BaseLine一樣，我們block大小還是設(shè)置為128個(gè)線程，一個(gè)warp是32個(gè)線程，所以我們一個(gè)block可以組織成(32, 4)，包含4個(gè)warp。

根據(jù)這個(gè)層級(jí)劃分，我們可以計(jì)算出：

• global_thread_group_id 當(dāng)前warp的全局index
• num_total_thread_group warp的總數(shù)量
• lane_id 線程束內(nèi)的線程id
• num_packs pack的數(shù)目，即每個(gè)線程需要處理的元素個(gè)數(shù) / pack_size

const int32_t global_thread_group_id = blockIdx.x * blockDim.y + threadIdx.y;   
    const int32_t num_total_thread_group = gridDim.x * blockDim.y;   
    const int32_t lane_id = threadIdx.x;   
    using LoadStoreType = PackType<T, pack_size>;  
    using LoadStorePack = Pack<T, pack_size>;  
    T buf[cols_per_thread];   
    constexpr int num_packs = cols_per_thread / pack_size;  

由于存在啟動(dòng)的warp的數(shù)量小于行的數(shù)量，所以我們要引入一個(gè) for 循環(huán)。假設(shè)我們 cols = 256，那么線程束里的每個(gè)線程需要處理 256 /32 = 8個(gè)元素，而4個(gè)float可以pack到一起，所以我們線程束里的每個(gè)線程要處理2個(gè)pack，因此也要引入一個(gè)關(guān)于 num_packs 的 for 循環(huán)，以保證整一行都有被讀取到：

一次性讀取到一個(gè) pack 后，我們?cè)僖粋€(gè)個(gè)放到寄存器當(dāng)中緩存起來(lái)，并計(jì)算線程上的 AbsMaxVal。

for(IDX row_idx = global_thread_group_id; row_idx < row_size; row_idx += num_total_thread_group){  
        T thread_abs_max_val = 0.0;   
        for(int pack_idx = 0; pack_idx < num_packs; pack_idx++){  
            const int32_t pack_offset = pack_idx * pack_size;   
            const int32_t col_offset = pack_idx * kWarpSize * pack_size + lane_id * pack_size;   
            const int32_t load_offset = (row_idx * col_size + col_offset) / pack_size;   
            LoadStorePack load_pack;   
            load_pack.storage = *(reinterpret_cast<LoadStoreType*>(x)+ load_offset);   
            #pragma unroll   
            for(int i = 0; i < pack_size; i++){  
                buf[pack_offset] = load_pack.elem[i];   
                thread_abs_max_val = max_func(thread_abs_max_val, abs_func(buf[pack_offset]));  
            }   
        }  

接著我們調(diào)用 WarpAbsMaxAllReduce 進(jìn)行reduce，獲得線程束中的 AbsMaxVal，并對(duì)緩存的數(shù)據(jù)進(jìn)行數(shù)值縮放。

T warp_max_val = WarpAbsMaxAllReduce<T>(thread_abs_max_val);   
        #pragma unroll  
        for (int col = 0; col < cols_per_thread; col++) {  
            buf[col] = buf[col] / warp_max_val;  
        }  

最后跟一開(kāi)始讀取類(lèi)似，我們將寄存器里的值再寫(xiě)回去，相關(guān)索引的計(jì)算邏輯都是一致的：

for(int pack_idx = 0; pack_idx < num_packs; pack_idx++){  
            const int32_t pack_offset = pack_idx * pack_size;   
            const int32_t col_offset = pack_idx * pack_size * kWarpSize + lane_id * pack_size;   
            const int32_t store_offset = (row_idx * col_size + col_offset) / pack_size;   
            LoadStorePack store_pack;   
            #pragma unroll   
            for(int i = 0; i < pack_size; i++){  
                store_pack.elem[i] = buf[pack_offset + i];   
            }   
            *(reinterpret_cast<LoadStoreType*>(x)+ store_offset) = store_pack.storage;   
        }  

完整代碼如下：

template<typename T>  
__inline__ __device__ T WarpAbsMaxAllReduce(T val){  
    for(int lane_mask = kWarpSize/2; lane_mask > 0; lane_mask /= 2){  
        val = AbsMaxOp<T>()(val, __shfl_xor_sync(0xffffffff, val, lane_mask));   
    }  
    return val;   
}  
  
template<typename T, typename IDX, int pack_size, int cols_per_thread>  
__global__ void ReduceScaleWarpKernel(T* x, IDX row_size, IDX col_size) {  
    const int32_t global_thread_group_id = blockIdx.x * blockDim.y + threadIdx.y;   
    const int32_t num_total_thread_group = gridDim.x * blockDim.y;   
    const int32_t lane_id = threadIdx.x;   
    using LoadStoreType = PackType<T, pack_size>;  
    using LoadStorePack = Pack<T, pack_size>;  
    T buf[cols_per_thread];   
    constexpr int num_packs = cols_per_thread / pack_size;  
    for(IDX row_idx = global_thread_group_id; row_idx < row_size; row_idx += num_total_thread_group){  
        T thread_abs_max_val = 0.0;   
        for(int pack_idx = 0; pack_idx < num_packs; pack_idx++){  
            const int32_t pack_offset = pack_idx * pack_size;   
            const int32_t col_offset = pack_idx * kWarpSize * pack_size + lane_id * pack_size;   
            const int32_t load_offset = (row_idx * col_size + col_offset) / pack_size;   
            LoadStorePack load_pack;   
            load_pack.storage = *(reinterpret_cast<LoadStoreType*>(x)+ load_offset);   
            #pragma unroll   
            for(int i = 0; i < pack_size; i++){  
                buf[pack_offset] = load_pack.elem[i];   
                thread_abs_max_val = max_func(thread_abs_max_val, abs_func(buf[pack_offset]));  
            }   
        }  
        T warp_max_val = WarpAbsMaxAllReduce<T>(thread_abs_max_val);   
        #pragma unroll  
        for (int col = 0; col < cols_per_thread; col++) {  
            buf[col] = buf[col] / warp_max_val;  
        }  
        for(int pack_idx = 0; pack_idx < num_packs; pack_idx++){  
            const int32_t pack_offset = pack_idx * pack_size;   
            const int32_t col_offset = pack_idx * pack_size * kWarpSize + lane_id * pack_size;   
            const int32_t store_offset = (row_idx * col_size + col_offset) / pack_size;   
            LoadStorePack store_pack;   
            #pragma unroll   
            for(int i = 0; i < pack_size; i++){  
                store_pack.elem[i] = buf[pack_offset + i];   
            }   
            *(reinterpret_cast<LoadStoreType*>(x)+ store_offset) = store_pack.storage;   
        }  
    }  
}  

這里我們方便測(cè)試，調(diào)用的時(shí)候就直接寫(xiě)死一些模板參數(shù)

constexpr int cols_per_thread = 128 / kWarpSize;   
ReduceScaleWarpKernel<float, int32_t, 4, cols_per_thread><<<55296, block_dim>>>(device_ptr, row_size, col_size);  

最后我們看一下 ncu 的結(jié)果：

吞吐量達(dá)到了1.3T，時(shí)間位333us，相比 BaseLine 快了 73 %。

還有更多特殊情況可以參考 Softmax 優(yōu)化的代碼，這里僅實(shí)現(xiàn)了第一個(gè) Warp 計(jì)算方式。我感覺(jué)看著還行，真自己寫(xiě)起來(lái)理解還是有點(diǎn)困難的，希望這篇博客能幫助讀者理解到一些 warp 的使用。

本文僅做學(xué)術(shù)分享，如有侵權(quán)，請(qǐng)聯(lián)系刪文。