企业官网建设流程全解析

CANN ops-nn 中 HardShrinkGrad 反向梯度算子的 Ascend C 实现与 ACLNN 调用指南

【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn

导读

本文以 CANN ops-nn 仓库中experimental/activation/hard_shrink_grad算子为研究对象，系统讲解 HardShrink 激活函数反向梯度算子（HardShrinkGrad）的数学语义、算子规格、ACLNN 调用方式、Tiling 切分策略与 Kernel 计算方案。读完本文，你将掌握如何在昇腾 NPU（Ascend 910B / arch22）上通过aclnnHardShrinkGrad两段式接口完成梯度计算，并理解 fp16/bf16 在 arch22 上"先升精度、再计算、后降精度"的实现动机与精度验收标准。

算子背景与数学语义

HardShrink 是一种阈值型激活函数：对输入中小于等于阈值lambd的部分进行收缩置零。在神经网络训练中，反向传播需要把上游梯度按前向的收缩结果进行"门控"，这就是 HardShrinkGrad 的职责。

HardShrinkGrad 根据前向输入self的绝对值与阈值lambd的比较结果，决定是否将上游梯度grad_output传递到下游，其逐元素计算规则为：

output_i = grad_output_i, if |self_i| > lambd 0, otherwise

等价表达为output = grad_output * (|self| > lambd ? 1 : 0)，其中严格大于lambd才透传梯度，等于或小于lambd时梯度被置零。对标框架为 PyTorch 的hardshrink_backward。

这是一个典型的 Elementwise（逐元素）条件赋值算子：输出要么是精确的grad_output原值，要么是精确的 0，不引入中间舍入误差，因此精度理论上只受Cast与比较边界的影响（详见"精度标准"一节）。

算子规格

输入

输出

属性

lambd即 Hardshrink 前向的收缩阈值。在算子原型注册中它被声明为可选属性并默认取0.5f（见 hard_shrink_grad_def.cpp），在 Tiling 阶段若属性指针为空则回退到该默认值。

约束

• grad_output和self的 dtype 与 shape 必须完全一致，不支持广播；
• 最大支持 8 维张量（原型注册中开启DynamicRankSupportFlag(true)，同时输入均声明为FORMAT_ND）；
• 空 tensor（元素数量为 0）：直接返回空 tensor，不触发实际计算。

ACLNN 调用方法

两段式接口调用

与 CANN 其他算子一致，HardShrinkGrad 通过 ACLNN 的两段式接口完成调用。第一段aclnnHardShrinkGradGetWorkspaceSize用于查询 workspace 大小并创建 executor，第二段aclnnHardShrinkGrad真正下发计算：

#include "aclnnop/aclnn_hard_shrink_grad.h" // 第一段：计算 workspace 大小 uint64_t workspaceSize = 0; aclOpExecutor* executor = nullptr; aclnnStatus ret = aclnnHardShrinkGradGetWorkspaceSize( gradOutput, self, lambd, output, &workspaceSize, &executor); // 分配 workspace void* workspace = nullptr; if (workspaceSize > 0) { aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 第二段：执行计算 ret = aclnnHardShrinkGrad(workspace, workspaceSize, executor, stream);

需要说明的是，从 Tiling 实现看该算子在当前版本申请的系统 workspace 大小为 0（hard_shrink_grad_tiling.cpp 中WS_SYS_SIZE = 0），即计算过程不需要额外 GM workspace，但按 ACLNN 两段式接口的通用规范仍需保留上述分配逻辑，以兼容后续版本变化。

仓库提供了完整的可直接编译运行的单算子调用示例 examples/test_aclnn_hard_shrink_grad.cpp，其中包含完整的资源初始化（aclInit、aclrtSetDevice、aclrtCreateStream）、Host 数据到 Device 的搬运（aclrtMalloc+aclrtMemcpy）、aclTensor构造、两段式调用以及结果回拷校验流程，是学习 ACLNN 调用范式的最佳模板。

编译与安装

# 编译算子包 bash build.sh # 安装到 CANN OPP 目录 # 编译产物在 build/ 目录下，按自定义算子安装流程部署

算子子目录通过顶层 CMakeLists.txt 统一组织：当未开启ENABLE_TEST时，tests目录不会参与构建；对每个包含CMakeLists.txt的子目录递归执行add_subdirectory，从而把op_host、op_kernel、examples等模块纳入同一构建体系。

运行测试

单元测试（UT）与系统测试（ST）均按算子目录独立组织：

# 运行 UT 测试（Tiling + InferShape 逻辑，无 NPU 依赖） cd tests/ut bash run.sh # 运行 ST 测试（ACLNN 全流程精度验证，需要 NPU 设备） cd tests/st bash run.sh

其中 ST 测试主程序为test_aclnn_hard_shrink_grad.cpp，测试用例定义集中在tests/st/testcases/下，通过逐元素 allclose 与 golden 结果比对来验收精度。

精度标准

精度比对方法为逐元素 allclose，判定公式：

|golden - actual| <= atol + rtol * |golden|

HardShrinkGrad 为条件赋值算子，输出为精确的 0 或精确的grad_output值，理论上不引入额外计算误差。需要留意的是 fp16/bf16 路径涉及 Cast 升/降精度，但该算子输出的候选值只有grad_output原值和 0 两种，Cast 往返不会改变这两个数值本身（0 恒为 0，grad_output的 fp16 值经 fp32 再转回 fp16 保持原值），因此精度验收表中 fp16/bf16 采用双千分之一标准即可覆盖。

目录结构与工程组织

算子工程采用"Host 侧 + Kernel 侧 + 测试 + 文档"的标准划分：

hard_shrink_grad/ ├── CMakeLists.txt # 顶层构建配置 ├── README.md # 算子说明文档 ├── examples/ │ └── test_aclnn_hard_shrink_grad.cpp # ACLNN 调用示例 ├── op_host/ # Host 侧实现 │ ├── CMakeLists.txt # Host 构建配置 │ ├── hard_shrink_grad_def.cpp # 算子原型注册（输入/输出/属性/芯片配置） │ ├── hard_shrink_grad_infershape.cpp # Shape 推导（output shape = input shape） │ └── hard_shrink_grad_tiling.cpp # Tiling 实现（多核切分、UB 切分、TilingKey 选择） ├── op_kernel/ # Kernel 侧实现 │ ├── CMakeLists.txt # Kernel 构建配置 │ ├── hard_shrink_grad.cpp # Kernel 入口（模板实例化，双类分发） │ ├── hard_shrink_grad.h # Kernel 类定义与实现 │ ├── hard_shrink_grad_tiling_data.h # TilingData 结构体 │ └── hard_shrink_grad_tiling_key.h # TilingKey 定义（6 种模板组合） └── tests/ # 测试代码（UT/ST 用例）

关于工程演进：该算子还有一个面向稳定发布形态的版本目录 activation/hard_shrink_grad，其在 experimental 版本的基础上进一步补齐了op_api（ACLNN 接口封装层，对应aclnn_hardshrink_backward系列接口）、op_graph（图融合与算子原型）以及arch35架构的 Kernel 实现，并配套了接口级 UT 与接口说明文档 aclnnHardshrinkBackward.md。对比阅读两个目录，可以直观看到算子从实验形态走向正式发布所经历的工程化过程。

实现要点与源码级原理

计算方案

Kernel 侧实现了两套计算类，由 hard_shrink_grad.cpp 的入口函数依据模板参数sizeof(D_T)在编译期分发：

• fp32（HardShrinkGradDirect 直算路径）：直接在原生类型上执行Abs -> Compare(GT, tensor-tensor) -> Select三步计算；
• fp16/bf16（HardShrinkGradCastFp32 升精度路径）：先Cast到 fp32，再执行Abs -> Compare -> Select，最后将结果Cast回原 dtype 输出。

两条路径均采用 Ascend C 的TPipe流水线 +TQue队列（VECIN/VECOUT）+TBuf计算缓冲（VECCALC）结构，主循环按CopyIn -> Compute -> CopyOut三段推进（见 hard_shrink_grad.h）。

fp32 直算路径的 Kernel 级数据流为：

CopyIn: GM(grad_output, self) -> UB Compute: Abs(self) -> absLocal Compare(absLocal, lambdLocal, GT) -> cmpMask Select(cmpMask, grad, zero) -> output CopyOut: UB(output) -> GM

fp16/bf16 升精度路径则在 Compute 阶段前后各插入一次 Cast：

CopyIn: GM(grad_output_fp16, self_fp16) -> UB Compute: Cast(fp16 -> fp32) Abs(self_fp32) -> absLocal_fp32 Compare(absLocal_fp32, lambdLocal_fp32, GT) -> cmpMask Select(cmpMask, grad_fp32, zero_fp32) -> out_fp32 Cast(fp32 -> fp16) CopyOut: UB(output_fp16) -> GM

lambd与常数 0 在Init阶段通过Duplicate预填充为 UB 常驻缓冲，避免在逐轮循环中重复构造。

关键技术决策

源码注释与实现共同印证了以下三个关键决策（这也是该算子最具移植参考价值的部分）：

1. arch22 上 Compare API 对 fp16/bf16 不可靠：统一 Cast 到 fp32 计算，确保精度达标。这是 fp16/bf16 必须走升精度路径的根本原因，直接决定了两套 Kernel 类的存在。
2. bf16 不支持 Abs/Compare/Select API：复用 fp16 的CastFp32路径。由于 bf16 与 fp16 的元素宽度相同（sizeof(T) == 2），入口函数仅凭sizeof(D_T)即可完成统一分发，无需为 bf16 单独编写 Kernel。
3. 使用 Compare(tensor-tensor) 替代 CompareScalar：lambd通过Duplicate广播为与数据同长度的 tensor 再参与比较，相比标量比较更通用、更可靠，同时规避了标量比较在边界条件下的潜在行为差异。

TilingKey 组合

算子在 Tiling 阶段通过ASCENDC_TPL_ARGS_DECL/ASCENDC_TPL_SEL声明 6 种模板组合（见 hard_shrink_grad_tiling_key.h），模板参数为数据类型D_T与缓冲模式BUFFER_MODE（0=单缓冲，1=双缓冲）：

双缓冲阈值MIN_SPLIT_THRESHOLD = 1024（见 hard_shrink_grad_tiling.cpp）：当totalNum > 1024时启用双缓冲以隐藏 GM->UB 拷贝与向量计算之间的等待，数据量小时则用单缓冲以节省 UB 空间。

Tiling 切分流程

Tiling 函数按五步完成切分决策（对应 hard_shrink_grad_tiling.cpp）：

1. 获取平台信息：通过GetCoreNumAiv()获取可用 AI Core 数coreNum，通过GetCoreMemSize(UB)获取 UB 容量；
2. 获取输入信息：读取grad_output与self的 storage shape，校验二者元素总数一致（否则报错），并校验 dtype 属于{DT_FLOAT16, DT_FLOAT, DT_BF16}；
3. 获取属性：读取lambd（缺省 0.5f）；
4. 计算切分参数：
   • 多核切分：blockFactor = CeilDiv(totalNum, coreNum)，实际使用核数usedCoreNum = CeilDiv(totalNum, blockFactor)；
   • UB 切分：先按数据类型估算单元素占用字节数（fp32 直算路径为(3 * bufNum + 3) * 4，fp16/bf16 升精度路径为3 * bufNum * 2 + 6 * 4），再从 UB 容量中预留 256 字节cmpMask位图空间，计算出ubFactor并分别按ubBlockSize与 Compare 的 256 字节对齐约束（即 64 个 fp32 元素）向下取整对齐；
5. 设置 BlockDim 与 TilingKey：BlockDim = usedCoreNum，TilingKey 由dTypeVal与useDoubleBuffer组合选择。

空 tensor 与边界处理

当totalNum == 0时，Tiling 直接设置BlockDim = 1并返回，Kernel 因blockLength_ <= 0直接跳过计算（hard_shrink_grad.h），从而保证空 tensor 输入返回空 tensor。TilingData 结构仅包含 4 个字段（hard_shrink_grad_tiling_data.h）：

struct HardShrinkGradTilingData { int64_t totalNum = 0; // 总元素数 int64_t blockFactor = 0; // 每个 AI Core 分到的元素数 int64_t ubFactor = 0; // 每轮 UB 迭代处理的元素数 float lambd = 0.5f; // 收缩阈值（来自属性） };

测试结果与验收

UT 覆盖 Tiling 逻辑（48 条）与 InferShape 逻辑（6 条），ST 为 20 条 L0 用例，在真实 NPU 上验证aclnnHardShrinkGrad全链路精度；穿刺验证工程覆盖 fp16 单缓冲、fp32 单/双缓冲、bf16 支持、空 tensor、lambd=0边界共 6 个专项场景。

开发环境

小结

HardShrinkGrad 是典型的逐元素反向梯度算子，其工程实现完整覆盖了"原型注册 -> InferShape -> Tiling -> Kernel 计算"的自定义算子开发全流程。对读者而言，本算子最具借鉴价值的三个点在于：fp16/bf16 升精度计算的规避策略（应对 arch22 Compare API 的可靠性问题）、tensor-tensor Compare + Select 的条件赋值范式（替代标量比较的通用写法），以及按数据量阈值在单/双缓冲间切换的 Tiling 模板设计。如需将同类激活梯度算子移植到其他昇腾架构，可对照本算子的实现要点逐项核验 API 支持度后复用其整体框架。

【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn