CANN ops-nn 算子解析：aclnnHardshrinkBackward 两段式接口的梯度传播实现与调用指南

CANN ops-nn 算子解析aclnnHardshrinkBackward 两段式接口的梯度传播实现与调用指南【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn导读aclnnHardshrinkBackward是 CANN 神经网络算子库ops-nn中 HardShrink 激活函数的反向算子接口用于计算 HardShrink 前向传播过程中输入x的梯度gradInput。本文基于 aclnnHardshrinkBackward 接口文档结合仓库中 op_api 调用实现、算子定义、内核实现 与 单元测试系统讲解该接口的产品支持范围、数学原理、两段式调用范式、参数约束、返回码语义并给出可直接编译运行的完整调用示例帮助读者在昇腾 NPU 上正确接入 HardShrink 反向梯度计算。产品支持情况aclnnHardshrinkBackward在不同昇腾硬件平台上的支持情况如下表所示产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品支持注意在 Atlas 训练系列产品上数据类型仅支持 FLOAT16、FLOAT即不支持 BFLOAT16。这一点与源码中 aclnn_hardshrink_backward.cpp 的分架构数据类型支持列表实现相互印证Ascend 910Atlas 训练系列产品对应的支持列表为{DT_FLOAT, DT_FLOAT16}而 A2/A3/950 等架构DAV_2201及 Regbase 架构支持{DT_BF16, DT_FLOAT, DT_FLOAT16}三种类型。功能说明与数学原理与前向算子的关系aclnnHardshrinkBackward是 aclnnHardshrinkHardShrink 前向算子的反向接口接收反向传播过程中上一环节传来的梯度grad结合前向输入x与阈值λ计算出当前算子对x的梯度gradInput。整个算子的梯度语义在 hard_shrink_grad_proto.h 中有明确注释与 PyTorch 的hardshrink_backward算子兼容。计算公式HardShrink 前向将绝对值不超过阈值λ的元素收缩为 0即Hardshrink(x) x if |x| λ else 0因此其反向传播的梯度遵循如下分段规则$$ HardshrinkBackward(x,grad) \begin{cases} grad, if(x \lambda) \ grad, if(x -\lambda) \ 0, otherwise \end{cases} $$即当输入x的绝对值严格大于阈值λ时梯度原样透传为grad当x落在[-λ, λ]区间内时前向输出为 0梯度被置为 0。这也是“硬收缩”hard shrink名称的由来——只有越过阈值区间的元素才参与梯度回传。从源码结构看该规则最终在 l0 算子层 hardshrink_grad.cpp 中以l0op::HardShrinkGrad(gradOutput, self, lambd, executor)的形式下发给 AI Core 执行算子类型注册为HardShrinkGrad通过ADD_TO_LAUNCHER_LIST_AICORE宏将输入、输出与阈值属性lambd加入任务队列。函数原型两段式接口与 CANN 单算子 API 的通用约定一致参见 两段式接口说明aclnnHardshrinkBackward采用“先查询、后执行”的两段式调用范式必须先调用aclnnHardshrinkBackwardGetWorkspaceSize获取计算所需的 workspace 大小以及封装了算子计算流程的执行器executor再调用aclnnHardshrinkBackward真正下发计算。aclnnStatus aclnnHardshrinkBackwardGetWorkspaceSize( const aclTensor* gradOutput, const aclTensor* self, const aclScalar* lambd, aclTensor* gradInput, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnHardshrinkBackward( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)workspace 是指除输入/输出之外算子在 NPU 上完成计算所需的临时内存workspaceSize即该临时内存的大小由第一段接口计算得出。第二段接口aclnnHardshrinkBackward(...)不可重复调用一次GetWorkspaceSize对应一次执行。aclnnHardshrinkBackwardGetWorkspaceSize 参数详解参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorgradOutputaclTensor*输入反向传播过程中上一步输出的梯度作为本反向算子的输入即公式中的 gradgradOutput 与 self 的 shape 满足 broadcast 关系gradOutput 与 self、gradInput 的数据类型一致BFLOAT16、FLOAT16、FLOATND0-8√selfaclTensor*输入表示输入的 Tensor即公式中的 xgradOutput 与 self 的 shape 满足 broadcast 关系gradOutput 与 self、gradInput 的数据类型一致BFLOAT16、FLOAT16、FLOATND0-8√lambdaclScalar*输入公式中的 λ当 λ 小于 0 时按 0 计算FLOAT---gradInputaclTensor*输出公式中的 HardshrinkBackward(x,grad)gradInput 与 self 和 gradOutput broadcast 后的 Tensor 的 shape 一致支持空 Tensor 传入gradOutput 与 self、gradInput 的数据类型一致BFLOAT16、FLOAT16、FLOATND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----补充说明Atlas 训练系列产品数据类型仅支持 FLOAT16、FLOAT。关键约束的源码级解读shape 维度上限gradOutput、self、gradInput的维度均为 0-8 维。对应源码中的 CheckShape 通过OP_CHECK_MAX_DIM(..., 8, ...)逐一校验三个 Tensor 的维度不超过 8。broadcast 关系gradOutput与self的 shape 必须满足 broadcast 条件且输出gradInput的 shape 必须与二者 broadcast 后的结果一致。对应源码中的 CheckShapeBroadcast先用OP_CHECK_BROADCAST_AND_INFER_SHAPE推导 broadcast 结果 shape再用OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE校验gradInput与之相等。在 BroadcastTo 中当gradOutput与self的 viewShape 不一致时两者会被统一 broadcast 到同一 shape 后再参与计算。空 TensorgradInput支持空 Tensor 传入当gradOutput或self为空 Tensor 时第一段接口直接返回workspaceSize 0并成功结束见 aclnn_hardshrink_backward.cpp。类型推导与转换第一段接口内部会调用 CheckPromoteType对gradOutput与self做数据类型提升promote推导推导结果必须落在算子支持的数据类型列表内且必须能够转换为gradInput的数据类型。实际计算链路中输入会先被l0op::Cast提升到统一计算类型计算完成后再 cast 回gradInput的数据类型参见 aclnn_hardshrink_backward.cpp。非连续 Tensor三个 Tensor 均支持非连续存储。源码在处理时先通过l0op::Contiguous将输入转为连续 Tensoraclnn_hardshrink_backward.cpp计算完成后若输出为非连续 Tensor再通过l0op::ViewCopy将连续结果写回aclnn_hardshrink_backward.cpp。负阈值处理lambd以aclScalarFLOAT 类型传入当 λ 小于 0 时按 0 计算。该规则在源码中有直接体现auto lambdFloat lambd-ToFloat(); if (lambdFloat 0) { lambdFloat 0; }aclnn_hardshrink_backward.cpp。需要说明的是源码中算子属性lambd的默认值为 0.5见 hard_shrink_grad_def.cpp 与 hard_shrink_grad_proto.h但通过 aclnn 接口显式传入aclScalar时调用方应自行给出确定的阈值。返回值aclnnStatus返回状态码具体定义参见 aclnn 返回码说明。第一段接口会完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 gradOutput、self、lambd 或 gradInput 是空指针ACLNN_ERR_PARAM_INVALID161002self、lambd 和 gradInput 的数据类型和数据格式不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self 和 gradOutput 的 shape 不满足 broadcast 条件ACLNN_ERR_PARAM_INVALID161002gradInput 的 shape 与 self 和 gradOutput 的 broadcast 结果不一致这些错误路径在 CheckParams 中按序执行先校验空指针返回ACLNN_ERR_PARAM_NULLPTR再校验数据类型支持范围与 shape 约束返回ACLNN_ERR_PARAM_INVALID与文档中的错误码表完全对应。仓库单元测试 test_aclnn_hardshrink_backward.cpp 中的case_nullptr用例即验证了任一参数为空指针时返回ACLNN_ERR_PARAM_NULLPTRcase_dtype_double用例验证了传入不支持的ACL_DOUBLE类型时返回ACLNN_ERR_PARAM_INVALID。aclnnHardshrinkBackward 参数详解参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnHardshrinkBackwardGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值aclnnStatus返回状态码具体参见 aclnn 返回码说明。第二段接口内部通过CommonOpExecutorRun(workspace, workspaceSize, executor, stream)统一完成执行器调度见 aclnn_hardshrink_backward.cpp。约束说明确定性计算aclnnHardshrinkBackward默认采用确定性实现即相同输入与相同硬件环境下多次运行的计算结果保持一致。调用示例以下示例代码展示了aclnnHardshrinkBackward的完整调用流程与仓库中的 examples/test_aclnn_hard_shrink_backward.cpp 保持一致。编译与执行的具体过程可参考 编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_hardshrink_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 gradOutputShape {2, 2}; std::vectorint64_t selfShape {2, 2}; std::vectorint64_t gradInputShape {2, 2}; void* gradOutputDeviceAddr nullptr; void* selfDeviceAddr nullptr; void* gradInputDeviceAddr nullptr; aclTensor* gradOutput nullptr; aclTensor* self nullptr; aclScalar* lambda nullptr; aclTensor* gradInput nullptr; std::vectorfloat gradOutputHostData {0, 0, 1, 1}; std::vectorfloat selfHostData {1, 2, 3, 4}; std::vectorfloat gradInputHostData(4, 0); float lambdaValue 2.0f; // 创建gradOutput aclTensor ret CreateAclTensor(gradOutputHostData, gradOutputShape, gradOutputDeviceAddr, aclDataType::ACL_FLOAT, gradOutput); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建gradInput aclTensor ret CreateAclTensor(gradInputHostData, gradInputShape, gradInputDeviceAddr, aclDataType::ACL_FLOAT, gradInput); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建lambda aclScalar lambda aclCreateScalar(lambdaValue, aclDataType::ACL_FLOAT); CHECK_RET(lambda ! nullptr, return ret); // 3. 调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnHardshrinkBackward第一段接口 ret aclnnHardshrinkBackwardGetWorkspaceSize(gradOutput, self, lambda, gradInput, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnHardshrinkBackwardGetWorkspaceSize 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); } // 调用aclnnHardshrinkBackward第二段接口 ret aclnnHardshrinkBackward(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnHardshrinkBackward 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侧需要根据具体API的接口定义修改 auto size GetShapeSize(gradInputShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), gradInputDeviceAddr, 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需要根据具体API的接口定义修改 aclDestroyTensor(gradOutput); aclDestroyTensor(self); aclDestroyTensor(gradInput); aclDestroyScalar(lambda); // 7. 释放device资源需要根据具体API的接口定义参数 aclrtFree(gradOutputDeviceAddr); aclrtFree(selfDeviceAddr); aclrtFree(gradInputDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例预期结果推演示例中取gradOutput {0, 0, 1, 1}、self {1, 2, 3, 4}、λ 2.0。根据公式逐元素计算x1与x2落在[-2, 2]区间内梯度置 0x3与x4大于 λ梯度原样透传。因此预期gradInput {0, 0, 1, 1}程序运行后应打印result[0] is: 0.000000、result[1] is: 0.000000、result[2] is: 1.000000、result[3] is: 1.000000。底层调用链与实现结构从仓库源码结构看aclnnHardshrinkBackward的完整实现横跨算子库的多个层次理解这条调用链有助于排查问题和二次开发aclnn 接口层Level 2op_api/aclnn_hardshrink_backward.cpp 负责入参校验空指针、数据类型、维度上限、broadcast、类型推导随后依次完成输入连续化Contiguous、broadcast 对齐、类型提升Cast、调用 l0 算子计算、结果类型回转与ViewCopy回写最后通过executor-GetWorkspaceSize()汇总 workspace 需求。l0 算子层op_api/hardshrink_grad.cpp 中l0op::HardShrinkGrad通过BroadcastInferShape推导结果 shape并以ADD_TO_LAUNCHER_LIST_AICORE(HardShrinkGrad, OP_INPUT(...), OP_OUTPUT(...), OP_ATTR(lambd))将计算任务加入执行器队列。算子定义与图层op_host/hard_shrink_grad_def.cpp 注册HardShrinkGrad的输入gradients、features、输出backprops与属性lambd默认 0.5fop_graph/hard_shrink_grad_proto.h 给出图协议定义op_host/hard_shrink_grad_infershape.cpp 实现输出 shape 与输入features一致op_graph/hard_shrink_grad_graph_infer.cpp 实现输出数据类型与输入一致。AI Core 内核层op_kernel/hard_shrink_grad_apt.cpp 是 arch35 架构的内核入口通过NsHardShrinkGrad::HardShrinkGradD_T_X, BUFFER_MODE模板实例化D_T_X为数据类型、BUFFER_MODE为单/双 buffer 模式Init后调用Process()执行逐元素的分段梯度计算对应的 tiling 计算位于 op_host/arch35/hard_shrink_grad_tiling.cpp。此外仓库在 examples/arch35/test_aclnn_hard_shrink_backward.cpp 提供了 arch35 平台专属示例在 tests/ut/op_host/op_api/test_aclnn_hardshrink_backward.cpp 提供了覆盖正常用例case_normal4D 输入 shape{1,1,4,4}、FLOAT16 用例、空指针用例case_nullptr、非法数据类型用例case_dtype_double与类型提升用例case_promote_dtype的单元测试可作为接入与验证的参考基线。总结aclnnHardshrinkBackward是 CANN ops-nn 中 HardShrink 激活反向传播的标准入口其核心要点可归纳为数学语义|x| λ时梯度透传|x| ≤ λ时梯度为 0与 PyTorchhardshrink_backward语义兼容调用范式严格遵守“GetWorkspaceSize查询 → 申请 workspace →aclnnHardshrinkBackward执行”的两段式流程第二段接口不可重复调用数据约束支持 BFLOAT16/FLOAT16/FLOATAtlas 训练系列产品仅 FLOAT16/FLOAT、ND 格式、0-8 维、非连续 Tensor 与空 Tensor输入间需满足 broadcast 关系λ 为负时按 0 处理错误处理空指针返回ACLNN_ERR_PARAM_NULLPTR161001数据类型、shape 不合法返回ACLNN_ERR_PARAM_INVALID161002。在接入前向 HardShrink 的训练或推理链路时可将本文的调用示例与仓库中的单元测试、arch35 示例结合使用快速完成 NPU 上的梯度计算验证。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考