解析OneFlow Element-Wise算子實(shí)現(xiàn)方法

0x0. 前言

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

__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激活函數(shù)為例子，分別測試 dtype = float32，不同shape下的前向耗時以及帶寬利用率（NVIDIA A100-PCIE-40GB）。性能情況如下圖所示：

在這里插入圖片描述

在這里插入圖片描述

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

0x2. 用法

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

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

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

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

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

然后我們就可以使用 cuda::Binary 這個模板函數(shù)來完成這個二元的 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 的元素個數(shù)，輸出張量的數(shù)據(jù)首地址 out->mut_dptr(), 輸入張量的數(shù)據(jù)首地址 x->dptr() && y->dptr() ，最后一個參數(shù)則是當(dāng)前 Kernel 運(yùn)行的 cuda Stream對象。

0x3. 原理&&代碼實(shí)現(xiàn)解析

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

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

下面這段代碼展示了 OneFlow 針對 Element-Wise 算子是如何設(shè)置 GridSize 和 BlockSize 的。對應(yīng)的源碼地址為：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 直接被設(shè)置為了 256 ，對應(yīng) constexpr int kBlockSize = 256; 這行代碼，也就是說每個 Block 有 256 個線程。為什么是 256 ？大家不妨讀一下俊丞大佬這篇經(jīng)典的 給CUDA Kernel設(shè)置合適的 GridSize 和 Block Size 的文章 。文章中通過對 SM 的資源分析確定在主流的GPU上將 BlockSize 設(shè)置為 128 或者 256 是比較合適，在這里直接設(shè)置為了 256 。

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

選自O(shè)neFlow CUDA Kernel 中 grid_size 和 block_size 應(yīng)該怎么設(shè)置 一文

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

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

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

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

0x3.2 向量化數(shù)據(jù)訪問提升帶寬

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

在 OneFlow 的 Element-Wise 系列算子中，為了更好的進(jìn)行向量化的數(shù)據(jù)訪問，俊丞設(shè)計了如下的 Pack 數(shù)據(jù)結(jié)構(gòu)（代碼位置：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 先聲明了一個內(nèi)存對齊的數(shù)據(jù)類型 type ，注意這個 type 的內(nèi)存長度為 pack_size * sizeof(T) 。然后這里的 T 是我們需要進(jìn)行 Pack 的數(shù)據(jù)類型，而 pack_size 則表示我們需要 Pack 的元素個數(shù)。接下來我們看到 Pack 聯(lián)合體中聲明了 storage 和 elem 兩個數(shù)組，它們公用同一段對齊的內(nèi)存。然后 Pack 聯(lián)合體的入口有一個檢查: static_assert(sizeof(PackType) == sizeof(T) * pack_size, ""); 這是用來判斷我們之前聲明的 type 的內(nèi)存長度是否符合預(yù)期。

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

接下來，OneFlow 定義了真正要執(zhí)行數(shù)據(jù) Pack 的數(shù)據(jù)結(jié)構(gòu) Packed 并且定義了計算 PackSize 的工具函數(shù)。代碼位置為：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。然后對于各種數(shù)據(jù)類型來說，Half 數(shù)據(jù)類型的 bit 數(shù)是最少的即 16，所以一次性可以支持 Pack 8個half類型的數(shù)據(jù)，4個float32的數(shù)據(jù)，以此類推。所以這里的定義的 kMaxPackSize 表示 128/16=8 ，然后 kMaxPackBytes 則表示最大可以 Pack 的 byte 數(shù) 。

請注意區(qū)分 bit 和 byte 。

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

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

在這里插入圖片描述

首先這里定義了一個 HasApply2 類用來判斷是否可以支持一次性Pack 2個 char/int8/half2 類型的元素，這個地方是一個針對 int8/half2/char 數(shù)據(jù)類型的特殊處理，某些 Element-Wise 算子 Kernel 確實(shí)需要支持這種數(shù)據(jù)類型的計算。也就是說對于 half2 的話，在一個內(nèi)存訪問粒度里我們其實(shí)是可以 Pack 128 / 8 = 16個的。然后用了C++模板元編程的 std::enable_if 來控制針對 half2 類型的特殊 Pack 處理，也就是上圖代碼中的兩個 ApplyPack 函數(shù)?？梢钥吹綄τ?half2 類型的 Element-Wise 操作我們需要給對應(yīng)的 Functor 定義一個 Apply2 函數(shù)，比如對于 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 實(shí)現(xiàn)：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh#L133-L144 。

在這里插入圖片描述

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

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

最后，在循環(huán)之外，我們還需要根據(jù)傳入的 n_tail 參數(shù)，看一下還有沒有因?yàn)闆]有被 pack_size 整除的剩余元素，如果有的話就單獨(dú)調(diào)用 functor 進(jìn)行處理。

0x3.4 unroll

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

0x3.5 Kernel Launch的細(xì)節(jié)

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

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

0x4. 總結(jié)

以上就是我對 OneFlow Element-Wise 系列 CUDA 算子實(shí)現(xiàn)的解析，后續(xù)有空會持續(xù)更新學(xué)習(xí)到的新知識。

審核編輯：郭婷