ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

aclnnThresholdBackward:CANN ops-nn 中 Threshold 反向传播算子的两段式接口与源码级解析

aclnnThresholdBackward:CANN ops-nn 中 Threshold 反向传播算子的两段式接口与源码级解析 人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载aclnnThresholdBackward是 CANN ops-nn 开源算子库中 aclnnThreshold 前向算子的反向实现用于在 NPU 上计算 Threshold 激活函数的梯度。本文基于 aclnnThresholdBackward.md 展开完整讲解其计算公式、两段式接口原型、全部参数约束、返回码校验逻辑与可直接编译运行的调用示例并结合threshold_grad_v2_d目录下的 host 侧tiling/算子定义与 kernel 侧AscendC 实现源码揭示该接口在底层如何被映射到ThresholdGradV2D与ReluGrad两个核函数并完成多核切分计算。一、功能说明与计算公式aclnnThresholdBackward完成 Threshold 前向函数的反向传播给定反向梯度gradOutput与前向输入self输出out仅在self元素严格大于阈值threshold时保留对应位置的梯度其余位置置 0。其数学定义为$$ output \begin{cases} gradOutput(i) \text{if } self(i) threshold \ 0 \text{otherwise} \end{cases} $$从计算语义上看该接口与同目录下的核算子ThresholdGradV2D见 README.md其公式为input_feature threshold时输出input_gradient否则输出 0完全一致是这一核算子的上层 Aclnn 封装。值得注意的一个工程细节是当threshold取值为 0.0 时该函数退化为self(i) 0的梯度门控与 ReLU 反向gradOutput(i) * (self(i) 0)在语义上完全等价因此 host 侧实现会直接复用ReluGrad核算子来消除冗余计算详见第五节源码解析。二、产品支持情况产品是否支持Atlas A2 训练系列产品 / Atlas 800I A2 推理产品√该支持范围与底层核算子ThresholdGradV2D保持一致。需要特别说明的是不同产品代际在数据类型支持上存在差异在 aclnn_threshold_backward.cpp 中源码按 NPU 架构维护了三个 dtype 支持列表ASCEND910_DTYPE_DTYPE_SUPPORT_LISTFLOAT、INT32、FLOAT16、INT8、UINT8ASCEND910B_DTYPE_DTYPE_SUPPORT_LIST在上一列表基础上增加 BF16REGBASE_DTYPE_DTYPE_SUPPORT_LIST在 910B 列表基础上再增加 INT64仅在threshold 0.0走 ReluGrad 路径时启用。由此可以推断文档中列出的 INT64 支持是有条件的依赖 threshold 取值与产品架构而 threshold_grad_v2_d_def.cpp 中核算子自身声明的数据类型仅为 FLOAT16、FLOAT、BF16、INT32、INT8、UINT8 六种。在 Atlas A2/800I A2 上实际支持 FLOAT、BFLOAT16、FLOAT16、INT32、INT8、UINT8 六种数据类型。三、函数原型两段式接口与 CANN 算子库的通用约定一致详见 两段式接口说明该 API 由“获取 workspace 大小”与“执行计算”两段组成必须先调用第一段获取入参校验结果、workspace 大小与执行器再调用第二段提交计算任务。第一段接口获取 workspace 大小并完成入参校验aclnnStatus aclnnThresholdBackwardGetWorkspaceSize( const aclTensor *gradOutput, const aclTensor *self, const aclScalar *threshold, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)第二段接口执行计算aclnnStatus aclnnThresholdBackward( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)其中第一段接口内部会依次完成创建OpExecutor→ 空指针校验 → 数据类型校验 → 维度校验 → 输入 broadcast 校验 → 输出 shape 校验 → 空 Tensor 短路处理 → 构图生成Contiguous、ReluGrad/ThresholdGradV2D、ViewCopy算子节点→ 计算并返回 workspace 大小。四、aclnnThresholdBackwardGetWorkspaceSize 参数详解第一段接口共 6 个参数其中 4 个是计算相关的张量/标量参数2 个是框架返回参数参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorgradOutput输入公式中的 gradOutput支持空 Tensordtype 需与 self 一致shape 需与 self 满足 broadcast 关系FLOAT、BFLOAT16、FLOAT16、INT32、INT8、UINT8、INT64ND0-8√self输入前向输入张量公式中的比较基准数据类型与 gradOutput 满足互推导关系FLOAT、BFLOAT16、FLOAT16、INT32、INT8、UINT8、INT64ND0-8√threshold输入公式中的阈值dtype 需与 gradOutput 一致shape 需与 gradOutput 满足 broadcast 关系FLOAT、BFLOAT16、FLOAT16、INT32、INT8、UINT8、INT64ND0-8√out输出公式中的输出 outdtype 需与 self 相同shape 需等于 self 与 gradOutput broadcast 之后的 shapeFLOAT、BFLOAT16、FLOAT16、INT32、INT8、UINT8、INT64ND0-8√workspaceSize输出需要在 Device 侧申请的 workspace 大小-----executor输出op 执行器包含算子计算流程-----关于上表需要补充几点实现层面的说明支持空 Tensor在 aclnn_threshold_backward.cpp 中当self或gradOutput为空 Tensor 时接口直接返回workspaceSize 0并释放执行器不进入 kernel 计算路径。非连续 Tensor 支持源码在构图前会分别对gradOutput、self调用l0op::Contiguous转为连续内存最终结果通过l0op::ViewCopy写回可能非连续的out因此三个张量均允许非连续排布。broadcast 校验接口通过OP_CHECK_BROADCAST_AND_INFER_SHAPE推导self与gradOutput的广播后 shape并要求out的 view shape 与之严格相等否则返回ACLNN_ERR_PARAM_INVALID源码第 122-128 行。返回值与第一段接口报错场景第一段接口返回aclnnStatus状态码具体取值参见 aclnn返回码。在 op_api 单元测试 中可以看到这些错误场景的实测覆盖返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 gradOutput 或 self 是空指针ACLNN_ERR_PARAM_INVALID161002gradOutput 或 self 的数据类型不在支持范围之内ACLNN_ERR_PARAM_INVALID161002gradOutput 或 self 的 shape 超过 8 维ACLNN_ERR_PARAM_INVALID161002gradOutput、out 与 self 数据类型不一致ACLNN_ERR_PARAM_INVALID161002out 的 shape 与 self、gradOutput broadcast 后的 shape 不一致此外源码中threshold、out、workspaceSize指针为空也会分别触发ACLNN_ERR_PARAM_NULLPTR。对应测试用例如l2_test_unsupport_dtype、l2_test_unmatch_dtype、l2_test_invalid_out_shape、l2_test_nullptr验证了这些返回码的实际行为。五、aclnnThresholdBackward 参数详解第二段接口仅负责提交执行参数含义如下参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnThresholdBackwardGetWorkspaceSize获取executor输入op 执行器包含算子计算流程stream输入指定执行任务的 Stream其实现非常简洁在 aclnn_threshold_backward.cpp 中第二段接口直接调用框架统一入口CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成异步计算返回值同样为aclnnStatus。底层计算路径ThresholdGradV2D 与 ReluGrad 的动态选择这是本算子最值得关注的源码设计点。第一段接口在完成连续化与校验后会根据 threshold 的取值动态选择底层核算子aclnn_threshold_backward.cppconst aclTensor* opOut; if (IsFloatEqual(thresholdVal_, 0.0)) { opOut l0op::ReluGrad(gradOutputContiguous, selfContiguous, uniqueExecutor.get()); } else { opOut l0op::ThresholdGradV2D(gradOutputContiguous, selfContiguous, thresholdVal_, uniqueExecutor.get()); }即threshold 0.0用浮点 epsilon 判等见IsFloatEqual时复用 activation/relu_grad 的 ReluGrad 核此时可额外支持 INT64见前文 REGBASE 列表否则调用ThresholdGradV2D核算子tiling 入口位于 threshold_grad_v2_d_tiling.cpp。ThresholdGradV2D的算子定义见 threshold_grad_v2_d_def.cpp其中Attr(threshold)的默认值为 1.0输入输出均为 ND 格式并启用了动态 shape、动态 rank 与精度降低PrecisionReduceFlag(true)支持。其 shape 推导threshold_grad_v2_d_infershape.cpp将输出 shape 直接复制自输入input_gradient。Kernel 侧实现CompareScalar Select 的门控逻辑核函数实现位于 threshold_grad_v2_d.cpp 与 threshold_grad_v2_d.h。以 FLOAT 路径为例核心计算仅两步KernelThresholdGradV2D::ComputeCompareScalar(maskLocal, fLocal, threshold, CMPMODE::GT, ...)逐元素比较input_feature threshold生成 0/1 maskSelect(outLocal, maskLocal, gLocal, 0.0, VSEL_TENSOR_SCALAR_MODE, ...)mask 为 1 的位置取input_gradient否则取 0。对于 INT8/UINT8 等整型输入会先经Cast提升到 half/float 参与比较与选择再以CAST_TRUNC截断回原类型BF16 路径使用CAST_RINT四舍五入回写。数据流采用标准 AscendC 流水双缓冲TQueBUFFER_NUM 2实现CopyIn → Compute → CopyOut三段流水Process()中前tileNum - 1个 tile 以满tileDataNum处理最后一个 tile 按tailDataNum处理尾部数据。Tiling 切分策略tiling 计算位于 threshold_grad_v2_d_tiling.cpp整体流程为通过PlatformAscendC获取 UB 大小与 AIV 核数GetCoreNumAiv按数据类型选择单核可容纳的数据块数ubDataNumberF32/F16 为 7、INT32 为 8、BF16 为 10、INT8/UINT8 为 14再结合BLOCK_SIZE 256推导tileDataNum将输入按 256 字节对齐后均分到各核计算大核/小核数据量bigCoreDataNum/smallCoreDataNum、每核 tile 数及尾部数据量bigTailDataNum/smallTailDataNum、尾核数tailBlockNum实现核间负载均衡将全部切分参数写入ThresholdGradV2DTilingData并设置 block dim 与 tiling key。六、约束说明确定性计算aclnnThresholdBackward默认采用确定性实现即相同输入与配置下多次运行结果一致。七、调用示例可直接编译运行以下示例来自 examples/test_aclnn_threshold_backward.cpp使用{2, 2}的 FLOAT 张量演示完整调用流程。整体编译与执行步骤请参考 编译与运行样例。#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) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); 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); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides 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]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. 固定写法device/stream初始化参考acl API手册 // 根据自己的实际device填写deviceId int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 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.2, 1.2, 2.2, 3.2}; std::vectorfloat gradOutputHostData {4.5, 4.4, 4.3, 4.2}; std::vectorfloat outHostData {0.0, 0.0, 0.0, 0.0}; float thresholdValue 1.0f; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建gradOutput aclTensor ret CreateAclTensor(gradOutputHostData, gradOutputShape, gradOutputDeviceAddr, aclDataType::ACL_FLOAT, gradOutput); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建threshold aclScalar threshold aclCreateScalar(thresholdValue, aclDataType::ACL_FLOAT); CHECK_RET(threshold ! nullptr, return ret); // 创建out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnThresholdBackward第一段接口 ret aclnnThresholdBackwardGetWorkspaceSize(gradOutput, self, threshold, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnThresholdBackwardGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 调用aclnnThresholdBackward第二段接口 ret aclnnThresholdBackward(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnThresholdBackward failed. ERROR: %d\n, ret); return ret); // 4. 固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至host侧 auto size GetShapeSize(outShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] 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、gradOutput {4.5, 4.4, 4.3, 4.2}。按公式逐元素计算0.2 1.0为假 → 输出 01.2 1.0为真 → 输出 4.42.2 1.0为真 → 输出 4.33.2 1.0为真 → 输出 4.2。即最终输出约为{0.0, 4.4, 4.3, 4.2}。若将thresholdValue改为 0.0则接口内部会切换到 ReluGrad 路径输出变为{0.0, 4.4, 4.3, 4.2}与 0 阈值门控一致。该示例可直接用于验证接口的正确性与两段式调用的完整性。八、关联测试与工程入口汇总用途仓库路径API 文档本文主体experimental/activation/threshold_grad_v2_d/docs/aclnnThresholdBackward.md可运行示例examples/test_aclnn_threshold_backward.cpp两段式接口实现op_host/op_api/aclnn_threshold_backward.cpp算子定义IR 注册op_host/threshold_grad_v2_d_def.cppshape 推导op_host/threshold_grad_v2_d_infershape.cpptiling 切分op_host/threshold_grad_v2_d_tiling.cppAscendC 核函数op_kernel/threshold_grad_v2_d.cpp、op_kernel/threshold_grad_v2_d.hop_api 单元测试tests/ut/op_api/test_aclnn_threshold_grad_v2_d.cpp前向算子文档activation/threshold/docs/aclnnThresholdaclnnInplaceThreshold.md九、总结aclnnThresholdBackward是一个典型的“薄封装、强校验”的两段式 Aclnn 接口对外提供标准的 workspace/executor 编程模型与完善的参数校验对内则通过 threshold 取值的分支判断灵活复用ThresholdGradV2D与ReluGrad两个经过多核 tiling 优化的核算子兼顾了代码复用与计算效率。开发者在使用时只需牢记三点第一段接口必须先于第二段调用、out的 shape 必须等于self与gradOutput广播后的 shape、gradOutput与self的数据类型必须一致其余的内存管理与异步执行细节均可参照本文示例完整落地。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐CANN ops-nn aclnnThresholdBackward 算子接口详解Threshold 反向传播的两段式调用与源码实现CANN ops nn aclnnThresholdBackward 算子接口详解Threshold 反向传播的两段式调用与源码实现 aclnnThresho人工智能算子库深度学习CANNAscendaclnnLogSigmoidBackwardCANN ops-nn 中 LogSigmoid 反向传播算子的两段式接口与 NPU 实现解析aclnnLogSigmoidBackwardCANN ops nn 中 LogSigmoid 反向传播算子的两段式接口与 NPU 实现解析 本篇技术指南面向人工智能算子库深度学习CANNAscendCANN ops-nn 中 aclnnSoftplusBackward 算子解析Softplus 反向传播的两段式接口与 NPU 实现CANN ops nn 中 aclnnSoftplusBackward 算子解析Softplus 反向传播的两段式接口与 NPU 实现 本文围绕 CANN o人工智能算子库深度学习CANNAscend上一篇大数据管道监控终极指南吞吐量、延迟与错误率三大核心指标详解下一篇Tuigreet 项目常见问题解决方案创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进