0x0. 前言

由于CUDA水平太菜，所以一直沒寫過這方面的筆記。現在日常的工作中已經不能離開寫CUDA代碼，所以準備學習ZZK隨緣做一做CUDA的筆記記錄一下學習到的知識和技巧。這篇文章記錄的是閱讀OneFlow的Element-Wise系列CUDA算子實現方案學習到的技巧，希望可以幫助到一起入門CUDA的小伙伴們。Elemet-Wise算子指的是針對輸入Tensor進行逐元素操作，比如ReLU就是針對輸入Tensor的每個值進行判斷是否大于0，大于0的話輸出就是輸入否則就是0。用CUDA來表達最簡單的寫法就是：

__global__voidrelu_kernel(float*input,float*output){
int32_tidx=blockIdx.x*blockDim.x+threadIdx.x;
output[idx]=input[idx]>>(src,dst);

cudaDeviceSynchronize();
cudaFree(src);
cudaFree(dst);
return0;
}

雖然這種寫法非常簡單明了，但卻存在明顯的性能問題。所以這篇文章將基于OneFlow開源的Element-Wise CUDA算子方案來解釋如何寫一個高性能的Element-Wise CUDA算子。

0x1. 性能

以GELU激活函數為例子，分別測試 dtype = float32，不同shape下的前向耗時以及帶寬利用率（NVIDIA A100-PCIE-40GB）。性能情況如下圖所示：

在這里插入圖片描述

在這里插入圖片描述

可以看到對于 GeLU 來說，無論是性能還是帶寬 OneFlow 的實現都是更優的，接下來我們就來了解一下為什么 OneFlow 的 Element-Wise 算子性能可以做到更優。

0x2. 用法

OneFlow在 elementwise.cuh 文件中分別針對一元，二元，三元運算的 Element-Wise 操作實現了模板函數。在包含這個頭文件之后我們可以使用 cuda::Unary/Binary/Ternary 這幾個模板函數來針對我們自己定義的 Element-Wise 操作進行計算。注意，這里說的一元，二元，三元代表的是這個 Element-Wise 操作有幾個輸入 Tensor。

我們舉個例子，假設我們要做的 Element-Wise 操作是逐點乘法，也即有 2 個輸入Tensor x 和 y，然后 x 和 y的形狀和數據類型都是一致的。那么我們可以定義一個模板類：

template
structMultiplyFunctor{
OF_DEVICE_FUNCToperator()(Tx,Ty)const{
returnx*y;
}
};

這里 OF_DEVICE_FUNC 表示我們定義的這個函數既可以運行在 CPU 又可以運行在 GPU 上，它的定義是：

#ifdefined(__CUDACC__)
#defineOF_DEVICE_FUNCTION__device____host____forceinline__
#else
#defineOF_DEVICE_FUNCTIONinline
#endif

然后我們就可以使用 cuda::Binary 這個模板函數來完成這個二元的 Element-Wise 算子了。示例代碼如下：

constuser_op::Tensor*x=ctx->Tensor4ArgNameAndIndex("x",0);
constuser_op::Tensor*y=ctx->Tensor4ArgNameAndIndex("y",0);
user_op::Tensor*out=ctx->Tensor4ArgNameAndIndex("out",0);
constint64_telem_cnt=x->shape().elem_cnt();
OF_CUDA_CHECK（cuda::Binary(MultiplyFunctor(),elem_cnt,out->mut_dptr(),
x->dptr(),
y->dptr(),
ctx->device_ctx()->cuda_stream()));

這里的 x, y, out 分別代表這個 Element-Wise 操作的輸入輸出 Tensor，然后 element_cnt 表示 Tensor 的元素個數，輸出張量的數據首地址 out->mut_dptr(), 輸入張量的數據首地址 x->dptr() && y->dptr() ，最后一個參數則是當前 Kernel 運行的 cuda Stream對象。

0x3. 原理&&代碼實現解析

我個人認為這里有幾個要點，分別是一個線程處理多個數據，向量化數據訪問提升帶寬，設置合理的Block數量（GridSize）和線程數量（BlockSize）以及在合適的地方進行循環展開（unrool）以及一些編程上的技巧。

0x3.1 給 Element-Wise 操作設置合理的 GridSize 和 BlockSize

下面這段代碼展示了 OneFlow 針對 Element-Wise 算子是如何設置 GridSize 和 BlockSize 的。對應的源碼地址為：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh#L30-L52 。

constexprintkBlockSize=256;
constexprintkNumWaves=32;

inlinecudaError_tGetNumBlocks(int64_tn,int*num_blocks){
intdev;
{
cudaError_terr=cudaGetDevice(&dev);
if(err!=cudaSuccess){returnerr;}
}
intsm_count;
{
cudaError_terr=cudaDeviceGetAttribute(&sm_count,cudaDevAttrMultiProcessorCount,dev);
if(err!=cudaSuccess){returnerr;}
}
inttpm;
{
cudaError_terr=cudaDeviceGetAttribute(&tpm,cudaDevAttrMaxThreadsPerMultiProcessor,dev);
if(err!=cudaSuccess){returnerr;}
}
*num_blocks=std::max(1,std::min((n+kBlockSize-1)/kBlockSize,
sm_count*tpm/kBlockSize*kNumWaves));
returncudaSuccess;
}

這個地方 BlockSize 直接被設置為了 256 ，對應 constexpr int kBlockSize = 256; 這行代碼，也就是說每個 Block 有 256 個線程。為什么是 256 ？大家不妨讀一下俊丞大佬這篇經典的 給CUDA Kernel設置合適的 GridSize 和 Block Size 的文章 。文章中通過對 SM 的資源分析確定在主流的GPU上將 BlockSize 設置為 128 或者 256 是比較合適，在這里直接設置為了 256 。

確定了 BlockSize 之后需要確定 Kernel 啟動線程塊的數量，我一直覺得上述文章中對這一段的分析是尤其精彩的，這里再截圖展示一下：

選自OneFlow CUDA Kernel 中 grid_size 和 block_size 應該怎么設置 一文

根據這里的分析，對于 Element-Wise 操作要設置合適的 GridSize 不僅需要考慮元素的數量還要考慮由于 SM 硬件本身帶來的限制。如下公式所述：

*num_blocks=std::max(1,std::min((n+kBlockSize-1)/kBlockSize,
sm_count*tpm/kBlockSize*kNumWaves));

這里的 (n + kBlockSize - 1) / kBlockSize 就是根據 Element-Wise 操作的元素個數來計算需要啟動多少個線程塊，比如在文章開頭的例子中有 = 個元素，那么就一共需要 個線程塊。然后這里以GTX 3080Ti為例，它的SM個數也就是sm_count=80，每個SM最多調度的線程數tpm=1536，那么sm_count * tpm / kBlockSize * kNumWaves = 80 * 1536 / 256 * 32 = 15360，所以在這個例子中我們最終設置的線程塊個數為 588 個。

通過上述講解和分析我們已經確定了啟動 Element-Wise CUDA Kernel 的 GridSize 和 BlockSize。

0x3.2 向量化數據訪問提升帶寬

對于大多數 Element-Wise 算子來說，一般它們的計算量不會太大，所以它們的瓶頸一般在GPU的帶寬上。在 NVIDIA 的性能優化博客 https://developer.nvidia.com/blog/cuda-pro-tip-increase-performance-with-vectorized-memory-access/ 中提到，對于很多 CUDA 核函數我們都可以通過向量化數據訪問的方式來提升帶寬受限的 Kernel 的性能，特別是對于架構比較新的 GPU 向量化數據訪問的效果會更加明顯。

在 OneFlow 的 Element-Wise 系列算子中，為了更好的進行向量化的數據訪問，俊丞設計了如下的 Pack 數據結構（代碼位置：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh#L54-L70）：

template
structGetPackType{
usingtype=typenamestd::aligned_storage::type;
};

template
usingPackType=typenameGetPackType::type;

template
unionPack{
static_assert(sizeof(PackType)==sizeof(T)*pack_size,"");
__device__Pack(){
//donothing
}
PackTypestorage;
Telem[pack_size];
};

對GetPackType理解有誤請看知乎的修改后正確版本用了 std::aligned_storage 先聲明了一個內存對齊的數據類型 type ，注意這個 type 的內存長度為 pack_size * sizeof(T) 。然后這里的 T 是我們需要進行 Pack 的數據類型，而 pack_size 則表示我們需要 Pack 的元素個數。接下來我們看到 Pack 聯合體中聲明了 storage 和 elem 兩個數組，它們公用同一段對齊的內存。然后 Pack 聯合體的入口有一個檢查: static_assert(sizeof(PackType) == sizeof(T) * pack_size, ""); 這是用來判斷我們之前聲明的 type 的內存長度是否符合預期。

接下來我們從 https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh#L155-L194 這里可以看到這個 Pack 聯合體主要是用在 Kernel 啟動之前判斷 Element-Wise 操作的輸入輸出 Tensor 對應的數據指針地址是否滿足內存對齊的條件，如果不滿足則這個 Element-Wise 操作無法執行數據 Pack 。對應下圖2個畫紅色框的地方。

接下來，OneFlow 定義了真正要執行數據 Pack 的數據結構 Packed 并且定義了計算 PackSize 的工具函數。代碼位置為：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh#L72-L95 。

template
structalignas(sizeof(T)*pack_size)Packed{
__device__Packed(){
//donothing
}
union{
Telem[pack_size];
};
};

constexprintkMaxPackBytes=128/8;
constexprintkMaxPackSize=8;

constexprintMin(inta,intb){returna
constexprintPackSize(){
returnMin(kMaxPackBytes/sizeof(T),kMaxPackSize);
}

template
constexprintPackSize(){
returnMin(PackSize(),PackSize());
}

這里需要注意的是對于 CUDA 來說，最多支持 128 個 bit 的訪問粒度，也就是說 PackSize 的大小不能超過 128 個bit。然后對于各種數據類型來說，Half 數據類型的 bit 數是最少的即 16，所以一次性可以支持 Pack 8個half類型的數據，4個float32的數據，以此類推。所以這里的定義的 kMaxPackSize 表示 128/16=8 ，然后 kMaxPackBytes 則表示最大可以 Pack 的 byte 數 。

請注意區分 bit 和 byte 。

接下來 https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh#L97-L144 則是真正的為 Element-Wise 操作完成數據 Pack 并執行計算。

首先來看這段充滿技巧的代碼：

在這里插入圖片描述

首先這里定義了一個 HasApply2 類用來判斷是否可以支持一次性Pack 2個 char/int8/half2 類型的元素，這個地方是一個針對 int8/half2/char 數據類型的特殊處理，某些 Element-Wise 算子 Kernel 確實需要支持這種數據類型的計算。也就是說對于 half2 的話，在一個內存訪問粒度里我們其實是可以 Pack 128 / 8 = 16個的。然后用了C++模板元編程的 std::enable_if 來控制針對 half2 類型的特殊 Pack 處理，也就是上圖代碼中的兩個 ApplyPack 函數。可以看到對于 half2 類型的 Element-Wise 操作我們需要給對應的 Functor 定義一個 Apply2 函數，比如對于 Cast 操作的 Functor 定義如下：

template
structCastFunctor{
__device__Tooperator()(Fromfrom)const{returnstatic_cast(from);}
};

template
structCastFunctor::value>::type>{
__device__Tooperator()(halffrom)const{returnstatic_cast(static_cast(from));}

__device__voidApply2(To*to,consthalf*from)const{
constfloat2f2=__half22float2(*reinterpret_cast(from));
to[0]=static_cast(f2.x);
to[1]=static_cast(f2.y);
}
};

0x3.3 啟動 Kernel

我們接下來看一下 Element-Wise 的 Kernel 實現：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh#L133-L144 。

在這里插入圖片描述

在 Kernel 中我們發現每一個線程實際上處理了多個 Pack 后的數據，也即：for (int64_t i = global_tid; i < n_pack; i += blockDim.x * gridDim.x) 。初學者看到這個循環也許會比較疑惑，為什么它的步幅是 blockDim.x * gridDim.x ?? 這個 blockDim.x * gridDim.x 表示的是 CUDA 線程網格中的線程總數。假設線程網格中有 1280 個線程，線程 0 將計算元素 0、1280、2560 等。通過使用步幅等于網格大小的循環，確保了 warp 中的所有尋址都是單位步幅，可以獲得最大的內存合并。想了解更多細節可以查看：https://zhuanlan.zhihu.com/p/571320529 。

除此之外，使用這種技巧的還有個好處就是如果對于 Kernel 中存在每個線程都包含一個公共的操作，那么線程數的增多，也代表著這部分的開銷變大。這個時候我們減少線程的數量并循環進行處理的話那么這個公共操作的開銷就會更低。

最后，在循環之外，我們還需要根據傳入的 n_tail 參數，看一下還有沒有因為沒有被 pack_size 整除的剩余元素，如果有的話就單獨調用 functor 進行處理。

0x3.4 unroll

實際上就是代碼中的 #pragma unroll ，這個宏會對我們的 for 循環做循環展開，讓更多的指令可以并行執行。但容易想到，只有處理的數據沒有前后依賴關系的時候我們可以做。對于大多數的 ElementWise 算子來說一般是滿足這個條件的。

0x3.5 Kernel Launch的細節

在 https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh#L166-L181 這個位置 OneFlow 展示了 Element-Wise Kernel 的啟動細節，我們簡單注釋一下：

template
cudaError_tLaunchKernel(FactoryTfactory,int64_tn,R*r,constIN*...in,cudaStream_tstream){
constint64_tn_pack=n/pack_size;//根據元素個數和pack_size，計算pack數目，比如1026/4=256。
constint64_ttail_offset=n_pack*pack_size;//如果存在不被整除的情況，我們計算使用pack的偏移量：256*4；
constint64_tn_tail=n-tail_offset;////元素數目-偏移量=剩下的元素個數->1026-1024=2
intnum_blocks;
{
cudaError_terr=GetNumBlocks(n_pack,&num_blocks);//計算線程塊數目
if(err!=cudaSuccess){returnerr;}
}
ApplyGeneric<<>>(
factory,n_pack,reinterpret_cast*>(r),
(reinterpret_cast*>(in))...,n_tail,r+tail_offset,
(in+tail_offset)...);
returncudaPeekAtLastError();
}

0x4. 總結

以上就是我對 OneFlow Element-Wise 系列 CUDA 算子實現的解析，后續有空會持續更新學習到的新知識。

審核編輯：郭婷