ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

CANN ops-nn MishGrad 算子深度指南:NPU 上 Mish 反向传播梯度的原理与 aclnn 调用实战

2026/9/20 6:25:48 拓冰建站 浏览量
CANN ops-nn MishGrad 算子深度指南:NPU 上 Mish 反向传播梯度的原理与 aclnn 调用实战 CANN ops-nn MishGrad 算子深度指南NPU 上 Mish 反向传播梯度的原理与 aclnn 调用实战【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn本文围绕 CANN ops-nn 开源仓库中 MishGrad 算子文档 展开系统讲解该算子在 NPU 上计算 Mish 激活函数梯度的数学原理、算子定义与 tiling/kernel 实现结构并结合 aclnnMishBackward 接口文档 与 完整调用样例给出可直接参考的两段式 aclnn 调用流程。读完本文你将掌握 MishGrad 的两种梯度计算分支、其 Host 侧算子注册与 Device 侧 Kernel 的协作方式并能在 Atlas A2 训练系列产品上编写、编译与运行基于 aclnnMishBackward 的反向梯度计算程序。产品支持情况MishGrad 算子的产品支持情况如下表所示引自 README产品是否支持Atlas A2 训练系列产品 / Atlas 800I A2 推理产品√从算子注册代码 mish_grad_def.cpp 可以看到该算子在AICore()配置中为 SoC 版本ascend910b注册了 AICore 实现与文档声明的产品支持范围一致。需要说明的是文档中的支持矩阵以当前仓库发布时的适配状态为准其他 SoC 版本如要使用需参照该注册方式补充对应配置。功能说明计算 Mish 激活函数的梯度MishGrad 算子的功能是计算 Mish 激活函数的反向传播梯度它是正向 Mish 算子self × tanh(softplus(self))详见 Mish 算子文档的反向配套算子训练过程中由上层框架在反向传播阶段调用。根据是否传入正向阶段保存的中间量tanhx算子存在两种计算分支对应两套梯度公式分支一传入 tanhx当调用方同时传入正向计算中得到的tanhx即tanh(softplus(x))的值时梯度按下式计算$$ xgrad grad \times \frac{tanhx - x \times (1 - tanh^2x) \times e^{x}}{1 e^{x}} $$该分支利用正向阶段缓存的tanhx避免重复计算tanh在反向计算中只需一次Exp与少量乘加运算即可完成计算开销更小适合在训练过程中保存中间激活的场景。分支二不传入 tanhx当未传入tanhx时算子需要自行构造等价表达式先计算辅助量y1$$ y_1 \frac{1}{(1 e^{x})^{2} 1} $$再计算梯度$$ xgrad grad \times (4 \times x \times e^{x} \times (e^{x} 1) \times y_1^2 - 2 \times y_1 1) $$从 Kernel 源码看数值稳定实现从 mish_grad.h 的Compute实现可以看出无tanhx分支在实际计算时并非直接套用上述统一公式而是按x的符号分段、使用exp(-x)/exp(-2x)或exp(x)/exp(2x)构造数值稳定的等价形式当x 0时通过CompareScalar(x, 0, CMPMODE::GT)生成掩码mask1使用exp(-x)与exp(-2x)构造梯度表达式避免大正数exp(x)溢出当x ≤ 0时通过CompareScalar(x, 0, CMPMODE::LE)生成掩码mask2改用exp(x)与exp(2x)避免exp(-x)在强负输入下的上溢两段结果分别用Select按掩码合并后再求和最后乘以grad得到xgrad。在 mish_grad.cpp 中mish_grad内核入口通过REGISTER_TILING_DEFAULT/GET_TILING_DATA_WITH_STRUCT读取 tiling 数据实例化NsMishGrad::KernelMishGradDTYPE_X后调用Init与Process完成流水计算。Kernel 内部使用 AscendC 的TPipe流水VECIN/VECOUT双缓冲队列BUFFER_NUM 2对grad、x以及可选的tanhx执行DataCopy搬入、向量计算、结果搬出当输入为 FLOAT16/BFLOAT16 时先Cast到 float 做中间计算再Cast回原类型以保证精度见 mish_grad.h。参数说明MishGrad 算子的参数定义如下引自 README 参数说明表格参数名输入/输出/属性描述数据类型数据格式grad输入待进行 mishgrad 计算的入参公式中的 grad。FLOAT、FLOAT16、BFLOAT16ND、NCHW、NHWC、NC1HWC0、FRACTAL_Z、FRACTAL_NZx输入待进行 mishgrad 计算的入参公式中的 x。FLOAT、FLOAT16、BFLOAT16ND、NCHW、NHWC、NC1HWC0、FRACTAL_Z、FRACTAL_NZtanhx输入待进行 mishgrad 计算的可选入参公式中的 tanhx。FLOAT、FLOAT16、BFLOAT16ND、NCHW、NHWC、NC1HWC0、FRACTAL_Z、FRACTAL_NZx_grad输出待进行 mishgrad 计算的出参公式中的 xgrad。FLOAT、FLOAT16、BFLOAT16ND、NCHW、NHWC、NC1HWC0、FRACTAL_Z、FRACTAL_NZ与算子定义源码的对应关系上述参数约束在 mish_grad_def.cpp 的算子信息库注册中逐一落地grad、x、x_grad均为ParamType(REQUIRED)必选tanhx为ParamType(OPTIONAL)可选与文档中可选入参的描述一致数据类型枚举覆盖ge::DT_FLOAT16、ge::DT_FLOAT、ge::DT_BF16三种对应 FLOAT16、FLOAT、BFLOAT16每种数据类型都显式列出了 ND、NCHW、NHWC、NC1HWC0、FRACTAL_Z、FRACTAL_NZ 六种 format与文档数据格式列完全对应配置项DynamicShapeSupportFlag(true)、DynamicRankSupportFlag(true)表明算子支持动态 shape 与动态 rankDynamicCompileStaticFlag(true)支持动态编译静态标志PrecisionReduceFlag(true)允许精度降低如低精度计算详见 mish_grad_def.cpp。输出 shape 推导在 mish_grad_infershape.cpp 中InferShapeMishGrad将输出x_grad的 shape 直接置为输入x的 shape*yShape *xShape即梯度输出与正向输入逐元素对应、shape 完全一致这是逐元素激活函数反向算子的典型 shape 推导规则。约束说明按 README 说明本算子无额外约束。需要补充说明的是通过 aclnn 接口调用时仍需遵守接口层对 Tensor 合法性、数据类型与广播关系的通用约束详见下文。调用说明aclnn 两段式接口MishGrad 算子通过 CANN 的aclnn 两段式接口对外暴露接口名为aclnnMishBackward。完整接口说明见 aclnnMishBackward.md该接口的实现位于 aclnn_mish_backward.cpp。两段式接口机制根据 两段式接口说明aclnn 算子 API 必须先调用GetWorkspaceSize 接口获取入参校验结果与所需 workspace 大小再调用执行接口真正下发计算任务。MishGrad 对应的两个函数原型如下aclnnStatus aclnnMishBackwardGetWorkspaceSize( const aclTensor* gradOutput, // 输入反向传播上一步输出的梯度公式中的 grad const aclTensor* self, // 输入正向输入数据公式中的 x aclTensor* gradInput, // 输出计算得到的梯度公式中的 xgrad uint64_t* workspaceSize, // 输出Device 侧需要申请的 workspace 大小 aclOpExecutor** executor) // 输出算子执行器包含计算流程aclnnStatus aclnnMishBackward( void* workspace, // 输入Device 侧 workspace 内存地址 uint64_t workspaceSize, // 输入由第一段接口返回的 workspace 大小 aclOpExecutor* executor, // 输入第一段接口返回的执行器 aclrtStream stream) // 输入执行任务的 Stream第一段接口参数与返回值aclnnMishBackwardGetWorkspaceSize的入参/出参要点详见 aclnnMishBackward.md参数名输入/输出描述数据类型数据格式维度(shape)非连续 TensorgradOutput输入反向传播过程中上一步输出的梯度公式中的 gradOutput。BFLOAT16、FLOAT16、FLOAT32ND0-8√self输入正向的输入数据公式中的 self。BFLOAT16、FLOAT16、FLOAT32ND0-8√gradInput输出计算得到梯度作为反向传播下一步反向算子的计算输入。BFLOAT16、FLOAT16、FLOAT32ND0-8√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小。----executor输出返回 op 执行器包含了算子计算流程。----使用说明要点三个 Tensor 均不支持空 TensorgradOutput与self的数据类型需满足互推导关系gradOutput与self的 shape 需满足 broadcast 关系gradInput的 shape 需与广播后的 shape 一致返回的aclnnStatus状态码定义参见 aclnn 返回码。第一段接口会完成入参校验出现以下场景时返回对应错误码返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 gradOutput、self 或 gradInput 是空指针时。ACLNN_ERR_PARAM_INVALID161002gradOutput、self 或 gradInput 的数据类型不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002gradOutput 的 shape 不能向 self broadcast。ACLNN_ERR_PARAM_INVALID161002gradOutput 或 self 的 shape dim 大于 8。这些校验逻辑在 aclnn_mish_backward.cpp 中由CheckNotNull、CheckDtypeValid、CheckShape三个函数实现分别检查空指针、按 SoC 版本匹配的数据类型支持列表ASCEND910B/ASCEND910_93支持DT_FLOAT/DT_FLOAT16/DT_BF16其余支持DT_FLOAT/DT_FLOAT16见 aclnn_mish_backward.cpp、以及最大维度与广播推导此外还会对 shape 不匹配的输入执行BroadcastTo后再进入底层计算。第二段接口参数aclnnMishBackward的四个入参均为执行阶段所需运行时信息详见 aclnnMishBackward.mdworkspace在 Device 侧申请的 workspace 内存地址workspaceSize第一段接口计算出的 workspace 大小executor第一段接口返回的算子执行器stream指定执行任务的 Stream。关于 workspace 的来源在 Host 侧 tiling 阶段mish_grad_tiling.cpp 的GetWorkspaceSize通过GetLibApiWorkSpaceSize()获取系统库 API 所需 workspace 大小并写入context-GetWorkspaceSizes(1)当前使用一块 workspace。因此调用方必须严格按第一段接口返回的workspaceSize申请内存否则可能因缓冲区不足导致执行失败。完整调用示例与逐步讲解test_aclnn_mish_grad.cpp 给出了完整的 aclnn 调用样例aclnnMishBackward.md 中亦包含等价示例代码其编译与执行流程参考编译与运行样例。整个调用过程分为七个步骤1. 设备与 Stream 初始化int Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); // 初始化 ACL 运行环境 CHECK_RET(ret ACL_SUCCESS, ...); ret aclrtSetDevice(deviceId); // 设置计算设备 CHECK_RET(ret ACL_SUCCESS, ...); ret aclrtCreateStream(stream); // 创建 Stream CHECK_RET(ret ACL_SUCCESS, ...); return 0; }2. 构造输入输出 aclTensor样例中定义 shape 均为{4, 2}输入数据为gradOutput {0..7}、self {1,1,1,2,2,2,3,3}std::vectorint64_t gradOutputShape {4, 2}; std::vectorint64_t selfShape {4, 2}; std::vectorint64_t gradInputShape {4, 2}; // ... aclrtMalloc 申请 Device 内存、aclrtMemcpy 拷贝 Host-Device // ... 按连续 tensor 规则计算 strides *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr);其中CreateAclTensor模板函数封装了申请 Device 内存 → 拷贝数据 → 计算 strides → 创建 aclTensor的完整流程数据类型使用aclDataType::ACL_FLOAT对应 FLOAT32。3. 两段式调用算子 APIuint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段获取 workspace 大小与执行器 ret aclnnMishBackwardGetWorkspaceSize(gradOutput, self, gradInput, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, ...); // 按返回的 workspaceSize 申请 device 内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, ...); } // 第二段执行计算 ret aclnnMishBackward(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, ...);注意workspaceSize 0时才需要申请 workspace 内存且申请时必须使用第一段接口返回的确切大小。4~7. 同步、取数、释放资源ret aclrtSynchronizeStream(stream); // 4. 同步等待任务执行结束 // 5. aclrtMemcpy 将 gradInputDeviceAddr 上的结果拷回 Host 并打印 aclDestroyTensor(gradOutput); // 6. 释放 aclTensor aclDestroyTensor(self); aclDestroyTensor(gradInput); aclrtFree(gradOutputDeviceAddr); // 7. 释放 Device 内存与运行资源 aclrtFree(selfDeviceAddr); aclrtFree(gradInputDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize();需要特别说明的是样例中tanhx分支并未启用接口层aclnnMishBackward的参数列表本身不含tanhx对应 README 中不传入 tanhx的计算分支若需要复用正向保存的tanhx走梯度计算的分支一需在框架图层面直接下发MishGrad算子携带tanhx输入而非走此 aclnn 接口。单元测试与验证仓库为 MishGrad 提供了 Host 侧、Kernel 侧与 op_api 三套单元测试可作为算子行为与数值正确性的参考Kernel 侧test_mish_grad.cpp 使用tikicpulib的 CPU 仿真环境直接运行mish_grad0内核手动构造MishGradTilingData如smallCoreDataNum 2048、tileDataNum 4032、haveTanhx 0以blockDim 1、AIV_MODE方式验证无tanhx分支的计算Host 侧test_mish_grad_tiling.cpp 验证 tiling 计算逻辑op_api 侧test_aclnn_mish_backward.cpp 验证aclnnMishBackward两段式接口的端到端调用。从 tiling 结构体 mish_grad_tiling_data.h 可以看出Host 侧 MishGradTilingFunc 会依据平台 UB 大小与 core 数量计算每个 core 处理的数据量smallCoreDataNum/bigCoreDataNum、每个 tile 的数据量tileDataNum、尾块数据量smallTailDataNum/bigTailDataNum以及是否携带tanhxhaveTanhx并通过context-SetBlockDim(coreNum)与SetTilingKey(ELEMENTWISE_TPL_SCH_MODE_0)将切分策略下发给 KernelKernel 侧Init根据coreId与tailBlockNum的关系确定各 core 的数据范围与全局缓冲区偏移Process则按tileNum循环执行CopyIn → Compute → CopyOut流水见 mish_grad.h。tiling 还根据是否传入tanhx使用不同的 UB 缓冲策略传入时每 tile 至多容纳 10FLOAT/16其他类型个 32B 块未传入时对应 12/22 个 256B 块并据此调整BLOCK_SIZE与tileDataNum见 mish_grad_tiling.cpp。贡献说明MishGrad 算子的贡献信息如下引自 README 贡献说明表格贡献者贡献方贡献算子贡献时间贡献内容ilovescrapy个人开发者MishGrad2025/12/23MishGrad 算子适配开源仓总结MishGrad 是 CANN ops-nn 中为 Mish 激活函数配套的反向梯度算子支持 FLOAT、FLOAT16、BFLOAT16 三种数据类型与 ND、NCHW、NHWC、NC1HWC0、FRACTAL_Z、FRACTAL_NZ 六种数据格式通过可选输入tanhx提供复用正向中间量与自行构造等价式两种计算分支Kernel 内部按x符号分段使用数值稳定形式并对低精度输入做 float 中间计算。对外通过aclnnMishBackward两段式接口暴露调用时需严格遵循GetWorkspaceSize 取 workspace → 申请内存 → 执行 → 同步 → 释放的流程具体代码可直接复用 test_aclnn_mish_grad.cpp并参考 编译与运行样例 完成构建与验证。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考