0x0. 前言

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

__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;
}

雖然這種寫法非常簡(jiǎn)單明了，但卻存在明顯的性能問(wèn)題。所以這篇文章將基于OneFlow開(kāi)源的Element-Wise CUDA算子方案來(lái)解釋如何寫一個(gè)高性能的Element-Wise CUDA算子。

0x1. 性能

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

在這里插入圖片描述

在這里插入圖片描述

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

0x2. 用法

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

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

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

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

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

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

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

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

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

下面這段代碼展示了 OneFlow 針對(duì) Element-Wise 算子是如何設(shè)置 GridSize 和 BlockSize 的。對(duì)應(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;
}

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

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

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

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

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

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

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

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

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

在 OneFlow 的 Element-Wise 系列算子中，為了更好的進(jìn)行向量化的數(shù)據(jù)訪問(wèn)，俊丞設(shè)計(jì)了如下的 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];
};

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

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

接下來(lái)，OneFlow 定義了真正要執(zhí)行數(shù)據(jù) Pack 的數(shù)據(jù)結(jié)構(gòu) Packed 并且定義了計(jì)算 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());
}

這里需要注意的是對(duì)于 CUDA 來(lái)說(shuō)，最多支持 128 個(gè) bit 的訪問(wèn)粒度，也就是說(shuō) PackSize 的大小不能超過(guò) 128 個(gè)bit。然后對(duì)于各種數(shù)據(jù)類型來(lái)說(shuō)，Half 數(shù)據(jù)類型的 bit 數(shù)是最少的即 16，所以一次性可以支持 Pack 8個(gè)half類型的數(shù)據(jù)，4個(gè)float32的數(shù)據(jù)，以此類推。所以這里的定義的 kMaxPackSize 表示 128/16=8 ，然后 kMaxPackBytes 則表示最大可以 Pack 的 byte 數(shù) 。

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

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

首先來(lái)看這段充滿技巧的代碼：

在這里插入圖片描述

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

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

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

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

0x3.4 unroll

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

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

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

template
cudaError_tLaunchKernel(FactoryTfactory,int64_tn,R*r,constIN*...in,cudaStream_tstream){
constint64_tn_pack=n/pack_size;//根據(jù)元素個(gè)數(shù)和pack_size，計(jì)算pack數(shù)目，比如1026/4=256。
constint64_ttail_offset=n_pack*pack_size;//如果存在不被整除的情況，我們計(jì)算使用pack的偏移量：256*4；
constint64_tn_tail=n-tail_offset;////元素?cái)?shù)目-偏移量=剩下的元素個(gè)數(shù)->1026-1024=2
intnum_blocks;
{
cudaError_terr=GetNumBlocks(n_pack,&num_blocks);//計(jì)算線程塊數(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é)

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

審核編輯：郭婷