人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载本篇技术指南以 CANN 神经网络算子库 ops-nn 中的experimental/activation/threshold_grad_v2_d算子为对象系统讲解 ThresholdGradV2D 算子的数学定义、参数约束、Host 侧算子注册与 Shape 推导、Tiling 数据切分、NPU Kernel 实现原理并给出基于aclnnThresholdBackward两段式接口的完整调用示例与单元测试验证路径。读完本文你将掌握如何阅读一个 Ascend 自定义算子从定义、编译到调用的完整链路并能够直接复用文中的示例代码在 Atlas A2 训练系列产品/Atlas 800I A2 推理产品上运行该算子。一、算子概览ThresholdGradV2D 解决什么问题ThresholdGradV2D 是threshold_v2d激活函数前向算子ThresholdV2D的反向传播梯度算子核心功能是求 threshold 激活函数的梯度。其计算逻辑可以用如下分段函数描述$$ y \begin{cases} input_gradient, input_feature threshold \ 0, input_feature threshold \end{cases} $$从语义上理解前向时ThresholdV2D会把输入中小于等于threshold的元素置为threshold而大于threshold的元素保持不变因此反向传播时只有input_feature threshold的位置梯度才能正常回传值为input_gradient其余位置的梯度被置零。这与 PyTorch 中torch.nn.functional.threshold的反向行为一致。在 ops-nn 仓库中该算子的正向对应接口为aclnnThreshold可参考 activation/threshold 目录 下的算子文档而本算子aclnnThresholdBackward专门完成其反向计算。算子整体位于 experimental/activation/threshold_grad_v2_d 目录遵循 CANN 算子开发的标准目录结构op_host负责算子定义/推导/Tilingop_kernel负责 NPU 侧 Kernel 实现op_api提供上层 aclnn 接口tests/ut提供单元测试examples提供可直接运行的调用样例。产品支持情况产品是否支持Atlas A2 训练系列产品 / Atlas 800I A2 推理产品√需要说明的是算子定义文件 op_host/threshold_grad_v2_d_def.cpp 中通过AICore().AddConfig(ascend910b, aicoreConfig)将 AICore 配置挂载到 ascend910b 架构上即当前算子面向 910B 及同架构系列产品编译部署。二、参数说明输入、输出与属性根据 README 与算子注册源码op_host/threshold_grad_v2_d_def.cppThresholdGradV2D 共包含 2 个输入、1 个属性、1 个输出参数名输入/输出/属性描述数据类型数据格式input_gradient输入反向传播梯度即公式中的input_gradientFLOAT、FLOAT16、BFLOAT16、INT32、INT8、UINT8NDinput_feature输入前向激活的输入特征与input_gradient具有相同的类型、格式和形状FLOAT、FLOAT16、BFLOAT16、INT32、INT8、UINT8NDthreshold属性阈值标量参与input_feature threshold的比较FLOAT—output_backprops输出公式中的输出张量即回传的梯度FLOAT、FLOAT16、BFLOAT16、INT32、INT8、UINT8ND以下几点需要特别关注类型一致性input_gradient、input_feature、output_backprops三者数据类型必须一致注册源码中对三个张量使用了相同的DataType与Format集合threshold_grad_v2_d_def.cppthreshold 是属性而非输入它是编译期常量在算子定义中以this-Attr(threshold).AttrType(REQUIRED).Float(1.0)声明默认值为 1.0形状约束Shape 推导逻辑op_host/threshold_grad_v2_d_infershape.cpp直接将输出 shape 赋值为输入 shape即*yShape *xShape输出与输入同形状属于逐元素elementwise语义。与 aclnn 层参数的对应关系在 aclnn 接口层op_host/op_api/aclnn_threshold_backward.cpp上述 4 个参数映射为aclnnThresholdBackward的 4 个核心参数gradOutput对应 input_gradient、self对应 input_feature、threshold对应属性但接口层以aclScalar传入、out对应 output_backprops。由于接口层支持 broadcastgradOutput与self的 shape 可以满足广播关系最终输出 shape 为二者广播后的形状详见 docs/zh/context/broadcast_relationship.md。三、实现链路从算子定义到 NPU Kernel一个 Ascend 算子的完整实现由 Host 侧的定义—推导—Tiling与 Device 侧的Kernel 计算两部分组成。下面结合源码逐层剖析。3.1 算子注册与属性声明op_host/threshold_grad_v2_d_def.cpp 通过继承OpDef完成算子注册this-Input(input_gradient) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_INT32, ge::DT_INT8, ge::DT_UINT8}) .Format({ge::FORMAT_ND, ...}) ... this-Attr(threshold).AttrType(REQUIRED).Float(1.0); OpAICoreConfig aicoreConfig; aicoreConfig.DynamicCompileStaticFlag(true) .DynamicRankSupportFlag(true) .DynamicShapeSupportFlag(true) .PrecisionReduceFlag(true) .ExtendCfgInfo(opFile.value, threshold_grad_v2_d); this-AICore().AddConfig(ascend910b, aicoreConfig); OP_ADD(ThresholdGradV2D);几个关键配置的含义DynamicShapeSupportFlag(true)与DynamicRankSupportFlag(true)支持动态 shape 与动态 rank算子运行时根据实际输入形状动态推导输出并动态分配 TilingPrecisionReduceFlag(true)允许在满足精度的前提下采用低精度计算路径Kernel 中对 BF16 使用 float 中间量计算即与此相关ExtendCfgInfo(opFile.value, threshold_grad_v2_d)指定 Kernel 入口文件名即op_kernel/threshold_grad_v2_d.cppAddConfig(ascend910b, ...)为 ascend910b 架构注册该算子。3.2 Shape 推导输出与输入同形op_host/threshold_grad_v2_d_infershape.cpp 中InferShapeThresholdGradV2D的实现非常简洁读取输入 0 的 shape 后整体赋给输出 0即*yShape *xShape随后通过IMPL_OP_INFERSHAPE(ThresholdGradV2D).InferShape(...)完成注册。这从图编译层面保证了输出张量的形状在编译期即可确定。3.3 Tiling 计算多核并行与数据切分由于 NPU 的片上统一缓冲区UB容量有限大张量必须切分为多个 tile 依次搬运计算。Tiling 逻辑位于 op_host/threshold_grad_v2_d_tiling.cpp核心流程如下获取平台信息GetPlatformInfo通过PlatformAscendC获取 UB 大小ubSize与 AIV 核数coreNum计算 tile 粒度GetShapeAttrsInfo依据数据类型选择不同的 UB 切分数宏定义UB_NUM_F32_F16 7、UB_NUM_BF16 10、UB_NUM_INT8_UINT8 14、UB_NUM_INT32 8结合BLOCK_SIZE 256字节块对齐算出单个 tile 的数据量tileDataNum核间负载均衡CalculateCoreBlockNums将数据按 256 字节块均分到各核余数由前tailBlockNum个核多承担一块bigCoreDataNum其余核承担smallCoreDataNum并分别算出大/小核的 tile 数finalBigTileNum/finalSmallTileNum与尾部数据量bigTailDataNum/smallTailDataNum写入 tiling 数据以上字段连同threshold一起写入ThresholdGradV2DTilingData结构体定义见 op_kernel/threshold_grad_v2_d_tiling_data.h并通过context-SetBlockDim(coreNum)设置核数、context-SetTilingKey(...)设置 tiling key。从源码结构看这是一种典型的双核型负载均衡策略big/small core 搭配当单个 tile 数据量大于等于总数据量时coreNum会被收敛为 1退化为单核执行。3.4 Kernel 实现CompareScalar Select 的掩码选择Device 侧 Kernel 入口在 op_kernel/threshold_grad_v2_d.cpp核心计算类KernelThresholdGradV2D定义于 op_kernel/threshold_grad_v2_d.h采用 AscendC 编程模型使用双缓冲流水线BUFFER_NUM 2VECIN/VECOUT队列重叠搬入、计算、搬出三个阶段。Compute针对不同数据类型做了分支处理但核心思想一致先比较生成掩码再按掩码选择。以 float 类型为例计算主体仅两条指令AscendC::CompareScalar(maskLocal, fLocal, static_castfloat(this-thresholdValue), AscendC::CMPMODE::GT, this-processDataNum); AscendC::Select(outLocal, maskLocal, gLocal, static_castfloat(0.0), AscendC::SELMODE::VSEL_TENSOR_SCALAR_MODE, this-processDataNum);即用CompareScalar将input_feature与threshold逐元素比较CMPMODE::GT大于生成 0/1 掩码再用Select依据掩码从input_gradient与常量 0 中选取结果。halfFLOAT16类型走相同路径仅常量类型改为half。对于需要提升精度计算的类型BF16、INT32、INT8、UINT8Kernel 先通过Cast转换到 floatINT8/UINT8 经 half 两级转换完成比较与选择输出时再转回原类型其中整数类型使用CAST_TRUNC截断舍入、BF16 使用CAST_RINT四舍五入这与算子定义中的PrecisionReduceFlag设置相互呼应。Process()则按 tile 循环执行CopyIn → Compute → CopyOut最后一轮用tailDataNum处理尾部不足一个 tile 的数据for (int32_t i 0; i loopCount - 1; i) { CopyIn(i); Compute(i); CopyOut(i); } this-processDataNum this-tailDataNum; CopyIn(loopCount - 1); Compute(loopCount - 1); CopyOut(loopCount - 1);四、调用方式aclnnThresholdBackward 两段式接口CANN 为每个算子封装了统一的 aclnn 两层接口必须先调用第一段接口获取 workspace 大小再调用第二段接口执行计算参见 docs/zh/context/two_phase_api.md。本算子的 aclnn 接口文档位于 experimental/activation/threshold_grad_v2_d/docs/aclnnThresholdBackward.md头文件实现位于 op_host/op_api/aclnn_threshold_backward.h。4.1 第一段接口aclnnThresholdBackwardGetWorkspaceSizeaclnnStatus aclnnThresholdBackwardGetWorkspaceSize( const aclTensor *gradOutput, // 输入反向梯度支持空 Tensor const aclTensor *self, // 输入前向特征与 gradOutput 满足 broadcast const aclScalar *threshold, // 输入阈值标量 aclTensor *out, // 输出与 broadcast 后的 shape 一致 uint64_t *workspaceSize,// 输出需要在 Device 侧申请的 workspace 大小 aclOpExecutor **executor) // 输出算子执行器封装了计算流程各输入参数的约束如下参数类型关键约束gradOutputFLOAT/BFLOAT16/FLOAT16/INT32/INT8/UINT8/INT64dtype 与 self 一致shape 与 self 满足 broadcast支持非连续 Tensor0-8 维self同上数据类型遵循互推导规则docs/zh/context/deduction_relationship.mdthresholdFLOAT接口层为aclScalardtype 需可转换为推导后的类型out同上dtype 与 self 相同shape 等于 self 与 gradOutput broadcast 后的 shapeworkspaceSize / executor—输出参数由框架计算并返回注意接口文档中的 dtype 列表包含 INT64但这是 REGBASE 架构且threshold 0.0时的特例此时内部可复用ReluGrad路径源码见 op_host/op_api/aclnn_threshold_backward.cpp对于本文面向的Atlas A2 训练/Atlas 800I A2 推理产品实际支持的数据类型为 FLOAT、BFLOAT16、FLOAT16、INT32、INT8、UINT8与算子注册的 dtype 集合一致。第一段接口会完成入参校验常见错误返回码如下返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001gradOutput 或 self 为空指针ACLNN_ERR_PARAM_INVALID161002dtype 不在支持范围 / shape 超过 8 维 / gradOutput、out 与 self 数据类型不一致校验逻辑与上述返回码一一对应aclnn_threshold_backward.cpp 中的CheckParams先查空指针再查 dtype 支持列表与一致性最后查最大维度完整返回码说明可参考 docs/zh/context/aclnn_return_code.md。4.2 第二段接口aclnnThresholdBackwardaclnnStatus aclnnThresholdBackward( void *workspace, // 输入Device 侧 workspace 内存地址 uint64_t workspaceSize, // 输入由第一段接口计算得到 aclOpExecutor *executor, // 输入第一段接口返回的执行器 const aclrtStream stream) // 输入执行任务所在的 Stream第二段接口不重复校验参数直接基于执行器在指定 Stream 上异步下发计算任务。4.3 确定性计算本算子默认采用确定性实现即相同输入与参数多次运行结果完全一致这在 docs/aclnnThresholdBackward.md 的约束说明中有明确说明适合用于训练反向等对可复现性有要求的场景。五、完整调用示例从设备初始化到结果回拷仓库在 examples/test_aclnn_threshold_backward.cpp 提供了可直接参考的完整样例。以下按七个步骤还原其主流程以{2, 2}形状、FLOAT 类型、threshold 1.0f为例#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_threshold_backward.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) int Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); // GetShapeSize 各维度乘积 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, return ret); ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, return ret); std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; // 按连续排布计算 strides } *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. 设备/Stream 初始化 int32_t deviceId 0; // 按实际设备填写 aclrtStream stream; CHECK_RET(Init(deviceId, stream) ACL_SUCCESS, return 1); // 2. 构造输入与输出shape 均为 {2, 2} std::vectorint64_t selfShape {2, 2}; std::vectorint64_t gradOutputShape {2, 2}; std::vectorint64_t outShape {2, 2}; void* selfDeviceAddr nullptr; void* gradOutputDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* gradOutput nullptr; aclScalar* threshold nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0.2f, 1.2f, 2.2f, 3.2f}; // input_feature std::vectorfloat gradOutputHostData {4.5f, 4.4f, 4.3f, 4.2f}; // input_gradient std::vectorfloat outHostData {0.0f, 0.0f, 0.0f, 0.0f}; float thresholdValue 1.0f; CHECK_RET(CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, ACL_FLOAT, self) ACL_SUCCESS, return 1); CHECK_RET(CreateAclTensor(gradOutputHostData, gradOutputShape, gradOutputDeviceAddr, ACL_FLOAT, gradOutput) ACL_SUCCESS, return 1); threshold aclCreateScalar(thresholdValue, aclDataType::ACL_FLOAT); CHECK_RET(threshold ! nullptr, return 1); CHECK_RET(CreateAclTensor(outHostData, outShape, outDeviceAddr, ACL_FLOAT, out) ACL_SUCCESS, return 1); // 3. 两段式接口调用先算 workspace再执行 uint64_t workspaceSize 0; aclOpExecutor* executor; auto ret aclnnThresholdBackwardGetWorkspaceSize(gradOutput, self, threshold, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, return 1); void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, return 1); } ret aclnnThresholdBackward(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, return 1); // 4. 同步等待任务执行完成 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, return 1); // 5. 结果回拷并打印 std::vectorfloat resultData(4, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(float), outDeviceAddr, resultData.size() * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, return 1); for (size_t i 0; i resultData.size(); i) { printf(result[%zu] is: %f\n, i, resultData[i]); } // 6. 释放 aclTensor / aclScalar aclDestroyTensor(self); aclDestroyTensor(gradOutput); aclDestroyScalar(threshold); aclDestroyTensor(out); // 7. 释放 Device 资源 aclrtFree(selfDeviceAddr); aclrtFree(gradOutputDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }对照公式可以手动验证该示例的结果self {0.2, 1.2, 2.2, 3.2}、threshold 1.0只有第一个元素0.2 1.0因此result[0] 0其余三个位置梯度原样通过4.4 / 4.3 / 4.2。示例的编译与运行流程可参考 docs/zh/context/compile_and_run_sample.md。六、单元测试接口层校验与 Kernel 行为验证仓库在 tests/ut 目录下按层组织了三个测试维度op_api 层tests/ut/op_api/test_aclnn_threshold_grad_v2_d.cpp基于OP_API_UT框架覆盖接口层校验行为例如非法 dtypeINT64 在非 REGBASE 场景返回ACLNN_ERR_PARAM_INVALIDgradOutput与selfdtype 不一致返回ACLNN_ERR_PARAM_INVALID输入 shape 满足 broadcast 关系时返回ACLNN_SUCCESS。 这些用例与 4.1 节的返回码表格相互印证op_host 层tests/ut/op_host/test_threshold_grad_v2_d_tiling.cpp验证 Tiling 计算在不同 shape、不同数据类型下的 tile 切分与核数分配结果op_kernel 层tests/ut/op_kernel/test_threshold_grad_v2_d.cpp在核函数级验证CompareScalar Select的计算结果与公式一致。七、总结ThresholdGradV2D 是 ops-nn 仓库中一个结构清晰、可作为 Ascend 自定义算子学习范本的激活反向算子数学上它仅需一次阈值比较与掩码选择工程上它完整覆盖了算子注册op_def、Shape 推导infershape、多核 Tiling 切分、AscendC 双缓冲 Kernel 实现以及 aclnn 两段式接口封装。通过本文你既可以快速理解梯度如何按阈值做门控这一核心计算语义也可以按图索骥沿 examples/test_aclnn_threshold_backward.cpp 与 tests/ut 完成从调用到验证的完整闭环为在 CANN 平台上开发和移植同类逐元素反向算子打下基础。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐CANN ops-nn 算子解析threshold_grad_v2_daclnnThresholdBackward阈值反向梯度算子实战指南CANN ops nn 算子解析threshold_grad_v2_daclnnThresholdBackward阈值反向梯度算子实战指南 本文以 CAN人工智能算子库深度学习CANNAscendCANN ops-nn 算子解析aclnnSeluBackward 两段式接口实现 SELU 反向梯度计算CANN ops nn 算子解析aclnnSeluBackward 两段式接口实现 SELU 反向梯度计算 aclnnSeluBackward 是 CANN人工智能算子库深度学习CANNAscendManimGL 如何在本地构建官方 Sphinx 文档并查看构建产物ManimGL 如何在本地构建官方 Sphinx 文档并查看构建产物 ManimGL本仓库3Blue1Brown 使用的 OpenGL 版本 manim把人工智能算子库深度学习CANNAscend上一篇如何选择最佳语音转文字方案Glass21中Whisper与Deepgram的终极对比指南下一篇awesome-tauri性能监控实时追踪应用运行状态终极指南创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考