解析OneFlow BatchNorm相關(guān)算子實(shí)現(xiàn)

0x1. 前言

在ResNet中（https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py），關(guān)于BatchNorm的調(diào)用一共有兩種模式，第一種是ReLU接在BN之后：

out=self.bn1(out)
out=self.relu(out)

另外一種模式是殘差結(jié)構(gòu)引入的 BNAddReLU 的模式：

out=self.bn2(out)

ifself.downsampleisnotNone:
identity=self.downsample(x)

out+=identity
out=self.relu(out)

我們知道在 CUDA 優(yōu)化中常見的一個(gè)技巧是將一些ElementWise的算子融合到之前的計(jì)算密集型算子如卷積，矩陣乘等。在OneFlow中針對(duì)上述兩種情況并且cudnn無法fuse時(shí)分別進(jìn)行了fuse和優(yōu)化，本篇文章就來解析一下這里的代碼實(shí)現(xiàn)，體會(huì)其中的CUDA優(yōu)化技巧。這里的源碼開源在OneFlow的github倉庫：「https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/normalization_kernel.cu」 。如果本文對(duì)你產(chǎn)生了啟發(fā)，不妨為OneFlow投個(gè)star。

0x2. 代碼解析

0x2.1 CUDNN BatchNorm算子的實(shí)現(xiàn)和局限

我們先來看一下OneFlow中是如何使用CUDNN庫實(shí)現(xiàn)BatchNorm算子的。代碼見：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/normalization_kernel.cu#L31-L244 。這段代碼中首先實(shí)現(xiàn)了一個(gè)getCudnnBatchNormMode工具函數(shù)：

cudnnBatchNormMode_tgetCudnnBatchNormMode(constint64_tdim){
if(dim==2){
returnCUDNN_BATCHNORM_PER_ACTIVATION;
}elseif(ParseBooleanFromEnv("ONEFLOW_ENABLE_NHWC",false)){
returnCUDNN_BATCHNORM_SPATIAL_PERSISTENT;
}else{
//NOTE(LiangDepeng):ThenewCUDNN_BATCHNORM_SPATIAL_PERSISTENTmodewas
//introducedinCuDNN7forperformanceoptimization,butitresultsin
//accuracylossesinconvolutionmodelssuchasResNeXt-101and
//videoR(2+1)D.WewillfallbacktothenormalCUDNN_BATCHNORM_SPATIAL
returnCUDNN_BATCHNORM_SPATIAL;
}
}

這里的dim表示輸入Tensor的維度，比如形狀為的輸入Tensor，這里的維度就是4。然后這里涉及到三種不同的cudnnBatchNormMode_t，我們看一下CUDNN的文檔（https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnBatchNormMode_t）：

可以看到 CUDNN_BATCHNORM_PER_ACTIVATION 被用于非卷積層，在OneFlow中只有當(dāng)輸入Tensor的維度為2時(shí)才選取這種模式。而CUDNN_BATCHNORM_SPATIAL_PERSISTENT這種模式只有當(dāng)輸入Tensor的數(shù)據(jù)排布為NHWC方式時(shí)才會(huì)啟用。而對(duì)于其它的模式，在OneFlow中一律選取CUDNN_BATCHNORM_SPATIAL模式。

接下來閱讀一下 InferDimSizeAndDataFormat 函數(shù)：

voidInferDimSizeAndDataFormat(constShapeView&x_shape,constint32_taxis,int32_t*n,int32_t*c,
int32_t*h,int32_t*w,cudnnTensorFormat_t*format){
if(x_shape.Count(axis+1)==1){
if(axis==0){
*n=1;
*h=1;
}else{
*n=x_shape.At(0);
*h=x_shape.Count(1,axis);
}
*w=1;
*c=x_shape.At(axis);
*format=CUDNN_TENSOR_NHWC;
}else{
*n=x_shape.Count(0,axis);
*c=x_shape.At(axis);
*h=x_shape.Count(axis+1);
*w=1;
*format=CUDNN_TENSOR_NCHW;
}
}

這個(gè)函數(shù)會(huì)根據(jù)輸入Tensor的shape以及axis推斷這個(gè)Tensor的內(nèi)存排布是NCHW還是NHWC模式，并設(shè)置對(duì)應(yīng)的n, c, h, w變量。

//推斷和設(shè)置cudnn中的Tensor描述符
voidInferXYCudnnTensorDesc(constShapeView&xy_shape,constDataType&data_type,
constint32_taxis,cudnnTensorDescriptor_txy_desc){
int32_tn,c,h,w;
cudnnTensorFormat_tformat;
//根據(jù)輸入Tensor的shape推斷format和n,c,h,w
InferDimSizeAndDataFormat(xy_shape,axis,&n,&c,&h,&w,&format);
//根據(jù)上述的推斷結(jié)果，設(shè)置Tensor的描述符
OF_CUDNN_CHECK(
cudnnSetTensor4dDescriptor(xy_desc,format,GetCudnnDataType(data_type),n,c,h,w));
}
//根據(jù)輸入Tensor的描述符xy_desc和cudnnBatchNormMode_t模式設(shè)置參數(shù)的描述符param_desc
voidInferParamCudnnTensorDesc(constcudnnTensorDescriptor_txy_desc,cudnnBatchNormMode_tmode,
cudnnTensorDescriptor_tparam_desc){
OF_CUDNN_CHECK(cudnnDeriveBNTensorDescriptor(param_desc,xy_desc,mode));
}
//這個(gè)類就是完整使用上述的工具函數(shù)的工具類，負(fù)責(zé)推斷cudnnBatchNorm接口需要的各種描述信息比如這里的xy_desc_，param_desc_，param_data_type_和param_size_
classCudnnTensorDescHelperfinal{
public:
OF_DISALLOW_COPY_AND_MOVE(CudnnTensorDescHelper);
CudnnTensorDescHelper(constShapeView&xy_shape,constDataType&data_type,constint32_taxis,
cudnnBatchNormMode_tmode){
OF_CUDNN_CHECK(cudnnCreateTensorDescriptor(&xy_desc_));
InferXYCudnnTensorDesc(xy_shape,data_type,axis,xy_desc_);
OF_CUDNN_CHECK(cudnnCreateTensorDescriptor(¶m_desc_));
InferParamCudnnTensorDesc(xy_desc_,mode,param_desc_);
intn,c,h,w,n_stride,c_stride,h_stride,w_stride;
OF_CUDNN_CHECK(cudnnGetTensor4dDescriptor(param_desc_,¶m_data_type_,&n,&c,&h,&w,
&n_stride,&c_stride,&h_stride,&w_stride));
param_size_=c;
}
~CudnnTensorDescHelper(){
OF_CUDNN_CHECK(cudnnDestroyTensorDescriptor(param_desc_));
OF_CUDNN_CHECK(cudnnDestroyTensorDescriptor(xy_desc_));
}

cudnnTensorDescriptor_txy_desc()const{returnxy_desc_;}

cudnnTensorDescriptor_tparam_desc()const{returnparam_desc_;}

voidCheckParamTensor(constuser_op::Tensor*tensor)const{
CHECK_NOTNULL(tensor);
CHECK_EQ(tensor->shape_view().NumAxes(),1);
CHECK_EQ(tensor->shape_view().At(0),param_size_);
CHECK_EQ(GetCudnnDataType(tensor->data_type()),param_data_type_);
}

private:
cudnnTensorDescriptor_txy_desc_=nullptr;
cudnnTensorDescriptor_tparam_desc_=nullptr;
cudnnDataType_tparam_data_type_;
int32_tparam_size_=0;
};

除了這些描述信息之外，我們還可以在cudnn提供的文檔中查看BatchNorm相關(guān)的算子一般還需要什么特殊的輸入信息。我們來看 cudnnBatchNormalizationForwardTrainingEx() 這個(gè)API ：https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnBatchNormalizationForwardTrainingEx 。

可以看到這個(gè)算子是 cudnnBatchNormalizationForwardTraining() 這個(gè)算子的擴(kuò)展，擴(kuò)展的內(nèi)容就是可以我們可以傳入額外的一個(gè)Activation的算子比如ReLU以及一個(gè)Add算子分別對(duì)應(yīng)我們?cè)谇把灾薪榻B的 ResNet 中的 BNReLU 和 BNAddReLU 模式?？梢钥吹皆谶@個(gè)算子接口中除了對(duì)輸入Tensor x，BN后需要add的輸入Tensor z以及輸出Tensor y的描述信息外，還需要指定workspace和reserveSpace，這個(gè)workspace是cudnn的BatchNorm以NHWC模式計(jì)算時(shí)需要的GPU內(nèi)存buffer，而reserveSpace則表示當(dāng)前這個(gè)配置的BN算子至少還需要多少可以申請(qǐng)的GPU顯存（從文檔猜測(cè)應(yīng)該是和BNReLU/BNAddReLU這倆Pattern相關(guān)）。

在OneFlow中， https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/normalization_kernel.cu#L126-L175 以及 https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/normalization_kernel.cu#L637-L684 就是為了推斷BN算子以及BN擴(kuò)展的算子需要的額外GPU內(nèi)存大小，然后在OneFlow的內(nèi)存池中開辟一塊顯存供調(diào)用cudnn的 cudnnBatchNormalizationForwardTrainingEx() 和 cudnnBatchNormalizationBackwardEx() 接口時(shí)使用。

關(guān)于調(diào)用cudnn的BatchNorm相關(guān)的算子api，我們還需要注意一點(diǎn)，那就是要使用cudnn提供的擴(kuò)展接口cudnnBatchNormalizationForwardTrainingEx() 和 cudnnBatchNormalizationBackwardEx() 還存在一些限制：

首先是cudnn版本的限制，然后對(duì)于CUDNN_BATCHNORM_OPS_BN_ADD_ACTIVATION的Op模式，輸入Tensor的通道數(shù)必須是4的倍數(shù)，最后這個(gè)擴(kuò)展Op必須在輸入Tensor的數(shù)據(jù)排布模式是NHWC時(shí)才能啟動(dòng)。這些限制對(duì)應(yīng)到OneFlow的代碼在：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/job_rewriter/cudnn_fused_normalization_add_relu_pass.cpp#L79-L86 。

0x2.2 善用CUDA優(yōu)化打破cudnn的限制

上面提到要使用CUDNN的擴(kuò)展算子有一系列限制，我們有沒有辦法打破這限制呢？有的。以ResNet為例，針對(duì)BNReLu和BNAddReLU這兩種Pattern，我們可以分別針對(duì)ReLU和AddReLU實(shí)現(xiàn)一個(gè)CUDA Kernel，相信入門CUDA的小伙伴寫這兩個(gè)算子都沒什么問題。但如何在考慮到Backward的時(shí)候把這兩個(gè)算子優(yōu)化到位呢？OneFlow給出了一個(gè)解決方案。

前向的CUDA實(shí)現(xiàn)：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/normalization_kernel.cu#L246-L272

反向的CUDA實(shí)現(xiàn)：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/normalization_kernel.cu#L246-L272

以 ReLU 算子為例，前向的輸入為x，輸出為y，后向的輸入為dy和y，輸出dx。后向計(jì)算中的y僅用來判斷對(duì)應(yīng)元素是否大于0，因此可以將y替換為由前向生成的bitset（對(duì)應(yīng)上述代碼中的mask），理論上可以省掉ReLU的后向算子對(duì)冗余的y的訪問操作，減少約y大小的讀取，也對(duì)應(yīng)約1/3的global memory訪問。對(duì)于ReLU/ReLUAdd這種ElementWise算子來說，GPU的帶寬是極容易成為瓶頸的，通過這種優(yōu)化可以大大提升ReLU和ReLUAdd算子的帶寬。

在 《OneFlow是如何做到世界上最快的深度學(xué)習(xí)框架》(https://zhuanlan.zhihu.com/p/271740706) 文章中已經(jīng)介紹到了這種基于bitmask優(yōu)化后向算子的方案。并且文章中給出了3種方案，但沒有給出對(duì)應(yīng)的代碼實(shí)現(xiàn)，實(shí)際上我只讀懂了第一種和第三種方案，接下來我們描述一下這兩種方案。

Bitset mask生成方案一：順序遍歷法

這種方法是讓每個(gè)CUDA線程連續(xù)讀取內(nèi)存中的8個(gè)元素，并根據(jù)每個(gè)元素是否大于0生成一個(gè)int8類型的mask，并寫入到最終的bitset mask中。這種訪問對(duì)于寫全局內(nèi)存是連續(xù)訪問的，但對(duì)于讀（Read）全局內(nèi)存，線程間內(nèi)存訪問不連續(xù)，所以沒有充分合并內(nèi)存事務(wù)。下圖展示了這種方案讀寫內(nèi)存的示例：

以ReLU為例子，這種方案的代碼實(shí)現(xiàn)如下：

template
__global__voidReluGpu(int64_tn,constT*x,T*y,int8_t*mask){
CUDA_1D_KERNEL_LOOP(i,n){
int8_tmask_val=0;
for(int32_tj=0;j0);
if(is_positive){
y[offset]=sum;
mask_val|=(1<

	

	在這種方案中，每個(gè)thread需要連續(xù)讀的8個(gè)float32數(shù)據(jù)，則相鄰線程每次加載數(shù)據(jù)的間隔為32 bytes = 4 bytes * 8。所以每個(gè)線程一次加載指令就要執(zhí)行一個(gè)32字節(jié)的內(nèi)存事務(wù)。故warp內(nèi)的線程間全局內(nèi)存訪問完全沒有合并，實(shí)際有效訪存帶寬僅為 1/8，訪存效率十分低下，性能很差。

	Bitset mask生成方案三：warp同步法

	我們可以采用warp級(jí)別的同步原語：__ballot_sync(unsigned mask, predicate)，這個(gè)函數(shù)接收兩個(gè)參數(shù)，第一個(gè)參數(shù)是warp中參與計(jì)算的線程掩碼，第二個(gè)參數(shù)是要參與判斷的bool值，返回一個(gè)32bit的mask，每個(gè)bit代表warp中各個(gè)線程傳入的元素是否大于0，最后由每個(gè)warp中的0號(hào)線程將生成的mask寫入global memory中。(idea可以參考NVIDIA的性能優(yōu)化博客：https://developer.nvidia.com/blog/using-cuda-warp-level-primitives/)

	這種方案的示意圖如下：

	


	以ReLU為例，代碼實(shí)現(xiàn)如下：

	
template
__global__voidReluGpu(int64_tn,constT*x,T*y,int32_t*mask){
constint32_tlane_id=threadIdx.x%kCudaWarpSize;//如果lane_id=0，表示當(dāng)前線程是一個(gè)warp的0號(hào)線程
CUDA_1D_KERNEL_LOOP(i,n){
constboolis_positive=(x[i]>0);
int32_twarp_mask=__ballot_sync(__activemask(),static_cast(is_positive));
if(lane_id==0){mask[i/kCudaWarpSize]=warp_mask;}//0號(hào)線程將生成的mask寫入globalmemory
y[i]=is_positive?sum:0;
}
}


	

	0x3. 性能

	我們這里對(duì)比一下BNReLU這個(gè)Pattern在優(yōu)化前后的后向Kernel（也就是ReLU Grad Kernel）的性能和帶寬表現(xiàn)，本次測(cè)試的環(huán)境為A100 PCIE 40G，使用Nsight Compute工具進(jìn)行Profile。Profile的腳本為：

	
importoneflowasflow
bn=flow.nn.BatchNorm2d(num_features=32,eps=1e-5,momentum=0.1).to("cuda")
fused_bn=flow.nn.FusedBatchNorm2d(32).to("cuda")
bn.train()
fused_bn.train()

x=flow.randn(16,32,112,112).to("cuda").requires_grad_()

y=flow.relu(bn(x))#這個(gè)是未優(yōu)化的實(shí)現(xiàn)
#y=fused_bn(x)#打開這行代表啟用上述介紹的優(yōu)化
res=y.sum()
res.backward()
res_scalar=res.detach().cpu().numpy()


	

	經(jīng)過多次測(cè)試，flow.relu(bn(x))中對(duì)應(yīng)的ReLU的反向Kernel耗時(shí)大概為 「48.3us」，而fused_bn(x)中對(duì)應(yīng)的ReLU的反向Kernel耗時(shí)大概為 「42.8us」 ，可以說明上述基于mask掩碼降低全局內(nèi)存訪問的優(yōu)化方法是有效的。而對(duì)于BNAddReLU的Pattern來說，則可以獲得更好的性能提升，因?yàn)镽eluBackward相當(dāng)于將這兩個(gè)ElementWise操作給fuse了。

	0x4. 總結(jié)

	這里暫時(shí)寫了一下個(gè)人看OneFlow Normalization 系列算子實(shí)現(xiàn)的理解。實(shí)際上我個(gè)人還是有一些疑問在，如果后續(xù)能搞清楚的話，會(huì)繼續(xù)補(bǔ)充和修改。

	0x5. 相關(guān)鏈接

	cudnn文檔：https://docs.nvidia.com/deeplearning/cudnn/api/index.html

	oneflow代碼實(shí)現(xiàn)：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/normalization_kernel.cu

	審核編輯：湯梓紅