0
  • 聊天消息
  • 系统消息
  • 评论与回复
登录后你可以
  • 下载海量资料
  • 学习在线课程
  • 观看技术视频
  • 写文章/发帖/加入社区
会员中心
创作中心

完善资料让更多小伙伴认识你,还能领取20积分哦,立即完善>

3天内不再提示

解析OneFlow Element-Wise算子实现方法

jf_pmFSk4VX 来源:GiantPandaCV 作者:GiantPandaCV 2022-12-12 10:54 次阅读

0x0. 前言

由于CUDA水平太菜,所以一直没写过这方面的笔记。现在日常的工作中已经不能离开写CUDA代码,所以准备学习ZZK随缘做一做CUDA的笔记记录一下学习到的知识和技巧。这篇文章记录的是阅读OneFlow的Element-Wise系列CUDA算子实现方案学习到的技巧,希望可以帮助到一起入门CUDA的小伙伴们。Elemet-Wise算子指的是针对输入Tensor进行逐元素操作,比如ReLU就是针对输入Tensor的每个值进行判断是否大于0,大于0的话输出就是输入否则就是0。用CUDA来表达最简单的写法就是:

__global__voidrelu_kernel(float*input,float*output){
int32_tidx=blockIdx.x*blockDim.x+threadIdx.x;
output[idx]=input[idx]< 0 ? 0 : input[idx];
}

int main(){
  float* input;
  float* output;
  int32_t elem_cnt = 3*224*224;
  
  cudaMalloc(&input, sizeof(float)*elem_cnt);
  cudaMalloc(&output, sizeof(float)*elem_cnt);
  int32_t thread_num = 256;
  int32_t grid_size = (elem_cnt + thread_num -1) / thread_num;
  relu_kernel<<>>(src,dst);

cudaDeviceSynchronize();
cudaFree(src);
cudaFree(dst);
return0;
}

虽然这种写法非常简单明了,但却存在明显的性能问题。所以这篇文章将基于OneFlow开源的Element-Wise CUDA算子方案来解释如何写一个高性能的Element-Wise CUDA算子。

0x1. 性能

以GELU激活函数为例子,分别测试 dtype = float32,不同shape下的前向耗时以及带宽利用率(NVIDIA A100-PCIE-40GB)。性能情况如下图所示:

9f2cb390-7987-11ed-8abf-dac502259ad0.png

在这里插入图片描述

9f2cb390-7987-11ed-8abf-dac502259ad0.png

在这里插入图片描述

可以看到对于 GeLU 来说,无论是性能还是带宽 OneFlow 的实现都是更优的,接下来我们就来了解一下为什么 OneFlow 的 Element-Wise 算子性能可以做到更优。

0x2. 用法

OneFlow在 elementwise.cuh 文件中分别针对一元,二元,三元运算的 Element-Wise 操作实现了模板函数。在包含这个头文件之后我们可以使用 cuda::Unary/Binary/Ternary 这几个模板函数来针对我们自己定义的 Element-Wise 操作进行计算。注意,这里说的一元,二元,三元代表的是这个 Element-Wise 操作有几个输入 Tensor。

我们举个例子,假设我们要做的 Element-Wise 操作是逐点乘法,也即有 2 个输入Tensor x 和 y,然后 x 和 y的形状和数据类型都是一致的。那么我们可以定义一个模板类:

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

这里 OF_DEVICE_FUNC 表示我们定义的这个函数既可以运行在 CPU 又可以运行在 GPU 上,它的定义是:

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

然后我们就可以使用 cuda::Binary 这个模板函数来完成这个二元的 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 的元素个数,输出张量的数据首地址 out->mut_dptr(), 输入张量的数据首地址 x->dptr() && y->dptr() ,最后一个参数则是当前 Kernel 运行的 cuda Stream对象。

0x3. 原理&&代码实现解析

我个人认为这里有几个要点,分别是一个线程处理多个数据,向量化数据访问提升带宽,设置合理的Block数量(GridSize)和线程数量(BlockSize)以及在合适的地方进行循环展开(unrool)以及一些编程上的技巧。

0x3.1 给 Element-Wise 操作设置合理的 GridSize 和 BlockSize

下面这段代码展示了 OneFlow 针对 Element-Wise 算子是如何设置 GridSize 和 BlockSize 的。对应的源码地址为: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 直接被设置为了 256 ,对应 constexpr int kBlockSize = 256; 这行代码,也就是说每个 Block 有 256 个线程。为什么是 256 ?大家不妨读一下俊丞大佬这篇经典的 给CUDA Kernel设置合适的 GridSize 和 Block Size 的文章 。文章中通过对 SM 的资源分析确定在主流的GPU上将 BlockSize 设置为 128 或者 256 是比较合适,在这里直接设置为了 256 。

确定了 BlockSize 之后需要确定 Kernel 启动线程块的数量,我一直觉得上述文章中对这一段的分析是尤其精彩的,这里再截图展示一下:

9f4990fa-7987-11ed-8abf-dac502259ad0.png

选自OneFlow CUDA Kernel 中 grid_size 和 block_size 应该怎么设置 一文

根据这里的分析,对于 Element-Wise 操作要设置合适的 GridSize 不仅需要考虑元素的数量还要考虑由于 SM 硬件本身带来的限制。如下公式所述:

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

这里的 (n + kBlockSize - 1) / kBlockSize 就是根据 Element-Wise 操作的元素个数来计算需要启动多少个线程块,比如在文章开头的例子中有 = 个元素,那么就一共需要 个线程块。然后这里以GTX 3080Ti为例,它的SM个数也就是sm_count=80,每个SM最多调度的线程数tpm=1536,那么sm_count * tpm / kBlockSize * kNumWaves = 80 * 1536 / 256 * 32 = 15360,所以在这个例子中我们最终设置的线程块个数为 588 个。

通过上述讲解和分析我们已经确定了启动 Element-Wise CUDA Kernel 的 GridSize 和 BlockSize。

0x3.2 向量化数据访问提升带宽

对于大多数 Element-Wise 算子来说,一般它们的计算量不会太大,所以它们的瓶颈一般在GPU的带宽上。在 NVIDIA 的性能优化博客 https://developer.nvidia.com/blog/cuda-pro-tip-increase-performance-with-vectorized-memory-access/ 中提到,对于很多 CUDA 核函数我们都可以通过向量化数据访问的方式来提升带宽受限的 Kernel 的性能,特别是对于架构比较新的 GPU 向量化数据访问的效果会更加明显。

在 OneFlow 的 Element-Wise 系列算子中,为了更好的进行向量化的数据访问,俊丞设计了如下的 Pack 数据结构(代码位置: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 先声明了一个内存对齐的数据类型 type ,注意这个 type 的内存长度为 pack_size * sizeof(T) 。然后这里的 T 是我们需要进行 Pack 的数据类型,而 pack_size 则表示我们需要 Pack 的元素个数。接下来我们看到 Pack 联合体中声明了 storage 和 elem 两个数组,它们公用同一段对齐的内存。然后 Pack 联合体的入口有一个检查: static_assert(sizeof(PackType) == sizeof(T) * pack_size, ""); 这是用来判断我们之前声明的 type 的内存长度是否符合预期。

接下来我们从 https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh#L155-L194 这里可以看到这个 Pack 联合体主要是用在 Kernel 启动之前判断 Element-Wise 操作的输入输出 Tensor 对应的数据指针地址是否满足内存对齐的条件,如果不满足则这个 Element-Wise 操作无法执行数据 Pack 。对应下图2个画红色框的地方。

9f77468a-7987-11ed-8abf-dac502259ad0.png

接下来,OneFlow 定义了真正要执行数据 Pack 的数据结构 Packed 并且定义了计算 PackSize 的工具函数。代码位置为: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< b ? a : b; }

template
constexprintPackSize(){
returnMin(kMaxPackBytes/sizeof(T),kMaxPackSize);
}

template
constexprintPackSize(){
returnMin(PackSize(),PackSize());
}

这里需要注意的是对于 CUDA 来说,最多支持 128 个 bit 的访问粒度,也就是说 PackSize 的大小不能超过 128 个bit。然后对于各种数据类型来说,Half 数据类型的 bit 数是最少的即 16,所以一次性可以支持 Pack 8个half类型的数据,4个float32的数据,以此类推。所以这里的定义的 kMaxPackSize 表示 128/16=8 ,然后 kMaxPackBytes 则表示最大可以 Pack 的 byte 数 。

请注意区分 bit 和 byte 。

接下来 https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh#L97-L144 则是真正的为 Element-Wise 操作完成数据 Pack 并执行计算。

首先来看这段充满技巧的代码:

9f848cbe-7987-11ed-8abf-dac502259ad0.png

在这里插入图片描述

首先这里定义了一个 HasApply2 类用来判断是否可以支持一次性Pack 2个 char/int8/half2 类型的元素,这个地方是一个针对 int8/half2/char 数据类型的特殊处理,某些 Element-Wise 算子 Kernel 确实需要支持这种数据类型的计算。也就是说对于 half2 的话,在一个内存访问粒度里我们其实是可以 Pack 128 / 8 = 16个的。然后用了C++模板元编程的 std::enable_if 来控制针对 half2 类型的特殊 Pack 处理,也就是上图代码中的两个 ApplyPack 函数。可以看到对于 half2 类型的 Element-Wise 操作我们需要给对应的 Functor 定义一个 Apply2 函数,比如对于 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 实现:https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh#L133-L144 。

9f98a0b4-7987-11ed-8abf-dac502259ad0.png

在这里插入图片描述

在 Kernel 中我们发现每一个线程实际上处理了多个 Pack 后的数据,也即:for (int64_t i = global_tid; i < n_pack; i += blockDim.x * gridDim.x) 。初学者看到这个循环也许会比较疑惑,为什么它的步幅是 blockDim.x * gridDim.x  ? 这个 blockDim.x * gridDim.x 表示的是 CUDA 线程网格中的线程总数。假设线程网格中有 1280 个线程,线程 0 将计算元素 0、1280、2560 等。通过使用步幅等于网格大小的循环,确保了 warp 中的所有寻址都是单位步幅,可以获得最大的内存合并。想了解更多细节可以查看:https://zhuanlan.zhihu.com/p/571320529 。

除此之外,使用这种技巧的还有个好处就是如果对于 Kernel 中存在每个线程都包含一个公共的操作,那么线程数的增多,也代表着这部分的开销变大。这个时候我们减少线程的数量并循环进行处理的话那么这个公共操作的开销就会更低。

最后,在循环之外,我们还需要根据传入的 n_tail 参数,看一下还有没有因为没有被 pack_size 整除的剩余元素,如果有的话就单独调用 functor 进行处理。

0x3.4 unroll

实际上就是代码中的 #pragma unroll ,这个宏会对我们的 for 循环做循环展开,让更多的指令可以并行执行。但容易想到,只有处理的数据没有前后依赖关系的时候我们可以做。对于大多数的 ElementWise 算子来说一般是满足这个条件的。

0x3.5 Kernel Launch的细节

在 https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/core/cuda/elementwise.cuh#L166-L181 这个位置 OneFlow 展示了 Element-Wise Kernel 的启动细节,我们简单注释一下:

template
cudaError_tLaunchKernel(FactoryTfactory,int64_tn,R*r,constIN*...in,cudaStream_tstream){
constint64_tn_pack=n/pack_size;//根据元素个数和pack_size,计算pack数目,比如1026/4=256。
constint64_ttail_offset=n_pack*pack_size;//如果存在不被整除的情况,我们计算使用pack的偏移量:256*4;
constint64_tn_tail=n-tail_offset;////元素数目-偏移量=剩下的元素个数->1026-1024=2
intnum_blocks;
{
cudaError_terr=GetNumBlocks(n_pack,&num_blocks);//计算线程块数目
if(err!=cudaSuccess){returnerr;}
}
ApplyGeneric<<>>(
factory,n_pack,reinterpret_cast*>(r),
(reinterpret_cast*>(in))...,n_tail,r+tail_offset,
(in+tail_offset)...);
returncudaPeekAtLastError();
}

0x4. 总结

以上就是我对 OneFlow Element-Wise 系列 CUDA 算子实现的解析,后续有空会持续更新学习到的新知识。

审核编辑:郭婷

声明:本文内容及配图由入驻作者撰写或者入驻合作网站授权转载。文章观点仅代表作者本人,不代表电子发烧友网立场。文章及其配图仅供工程师学习之用,如有内容侵权或者其他违规问题,请联系本站处理。 举报投诉
  • 代码
    +关注

    关注

    30

    文章

    4708

    浏览量

    68174
  • CUDA
    +关注

    关注

    0

    文章

    121

    浏览量

    13581

原文标题:【BBuf 的CUDA笔记】一,解析OneFlow Element-Wise 算子实现

文章出处:【微信号:GiantPandaCV,微信公众号:GiantPandaCV】欢迎添加关注!文章转载请注明出处。

收藏 人收藏

    评论

    相关推荐

    电源常用ic脚位解析方法 7脚电源芯片怎么看型号

    电源常用IC脚位解析方法 电源常用IC(集成电路)的脚位解析方法主要依赖于对IC引脚功能的理解,以及参考相关的技术手册或数据手册。以下是一些通用的
    的头像 发表于 10-07 17:10 842次阅读

    基于 DSP5509 进行数字图像处理中 Sobel 算子边缘检测的硬件连接电路图

    以下是基于 DSP5509 进行数字图像处理中 Sobel 算子边缘检测的硬件设计方案: 一、总体架构 图像采集:使用合适的图像传感器,如 CMOS 传感器,通过相应的接口(如 SPI、I2C 等
    发表于 09-25 15:25

    摩尔线程携手智源研究院完成基于Triton的大模型算子库适配

    里,即成功完成了近60个算子的功能验证,精度符合交付标准,并实现对Bert-large模型的全面支持。FlagGems算子库在摩尔线程MUSA架构上展现出了接近手写算子的计算性能,且性
    的头像 发表于 08-02 11:06 692次阅读

    微创软件推出AI大模型应用平台WISE

    微创软件在“2024微创人工智能战略发布会”上,正式推出了企业级AI大模型应用平台WISE。该平台以其独特的技术架构和卓越性能,为企业开发AI应用提供了全新的解决方案。
    的头像 发表于 05-31 11:31 759次阅读

    微创软件正式发布AI大模型应用平台WISE

    上海2024年5月28日 /美通社/ -- 5月20日,微创软件召开“2024微创人工智能战略发布会”,并正式推出企业级AI大模型应用平台WISE(Wicresoft Intelligence
    的头像 发表于 05-28 17:18 512次阅读
    微创软件正式发布AI大模型应用平台<b class='flag-5'>WISE</b>

    基于TPU-MLIR:详解EinSum的完整处理过程!

    、Reduce。EinSum支持任意多的输入,只要计算中只包含点乘(element-wise)、广播(broadcast)、归约求和(reductionsum)都可以使
    的头像 发表于 02-19 13:08 585次阅读
    基于TPU-MLIR:详解EinSum的完整处理过程!

    OneFlow Softmax算子源码解读之BlockSoftmax

    写在前面:笔者这段时间工作太忙,身心俱疲,博客停更了一段时间,现在重新捡起来。本文主要解读 OneFlow 框架的第二种 Softmax 源码实现细节,即 block 级别的 Softmax。
    的头像 发表于 01-08 09:26 626次阅读
    <b class='flag-5'>OneFlow</b> Softmax<b class='flag-5'>算子</b>源码解读之BlockSoftmax

    OneFlow Softmax算子源码解读之WarpSoftmax

    写在前面:近来笔者偶然间接触了一个深度学习框架 OneFlow,所以这段时间主要在阅读 OneFlow 框架的 cuda 源码。官方源码基于不同场景分三种方式实现 Softmax,本文主要介绍其中一种的
    的头像 发表于 01-08 09:24 682次阅读
    <b class='flag-5'>OneFlow</b> Softmax<b class='flag-5'>算子</b>源码解读之WarpSoftmax

    异步电机主要的三种调速方法解析

    异步电机主要的三种调速方法解析
    的头像 发表于 01-07 17:50 2148次阅读
    异步电机主要的三种调速<b class='flag-5'>方法</b><b class='flag-5'>解析</b>

    单片机解析g代码的方法

    的运动。 解析G代码是将其转化为单片机能够理解和执行的指令集。单片机解析G代码的方法主要包括以下几个方面:G代码的格式解析、指令的解析和执行
    的头像 发表于 12-22 14:15 1523次阅读

    OpenCV边缘检测算子Laplace、LoG详解

    一阶导数算子(例如 Sobel 算子)通过对图像求导来确定图像的边缘,数值绝对值较高的点对应了图像的边缘。如果继续求二阶导,原先数值绝对值较高的点对应了过零点。因此,也可以通过找到二阶导数的过零点来检测边缘。在某些情况下,找二阶导数的过零点可能更容易。
    的头像 发表于 12-21 16:34 1889次阅读
    OpenCV边缘检测<b class='flag-5'>算子</b>Laplace、LoG详解

    柠檬光子获选36氪WISE 2023未来商业潜力之王年度企业

    WISE2023 商业之王 年度企业系列名册 是知名科技媒体36氪的一项重要调研活动,也是中国商业领域内具有极高影响力和权威性的榜单之一。
    的头像 发表于 12-04 17:35 416次阅读

    时擎科技荣膺「WISE2023 未来商业之王 年度企业」

    11月28日,知名科技媒体36氪在北京国际会议中心发布了备受瞩目的「WISE2023未来商业之王」榜单。经由专业分析师的数轮遴选和专业评定,时擎科技凭借为万物智联的AIoT时代量身打造的多模态智能
    的头像 发表于 12-01 08:16 317次阅读
    时擎科技荣膺「<b class='flag-5'>WISE</b>2023 未来商业之王 年度企业」

    创新驱动商业增长 | 踏歌智行荣获“WISE2023未来商业之王年度企业”称号

    11月29日,36氪重磅发布「WISE2023商业之王年度企业系列名册」,踏歌智行凭借在 商业模式、主创团队、技术创新、财务指标、行业竞争力方面的杰出表现 ,荣获汽车出行领域“未来商业之王年度企业
    的头像 发表于 11-29 17:10 477次阅读
    创新驱动商业增长 | 踏歌智行荣获“<b class='flag-5'>WISE</b>2023未来商业之王年度企业”称号

    实力认证!望繁信科技入选2023WISE未来商业之王年度企业

    11月28日,36氪年度重磅IP——WISE2023商业之王大会在北京国际会议中心盛大开幕,备受瞩目的「WISE2023未来商业之王年度企业」榜单正式揭晓。望繁信科技凭借技术创新和在企业数字化转型
    的头像 发表于 11-29 14:55 704次阅读