CANN ops-nn SoftShrinkGrad 反向算子实战指南从公式推导到 aclnnSoftshrinkBackward 双段式调用与源码实现剖析【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn导读SoftShrinkGrad 是 CANN 神经网络算子库ops-nn中 Softshrink 激活函数的反向算子负责在反向传播阶段根据上游梯度与阈值lambd计算输入梯度并内置与 PyTorch 对齐的 NaN 保护逻辑。本文以仓库中的 README.md 与接口文档 aclnnSoftshrinkBackward.md 为主线结合 算子 API 实现、算子原型定义、shape 推导 与 AscendC 内核 等源码系统讲解其数学原理、参数约束、aclnn 两段式调用流程、GE 图模式注册方式以及底层计算实现帮助开发者在 Atlas/NPU 平台上快速完成该算子的接入、调用与问题定位。一、算子概述与产品支持情况1.1 算子定位SoftShrinkGrad 是 Softshrink 的反向算子。Softshrink软阈值收缩是一种常用于去噪、稀疏化场景的激活函数其正向公式为$$ Softshrink(x) \begin{cases} x - \lambda, if \ x \lambda\ x \lambda, if \ x -\lambda\ 0, otherwise \end{cases} $$在反向传播中SoftShrinkGrad 接收上游梯度input_grad接口文档中记为gradOutput与前向输入input_x接口文档中记为self输出对输入的梯度output_y接口文档中记为gradInput。该算子与 PyTorch 的SoftShrinkBackward算子对齐见 soft_shrink_grad_proto.h 中的框架兼容性注释可无缝对接 PyTorch 训练模型在 NPU 上的迁移。1.2 产品支持矩阵根据 README.md 与 aclnnSoftshrinkBackward.md 的声明该算子支持以下产品产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品√Atlas 推理系列产品√Atlas 训练系列产品√不同产品的数据类型支持范围存在差异Ascend 950PR/Ascend 950DT、Atlas A3、Atlas A2 系列input_grad、input_x、output_y支持 fp16、fp32、bf16 三种类型Atlas 200I/500 A2 推理、Atlas 推理系列、Atlas 训练系列仅支持 fp16、fp32。这一产品差异在 API 源码中也有对应实现aclnn_softshrink_backward.cpp 通过GetDtypeSupportList()依据当前 SoC 版本区分支持列表——ASCEND910、ASCEND310P、ASCEND310B平台仅支持DT_FLOAT与DT_FLOAT16其余平台额外支持DT_BF16。二、数学原理计算公式与 NaN 保护2.1 反向计算公式SoftShrinkGrad 的计算公式如下$$ output_y \begin{cases} input_grad, \lvert input_x \rvert lambd \text{ 或 } input_x \text{ 为NaN} \ 0, \text{其他情况} \end{cases} $$即当输入input_x的绝对值严格大于阈值lambd时梯度原样透传否则梯度置 0。这一规则正是 Softshrink 正向函数大于阈值区间内保持线性、阈值区间内输出恒 0的导数形式——阈值区间内正向输出为常量 0故梯度为 0。2.2 NaN 保护语义公式中的关键设计是NaN 保护当input_x为 NaN 时梯度直接透传output_y input_grad这与 PyTorch 的SoftShrinkBackward行为对齐。由于 IEEE 754 标准规定任何与 NaN 的比较结果均为 false|NaN| lambd恒为 false若不做特殊处理NaN 位置会被误置为 0从而污染梯度。源码层面通过x ! x的 NaN 自反检测实现该保护详见本文第六节内核分析。三、参数说明与约束3.1 算子参数总览GE 图模式视角在 soft_shrink_grad_def.cpp 与 soft_shrink_grad_proto.h 中算子定义了如下输入输出与属性参数名输入/输出/属性描述数据类型数据格式input_grad输入反向传播过程中上一步输出的梯度FLOAT16、FLOAT、BFLOAT16NDinput_x输入前向 Softshrink 的输入张量FLOAT16、FLOAT、BFLOAT16NDoutput_y输出反向传播梯度输出FLOAT16、FLOAT、BFLOAT16NDlambd可选属性Softshrink 的阈值参数默认值 0.5需满足 lambd≥0FLOAT-其中lambd在 GE 算子原型中注册为可选属性OPTIONAL默认值0.5f见 soft_shrink_grad_def.cpp 与 soft_shrink_grad_proto.h。算子定义还声明了AutoContiguous()输入输出自动转为连续张量DynamicRankSupportFlag(true)与DynamicShapeSupportFlag(true)支持动态 rank 与动态 shapeExtendCfgInfo(opFile.value, soft_shrink_grad)指定对应的 kernel 实现文件。3.2 约束说明根据 README.md 的约束章节dtype/shape 一致性input_grad与input_x的 dtype 和 shape 必须一致GE 图模式硬性要求维度范围支持0–8 维ND 格式阈值合法性lambd必须满足lambd ≥ 0。3.3 形状推导与数据类型推导soft_shrink_grad_infershape.cpp 中注册了算子的推断逻辑InferShape4SoftShrinkGradoutput_y.shape input_grad.shape输出形状直接继承输入梯度InferDataType4SoftShrinkGradoutput_y的数据类型继承input_grad的数据类型。四、aclnn 两段式接口调用详解4.1 接口背景两段式 API 设计aclnnSoftshrinkBackward遵循 CANN aclnn 统一的两段式接口设计详见 两段式接口说明第一段aclnnSoftshrinkBackwardGetWorkspaceSize完成入参校验与计算流程编排返回执行所需 workspace 大小及aclOpExecutor执行器第二段aclnnSoftshrinkBackward传入 workspace 与 executor在指定 stream 上真正执行算子计算。4.2 函数原型aclnnStatus aclnnSoftshrinkBackwardGetWorkspaceSize( const aclTensor* gradOutput, // 上游梯度 const aclTensor* self, // Softshrink 正向输入 const aclScalar* lambda, // 阈值 λ非负 aclTensor* gradInput, // 输出梯度 uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnSoftshrinkBackward( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)4.3 第一段接口参数说明aclnnSoftshrinkBackwardGetWorkspaceSize的完整参数如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorgradOutputaclTensor*输入反向传播的上游梯度不能传空指针支持空 Tensorshape 与 self 满足 broadcast 关系数据类型可与 self 不同不同时内部转换为 FLOAT 计算FLOAT16、FLOAT、BFLOAT16ND0-8√selfaclTensor*输入Softshrink 正向计算的输入不能传空指针支持空 Tensorshape 与 gradOutput 满足 broadcast 关系数据类型可与 gradOutput 不同不同时内部转换为 FLOAT 计算FLOAT16、FLOAT、BFLOAT16ND0-8√lambdaaclScalar*输入Softshrink 计算的阈值 λ不能传空指针取值必须 ≥ 0FLOAT---gradInputaclTensor*输出Softshrink 反向传播的输出梯度不能传空指针支持空 Tensorshape 必须与 self 和 gradOutput 的 broadcast 结果一致FLOAT16、FLOAT、BFLOAT16ND0-8√workspaceSizeuint64_t*输出需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出包含算子计算流程的执行器-----要点解读支持非连续 Tensor两个输入与输出均允许传入非连续含 view张量API 内部会通过l0op::Contiguous转为连续张量后计算再通过l0op::ViewCopy写回非连续输出见 aclnn_softshrink_backward.cpp支持 broadcastgradOutput与self的 shape 只需满足 broadcast 关系API 内部先BroadcastInferShape推断广播后形状再通过l0op::BroadcastTo将两者广播到一致 shape见 aclnn_softshrink_backward.cpp混合精度自动提升当gradOutput与self数据类型不同时非 FLOAT 的输入会被l0op::Cast提升为 FLOAT 计算输出再 Cast 回gradInput声明的 dtype见 aclnn_softshrink_backward.cpp支持空 TensorgradOutput或gradInput为空时第一段接口直接返回workspaceSize 0不执行实际计算见 aclnn_softshrink_backward.cpp。4.4 返回值与错误码两段接口均返回aclnnStatus状态码具体定义参见 aclnn 返回码说明。第一段接口完成入参校验以下场景会报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001gradOutput、self、lambda、gradInput、workspaceSize 或 executor 存在空指针ACLNN_ERR_PARAM_INVALID161002gradOutput、self 或 gradInput 的数据类型不在当前产品支持范围内ACLNN_ERR_PARAM_INVALID161002self 和 gradOutput 的 shape 不满足 broadcast 规则ACLNN_ERR_PARAM_INVALID161002lambda 为 NaN 或小于 0ACLNN_ERR_PARAM_INVALID161002gradInput 的 shape 与 self 和 gradOutput 的 broadcast 结果不一致ACLNN_ERR_PARAM_INVALID161002gradOutput 或 self 的维度大于 8这些校验逻辑在 aclnn_softshrink_backward.cpp 的CheckParams中依次执行先检查空指针CheckNotNull再检查数据类型是否在支持列表内CheckDtypeValid随后校验维度上限MAX_SUPPORT_DIMS_NUMS、broadcast 关系与输出 shape 一致性CheckShapeValid最后检查lambda是否为 NaN 且不小于 0。4.5 第二段接口参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取executor输入op 执行器包含算子计算流程stream输入指定执行任务的 Stream第二段接口的实现非常简洁直接调用框架统一的CommonOpExecutorRun完成计算见 aclnn_softshrink_backward.cpp这也是所有 aclnn 算子的固定收尾写法。4.6 确定性计算aclnnSoftshrinkBackward.md 明确说明aclnnSoftshrinkBackward 默认确定性实现同一输入在多次执行中计算结果完全一致适用于对可复现性有严格要求的训练场景。五、完整调用示例基于 aclnn API仓库提供了可直接编译运行的完整样例 examples/test_aclnn_soft_shrink_grad.cpp同时存在 arch35 变体 examples/arch35/test_aclnn_soft_shrink_grad.cpp。编译与执行方法可参考 编译与运行样例。下面按步骤拆解其调用流程5.1 头文件与宏定义#include cstdio #include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_softshrink_backward.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0)5.2 初始化 device 与 streamint Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, ...); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, ...); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, ...); return 0; }5.3 构造输入输出 aclTensor样例使用 shape 为{4, 2}的数据构造三个张量gradOutput上游梯度{0,1,2,3,4,5,6,7}、self前向输入{1,1,1,2,1,2,3,3}、gradInput输出占位初始全 0并创建阈值lambda 1.2f的aclScalarstd::vectorint64_t gradOutputShape {4, 2}; std::vectorint64_t selfShape {4, 2}; std::vectorint64_t gradInputShape {4, 2}; std::vectorfloat gradOutputHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectorfloat selfHostData {1, 1, 1, 2, 1, 2, 3, 3}; std::vectorfloat gradInputHostData {0, 0, 0, 0, 0, 0, 0, 0}; float lambdaValue 1.2f; // 通过 CreateAclTensor 模板函数完成aclrtMalloc 申请 device 内存、 // aclrtMemcpy 拷贝 host 数据到 device、计算 strides、aclCreateTensor 创建 aclTensor ret CreateAclTensor(gradOutputHostData, gradOutputShape, gradOutputDeviceAddr, aclDataType::ACL_FLOAT, gradOutput); ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); lambda aclCreateScalar(lambdaValue, aclDataType::ACL_FLOAT); ret CreateAclTensor(gradInputHostData, gradInputShape, gradInputDeviceAddr, aclDataType::ACL_FLOAT, gradInput);5.4 两段式调用// 第一段获取 workspace 大小与执行器 uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; ret aclnnSoftshrinkBackwardGetWorkspaceSize(gradOutput, self, lambda, 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 aclnnSoftshrinkBackward(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, ...);5.5 同步、取回结果与资源释放// 同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, ...); // 将 device 侧结果拷贝回 host 并打印 ret aclrtMemcpy(resultData.data(), ..., gradInputDeviceAddr, ..., ACL_MEMCPY_DEVICE_TO_HOST); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %f\n, i, resultData[i]); } // 释放 aclTensor / aclScalar 与 device 资源 aclDestroyTensor(gradOutput); aclDestroyTensor(self); aclDestroyScalar(lambda); aclDestroyTensor(gradInput); aclrtFree(gradOutputDeviceAddr); aclrtFree(selfDeviceAddr); aclrtFree(gradInputDeviceAddr); if (workspaceSize 0) aclrtFree(workspaceAddr); aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize();5.6 预期结果验证以样例数据手动推导lambda 1.2|self| 1.2的元素位置梯度透传。self {1,1,1,2,1,2,3,3}的绝对值为{1,1,1,2,1,2,3,3}其中大于 1.2 的是第 3、5、6、7 个元素值为 2、2、3、3对应gradOutput中值为 3、5、6、7 的梯度透传其余位置输出 0。即预期输出为{0, 0, 0, 3, 0, 5, 6, 7}。读者可通过运行样例程序打印的result[i]进行对照验证仓库中的单元测试 tests/ut/op_api/test_aclnn_softshrink_backward.cpp 与 tests/ut/op_kernel/test_soft_shrink_grad.cpp 也覆盖了该算子的行为断言。六、底层实现剖析AscendC Kernel 与 tiling6.1 内核入口kernel 入口 是一个典型的 AICore 算子入口从 tiling 数据中读取SoftShrinkGradTilingData实例化模板类NsSoftShrinkGrad::SoftShrinkGradDTYPE_INPUT_GRAD, BUFFER_MODE并依次调用Init与Process。BUFFER_MODE模板参数控制双缓冲BUFFER_NUM 2或单缓冲BUFFER_NUM 1模式。6.2 数据流与 tiling 分块内核实现 采用经典的三段流水结构Init根据tilingData计算当前 block 负责的blockLength_最后一个 block 处理余量blockTail、UB 缓冲长度ubLength_与阈值lambd_并按元素偏移blockFormer * blockIdx定位 Global Memory 上的输入输出地址随后pipe.InitBuffer为输入梯度、输入 x、输出分别分配BUFFER_NUM个 UB 队列缓冲CopyIn通过DataCopyPad将 Global Memory 数据按currentNum * sizeof(T)的分片长度搬入 UBCompute执行逐元素比较与选择逻辑详见 6.3CopyOut将计算结果DataCopyPad写回 Global Memory。Process中以loopCount (blockLength_ ubLength_ - 1) / ubLength_循环切分处理尾块通过currentNum blockLength_ - ubLength_ * i处理余量配合双缓冲实现数据搬运与计算的重叠最大化 AICore 利用率。6.3 比较-选择计算核心与 NaN 保护实现Compute中针对FP16/BF16 提升路径与FP32 直算路径分别实现核心逻辑一致取绝对值Abs(absX, xFP32, alignedNum)计算|x|阈值比较CompareScalar(cmpMask, absX, lambd_, CMPMODE::GT, alignedNum)生成比较位图cmpMask|x| lambd为 1条件选择Select(outFP32, cmpMask, gradFP32, 0.0f, SELMODE::VSEL_TENSOR_SCALAR_MODE, alignedNum)cmpMask1时输出梯度否则输出 0NaN 保护利用 IEEE 754 中NaN ! NaN恒为真的性质Compare(nanMask, xFP32, xFP32, CMPMODE::NE, alignedNum)生成 NaN 位图再执行第二次Select(outFP32, nanMask, gradFP32, outFP32, SELMODE::VSEL_TENSOR_TENSOR_MODE, alignedNum)将 NaN 位置的结果覆盖为输入梯度——这正是与 PyTorch 行为对齐的关键实现见 soft_shrink_grad.h。FP16/BF16 路径会先将输入Cast到 FP32 参与计算以避免精度损失最终Cast回原 dtype向下转换使用CAST_ROUND同时计算缓冲通过GetWithOffset划分出[gradFP32 | xFP32 | absX | outFP32 | cmpMask | nanMask]六个分区。另外由于CompareScalar/Compare要求计数占用 256B 对齐跨度代码将currentNum向上对齐到 64 元素CMP_ALIGN_ELEMS 256 / sizeof(float)并对中间缓冲做零初始化防止未对齐尾元素产生脏数据见 soft_shrink_grad.h。6.4 平台相关的 tiling 与 kernel 适配算子针对新平台如 arch35对应 Ascend 950 系列提供了专门的 tiling 计算与 kernel 实现Host 侧 tiling op_host/arch35/soft_shrink_grad_tiling_arch35.cpp负责计算 block 划分、UB 分片与阈值等 tiling 参数Kernel 侧数据结构 op_kernel/arch35/soft_shrink_grad_tiling_data.h 定义SoftShrinkGradTilingData含blockNum、blockFormer、dim0、ubFormer、lambd等字段soft_shrink_grad_tiling_key.h 定义 tiling key对应的单测 tests/ut/op_host/arch35/test_soft_shrink_grad_tiling.cpp 覆盖 tiling 计算正确性。七、GE 图模式接入除 aclnn 单算子 API 外SoftShrinkGrad 还以GE 图算子的形式注册供计算图构建与图编译使用。其算子原型定义在 op_graph/soft_shrink_grad_proto.hREG_OP(SoftShrinkGrad) .INPUT(input_grad, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16})) .INPUT(input_x, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16})) .OUTPUT(output_y, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16})) .ATTR(lambd, Float, 0.5) .OP_END_FACTORY_REG(SoftShrinkGrad)同时 op_graph/soft_shrink_grad_graph_infer.cpp 提供图模式下的 shape 推导op_graph/CMakeLists.txt 负责该模块的构建。开发者可在训练框架的构图阶段以SoftShrinkGrad为节点名、传入input_grad/input_x与属性lambd构建计算图由 GE 完成图编译后下发 NPU 执行。两种调用方式aclnn API 与 GE 图模式的对应关系总结如下调用方式入口/示例适用场景aclnn APIexamples/test_aclnn_soft_shrink_grad.cpp aclnnSoftshrinkBackward 接口文档单算子调测、自定义训练脚本中显式调用GE 图模式op_graph/soft_shrink_grad_proto.h 中的算子原型框架构图、整图下发与图优化场景八、工程结构与测试验证8.1 目录结构从仓库目录结构看soft_shrink_grad算子遵循 ops-nn 标准的算子工程组织方式每类平台文件按arch35等目录分层便于跨平台扩展与独立编译op_api/aclnn 对外接口实现aclnn_softshrink_backward.cpp与内部封装 soft_shrink_grad.cpp、soft_shrink_grad.hop_host/算子定义、shape/dtype 推导与 tiling 计算op_kernel/AscendC kernel 实现op_graph/GE 算子原型注册与图推导docs/aclnn 接口文档examples/可编译运行的单算子样例tests/包含ut/op_api、ut/op_host、ut/op_kernel三层单元测试以及st/aclnnSoftshrinkBackward的 ATK 用例配置。8.2 测试覆盖仓库为算子提供了完善的测试保障API 层单测tests/ut/op_api/test_aclnn_softshrink_backward.cpp 验证两段式接口的正确性Kernel 层单测tests/ut/op_kernel/test_soft_shrink_grad.cpp 验证逐元素计算逻辑Tiling 单测tests/ut/op_host/arch35/test_soft_shrink_grad_tiling.cpp 验证分块参数ST 测试tests/st/aclnnSoftshrinkBackward/atk_aclnnSoftshrinkBackward.json 提供昇腾算子工具链ATK的端到端用例配置。九、小结与使用建议SoftShrinkGrad 作为 Softshrink 的反向算子通过绝对值与阈值比较 NaN 自反检测覆盖的简单而严谨的规则在 NPU 上实现了与 PyTorch 语义完全对齐的梯度计算。实际使用时建议注意以下几点阈值约束lambda必须为非负值源码强制校验lambd 0或为 NaN 时报ACLNN_ERR_PARAM_INVALID默认 0.5与 PyTorchnn.Softshrink默认lambd0.5保持一致数据类型规划新平台950/A3/A2支持 bf16老平台仅 fp16/fp32混合精度输入会被统一提升为 FP32 计算后再转回存在轻微额外开销建议保持gradOutput与self同 dtypeshape 一致性GE 模式下要求input_grad与input_x的 dtype 和 shape 完全一致aclnn 模式下两者允许满足 broadcast 关系但输出gradInput必须与 broadcast 结果一致两段式调用务必先调用GetWorkspaceSize获取 workspace 大小并申请 Device 内存再调用第二段接口若workspaceSize为 0 可跳过申请空 Tensor 场景确定性保证该算子默认确定性实现适合对结果可复现有要求的训练与调测流程。更多算子家族信息可进一步阅读前向算子 Softshrink README 及其接口文档以及 ops-nn 仓库的 算子开发总览 与 两段式接口说明。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/19 9:35:37 · 阅读 1.2k CANN ops-nn SoftShrinkGrad 反向算子实战指南：从公式推导到 aclnnSoftshrinkBackward 双段式调用与源码实现剖析
CANN ops-nn SoftShrinkGrad 反向算子实战指南从公式推导到 aclnnSoftshrinkBackward 双段式调用与源码实现剖析【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn导读SoftShrinkGrad 是 CANN 神经网络算子库ops-nn中 Softshrink 激活函数的反向算子负责在反向传播阶段根据上游梯度与阈值lambd计算输入梯度并内置与 PyTorch 对齐的 NaN 保护逻辑。本文以仓库中的 README.md 与接口文档 aclnnSoftshrinkBackward.md 为主线结合 算子 API 实现、算子原型定义、shape 推导 与 AscendC 内核 等源码系统讲解其数学原理、参数约束、aclnn 两段式调用流程、GE 图模式注册方式以及底层计算实现帮助开发者在 Atlas/NPU 平台上快速完成该算子的接入、调用与问题定位。一、算子概述与产品支持情况1.1 算子定位SoftShrinkGrad 是 Softshrink 的反向算子。Softshrink软阈值收缩是一种常用于去噪、稀疏化场景的激活函数其正向公式为$$ Softshrink(x) \begin{cases} x - \lambda, if \ x \lambda\ x \lambda, if \ x -\lambda\ 0, otherwise \end{cases} $$在反向传播中SoftShrinkGrad 接收上游梯度input_grad接口文档中记为gradOutput与前向输入input_x接口文档中记为self输出对输入的梯度output_y接口文档中记为gradInput。该算子与 PyTorch 的SoftShrinkBackward算子对齐见 soft_shrink_grad_proto.h 中的框架兼容性注释可无缝对接 PyTorch 训练模型在 NPU 上的迁移。1.2 产品支持矩阵根据 README.md 与 aclnnSoftshrinkBackward.md 的声明该算子支持以下产品产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品√Atlas 推理系列产品√Atlas 训练系列产品√不同产品的数据类型支持范围存在差异Ascend 950PR/Ascend 950DT、Atlas A3、Atlas A2 系列input_grad、input_x、output_y支持 fp16、fp32、bf16 三种类型Atlas 200I/500 A2 推理、Atlas 推理系列、Atlas 训练系列仅支持 fp16、fp32。这一产品差异在 API 源码中也有对应实现aclnn_softshrink_backward.cpp 通过GetDtypeSupportList()依据当前 SoC 版本区分支持列表——ASCEND910、ASCEND310P、ASCEND310B平台仅支持DT_FLOAT与DT_FLOAT16其余平台额外支持DT_BF16。二、数学原理计算公式与 NaN 保护2.1 反向计算公式SoftShrinkGrad 的计算公式如下$$ output_y \begin{cases} input_grad, \lvert input_x \rvert lambd \text{ 或 } input_x \text{ 为NaN} \ 0, \text{其他情况} \end{cases} $$即当输入input_x的绝对值严格大于阈值lambd时梯度原样透传否则梯度置 0。这一规则正是 Softshrink 正向函数大于阈值区间内保持线性、阈值区间内输出恒 0的导数形式——阈值区间内正向输出为常量 0故梯度为 0。2.2 NaN 保护语义公式中的关键设计是NaN 保护当input_x为 NaN 时梯度直接透传output_y input_grad这与 PyTorch 的SoftShrinkBackward行为对齐。由于 IEEE 754 标准规定任何与 NaN 的比较结果均为 false|NaN| lambd恒为 false若不做特殊处理NaN 位置会被误置为 0从而污染梯度。源码层面通过x ! x的 NaN 自反检测实现该保护详见本文第六节内核分析。三、参数说明与约束3.1 算子参数总览GE 图模式视角在 soft_shrink_grad_def.cpp 与 soft_shrink_grad_proto.h 中算子定义了如下输入输出与属性参数名输入/输出/属性描述数据类型数据格式input_grad输入反向传播过程中上一步输出的梯度FLOAT16、FLOAT、BFLOAT16NDinput_x输入前向 Softshrink 的输入张量FLOAT16、FLOAT、BFLOAT16NDoutput_y输出反向传播梯度输出FLOAT16、FLOAT、BFLOAT16NDlambd可选属性Softshrink 的阈值参数默认值 0.5需满足 lambd≥0FLOAT-其中lambd在 GE 算子原型中注册为可选属性OPTIONAL默认值0.5f见 soft_shrink_grad_def.cpp 与 soft_shrink_grad_proto.h。算子定义还声明了AutoContiguous()输入输出自动转为连续张量DynamicRankSupportFlag(true)与DynamicShapeSupportFlag(true)支持动态 rank 与动态 shapeExtendCfgInfo(opFile.value, soft_shrink_grad)指定对应的 kernel 实现文件。3.2 约束说明根据 README.md 的约束章节dtype/shape 一致性input_grad与input_x的 dtype 和 shape 必须一致GE 图模式硬性要求维度范围支持0–8 维ND 格式阈值合法性lambd必须满足lambd ≥ 0。3.3 形状推导与数据类型推导soft_shrink_grad_infershape.cpp 中注册了算子的推断逻辑InferShape4SoftShrinkGradoutput_y.shape input_grad.shape输出形状直接继承输入梯度InferDataType4SoftShrinkGradoutput_y的数据类型继承input_grad的数据类型。四、aclnn 两段式接口调用详解4.1 接口背景两段式 API 设计aclnnSoftshrinkBackward遵循 CANN aclnn 统一的两段式接口设计详见 两段式接口说明第一段aclnnSoftshrinkBackwardGetWorkspaceSize完成入参校验与计算流程编排返回执行所需 workspace 大小及aclOpExecutor执行器第二段aclnnSoftshrinkBackward传入 workspace 与 executor在指定 stream 上真正执行算子计算。4.2 函数原型aclnnStatus aclnnSoftshrinkBackwardGetWorkspaceSize( const aclTensor* gradOutput, // 上游梯度 const aclTensor* self, // Softshrink 正向输入 const aclScalar* lambda, // 阈值 λ非负 aclTensor* gradInput, // 输出梯度 uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnSoftshrinkBackward( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)4.3 第一段接口参数说明aclnnSoftshrinkBackwardGetWorkspaceSize的完整参数如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorgradOutputaclTensor*输入反向传播的上游梯度不能传空指针支持空 Tensorshape 与 self 满足 broadcast 关系数据类型可与 self 不同不同时内部转换为 FLOAT 计算FLOAT16、FLOAT、BFLOAT16ND0-8√selfaclTensor*输入Softshrink 正向计算的输入不能传空指针支持空 Tensorshape 与 gradOutput 满足 broadcast 关系数据类型可与 gradOutput 不同不同时内部转换为 FLOAT 计算FLOAT16、FLOAT、BFLOAT16ND0-8√lambdaaclScalar*输入Softshrink 计算的阈值 λ不能传空指针取值必须 ≥ 0FLOAT---gradInputaclTensor*输出Softshrink 反向传播的输出梯度不能传空指针支持空 Tensorshape 必须与 self 和 gradOutput 的 broadcast 结果一致FLOAT16、FLOAT、BFLOAT16ND0-8√workspaceSizeuint64_t*输出需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出包含算子计算流程的执行器-----要点解读支持非连续 Tensor两个输入与输出均允许传入非连续含 view张量API 内部会通过l0op::Contiguous转为连续张量后计算再通过l0op::ViewCopy写回非连续输出见 aclnn_softshrink_backward.cpp支持 broadcastgradOutput与self的 shape 只需满足 broadcast 关系API 内部先BroadcastInferShape推断广播后形状再通过l0op::BroadcastTo将两者广播到一致 shape见 aclnn_softshrink_backward.cpp混合精度自动提升当gradOutput与self数据类型不同时非 FLOAT 的输入会被l0op::Cast提升为 FLOAT 计算输出再 Cast 回gradInput声明的 dtype见 aclnn_softshrink_backward.cpp支持空 TensorgradOutput或gradInput为空时第一段接口直接返回workspaceSize 0不执行实际计算见 aclnn_softshrink_backward.cpp。4.4 返回值与错误码两段接口均返回aclnnStatus状态码具体定义参见 aclnn 返回码说明。第一段接口完成入参校验以下场景会报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001gradOutput、self、lambda、gradInput、workspaceSize 或 executor 存在空指针ACLNN_ERR_PARAM_INVALID161002gradOutput、self 或 gradInput 的数据类型不在当前产品支持范围内ACLNN_ERR_PARAM_INVALID161002self 和 gradOutput 的 shape 不满足 broadcast 规则ACLNN_ERR_PARAM_INVALID161002lambda 为 NaN 或小于 0ACLNN_ERR_PARAM_INVALID161002gradInput 的 shape 与 self 和 gradOutput 的 broadcast 结果不一致ACLNN_ERR_PARAM_INVALID161002gradOutput 或 self 的维度大于 8这些校验逻辑在 aclnn_softshrink_backward.cpp 的CheckParams中依次执行先检查空指针CheckNotNull再检查数据类型是否在支持列表内CheckDtypeValid随后校验维度上限MAX_SUPPORT_DIMS_NUMS、broadcast 关系与输出 shape 一致性CheckShapeValid最后检查lambda是否为 NaN 且不小于 0。4.5 第二段接口参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取executor输入op 执行器包含算子计算流程stream输入指定执行任务的 Stream第二段接口的实现非常简洁直接调用框架统一的CommonOpExecutorRun完成计算见 aclnn_softshrink_backward.cpp这也是所有 aclnn 算子的固定收尾写法。4.6 确定性计算aclnnSoftshrinkBackward.md 明确说明aclnnSoftshrinkBackward 默认确定性实现同一输入在多次执行中计算结果完全一致适用于对可复现性有严格要求的训练场景。五、完整调用示例基于 aclnn API仓库提供了可直接编译运行的完整样例 examples/test_aclnn_soft_shrink_grad.cpp同时存在 arch35 变体 examples/arch35/test_aclnn_soft_shrink_grad.cpp。编译与执行方法可参考 编译与运行样例。下面按步骤拆解其调用流程5.1 头文件与宏定义#include cstdio #include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_softshrink_backward.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0)5.2 初始化 device 与 streamint Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, ...); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, ...); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, ...); return 0; }5.3 构造输入输出 aclTensor样例使用 shape 为{4, 2}的数据构造三个张量gradOutput上游梯度{0,1,2,3,4,5,6,7}、self前向输入{1,1,1,2,1,2,3,3}、gradInput输出占位初始全 0并创建阈值lambda 1.2f的aclScalarstd::vectorint64_t gradOutputShape {4, 2}; std::vectorint64_t selfShape {4, 2}; std::vectorint64_t gradInputShape {4, 2}; std::vectorfloat gradOutputHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectorfloat selfHostData {1, 1, 1, 2, 1, 2, 3, 3}; std::vectorfloat gradInputHostData {0, 0, 0, 0, 0, 0, 0, 0}; float lambdaValue 1.2f; // 通过 CreateAclTensor 模板函数完成aclrtMalloc 申请 device 内存、 // aclrtMemcpy 拷贝 host 数据到 device、计算 strides、aclCreateTensor 创建 aclTensor ret CreateAclTensor(gradOutputHostData, gradOutputShape, gradOutputDeviceAddr, aclDataType::ACL_FLOAT, gradOutput); ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); lambda aclCreateScalar(lambdaValue, aclDataType::ACL_FLOAT); ret CreateAclTensor(gradInputHostData, gradInputShape, gradInputDeviceAddr, aclDataType::ACL_FLOAT, gradInput);5.4 两段式调用// 第一段获取 workspace 大小与执行器 uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; ret aclnnSoftshrinkBackwardGetWorkspaceSize(gradOutput, self, lambda, 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 aclnnSoftshrinkBackward(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, ...);5.5 同步、取回结果与资源释放// 同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, ...); // 将 device 侧结果拷贝回 host 并打印 ret aclrtMemcpy(resultData.data(), ..., gradInputDeviceAddr, ..., ACL_MEMCPY_DEVICE_TO_HOST); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %f\n, i, resultData[i]); } // 释放 aclTensor / aclScalar 与 device 资源 aclDestroyTensor(gradOutput); aclDestroyTensor(self); aclDestroyScalar(lambda); aclDestroyTensor(gradInput); aclrtFree(gradOutputDeviceAddr); aclrtFree(selfDeviceAddr); aclrtFree(gradInputDeviceAddr); if (workspaceSize 0) aclrtFree(workspaceAddr); aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize();5.6 预期结果验证以样例数据手动推导lambda 1.2|self| 1.2的元素位置梯度透传。self {1,1,1,2,1,2,3,3}的绝对值为{1,1,1,2,1,2,3,3}其中大于 1.2 的是第 3、5、6、7 个元素值为 2、2、3、3对应gradOutput中值为 3、5、6、7 的梯度透传其余位置输出 0。即预期输出为{0, 0, 0, 3, 0, 5, 6, 7}。读者可通过运行样例程序打印的result[i]进行对照验证仓库中的单元测试 tests/ut/op_api/test_aclnn_softshrink_backward.cpp 与 tests/ut/op_kernel/test_soft_shrink_grad.cpp 也覆盖了该算子的行为断言。六、底层实现剖析AscendC Kernel 与 tiling6.1 内核入口kernel 入口 是一个典型的 AICore 算子入口从 tiling 数据中读取SoftShrinkGradTilingData实例化模板类NsSoftShrinkGrad::SoftShrinkGradDTYPE_INPUT_GRAD, BUFFER_MODE并依次调用Init与Process。BUFFER_MODE模板参数控制双缓冲BUFFER_NUM 2或单缓冲BUFFER_NUM 1模式。6.2 数据流与 tiling 分块内核实现 采用经典的三段流水结构Init根据tilingData计算当前 block 负责的blockLength_最后一个 block 处理余量blockTail、UB 缓冲长度ubLength_与阈值lambd_并按元素偏移blockFormer * blockIdx定位 Global Memory 上的输入输出地址随后pipe.InitBuffer为输入梯度、输入 x、输出分别分配BUFFER_NUM个 UB 队列缓冲CopyIn通过DataCopyPad将 Global Memory 数据按currentNum * sizeof(T)的分片长度搬入 UBCompute执行逐元素比较与选择逻辑详见 6.3CopyOut将计算结果DataCopyPad写回 Global Memory。Process中以loopCount (blockLength_ ubLength_ - 1) / ubLength_循环切分处理尾块通过currentNum blockLength_ - ubLength_ * i处理余量配合双缓冲实现数据搬运与计算的重叠最大化 AICore 利用率。6.3 比较-选择计算核心与 NaN 保护实现Compute中针对FP16/BF16 提升路径与FP32 直算路径分别实现核心逻辑一致取绝对值Abs(absX, xFP32, alignedNum)计算|x|阈值比较CompareScalar(cmpMask, absX, lambd_, CMPMODE::GT, alignedNum)生成比较位图cmpMask|x| lambd为 1条件选择Select(outFP32, cmpMask, gradFP32, 0.0f, SELMODE::VSEL_TENSOR_SCALAR_MODE, alignedNum)cmpMask1时输出梯度否则输出 0NaN 保护利用 IEEE 754 中NaN ! NaN恒为真的性质Compare(nanMask, xFP32, xFP32, CMPMODE::NE, alignedNum)生成 NaN 位图再执行第二次Select(outFP32, nanMask, gradFP32, outFP32, SELMODE::VSEL_TENSOR_TENSOR_MODE, alignedNum)将 NaN 位置的结果覆盖为输入梯度——这正是与 PyTorch 行为对齐的关键实现见 soft_shrink_grad.h。FP16/BF16 路径会先将输入Cast到 FP32 参与计算以避免精度损失最终Cast回原 dtype向下转换使用CAST_ROUND同时计算缓冲通过GetWithOffset划分出[gradFP32 | xFP32 | absX | outFP32 | cmpMask | nanMask]六个分区。另外由于CompareScalar/Compare要求计数占用 256B 对齐跨度代码将currentNum向上对齐到 64 元素CMP_ALIGN_ELEMS 256 / sizeof(float)并对中间缓冲做零初始化防止未对齐尾元素产生脏数据见 soft_shrink_grad.h。6.4 平台相关的 tiling 与 kernel 适配算子针对新平台如 arch35对应 Ascend 950 系列提供了专门的 tiling 计算与 kernel 实现Host 侧 tiling op_host/arch35/soft_shrink_grad_tiling_arch35.cpp负责计算 block 划分、UB 分片与阈值等 tiling 参数Kernel 侧数据结构 op_kernel/arch35/soft_shrink_grad_tiling_data.h 定义SoftShrinkGradTilingData含blockNum、blockFormer、dim0、ubFormer、lambd等字段soft_shrink_grad_tiling_key.h 定义 tiling key对应的单测 tests/ut/op_host/arch35/test_soft_shrink_grad_tiling.cpp 覆盖 tiling 计算正确性。七、GE 图模式接入除 aclnn 单算子 API 外SoftShrinkGrad 还以GE 图算子的形式注册供计算图构建与图编译使用。其算子原型定义在 op_graph/soft_shrink_grad_proto.hREG_OP(SoftShrinkGrad) .INPUT(input_grad, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16})) .INPUT(input_x, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16})) .OUTPUT(output_y, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16})) .ATTR(lambd, Float, 0.5) .OP_END_FACTORY_REG(SoftShrinkGrad)同时 op_graph/soft_shrink_grad_graph_infer.cpp 提供图模式下的 shape 推导op_graph/CMakeLists.txt 负责该模块的构建。开发者可在训练框架的构图阶段以SoftShrinkGrad为节点名、传入input_grad/input_x与属性lambd构建计算图由 GE 完成图编译后下发 NPU 执行。两种调用方式aclnn API 与 GE 图模式的对应关系总结如下调用方式入口/示例适用场景aclnn APIexamples/test_aclnn_soft_shrink_grad.cpp aclnnSoftshrinkBackward 接口文档单算子调测、自定义训练脚本中显式调用GE 图模式op_graph/soft_shrink_grad_proto.h 中的算子原型框架构图、整图下发与图优化场景八、工程结构与测试验证8.1 目录结构从仓库目录结构看soft_shrink_grad算子遵循 ops-nn 标准的算子工程组织方式每类平台文件按arch35等目录分层便于跨平台扩展与独立编译op_api/aclnn 对外接口实现aclnn_softshrink_backward.cpp与内部封装 soft_shrink_grad.cpp、soft_shrink_grad.hop_host/算子定义、shape/dtype 推导与 tiling 计算op_kernel/AscendC kernel 实现op_graph/GE 算子原型注册与图推导docs/aclnn 接口文档examples/可编译运行的单算子样例tests/包含ut/op_api、ut/op_host、ut/op_kernel三层单元测试以及st/aclnnSoftshrinkBackward的 ATK 用例配置。8.2 测试覆盖仓库为算子提供了完善的测试保障API 层单测tests/ut/op_api/test_aclnn_softshrink_backward.cpp 验证两段式接口的正确性Kernel 层单测tests/ut/op_kernel/test_soft_shrink_grad.cpp 验证逐元素计算逻辑Tiling 单测tests/ut/op_host/arch35/test_soft_shrink_grad_tiling.cpp 验证分块参数ST 测试tests/st/aclnnSoftshrinkBackward/atk_aclnnSoftshrinkBackward.json 提供昇腾算子工具链ATK的端到端用例配置。九、小结与使用建议SoftShrinkGrad 作为 Softshrink 的反向算子通过绝对值与阈值比较 NaN 自反检测覆盖的简单而严谨的规则在 NPU 上实现了与 PyTorch 语义完全对齐的梯度计算。实际使用时建议注意以下几点阈值约束lambda必须为非负值源码强制校验lambd 0或为 NaN 时报ACLNN_ERR_PARAM_INVALID默认 0.5与 PyTorchnn.Softshrink默认lambd0.5保持一致数据类型规划新平台950/A3/A2支持 bf16老平台仅 fp16/fp32混合精度输入会被统一提升为 FP32 计算后再转回存在轻微额外开销建议保持gradOutput与self同 dtypeshape 一致性GE 模式下要求input_grad与input_x的 dtype 和 shape 完全一致aclnn 模式下两者允许满足 broadcast 关系但输出gradInput必须与 broadcast 结果一致两段式调用务必先调用GetWorkspaceSize获取 workspace 大小并申请 Device 内存再调用第二段接口若workspaceSize为 0 可跳过申请空 Tensor 场景确定性保证该算子默认确定性实现适合对结果可复现有要求的训练与调测流程。更多算子家族信息可进一步阅读前向算子 Softshrink README 及其接口文档以及 ops-nn 仓库的 算子开发总览 与 两段式接口说明。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考