玄幻小说完本,盛世嫡妃凤轻小说,我欲封天耳根小说零

0x1. 前言

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

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

另外一種模式是殘差結構引入的 BNAddReLU 的模式：

out=self.bn2(out)

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

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

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

0x2. 代碼解析

0x2.1 CUDNN BatchNorm算子的實現和局限

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

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中只有當輸入Tensor的維度為2時才選取這種模式。而CUDNN_BATCHNORM_SPATIAL_PERSISTENT這種模式只有當輸入Tensor的數據排布為NHWC方式時才會啟用。而對于其它的模式，在OneFlow中一律選取CUDNN_BATCHNORM_SPATIAL模式。

接下來閱讀一下 InferDimSizeAndDataFormat 函數：

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

這個函數會根據輸入Tensor的shape以及axis推斷這個Tensor的內存排布是NCHW還是NHWC模式，并設置對應的n, c, h, w變量。

//推斷和設置cudnn中的Tensor描述符
voidInferXYCudnnTensorDesc(constShapeView&xy_shape,constDataType&data_type,
constint32_taxis,cudnnTensorDescriptor_txy_desc){
int32_tn,c,h,w;
cudnnTensorFormat_tformat;
//根據輸入Tensor的shape推斷format和n,c,h,w
InferDimSizeAndDataFormat(xy_shape,axis,&n,&c,&h,&w,&format);
//根據上述的推斷結果，設置Tensor的描述符
OF_CUDNN_CHECK(
cudnnSetTensor4dDescriptor(xy_desc,format,GetCudnnDataType(data_type),n,c,h,w));
}
//根據輸入Tensor的描述符xy_desc和cudnnBatchNormMode_t模式設置參數的描述符param_desc
voidInferParamCudnnTensorDesc(constcudnnTensorDescriptor_txy_desc,cudnnBatchNormMode_tmode,
cudnnTensorDescriptor_tparam_desc){
OF_CUDNN_CHECK(cudnnDeriveBNTensorDescriptor(param_desc,xy_desc,mode));
}
//這個類就是完整使用上述的工具函數的工具類，負責推斷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相關的算子一般還需要什么特殊的輸入信息。我們來看 cudnnBatchNormalizationForwardTrainingEx() 這個API ：https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnBatchNormalizationForwardTrainingEx 。

可以看到這個算子是 cudnnBatchNormalizationForwardTraining() 這個算子的擴展，擴展的內容就是可以我們可以傳入額外的一個Activation的算子比如ReLU以及一個Add算子分別對應我們在前言中介紹的 ResNet 中的 BNReLU 和 BNAddReLU 模式。可以看到在這個算子接口中除了對輸入Tensor x，BN后需要add的輸入Tensor z以及輸出Tensor y的描述信息外，還需要指定workspace和reserveSpace，這個workspace是cudnn的BatchNorm以NHWC模式計算時需要的GPU內存buffer，而reserveSpace則表示當前這個配置的BN算子至少還需要多少可以申請的GPU顯存（從文檔猜測應該是和BNReLU/BNAddReLU這倆Pattern相關）。

在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擴展的算子需要的額外GPU內存大小，然后在OneFlow的內存池中開辟一塊顯存供調用cudnn的 cudnnBatchNormalizationForwardTrainingEx() 和 cudnnBatchNormalizationBackwardEx() 接口時使用。

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

首先是cudnn版本的限制，然后對于CUDNN_BATCHNORM_OPS_BN_ADD_ACTIVATION的Op模式，輸入Tensor的通道數必須是4的倍數，最后這個擴展Op必須在輸入Tensor的數據排布模式是NHWC時才能啟動。這些限制對應到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優化打破cudnn的限制

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

前向的CUDA實現：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/normalization_kernel.cu#L246-L272

反向的CUDA實現：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/normalization_kernel.cu#L246-L272

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

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

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

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

以ReLU為例子，這種方案的代碼實現如下：

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<

	

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

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

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

	這種方案的示意圖如下：

	


	以ReLU為例，代碼實現如下：

	
template
__global__voidReluGpu(int64_tn,constT*x,T*y,int32_t*mask){
constint32_tlane_id=threadIdx.x%kCudaWarpSize;//如果lane_id=0，表示當前線程是一個warp的0號線程
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號線程將生成的mask寫入globalmemory
y[i]=is_positive?sum:0;
}
}


	

	0x3. 性能

	我們這里對比一下BNReLU這個Pattern在優化前后的后向Kernel（也就是ReLU Grad Kernel）的性能和帶寬表現，本次測試的環境為A100 PCIE 40G，使用Nsight Compute工具進行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))#這個是未優化的實現
#y=fused_bn(x)#打開這行代表啟用上述介紹的優化
res=y.sum()
res.backward()
res_scalar=res.detach().cpu().numpy()


	

	經過多次測試，flow.relu(bn(x))中對應的ReLU的反向Kernel耗時大概為 「48.3us」，而fused_bn(x)中對應的ReLU的反向Kernel耗時大概為 「42.8us」 ，可以說明上述基于mask掩碼降低全局內存訪問的優化方法是有效的。而對于BNAddReLU的Pattern來說，則可以獲得更好的性能提升，因為ReluBackward相當于將這兩個ElementWise操作給fuse了。

	0x4. 總結

	這里暫時寫了一下個人看OneFlow Normalization 系列算子實現的理解。實際上我個人還是有一些疑問在，如果后續能搞清楚的話，會繼續補充和修改。

	0x5. 相關鏈接

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

	oneflow代碼實現：https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/normalization_kernel.cu

	審核編輯：湯梓紅