【BBuf的CUDA笔记】二,解析 OneFlow BatchNorm 相关算子实现

news2024/11/23 19:03:30

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)

if self.downsample is not None:
    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_t getCudnnBatchNormMode(const int64_t dim) {
  if (dim == 2) {
    return CUDNN_BATCHNORM_PER_ACTIVATION;
  } else if (ParseBooleanFromEnv("ONEFLOW_ENABLE_NHWC", false)) {
    return CUDNN_BATCHNORM_SPATIAL_PERSISTENT;
  } else {
    // NOTE(Liang Depeng): The new CUDNN_BATCHNORM_SPATIAL_PERSISTENT mode was
    // introduced in CuDNN 7 for performance optimization, but it results in
    // accuracy losses in convolution models such as ResNeXt-101 and
    // video R(2+1)D. We will fall back to the normal CUDNN_BATCHNORM_SPATIAL
    return CUDNN_BATCHNORM_SPATIAL;
  }
}

这里的dim表示输入Tensor的维度,比如形状为 ( 1 , 3 , 224 , 224 ) (1, 3, 224, 224) (1,3,224,224)的输入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 函数:

void InferDimSizeAndDataFormat(const ShapeView& x_shape, const int32_t axis, 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描述符
void InferXYCudnnTensorDesc(const ShapeView& xy_shape, const DataType& data_type,
                            const int32_t axis, cudnnTensorDescriptor_t xy_desc) {
  int32_t n, c, h, w;
  cudnnTensorFormat_t format;
  // 根据输入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
void InferParamCudnnTensorDesc(const cudnnTensorDescriptor_t xy_desc, cudnnBatchNormMode_t mode,
                               cudnnTensorDescriptor_t param_desc) {
  OF_CUDNN_CHECK(cudnnDeriveBNTensorDescriptor(param_desc, xy_desc, mode));
}
// 这个类就是完整使用上述的工具函数的工具类,负责推断cudnn BatchNorm接口需要的各种描述信息比如这里的xy_desc_,param_desc_,param_data_type_和param_size_
class CudnnTensorDescHelper final {
 public:
  OF_DISALLOW_COPY_AND_MOVE(CudnnTensorDescHelper);
  CudnnTensorDescHelper(const ShapeView& xy_shape, const DataType& data_type, const int32_t axis,
                        cudnnBatchNormMode_t mode) {
    OF_CUDNN_CHECK(cudnnCreateTensorDescriptor(&xy_desc_));
    InferXYCudnnTensorDesc(xy_shape, data_type, axis, xy_desc_);
    OF_CUDNN_CHECK(cudnnCreateTensorDescriptor(&param_desc_));
    InferParamCudnnTensorDesc(xy_desc_, mode, param_desc_);
    int n, c, h, w, n_stride, c_stride, h_stride, w_stride;
    OF_CUDNN_CHECK(cudnnGetTensor4dDescriptor(param_desc_, &param_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_t xy_desc() const { return xy_desc_; }

  cudnnTensorDescriptor_t param_desc() const { return param_desc_; }

  void CheckParamTensor(const user_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_t xy_desc_ = nullptr;
  cudnnTensorDescriptor_t param_desc_ = nullptr;
  cudnnDataType_t param_data_type_;
  int32_t param_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<typename T>
__global__ void ReluGpu(int64_t n, const T* x, T* y, int8_t* mask) {
  CUDA_1D_KERNEL_LOOP(i, n) {
    int8_t mask_val = 0;
    for(int32_t j = 0; j < 8; j++) {
      int32_t offset = i * 8 + j;
      const bool is_positive = (x[offset] > 0);
      if(is_positive) {
        y[offset] = sum;
        mask_val |= (1<<j);
      } else {
        y[offset] = 0;
        mask_val &= (~(1<<j));
      }
    }
    mask[i] = mask_val;
  }
}

在这种方案中,每个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<typename T>
__global__ void ReluGpu(int64_t n, const T* x, T* y, int32_t* mask) {
  const int32_t lane_id = threadIdx.x % kCudaWarpSize; // 如果lane_id=0,表示当前线程是一个warp的0号线程
  CUDA_1D_KERNEL_LOOP(i, n) {
    const bool is_positive = (x[i] > 0);
    int32_t warp_mask = __ballot_sync(__activemask(), static_cast<int>(is_positive));
    if (lane_id == 0) { mask[i / kCudaWarpSize] = warp_mask; } // 0号线程将生成的mask写入global memory
    y[i] = is_positive ? sum : 0;
  }
}

0x3. 性能

我们这里对比一下BNReLU这个Pattern在优化前后的后向Kernel(也就是ReLU Grad Kernel)的性能和带宽表现,本次测试的环境为A100 PCIE 40G,使用Nsight Compute工具进行Profile。Profile的脚本为:

import oneflow as flow
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

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.coloradmin.cn/o/108723.html

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈,一经查实,立即删除!

相关文章

MicFunPred——最新16S rRNA扩增子数据功能预测数据库

近年来&#xff0c;基于扩增子测序进行物种的功能预测是研究微生物群落功能的主要方面&#xff0c;目前最常用的软件包括Tax4Fun以及PICRUSt2。关于这两款软件的使用方法详可参见凌波微课|扩增子研究第十六讲&#xff1a;扩增子测序结果中的物种功能预测。 Tax4Fun使用最近邻匹…

2022年终总结-两年Androider的成长之路

金句分享 生活金句 1.可难道我们生命中做的每一件事不都是为了被爱得更多一点吗 2.这不只是一种对承诺的恐惧&#xff0c;也不是我缺乏关心和爱的能力&#xff0c;因为我做得到&#xff0c;只不过&#xff0c;老老实实讲。我想 我宁愿为了某件我擅长的事&#xff0c;我能表现…

一个select死锁问题

以下代码的输出结果&#xff1a; func main() {var wg sync.WaitGroupfoo : make(chan int)bar : make(chan int)wg.Add(1)go func() {defer wg.Done()select {case foo <- <-bar:default:println("default")}}()wg.Wait() }结果 解析 对于 select 语句&#…

【Linux】进程信号

目录 一、什么是信号 二、信号产生的条件 1、键盘产生 2、进程异常 3、命令产生 4、软件条件 三、信号保存的方式 四、信号处理的方式 1、信号处理接口 2、信号处理时机 3、进程为什么要切换成为用户态才进行信号的捕获方法&#xff1f; 4、sigaction 五、可重入函…

Java+MySQL基于ssm的会议交接平台

随着社会竞争压力的不断加强,企事业单位内部的会议都在不断的增加,有效的会议可以提高企事业内部的沟通,更好的做出符合战略目标的决策,但是传统的会议交接有一定的问题存在,首先就是必须面对面进行传达,其次就是对任务的安排和执行没有很好的记录,为了改变这些情况,于是我们提…

信贷产品年终总结之贷后逾期分析

自本月月初疫情全面放开后&#xff0c;身边的朋友基本都阳了一遍&#xff0c;希望正在浏览本篇文章的读者您是还没阳过的幸运儿。另外&#xff0c;今天也是冬至了&#xff0c;祝各位读者身边健康&#xff0c;远离羊群&#xff01; 最近我们分享了信贷业务年终总结系列的前2篇文…

Python中转义字符是个啥

文章目录前言一、转义字符是什么&#xff1f;二、常见的转义字符有哪些&#xff1f;总结前言 昨天有粉丝问了我这个代码问题&#xff0c;如下图&#xff1a; 他很好奇代码都没有错误&#xff0c;怎么运行就报错&#xff0c;不知道有咩有小伙伴能看出问题在哪呢&#xff1f; 其…

Cookie 和 Session 的工作流程

文章目录1.Cookie1.什么是Cookie2.Cookie可以干嘛3.Cookie实现登陆逻辑的流程2.session1.session是什么2.session有什么用3.session的工作流程3.Cookie 和 session的区别(重点)1.Cookie 1.什么是Cookie Cookie中存储的是字符串,是浏览器在本地持久化保存数据的一种方案 通过点…

2023春招:Javaweb面试锦囊

cookie 和 session 的区别&#xff1f;&#xff08;必会&#xff09; 存储位置不同 cookie 存放在客户端电脑&#xff0c;是一个磁盘文件。Ie 浏览器是可以从文件夹中找到。session 是存放在服务器内存中的一个对象。 chrome 浏览器进行安全处理&#xff0c;只能通过浏览器找…

圣诞 HTML 代码汇总

文章目录Part.I 音效圣诞树Part.II 圣诞树小球Part.III 简笔圣诞树圣诞节快到了&#xff0c;在网上找了一些 html 代码&#xff0c;觉得挺有意思的&#xff0c;顺带分享一下~ Part.I 音效圣诞树 来源&#xff1a;https://blog.csdn.net/m0_73309780/article/details/128176149…

面临项目失控?四个维度应对项目进度优化【洞见2】

常见的对进度的管理的流程是制定进度目标&#xff0c;WBS工作任务拆解&#xff0c;任务的时间估算&#xff0c;然后执行监督。 有时候这样的过程就会出现的进度延迟&#xff0c;而针对进度的延迟&#xff0c;往往企业多选择是通过加班赶工来完成。 项目进度优化方案 但是如…

【第一章 Linux特点,结构,网路连接模式,Linux目录结构】

第一章 Linux特点&#xff0c;结构&#xff0c;网路连接模式&#xff0c;Linux目录结构 1.操作系统&#xff1a; 用于管理和控制计算机所有软、硬件资源的一组程序。 2. Linux特点总结&#xff1a; ① 开放性&#xff1b; ② 多用户&#xff1b; ③ 多任务&#xff1b; ④ 良好…

喇叭天线设计

电磁喇叭天线是最简单而常用的微波天线。它的主要优点是结构简单&#xff0c;馈电简便&#xff0c;便于控制主面波束宽度和增益&#xff0c;频率特性好且损耗较小。它由波导逐渐张开来形成&#xff0c;其作用是加强方向性&#xff0c;这与声学喇叭的原理相似。若主模TE10的矩形…

python爬虫爬取网页上的图片

目录 一&#xff1a;爬虫基础 二&#xff1a;安装html解析的python工具 三&#xff1a;爬取网页图片 一&#xff1a;爬虫基础 爬虫基本过程&#xff1a; 1.请求标头 headers 2.创建一个会话 requests.Session 3.确定请求的路径 4.根据路径获取网页资源(HTML文件) 5.解析html…

UG/NX二次开发Siemens官方NXOPEN实例解析—2.6 CreateNote

列文章目录 UG/NX二次开发Siemens官方NXOPEN实例解析—2.1 AssemblyViewer UG/NX二次开发Siemens官方NXOPEN实例解析—2.2 Selection UG/NX二次开发Siemens官方NXOPEN实例解析—2.3 Selection_UIStyler UG/NX二次开发Siemens官方NXOPEN实例解析—2.4 File2Points UG/NX二次…

【java随笔】面向对象思维

1.怎么建立面向对象的思维方式 &#xff08;1&#xff09;先整体&#xff0c;再局部 &#xff08;2&#xff09;先抽象&#xff0c;再具体 &#xff08;3&#xff09;能做什么&#xff0c;再怎么做 2.如何学习面向对象 &#xff08;1&#xff09;语法 &#xff08;2&#xff09…

执行docker restart [CONTAINER ID]命令时会把之前的参数都带上吗

我在搭建rocketmq时候&#xff0c;运行了好几个容器命令也比较长如下&#xff1a; # 启动 namesrv docker run -d -p 9876:9876 -v "D:/Program Files/rocketmq/namesrv/logs:/root/logs" -v "D:/Program Files/rocketmq/namesrv/store:/root/store" …

JAVA手机网站销售

开发工具(eclipse/idea/vscode等)&#xff1a; 数据库(sqlite/mysql/sqlserver等)&#xff1a; 功能模块(请用文字描述&#xff0c;至少200字)&#xff1a;

Hazelcast 在springboot下的使用集成

一、Hazelcast简介 Hazelcast是一款由Hazelcast公司开发的基于jvm环境的为各种应用提供分布式集群服务的分布式缓存解决方案。可以嵌入到java、c、.net等开发的产品中使用。其主要功能有&#xff1a; 提供了 Map、Queue、MultiMap、Set、List、Semaphore、Atomic 等接口的分布…

RocketMQ消费者没有成功消费消息的问题排查

背景 今天下游同事反馈&#xff0c;有一些以取消的订单库存还原异常了&#xff0c;导致部分商品库存没有还原。查日志发现没有收到还原消息&#xff0c;但是查看发送方是可以确认消息是已经发了的&#xff0c;那么是什么原因导致消费者没有收到&#xff0c;或者收到后没有处理消…