已合并
【CANN训练营第二季社区任务】Mish算子开发-算子提交 #490
wenyidylan创建于 2025年12月24日
【CANN训练营第二季社区任务】Mish算子开发-算子提交 #490
已合并
wenyidylan创建于 2025年12月24日
wenyidylan
2025年12月24日

描述

背景信息
基于Mish算子历史TBE版本使用Ascend C编程语言进行优化。

通过对Mish算子TBE版本的功能分析,当前支持的能力如下:
① x,y支持float16,float32,bfloat16三种数据类型的输入。
② Mish算子逐元素进行运算,仅有一个输入,不涉及到对输入数据进行广播。
③ 计算方式根据数据类型和实现模式采用不同策略:

Mish算子TBE版本的整体流程图如下图所示:
flow2_20251117.jpg

SUPER_PERFORMANCE模式的计算公式为:

y=x(12/(1+(1+(1+x/64)64)2))y = x*(1-2/(1+(1+(1+x/64)^{64})^2))

其他模式的计算公式为:

y={x(2ex+1)/(2e2x+2ex+1),if x>0x(2ex+e2x)/(2+2ex+e2x),if x0y = \begin{cases} x*(2e^{-x} + 1) / (2e^{-2x} + 2e^{-x} + 1), & \text{if $x > 0$} \\ x*(2e^x + e^{2x}) / (2 + 2e^x + e^{2x}), & \text{if $x \leq 0$} \end{cases}

算子原型

输入输出参数

名称 类别 dtype format shape 介绍
x 输入 float16/float32/bfloat16 ND 任意合法形状 输入
y 输出 float16/float32/bfloat16 ND 与x形状相同 输出

算子支持型号

Atlas A2 训练系列产品/Atlas 800I A2推理产品

host侧设计方案

算子计算过程不涉及数据的维度信息,故在host侧将数据视为一维向量,仅考虑数据个数,不考虑数据维度信息。
任务均分:coreNum 根据输入长度和块大小动态调整,确保每个核心处理的数据块数均匀。
批量搬运:tileBlockNum 和 tileDataNum 计算单次搬运的数据量,通过 finalSmallTileNum 和 finalBigTileNum 确定小核/大核的搬运次数,将多次搬运合并为批量操作,减少冗余开销。尾块的处理逻辑确保不完整块也能被合并到计算流程中,避免数据碎片。

分核策略

优先使用满核的原则。
如果核间能均分,可视作无大小核区分,大核小核数据块一致;
如果核间不能均分,需要将余出的数据块分配到前几个核上。
输入数据大小计算:通过GetInputShape和GetDataTypeLength函数获取输入数据的大小和类型长度,计算出输入数据的总字节数。
UB内存大小和核心数量获取:通过平台信息获取UB内存大小和核心数量,并根据这些信息调整核心数量。

数据分块和内存优化策略

充分使用UB空间的原则。
需要考虑不同硬件的UB大小不同、是否开启double buffer、kernel侧API实现过程中是否需要临时数据的储存,综合考虑单核内切分的大小。
UB内存大小获取:通过GetCoreMemSize函数获取UB内存的大小,用于后续的数据切分计算。
Tile块计算:根据UB内存大小和预定义的BLOCK_SIZE及BUFFER_NUM和不同类型下的ubDataNum,计算出每个Tile块的数据数量。
数据切分:将输入数据按照计算出的Tile块大小进行切分,计算出每个core需要处理的数据块数量和最后一个block的剩余数据量。
设置切分参数:将计算出的切分参数(如每个core的数据量、Tile块大小等)设置到MishTilingData对象中。
这些策略确保了数据在多个核心之间的均匀分布,并且在单个核心内进行了合理的切分,以提高并行处理的效率。

tilingkey规划策略

不进行tilingkey划分,在kernel侧利用输入数据的类型来走不同的分支。

kernel侧设计方案

进行Init和Process两个阶段,其中Process包括数据搬入(CopyIn)、计算(Compute)、搬出(CopyOut)三个阶段。
CopyIn、Compute、CopyOut阶段都是以连续的数据切片(tiling块)为单位处理数据。根据每个vector核心的索引号计算数据切片的位置。
依照TBE实现,在HIGH_PRECISION模式将float16数据精度转换为float32计算。
bfloat16数据类型转为float32数据类型进行计算。
把原TBE算子中的tbe.vmul、tbe.vmuls、tbe.vadds、tbe.vrec、tbe.vexp、tbe.cast_to指令替换为AscendC的Mul、Muls、Adds、Reciprocal、Exp、Cast等矢量指令。
使用AscendC的Compare和Select指令替换原TBE算子的tbe.vcmp和tbe.vsel指令,实现数据的条件判断和数据选取。
按照TBE算子的计算逻辑完成Mish计算。
AscendC的Mish算子流程见下图:
mish_ascendc_20251117.jpg

关联的Issue

https://gitcode.com/cann/ops-nn/issues/175

测试

单测:测试各种数据类型(float16/float32/bfloat16)
功能测试:测试不同维度的输入
性能测试:对比TBE版本的性能提升

文档更新

新增mish文档aclnnMish.md

类型标签

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 wenyidylan 的贡献)
Wwenyidylan
2025年12月24日 关联了issue:[Requirement|需求建议]: Mish算子AscendC实现贡献
CANN-robot
CANN-robot成员
2025年12月24日 评论:

Thank your for your pull-request.

The full list of commands accepted by me can be found at here.

You can get sig-info at here

likedislike
CANN-robot
CANN-robot成员
2025年12月24日 评论:

以下是根据您提交的修改文件推荐的Reviewer和Committer序列,需各模块评审通过后方可合入

Module List Reviewers Committers
sig-ops-nn chaotang233, zhajianqing123, zhou-qilong, fanqirui, lileizheng tangweiwei2, chenqi317, liubo75, crystalhu
likedislike
CANN-robotCANN-robot成员
2025年12月24日 添加了label:cann-cla/yes
wenyidylan
2025年12月24日 评论:

compile

likedislike
此处折叠了249条消息 查看更多
Ffulltower成员
4月22日 解决了最后一个问题
liujie12345678成员
4月24日 评论:

/lgtm

likedislike
CANN-robotCANN-robot成员
4月24日 添加了label:lgtm
CANN-robot
CANN-robot成员
4月24日 评论:

Review Guide

This pull-request passes review.
Committers who wrote a comment of /approve are: 陈琦.
Reviewers who wrote a comment of /lgtm are: liujie12345678, 陈琦.

likedislike
CANN-robotCANN-robot成员
4月24日 合入了pull request