實操教程｜CUDA WarpReduce 學習筆記
加入技術交流群
來源丨GiantPandaCV編輯丨極市平臺前言
之前看我司的 如何實現一個高效的Softmax CUDA kernel？多少還是有些細節(jié)沒有理解，恰好最近要做一個類似的 Reduce+Scale Kernel，原理機制還是比較相似的，所以翻出來重新理解一下。
背景我們定義這么一個ReduceScale操作：假設Tensor是(N, C)，首先在C這個維度計算出 absMax 值，我們記作scale，然后將每一行除以各自 行的scale，并最終輸出。一段樸素的numpy代碼是這樣：
import numpy as npBaseLine
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是直接調用cub庫中的 BlockReduce，一個 threadBlock 處理一行數據，計算出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;
}
}
}
參數中 x 是輸入數據，row_size是行的數量，col_size是列的大小測試機器是在 A100 40GB，為了讓結果區(qū)別比較明顯，我們將行數設置的比較大，輸入形狀為(55296*8, 128)，啟動的線程塊數目根據 如何設置CUDA Kernel中的grid_size和block_size？這篇文章來指定，這里比較粗暴的設置為(55296, 128)，數據類型為 Float，然后我們看下ncu的結果：
主要有這幾個指標，耗時為577.95us，吞吐量為 748.78Gb/s下面我們就根據 Softmax 優(yōu)化那篇文章所提及的點來逐步分析：
優(yōu)化1 數據Pack在之前的 高效、易用、可拓展我全都要：OneFlow CUDA Elementwise 模板庫的設計優(yōu)化思路 里很詳細的描述了如何做向量化讀寫，cuda里最大支持 128bit的讀寫，那么在數據類型為 Float 時，我們即可以將連續(xù)的4個 Float 打包到一起，一次性讀寫，提升吞吐。有了解過這方面的讀者應該就反應過來，誒 CUDA 里 不是剛好有一個類型叫 float4 就是干這件事的么，沒錯，但是為了更靈活的支持其他數據類型的向量化，我們利用union共享空間的特性實現了一個 Pack 類：
template<typename T, int N>優(yōu)化2 數據緩存
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];
};
整個算子邏輯是需要讀取一遍數據，計算scale，然后再讀取一遍數據，用scale進行縮放。很顯然這里我們讀取了兩遍數據，而數據是放在 Global Memory，帶寬比較低，會帶來讀取耗時。
一個很自然的想法是緩存到寄存器/Shared Memory中。由于這里我們只實現 WarpReduce 版本，所以我們是緩存到寄存器（其他版本可以參考開頭的優(yōu)化 Softmax 文章）中，減少一次對 Global Memory 的讀取。
template<typename T, typename IDX, int pack_size, int cols_per_thread>優(yōu)化3 使用Warp處理一行數據
__global__ void ReduceScaleWarpKernel(T* x, IDX row_size, IDX col_size) {
// ...
T buf[cols_per_thread];
// ...
相較 BaseLine，我們這里使用 warp 作為 Reduce 的單位進行操作，首先我們簡單看下 WarpReduce 的實現。
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 也經常看到，他是借助 __shfl_xor_sync 來實現比較，shuffle 指令允許同一線程束的兩個線程直接讀取對方的寄存器。
T __shfl_xor_sync(unsigned mask, T var, int laneMask, int width=warpSize);
其中 mask 是對線程的一個掩碼，我們一般所有線程都要參與計算，所以 mask 是 0xffffffffvar 則是寄存器值，laneMask 則是用來做按位異或的掩碼
這里引入一個概念叫 Lane，它表示線程束中的第幾號線程
示意圖如下：
當 laneMask = 16 時，其二進制為 0001 0000，然后線程束每個線程與 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
以此類推，最終得到一個 Warp 中的 absmax 值。接下來我們開始寫Kernel，模板參數分別為：
- T 數據類型
- IDX 索引類型
- pack_size pack數，比如float可以pack成4個，那對應pack_size=4
- cols_per_thread 每個線程需要處理的元素個數，比如一行大小是128，而我們一個warp有32個線程，那么這里就是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大小還是設置為128個線程，一個warp是32個線程，所以我們一個block可以組織成(32, 4)，包含4個warp。
根據這個層級劃分，我們可以計算出：
- global_thread_group_id 當前warp的全局index
- num_total_thread_group warp的總數量
- lane_id 線程束內的線程id
- num_packs pack的數目，即每個線程需要處理的元素個數 / 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;
由于存在啟動的warp的數量小于行的數量，所以我們要引入一個 for 循環(huán)。假設我們 cols = 256，那么線程束里的每個線程需要處理 256 /32 = 8個元素，而4個float可以pack到一起，所以我們線程束里的每個線程要處理2個pack，因此也要引入一個關于 num_packs 的 for 循環(huán)，以保證整一行都有被讀取到：
一次性讀取到一個 pack 后，我們再一個個放到寄存器當中緩存起來，并計算線程上的 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]));
}
}
接著我們調用 WarpAbsMaxAllReduce 進行reduce，獲得線程束中的 AbsMaxVal，并對緩存的數據進行數值縮放。
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;
}
完整代碼如下：
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;
}
}
}
這里我們方便測試，調用的時候就直接寫死一些模板參數
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 的結果：
吞吐量達到了1.3T，時間位333us，相比 BaseLine 快了 73 %。
總結還有更多特殊情況可以參考 Softmax 優(yōu)化的代碼，這里僅實現了第一個 Warp 計算方式。我感覺看著還行，真自己寫起來理解還是有點困難的，希望這篇博客能幫助讀者理解到一些 warp 的使用。
本文僅做學術分享，如有侵權，請聯系刪文。
*博客內容為網友個人發(fā)布，僅代表博主個人觀點，如有侵權請聯系工作人員刪除。
超級電容器相關文章:超級電容器原理
