InplaceSigmoid 算子开发设计文档【免费下载链接】cann-ops-competitions本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-ops-competitions一、需求背景1.1 需求来源通过社区任务完成昇腾 CANN 算子开源仓贡献需求丰富昇腾硬件上对神经网络激活函数数据类型的泛化支持。1.2 背景介绍1.2.1 InplaceSigmoid 算子实现优化本次任务要求参考昇腾内置 TBE 算子Sigmoid在 Atlas A2/A3 训练系列产品上基于 Ascend C 编程语言实现功能一致的aclnnInplaceSigmoid算子。主要扩展原算子仅支持浮点类型本次需扩展支持int16,int8,uint8数据类型并采用就地Inplace运算直接更新原 Tensor。TBE 参考路径/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/impl/dynamic/sigmoid.py1.2.2 InplaceSigmoid 算子现状分析1.2.2.1 TBE 算子支持的数据类型和数据格式根据 TBE 源码及原型定义原Sigmoid算子仅支持float16,float32,bfloat16数据类型。支持的数据格式为ND。1.2.2.2 TBE 算子实现描述原 TBE 算子的实现具有高度的硬件分支特征针对不同计算平台应用了不同的算法Ascend310P (高性能模式)采用泰勒展开式逼近 $L(x) ax bx^3 cx^5 d$然后结合极大极小值截断。Ascend v200/v300 (Tiny/Nano)采用极值安全转换公式 $y \frac{e^{\min(0, x)}}{1 e^{-|x|}}$ 防溢出。Ascend910B / Ascend910_93 (本次目标硬件)执行标准计算公式 $y \frac{1}{1 e^{-x}}$。在高精度HIGH_PRECISION模式下如果存在特定的运算如非卷积会将float16/bfloat16转换为float32进行运算完成后转回原类型。1.2.2.3 TBE 算子实现流程图二、需求分析2.1 外部组件依赖需要适配算子调用框架 ACLNN。2.2 内部适配模块算子内部需完成 Ascend C Host 侧的 Tiling 切分模块设计op_host/sigmoid_tiling.cpp和 Kernel 侧运算模块op_kernel/sigmoid.cpp并更新相应算子信息库。2.3 需求模块设计2.3.1 Ascend C 算子原型参数名输入/输出/属性描述使用说明数据类型数据格式维度(shape)非连续Tensorself输入待进行Sigmoid计算的入参。支持空Tensorshape需要与out一致。FLOAT、FLOAT16、BFLOAT16、INT16、INT8、UINT8ND0-8√out输出计算的出参。shape需要与self一致Inplace实现实际操作同一内存。FLOAT、FLOAT16、BFLOAT16、INT16、INT8、UINT8ND0-8√2.3.2 Ascend C 算子相关约束算子默认使用确定性实现。对新增整型数据的处理需严格遵循指定转换逻辑int16转为float32计算int8 / uint8转为float16计算。转回时需考虑四舍五入或截断等常规 Cast 行为以对齐框架。三、需求详细设计3.1 使能方式根据实际开发任务使用 ACLNN 框架调用算子。3.2 需求总体设计3.2.1 Host 侧设计3.2.1.1 分核策略获取输入数据的总元素个数基于物理设备可用的 AI Core 数量和硬件 UB (Unified Buffer) 内存上限计算每个核处理的数据总量 (coreDataNum)。将任务均匀分配至多核末尾多出的零散数据分配给最后一个核作为tailDataNum。3.2.1.2 数据分块和内存优化策略在每个 AI Core 内部基于 Double Buffer 机制将当前核的数据拆解为tileDataNum。分别申请用于数据流入 (inQueueX) 和流出 (outQueueY) 的缓存同时针对需要数据类型转换的场景如int16、int8等需额外申请TBuf充当中间高精度计算类型缓存如float32、float16缓存区以避免直接覆盖源数据内存导致错乱。3.2.1.3 tilingKey 规划策略由于新增了整型不同数据类型的中间态 Cast 路径差异极大。依据输入 Tensor 的dtype生成对应的tilingKey以便 Kernel 侧直接实例化对应数据类型的计算模板。3.2.2 Kernel 侧设计3.2.2.1 Kernel 侧实现描述Kernel 层面的处理流程高度专注于Atlas A2/A3架构舍弃 TBE 中的其它历史遗留硬件分支聚焦于公式 $y \frac{1}{1 e^{-x}}$。 针对不同类型的 Cast 策略设计如下FLOAT32:原生直接计算。Muls(x, -1.0)$\rightarrow$Exp(x)$\rightarrow$Adds(x, 1.0)$\rightarrow$Div(1.0, x)。FLOAT16 / BFLOAT16:参照 TBE 高精度模式提取数据Cast到float32向量中完成上述计算再Cast返回原类型。INT16 (新增):Cast为float32后进行计算计算完成后Cast(使用RoundMode::CAST_RINT或框架要求的转整型模式) 回int16。INT8 / UINT8 (新增):Cast为float16后进行计算完成后转回对应整型。Inplace 操作:在 CopyOut 阶段直接覆盖最初传入的 Global Memory 地址selfGm。3.2.2.2 Ascend C 实现流程图3.2.2.3 Ascend C 实现流程图与 TBE 流程图存在的差异点和原因差异点1. 抛弃了基于 Ascend 310P/310 的vabs、min极值处理和泰勒展开式的复杂分支判断仅保留标准 Sigmoid 计算。引入了大量的源头数据到计算精度的转换路线 (int16-fp32,int8-fp16等)。原因本次开发明确指定适配硬件为 Atlas A2/A3对应的就是原 TBE 中的Ascend910B及其衍生型号底层指令自带Inf/NaN支持性能强劲无需老旧硬件的防溢出 hack 代码。同时由于泛化了多种整型必须通过 Ascend C 的Cast接口进行类型升降级以利用浮点指令管线。3.3 支持硬件适配Atlas A2 训练系列产品 / Atlas A3 系列产品。3.4 算子约束限制支持 0-8 维ND及非连续 Tensor。Inplace 特性要求调用链路确保输入输出共享同一块内存块算子内不额外申请外存。四、特性交叉分析该算子为纯 Element-wise 的原地更新激活函数运算输入输出的 Shape 与 Stride 特性在执行中保持不变。对于非连续 Tensor依靠 ACLNN 框架下发的一维展平偏移量或循环迭代进行读取不引起网络结构层次的其他特性交叉冲突。五、可维可测分析5.1 精度标准/性能标准精度需满足AscendOpTest的默认阈值。新增的int16,int8,uint8的输入输出结果需与 PyTorch 参考实现由于其转浮点计算后再强制转整型会有精度舍入完全对齐。性能在所有核参与计算的前提下性能需达到原 TBE 的 95% 以上。小于 10us 的场景需提供仿真图作为调优依据。5.2 兼容性分析本实现采用 Ascend C 进行标准重构作为原有 TBE 的功能超集增加了整数泛化支持向后完全兼容已有模型的 Sigmoid 调用场景对现有模型业务运行无破坏性影响。【免费下载链接】cann-ops-competitions本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-ops-competitions创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考