這篇文章給大家介紹怎么實現(xiàn)一個高效的Softmax CUDA kernel，內(nèi)容非常詳細(xì)，感興趣的小伙伴們可以參考借鑒，希望對大家能有所幫助。
專業(yè)從事成都網(wǎng)站建設(shè)、做網(wǎng)站,高端網(wǎng)站制作設(shè)計,微信小程序,網(wǎng)站推廣的成都做網(wǎng)站的公司。優(yōu)秀技術(shù)團(tuán)隊竭力真誠服務(wù)，采用H5頁面制作+CSS3前端渲染技術(shù)，響應(yīng)式網(wǎng)站設(shè)計，讓網(wǎng)站在手機(jī)、平板、PC、微信下都能呈現(xiàn)。建站過程建立專項小組，與您實時在線互動，隨時提供解決方案，暢聊想法和感受。
Softmax操作是深度學(xué)習(xí)模型中最常用的操作之一。在深度學(xué)習(xí)的分類任務(wù)中，網(wǎng)絡(luò)最后的分類器往往是Softmax + CrossEntropy的組合：
盡管當(dāng)Softmax和CrossEntropy聯(lián)合使用時，其數(shù)學(xué)推導(dǎo)可以約簡，但還是有很多場景會單獨(dú)使用Softmax Op。如BERT的Encoder每一層的attention layer中就單獨(dú)使用了Softmax求解attention的概率分布；GPT-2的attention的multi-head部分也單獨(dú)使用了Softmax 等等。
深度學(xué)習(xí)框架中的所有計算的算子都會轉(zhuǎn)化為GPU上的CUDA kernel function，Softmax操作也不例外。Softmax作為一個被大多數(shù)網(wǎng)絡(luò)廣泛使用的算子，其CUDA Kernel實現(xiàn)高效性會影響很多網(wǎng)絡(luò)最終的訓(xùn)練速度。那么如何實現(xiàn)一個高效的Softmax CUDA Kernel？本文將會介紹OneFlow中優(yōu)化的Softmax CUDA Kernel的技巧，并跟cuDNN中的Softmax操作進(jìn)行實驗對比，結(jié)果表明，OneFlow深度優(yōu)化后的Softmax對顯存帶寬的利用率可以接近理論上限，遠(yuǎn)高于cuDNN的實現(xiàn)。
GPU基礎(chǔ)知識與CUDA性能優(yōu)化原則：
對GPU基礎(chǔ)知識的介紹以及CUDA性能優(yōu)化的原則與目標(biāo)可以參考之前的文章：
https://zhuanlan.zhihu.com/p/271740706
其中簡單介紹了GPU的硬件結(jié)構(gòu)與執(zhí)行原理：
Kernel：即CUDA Kernel function，是GPU的基本計算任務(wù)描述單元。每個Kernel都會根據(jù)配置參數(shù)在GPU上由非常多個線程并行執(zhí)行，GPU計算高效就是因為同時可以由數(shù)千個core（thread）同時執(zhí)行，計算效率遠(yuǎn)超CPU。
GPU的線程在邏輯上分為Thread、Block、Grid三級，硬件上分為core、warp兩級；
GPU內(nèi)存分為Global memory、Shared memory、Local memory三級。
GPU最主要提供的是兩種資源：計算資源和 顯存帶寬資源。如果我們能將這兩種資源充分利用，且對資源的需求無法再降低，那么性能就優(yōu)化到了極限，執(zhí)行時間就會最短。在大多數(shù)情況下，深度學(xué)習(xí)訓(xùn)練中的GPU計算資源是被充分利用的，而一個GPU CUDA Kernel的優(yōu)化目標(biāo)就是盡可能充分利用顯存帶寬資源。
如何評估一個CUDA Kernel是否充分利用了顯存帶寬資源？
對于顯存帶寬資源來說，“充分利用“指的是Kernel的有效顯存讀寫帶寬達(dá)到了設(shè)備顯存帶寬的上限，其中設(shè)備顯存帶寬可以通過執(zhí)行 cuda中的的bandwidthTest得到。 Kernel的有效顯存帶寬通過Kernel讀寫數(shù)據(jù)量和Kernel執(zhí)行時間進(jìn)行評估：
當(dāng) 前 K e r n e l 的 有 效 顯 存 帶 寬 = 讀 寫 數(shù) 據(jù) 量 / 執(zhí) 行 時 間 當(dāng)前Kernel的有效顯存帶寬 = 讀寫數(shù)據(jù)量 / 執(zhí)行時間 當(dāng)前Kernel的有效顯存帶寬=讀寫數(shù)據(jù)量/執(zhí)行時間
Naive的Softmax實現(xiàn)：
在介紹優(yōu)化技巧之前，我們先看看一個未經(jīng)優(yōu)化的Softmax Kernel的最高理論帶寬是多少。如下圖所示，一個最簡單的Softmax計算實現(xiàn)中，分別調(diào)用幾個基礎(chǔ)的CUDA Kernel function來完成整體的計算：
假設(shè)輸入的數(shù)據(jù)大小為D，shape = (num_rows, num_cols), 即 D = num_rows * num_cols ，最Navie的操作會多次訪問Global memory，其中：
ReduceMax = D + num_rows （read 為 D， write 為 num_rows）
BroadcaseSub = 2 * D + num_rows （read 為 D + num_rows，write 為 D）
Exp = 2 * D （read 、write 均為D）
ReduceSum = D + num_rows（read 為 D， write 為 num_rows）
BroadcastDive = 2 * D + num_rows （read 為 D + num_rows， write 為 D）
總共需要8 * D + 4 * num_rows 的訪存開銷。由于num_rows相比于D可以忽略，則Navie版本的Softmax CUDA Kernel需要訪問至少8倍數(shù)據(jù)的顯存，即：
N a i v e S o f t m a x K e r n e l 有 效 顯 存 帶 寬 < 理 論 帶 寬 / 8 Naive Softmax Kernel 有效顯存帶寬 < 理論帶寬 / 8 NaiveSoftmaxKernel有效顯存帶寬<理論帶寬/8
對于GeForce RTX? 3090顯卡，其理論帶寬上限為936GB/s，那么Navie版本的Softmax CUDA Kernel利用顯存帶寬的上界就是936 / 8 = 117 GB/s。
在文章 https://zhuanlan.zhihu.com/p/271740706 里，我們在方法：借助shared memory合并帶有Reduce計算的Kernel中提到了對Softmax Kernel的訪存優(yōu)化到了 2 * D，但這仍然沒有優(yōu)化到極致。在本文的優(yōu)化后，大多數(shù)場景下OneFlow的Softmax CUDA Kernel的顯存帶寬利用可以逼近理論帶寬。
OneFlow 與 cuDNN 的對比結(jié)果
我們對OneFlow深度優(yōu)化后的Softmax CUDA Kernel 與 cuDNN中的Softmax Kernel的訪存帶寬進(jìn)行了對比測試，測試結(jié)果如下：
Softmax Kernel Bandwidth：
Softmax Grad Kernel Bandwidth：
其中測試環(huán)境是 GeForce RTX? 3090 GPU，數(shù)據(jù)類型是half，Softmax的Shape = (49152, num_cols)，其中49152 = 32 * 12 * 128，是BERT-base網(wǎng)絡(luò)中的attention Tensor的前三維，我們固定了前三維，將最后一維動態(tài)變化，測試了從32到32768不同大小的Softmax前向Kernel和反向Kernel的有效顯存帶寬。從上面兩張圖中可以看出OneFlow在多數(shù)情況下都可以逼近理論帶寬（800GB/s左右，與cudaMemcpy的訪存帶寬相當(dāng)。官方公布的理論帶寬為936GB/S）。并且在所有情況下，OneFlow的CUDA Kernel的有效訪存帶寬都優(yōu)于cuDNN的實現(xiàn)。
OneFlow深度優(yōu)化Softmax CUDA Kernel的技巧
Softmax 函數(shù)的輸入形狀為:(num_rows, num_cols)，num_cols的變化，會對有效帶寬產(chǎn)生影響；因為，沒有一種 通用的優(yōu)化方法可以實現(xiàn)在所有 num_cols的情況下都是傳輸最優(yōu)的。所以，在 OneFlow 中采用分段函數(shù)優(yōu)化SoftmaxKernel：對于不同 num_cols范圍，選擇不同的實現(xiàn)，以在所有情況下都能達(dá)到較高的有效帶寬。見 Optimize softmax cuda kernel
OneFlow 中分三種實現(xiàn)，分段對 softmax 進(jìn)行優(yōu)化：
(1) 一個 Warp處理一行的計算，適用于 num_cols <= 1024 情況
硬件上并行執(zhí)行的32個線程稱之為一個warp，同一個warp的32個thread執(zhí)行同一條指令。 warp是GPU調(diào)度執(zhí)行的基本單元
(2) 一個 Block 處理一行的計算，借助 Shared Memory 保存中間結(jié)果數(shù)據(jù)，適用于需要的 Shared Memory 資源滿足 Kernel Launch 的可啟動條件的情況，在本測試環(huán)境中是 1024 < num_cols <= 4096
(3) 一個 Block 處理一行的計算，不使用 Shared Memory，重復(fù)讀輸入 x，適用于不支持(1)、（2)的情況
下面以前向計算為例分別對三種實現(xiàn)進(jìn)行介紹：
實現(xiàn)1：每個Warp處理一行元素。
每個 Warp 處理一行元素，行 Reduce 需要做 Warp 內(nèi)的 Reduce 操作，這里實現(xiàn) WarpAllReduce 完成 Warp 內(nèi)各線程間的求 Global Max 和 Global Sum 操作，WarpAllReduce 是利用Warp級別原語 __shfl_xor_sync 實現(xiàn)的，代碼如下。
template typename ReductionOp, typename T>
__inline__ __device__ T WarpAllReduce(T val) {
for (int mask = kWarpSize / 2; mask > 0; mask /= 2) {
val = ReductionOp()(val, __shfl_xor_sync(0xffffffff, val, mask));
}
return val;
}
主體循環(huán)邏輯如下，首先根據(jù) num_cols信息算出每個線程要處理的cols_per_thread，每個線程分配cols_per_thread大小的寄存器，將輸入x讀到寄存器中，后續(xù)計算均從寄存器中讀取。
理論上來說，每個 Warp 處理一行元素的速度是最快的，但是由于需要使用寄存器將輸入x緩存起來，而寄存器資源是有限的，并且在 num_cols>1024 情況下，使用(2)的 shared memory 方法已經(jīng)足夠快了，因此僅在 num_cols<=1024 時采用 Warp 的實現(xiàn)。
template
__global__ void SoftmaxWarpImpl(const int64_t rows, const int64_t cols, const T* x, T* y) {
static_assert(cols_per_thread % pack_size == 0, "");
constexpr int num_packs = cols_per_thread / pack_size;
assert(cols <= cols_per_thread * kWarpSize);
using ComputeType = typename GetComputeType::type;
ComputeType buf[cols_per_thread];
const int global_warp_id = blockIdx.x * blockDim.y + threadIdx.y;
const int num_global_warp = gridDim.x * blockDim.y;
const int lane_id = threadIdx.x;
for (int64_t row = global_warp_id; row < rows; row += num_global_warp) {
const int64_t row_offset = row * cols;
const T* row_x = x + row_offset;
T* row_y = y + row_offset;
ComputeType thread_max = -Inf();
#pragma unroll
for (int pack_id = 0; pack_id < num_packs; ++pack_id) {
const int col = (pack_id * kWarpSize + lane_id) * pack_size;
if (!padding || col < cols) {
MultiFetch()(buf + pack_id * pack_size, row_x + col);
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
thread_max = max(thread_max, buf[pack_id * pack_size + i]);
}
} else {
#pragma unroll
for (int i = 0; i < pack_size; ++i) { buf[pack_id * pack_size + i] = -Inf(); }
}
}
const ComputeType warp_max = WarpAllReduce(thread_max);
ComputeType thread_sum = 0;
#pragma unroll
for (int i = 0; i < cols_per_thread; ++i) {
buf[i] = exp(buf[i] - warp_max);
thread_sum += buf[i];
}
const ComputeType warp_sum = WarpAllReduce(thread_sum);
#pragma unroll
for (int i = 0; i < cols_per_thread; ++i) { buf[i] = buf[i] / warp_sum; }
#pragma unroll
for (int i = 0; i < num_packs; ++i) {
const int col = (i * kWarpSize + lane_id) * pack_size;
if (!padding || col < cols) {
MultiStore()(row_y + col, buf + i * pack_size);
}
}
}
}
實現(xiàn)2： 一個Block處理一行元素
一個 Block 處理一行元素，行 Reduce 需要做 Block 內(nèi)的 Reduce 操作，需要做 Block 內(nèi)線程同步，利用 BlockAllReduce 完成 Warp 內(nèi)各線程間的求 Global Max 和 Global Sum 操作。BlockAllReduce 是借助 Cub 的 BlockReduce 方法實現(xiàn)的，代碼如下：
template typename ReductionOp, typename T, int block_size>
__inline__ __device__ T BlockAllReduce(T val) {
typedef cub::BlockReduce BlockReduce;
__shared__ typename BlockReduce::TempStorage temp_storage;
__shared__ T result_broadcast;
T result = BlockReduce(temp_storage).Reduce(val, ReductionOp());
if (threadIdx.x == 0) { result_broadcast = result; }
__syncthreads();
return result_broadcast;
}
主體循環(huán)邏輯如下，根據(jù) num_cols算出需要的 Shared Memory 大小作為 Launch Kernel 參數(shù)，借助 Shared Memory 保存輸入，后續(xù)計算直接從 Shared Memory 讀取。
由于 SM 內(nèi)的 Shared Memory 資源同樣有限，因此當(dāng) num_cols超過一定范圍，kernel 啟動時申請 Shared Memory 超過最大限制，就會出現(xiàn)無法啟動的問題，因此，僅在調(diào)用 cudaOccupancyMaxActiveBlocksPerMultiprocessor 返回值大于0時采用 Shared Memory 方案。
此外，需要注意的是，由于 Block 內(nèi)線程要做同步，當(dāng) SM 中正在調(diào)度執(zhí)行的一個 Block 到達(dá)同步點時，SM 內(nèi)可執(zhí)行 Warp 逐漸減少，若同時執(zhí)行的 Block 只有一個，則 SM 中可同時執(zhí)行的 Warp 會在此時逐漸降成0，會導(dǎo)致計算資源空閑，造成浪費(fèi)，若此時同時有其他 Block 在執(zhí)行，則在一個 Block 到達(dá)同步點時仍然有其他 Block 可以執(zhí)行。當(dāng) block_size 越小時，SM 可同時調(diào)度的 Block 越多，因此在這種情況下 block_size 越小越好。但是當(dāng)在調(diào)大 block_size，SM 能同時調(diào)度的 Block 數(shù)不變的情況下，block_size 應(yīng)該是越大越好，越大就有越好的并行度。因此代碼中在選擇 block_size 時，對不同 block_size 都計算了 cudaOccupancyMaxActiveBlocksPerMultiprocessor，若結(jié)果相同，使用較大的 block_size。
template
__global__ void SoftmaxBlockSMemImpl(const int64_t rows, const int64_t cols, const T* x, T* y) {
using ComputeType = typename GetComputeType::type;
extern __shared__ __align__(sizeof(ComputeType)) 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;
for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
const int64_t row_offset = row * cols;
const T* row_x = x + row_offset;
T* row_y = y + row_offset;
ComputeType thread_max = -Inf();
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
MultiFetch()(pack, row_x + 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]);
}
}
const ComputeType row_max = BlockAllReduce(thread_max);
ComputeType thread_sum = 0;
for (int col = tid; col < cols; col += block_size) {
const ComputeType exp_x = exp(buf[col] - row_max);
buf[col] = exp_x;
thread_sum += exp_x;
}
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) {
pack[i] = buf[i * num_packs + pack_id] / row_sum;
thread_max = max(thread_max, pack[i]);
}
MultiStore()(row_y + pack_id * pack_size, pack);
}
}
}
實現(xiàn)3：一個Block處理一行的元素，不使用Shared Memory，重復(fù)讀輸入x
和實現(xiàn)2一樣，仍然是一個 Block 處理一行元素，不同的是，不再用 Shared Memory 緩存x，而是在每次計算時重新讀輸入 x，這種實現(xiàn)沒有最大 num_cols的限制，可以支持任意大小。
此外，需要注意的是，在這種實現(xiàn)中，block_size 應(yīng)該設(shè)越大越好，block_size 越大，SM 中能同時并行執(zhí)行的 Block 數(shù)就越少，對 cache 的需求就越少，就有更多機(jī)會命中 Cache，多次讀x不會多次訪問 Global Memory，因此在實際測試中，在能利用 Cache 情況下，有效帶寬不會因為讀3次x而降低幾倍。
template
__global__ void SoftmaxBlockUncachedImpl(const int64_t rows, const int64_t cols, const T* x, T* y) {
using ComputeType = typename GetComputeType::type;
const int tid = threadIdx.x;
assert(cols % pack_size == 0);
const int num_packs = cols / pack_size;
for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
const int64_t row_offset = row * cols;
const T* row_x = x + row_offset;
T* row_y = y + row_offset;
ComputeType thread_max = -Inf();
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
MultiFetch()(pack, row_x + pack_id * pack_size);
#pragma unroll
for (int i = 0; i < pack_size; ++i) { thread_max = max(thread_max, pack[i]); }
}
const ComputeType row_max = BlockAllReduce(thread_max);
ComputeType thread_sum = 0;
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
MultiFetch()(pack, row_x + pack_id * pack_size);
#pragma unroll
for (int i = 0; i < pack_size; ++i) { thread_sum += exp(pack[i] - row_max); }
}
const ComputeType row_sum = BlockAllReduce(thread_sum);
for (int pack_id = tid; pack_id < num_packs; pack_id += block_size) {
ComputeType pack[pack_size];
MultiFetch()(pack, row_x + pack_id * pack_size);
#pragma unroll
for (int i = 0; i < pack_size; ++i) { pack[i] = exp(pack[i] - row_max) / row_sum; }
MultiStore()(row_y + pack_id * pack_size, pack);
}
}
}
公共優(yōu)化技巧：
除了以上介紹的針對 softmax 的分段優(yōu)化技巧外，OneFlow 在 softmax 實現(xiàn)中，還使用了一些通用的公共優(yōu)化技巧，他們不僅應(yīng)用于 softmax 中，也應(yīng)用在其它 kernel 實現(xiàn)中。在此介紹兩種：
1、將 Half 類型 pack 成 Half2 進(jìn)行存取，在不改變延遲情況下提高指令吞吐，類似 CUDA template for element-wise kernels 的優(yōu)化
2、Shared Memory 中的 Bank Conflicts
CUDA 將 Shared Memory 按照4字節(jié)或8字節(jié)大小劃分到32個 bank 中，對于 Volta 架構(gòu)是4字節(jié)，這里以4字節(jié)為例，如下圖所示，0-128 bytes地址分別在bank 0-31中，128-256也分別在 bank0-31中。
注：此圖及以下 Bank Conflicts 圖片來自VOLTA Architecture and performance optimization
當(dāng)在一個 Warp 內(nèi)的不同線程訪問同一個 bank 的不同地址，就會出現(xiàn) Bank Conflicts。當(dāng)發(fā)生 Bank Conflicts 時，線程間只能順序訪問，增加延遲，下圖是一個 Warp 中 n 個線程同時訪問同一個 bank 中的不同地址時的延遲情況。
注：圖來自Dissecting the NVIDIA Volta GPU Architecture via Microbenchmarking
Bank Conflicts 的幾種情況：
若一個 Warp 內(nèi)每個線程讀4個字節(jié)，順序訪問，則不會有 Bank Conflicts，若一個 Warp 內(nèi)每個線程讀8個字節(jié)，則 Warp 內(nèi)0號線程訪問的地址在第0和第1個 bank，1號線程訪問的地址在第2和第3個 bank，以此類推，16號線程訪問地址又在第0和第1個 bank內(nèi)，和0號線程訪問了同一個bank的不同地址，此時即產(chǎn)生了 Bank Conflicts。
在前文的實現(xiàn)(2)中，給 Shared memory 賦值過程中，若采用下面方法，當(dāng) pack size=2，每個線程寫連續(xù)兩個4 byte 地址，就會產(chǎn)生 Bank Conflicts。
#pragma unroll
for (int j = 0; j < pack_size; ++j) {
buf[pack_id * pack_size * j] = pack[j];
thread_max = max(thread_max, pack[j]);
}因此，在實現(xiàn)(2)中，對Shared memory采用了新的內(nèi)存布局，避免了同一個Warp訪問相同bank的不同地址，避免了Bank Conflicts。
#pragma unroll
for (int j = 0; j < pack_size; ++j) {
buf[num_packs * j + pack_id] = pack[j];
thread_max = max(thread_max, pack[j]);
}關(guān)于怎么實現(xiàn)一個高效的Softmax CUDA kernel就分享到這里了，希望以上內(nèi)容可以對大家有一定的幫助，可以學(xué)到更多知識。如果覺得文章不錯，可以把它分享出去讓更多的人看到。
本文題目：怎么實現(xiàn)一個高效的SoftmaxCUDAkernel
本文URL：
http://weahome.cn/article/jhocjs.html
重慶分公司
重慶分公司
