人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载HardShrink硬收缩是一种逐元素稀疏化激活函数将输入张量中绝对值小于等于阈值 λ 的元素置零、其余元素原样保留常用于特征稀疏化与去噪场景。本文以 CANN 开源算子库 ops-nn 中 experimental/activation/hard_shrink 模块为准完整讲解aclnnHardShrink两段式接口的参数语义、错误码、约束与调用示例并结合算子定义、形状推导、Tiling 与 Kernel 源码剖析其 NPU 侧的实现原理使读者能够直接照抄示例运行算子并理解其底层工作机制。产品支持情况HardShrink 算子当前的适配范围如下支持的产品在表中以 √ 标识产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品xAtlas A2 训练系列产品/Atlas A2 推理系列产品xAtlas 200I/500 A2 推理产品xAtlas 推理系列产品xAtlas 训练系列产品x从仓库源码看算子定义 hard_shrink_def.cpp 中仅注册了ascend950的 AICore 配置this-AICore().AddConfig(ascend950, aicoreConfig)与上述产品支持矩阵一致对应的调用示例也全部位于arch35Ascend 950 架构目录下。因此若要在其他芯片型号上运行需先确认对应版本的算子包是否提供实现。功能说明与计算公式接口功能完成 HardShrink 激活函数计算将输入张量中绝对值小于等于阈值 lambd 的元素置零大于阈值的元素保持不变。该算子在功能上对标 PyTorch 的torch.nn.functional.hardshrink。计算公式$$ \text{HardShrink}(x) \begin{cases} x, \text{if } x \lambda \ x, \text{if } x -\lambda \ 0, \text{otherwise} \end{cases} $$其中x 为输入张量 self 中的元素λ 为阈值参数 lambd默认值为 0.5。从公式可以看出HardShrink 的关键语义是只有当元素严格大于 λ 或严格小于 -λ即 |x| λ时才保留原值恰等于 ±λ 的元素例如 x 0.5 且 lambd 0.5同样被置零这体现了“硬”收缩hard thresholding的截断特性。该行为在 op_kernel/hard_shrink.h 的两步比较实现中得到了精确对应第一步用Compare(x lambd)选出大于 λ 的元素第二步用Compare(x -lambd)选出小于 -λ 的元素两者均不命中含等于 ±λ的元素在 Select 中被置 0。函数原型与两段式调用流程与 CANN 其他单算子 API 一致aclnnHardShrink采用两段式接口必须先调用第一段接口获取 workspace 大小和包含算子计算流程的执行器再调用第二段接口真正执行计算。两段式接口的通用约定可参见仓库文档 两段式接口说明。aclnnStatus aclnnHardShrinkGetWorkspaceSize( const aclTensor *self, const aclScalar *lambd, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnHardShrink( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)说明workspace 指除输入/输出外算子在 NPU 上完成计算所需的临时内存workspaceSize 表示其大小由第一段接口计算得出。第二段接口aclnnHardShrink(...)不能重复调用同一 executor 只支持执行一次重复调用会异常。aclnnHardShrinkGetWorkspaceSize 参数说明第一段接口完成入参校验并构建执行器各参数语义如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入输入张量对应公式中的 x。支持空 Tensor。FLOAT、FLOAT16、BFLOAT16ND0-8√lambdaclScalar*输入阈值参数对应公式中的 λ默认值为 0.5。不支持空指针。FLOAT---outaclTensor*输出输出张量与 self 同 shape 同 dtype。不支持空 Tensor数据类型需与 self 一致shape 需与 self 一致。FLOAT、FLOAT16、BFLOAT16ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----几点实战提示空 Tensor 支持当 self 为空 Tensor0 元素时第一段接口仍会正常返回此时workspaceSize为 0out同样为空 Tensor不执行计算。此逻辑在 hard_shrink_tiling.cpp 的HandleEmptyTensor中实现将 blockDim 置为 1 并清空 TilingData。lambd 取值lambd 为 float 类型标量理论上取值无限制。但 Tiling 侧 GetLambdAttr 会对 NaN/Inf 打出告警日志此时输出可能全零同时参考主仓 op_api 实现 aclnn_hardshrink.cpp当 lambd 为负数时会被截断为 0.0f 再传入 L0 算子。实际使用时建议传入常规正数阈值。非连续 Tensorself 与 out 均支持非连续 Tensor表中 √。在 op_api 实现中输入先经Contiguous转为连续计算完成后通过ViewCopy将结果写回可能非连续的 out保证语义正确。aclnnHardShrink 参数说明第二段接口在申请好 workspace 后执行实际计算参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnHardShrinkGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。返回值与错误码两段接口均返回aclnnStatus状态码完整的状态码定义可参考仓库文档 aclnn 返回码。其中与 aclnnHardShrink 直接相关的参数校验错误如下返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001self、lambd、out 存在空指针。ACLNN_ERR_PARAM_INVALID161002self 的数据类型不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002out 的数据类型与 self 不一致。ACLNN_ERR_PARAM_INVALID161002out 的 shape 与 self 不一致。从 aclnn_hardshrink.cpp 的CheckParams校验链可以看到第一段接口的完整校验流程先做空指针检查CheckNotNull再做数据类型范围检查CheckDtypeValid再核对 shape 一致性CheckShape最后检查私有格式CheckFormat。错误码 161001/161002 分别对应 NULLPTR 与 INVALID 两类参数问题其余内部异常如算子二进制包未安装导致的 561xxx 系列可结合aclGetRecentErrMsg接口获取详细报错信息排查。约束说明aclnnHardShrink 为默认确定性实现即相同输入在多次运行中结果确定。self 与 out 的数据类型必须一致支持 FLOAT、FLOAT16、BFLOAT16。self 与 out 的 shape 必须一致不涉及广播。self 支持 0-8 维支持空 Tensor0 元素此时 out 也为空 Tensor不执行计算。上述约束在源码中有三重印证算子定义 hard_shrink_def.cpp 中self/out均声明为REQUIRED且数据类型限定为DT_FLOAT16/DT_FLOAT/DT_BF16形状推导 hard_shrink_infershape.cpp 中输出 shape 直接等于输入 shape*outputShape *inputShapeTiling 侧 GetInputInfo 同样对数据类型做了白名单校验。调用示例下面给出完整的可运行示例源自接口文档与仓库 examples/arch35/test_aclnn_hard_shrink.cpp 示例保持一致。示例输入特意构造了正值、负值、0、恰好等于 ±lambd 的值以及紧邻阈值的 0.49/0.51 等边界数据便于验证 HardShrink 的截断语义。具体编译与执行过程请参考仓库文档 编译与运行样例。#include iostream #include vector #include cstring #include acl/acl.h #include aclnn_hard_shrink.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; } 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); 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); 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); 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]; } *tensor aclCreateTensor( shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } 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; } int main() { // 1. ACL 初始化 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. 构造输入和输出 // 输入 self: shape[4, 4], dtypeFLOAT, 包含正值、负值和接近阈值的值 std::vectorint64_t selfShape {4, 4}; std::vectorfloat selfHostData { 1.0f, -1.0f, 0.3f, -0.3f, 0.5f, -0.5f, 0.0f, 2.0f, -2.0f, 0.1f, -0.1f, 10.0f, 0.49f, -0.49f, 0.51f, -0.51f }; aclTensor* self nullptr; void* selfDeviceAddr nullptr; ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 输出 out: 与 self 同 shape 同 dtype aclTensor* out nullptr; void* outDeviceAddr nullptr; std::vectorfloat outHostData(16, 0.0f); ret CreateAclTensor(outHostData, selfShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // lambd 标量参数double 类型 double lambd 0.5; // 3. 调用 aclnnHardShrink 第一段接口 uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; ret aclnnHardShrinkGetWorkspaceSize(self, lambd, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnHardShrinkGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 4. 申请 workspace void* workspaceAddr nullptr; if (workspaceSize static_castuint64_t(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); } // 5. 调用 aclnnHardShrink 第二段接口 ret aclnnHardShrink(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnHardShrink failed. ERROR: %d\n, ret); return ret); // 6. 同步等待计算完成 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 7. 获取输出的值将device侧内存上的结果拷贝至host侧 auto size GetShapeSize(selfShape); 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]); } // 8. 释放资源 aclDestroyTensor(self); aclDestroyTensor(out); aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize static_castuint64_t(0)) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }以 lambd 0.5 运行上述示例可预期的输出为1.0、-1.0、2.0、-2.0、10.0、0.51、-0.51等满足 |x| 0.5 的元素原样保留0.3、-0.3、0.5、-0.5、0.0、0.1、-0.1、0.49、-0.49等元素被置为 0。其中 0.5/-0.5恰等于阈值与 0.49/-0.49略小于阈值被置零直观验证了“绝对值小于等于阈值置零”的边界语义。从源码看 NPU 侧实现原理Host 侧算子定义与形状推导算子定义hard_shrink_def.cpp声明输入selfREQUIRED支持 FP16/FP32/BF16ND 格式AutoContiguous、输出out约束同 self并将 lambd 声明为Attr(lambd).AttrType(OPTIONAL).Float(0.5f)——即 lambd 以 Attr 形式经 TilingData 传递到 Kernel而非 Kernel 的运行时参数。同时开启动态编译、动态 rank、动态 shape 支持注册到ascend950架构。形状推导hard_shrink_infershape.cppHardShrink 是逐元素算子输出 shape 与 dtype 直接继承输入*outputShape *inputShape无需广播。Host 侧Tiling 切分策略hard_shrink_tiling.cpp 实现了运行时切分决策核心策略可归纳为多核切分blockFactor CeilDiv(totalNum, coreNum)按 AIV 核数GetCoreNumAiv把总元素均分到各核实际使用核数由CeilDiv(totalNum, blockFactor)得出。UB 切分根据 UB 内存大小GetCoreMemSize(UB)与缓冲模式计算ubFactor每轮循环处理的元素数并做 256B 对齐alignElems 256 / computeTypeSize。双缓冲阈值当totalNum 1024时BUFFER_MODE1双缓冲BUFFER_NUM2否则单缓冲用于隐藏搬数与计算延迟。bf16 特殊处理bf16 输入统一以 float 计算TilingKey 中IS_BF161Buffer 预算额外包含两个 float 临时缓冲区lambdBuf/negLambdBuf/tmpBuf/tmp2Buf。空 Tensor 处理HandleEmptyTensor将 blockDim 置 1、TilingData 清零、workspace 置 0并仍通过ASCENDC_TPL_SEL_PARAM完成模板参数选择保证编译期路径完备。最终 Tiling 结果totalNum/blockFactor/ubFactor/lambd写入 hard_shrink_tiling_data.h 定义的HardShrinkTilingData结构体随 Tiling 一并下发给 Kernel。Kernel 侧两次 Compare Select 的向量化实现Kernel 实现位于 op_kernel/hard_shrink.h入口 hard_shrink.cpp 通过REGISTER_TILING_DEFAULT读取 TilingData 后实例化模板类NsHardShrink::HardShrinkD_T, BUFFER_MODE, IS_BF16并执行。值得关注的是实现刻意采用“两次 Select”而非一次 Or 组合以避免 Or API 的兼容性风险计算流水为预填常量Duplicate将lambd与-lambd填充进两个 UB 常驻缓冲区lambdBuf/negLambdBuf。Compare(x lambd)生成 mask1Select(mask1 ? x : 0)得到中间结果 tmp保留大于 λ 的元素。Compare(x -lambd)生成 mask2Select(mask2 ? x : tmp)得到最终输出保留小于 -λ 的元素其余沿用 tmp 中的 0 或已保留值。对于 BF16 路径数据需先Cast bf16 → float参与比较与选择计算完成后再Cast float → bf16输出舍入模式CAST_RINTFP16/FP32 路径则直接以原生类型计算。Compare/Select 对 count 有 256B 对齐要求因此 Kernel 内按alignElems 256 / sizeof(COMPUTE_T)将currentNum向上对齐后调用向量指令。模板参数组合FP16/FP32/BF16 × 单/双缓冲由 hard_shrink_tiling_key.h 中的ASCENDC_TPL_ARGS_DECL/ASCENDC_TPL_SEL静态展开共 6 个编译实例。示例与测试配套仓库为该算子提供了多种验证入口可用于快速回归与精度确认test_aclnn_hard_shrink.cppaaclnn 两段式调用示例采用std::unique_ptr 自定义 deleterRAII管理 ACL 资源任何路径 return 都能保证 tensor、device 内存、stream、device、acl 正确释放是学习资源管理的良好范本。test_aclnn_hard_shrink_fp16.cpp / test_aclnn_hard_shrink_bf16.cpp分别验证 FP16、BF16 数据类型。test_aclnn_hard_shrink_large.cpp大 Tensor 场景覆盖多核切分与 UB 分块路径。总结aclnnHardShrink是 CANN ops-nn 中实现 HardShrink 硬阈值激活的标准入口遵循“GetWorkspaceSize 执行”的两段式调用范式支持 FLOAT/FLOAT16/BFLOAT16 与 0-8 维 ND 张量含空 Tensor 与非连续 Tensor默认确定性实现。从源码链路看其 NPU 实现依次经历算子定义Attr 传递 lambd、形状推导输出继承输入、Tiling多核 UB 双缓冲切分与 Kernel两次 CompareSelect 的向量化逐元素运算四层BF16 路径额外引入 float 中间计算以保证精度。开发者可直接复用本文示例完成单算子调用并结合 接口文档 与 模块 README 深入定制。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐CANN ops-nn 算子实战aclnnHardshrink 两段式接口的完整开发与调用指南CANN ops nn 算子实战aclnnHardshrink 两段式接口的完整开发与调用指南 导读 本文围绕 CANN 神经网络算子库 ops nn 中的人工智能算子库深度学习CANNAscendCANN ops-nn 算子开发指南aclnnSwiGlu 两段式接口详解与 SwiGlu 激活算子 NPU 实现剖析CANN ops nn 算子开发指南aclnnSwiGlu 两段式接口详解与 SwiGlu 激活算子 NPU 实现剖析 SwiGluSwish Gated人工智能算子库深度学习CANNAscendCANN ops-nn 算子开发指南aclnnCelu 与 aclnnInplaceCelu 两段式接口详解与 NPU 实战调用CANN ops nn 算子开发指南aclnnCelu 与 aclnnInplaceCelu 两段式接口详解与 NPU 实战调用 本篇技术指南以 activa人工智能算子库深度学习CANNAscend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-nn 算子开发指南：aclnnHardShrink 两段式接口详解与 NPU 实现原理
人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载HardShrink硬收缩是一种逐元素稀疏化激活函数将输入张量中绝对值小于等于阈值 λ 的元素置零、其余元素原样保留常用于特征稀疏化与去噪场景。本文以 CANN 开源算子库 ops-nn 中 experimental/activation/hard_shrink 模块为准完整讲解aclnnHardShrink两段式接口的参数语义、错误码、约束与调用示例并结合算子定义、形状推导、Tiling 与 Kernel 源码剖析其 NPU 侧的实现原理使读者能够直接照抄示例运行算子并理解其底层工作机制。产品支持情况HardShrink 算子当前的适配范围如下支持的产品在表中以 √ 标识产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品xAtlas A2 训练系列产品/Atlas A2 推理系列产品xAtlas 200I/500 A2 推理产品xAtlas 推理系列产品xAtlas 训练系列产品x从仓库源码看算子定义 hard_shrink_def.cpp 中仅注册了ascend950的 AICore 配置this-AICore().AddConfig(ascend950, aicoreConfig)与上述产品支持矩阵一致对应的调用示例也全部位于arch35Ascend 950 架构目录下。因此若要在其他芯片型号上运行需先确认对应版本的算子包是否提供实现。功能说明与计算公式接口功能完成 HardShrink 激活函数计算将输入张量中绝对值小于等于阈值 lambd 的元素置零大于阈值的元素保持不变。该算子在功能上对标 PyTorch 的torch.nn.functional.hardshrink。计算公式$$ \text{HardShrink}(x) \begin{cases} x, \text{if } x \lambda \ x, \text{if } x -\lambda \ 0, \text{otherwise} \end{cases} $$其中x 为输入张量 self 中的元素λ 为阈值参数 lambd默认值为 0.5。从公式可以看出HardShrink 的关键语义是只有当元素严格大于 λ 或严格小于 -λ即 |x| λ时才保留原值恰等于 ±λ 的元素例如 x 0.5 且 lambd 0.5同样被置零这体现了“硬”收缩hard thresholding的截断特性。该行为在 op_kernel/hard_shrink.h 的两步比较实现中得到了精确对应第一步用Compare(x lambd)选出大于 λ 的元素第二步用Compare(x -lambd)选出小于 -λ 的元素两者均不命中含等于 ±λ的元素在 Select 中被置 0。函数原型与两段式调用流程与 CANN 其他单算子 API 一致aclnnHardShrink采用两段式接口必须先调用第一段接口获取 workspace 大小和包含算子计算流程的执行器再调用第二段接口真正执行计算。两段式接口的通用约定可参见仓库文档 两段式接口说明。aclnnStatus aclnnHardShrinkGetWorkspaceSize( const aclTensor *self, const aclScalar *lambd, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnHardShrink( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)说明workspace 指除输入/输出外算子在 NPU 上完成计算所需的临时内存workspaceSize 表示其大小由第一段接口计算得出。第二段接口aclnnHardShrink(...)不能重复调用同一 executor 只支持执行一次重复调用会异常。aclnnHardShrinkGetWorkspaceSize 参数说明第一段接口完成入参校验并构建执行器各参数语义如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入输入张量对应公式中的 x。支持空 Tensor。FLOAT、FLOAT16、BFLOAT16ND0-8√lambdaclScalar*输入阈值参数对应公式中的 λ默认值为 0.5。不支持空指针。FLOAT---outaclTensor*输出输出张量与 self 同 shape 同 dtype。不支持空 Tensor数据类型需与 self 一致shape 需与 self 一致。FLOAT、FLOAT16、BFLOAT16ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----几点实战提示空 Tensor 支持当 self 为空 Tensor0 元素时第一段接口仍会正常返回此时workspaceSize为 0out同样为空 Tensor不执行计算。此逻辑在 hard_shrink_tiling.cpp 的HandleEmptyTensor中实现将 blockDim 置为 1 并清空 TilingData。lambd 取值lambd 为 float 类型标量理论上取值无限制。但 Tiling 侧 GetLambdAttr 会对 NaN/Inf 打出告警日志此时输出可能全零同时参考主仓 op_api 实现 aclnn_hardshrink.cpp当 lambd 为负数时会被截断为 0.0f 再传入 L0 算子。实际使用时建议传入常规正数阈值。非连续 Tensorself 与 out 均支持非连续 Tensor表中 √。在 op_api 实现中输入先经Contiguous转为连续计算完成后通过ViewCopy将结果写回可能非连续的 out保证语义正确。aclnnHardShrink 参数说明第二段接口在申请好 workspace 后执行实际计算参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnHardShrinkGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。返回值与错误码两段接口均返回aclnnStatus状态码完整的状态码定义可参考仓库文档 aclnn 返回码。其中与 aclnnHardShrink 直接相关的参数校验错误如下返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001self、lambd、out 存在空指针。ACLNN_ERR_PARAM_INVALID161002self 的数据类型不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002out 的数据类型与 self 不一致。ACLNN_ERR_PARAM_INVALID161002out 的 shape 与 self 不一致。从 aclnn_hardshrink.cpp 的CheckParams校验链可以看到第一段接口的完整校验流程先做空指针检查CheckNotNull再做数据类型范围检查CheckDtypeValid再核对 shape 一致性CheckShape最后检查私有格式CheckFormat。错误码 161001/161002 分别对应 NULLPTR 与 INVALID 两类参数问题其余内部异常如算子二进制包未安装导致的 561xxx 系列可结合aclGetRecentErrMsg接口获取详细报错信息排查。约束说明aclnnHardShrink 为默认确定性实现即相同输入在多次运行中结果确定。self 与 out 的数据类型必须一致支持 FLOAT、FLOAT16、BFLOAT16。self 与 out 的 shape 必须一致不涉及广播。self 支持 0-8 维支持空 Tensor0 元素此时 out 也为空 Tensor不执行计算。上述约束在源码中有三重印证算子定义 hard_shrink_def.cpp 中self/out均声明为REQUIRED且数据类型限定为DT_FLOAT16/DT_FLOAT/DT_BF16形状推导 hard_shrink_infershape.cpp 中输出 shape 直接等于输入 shape*outputShape *inputShapeTiling 侧 GetInputInfo 同样对数据类型做了白名单校验。调用示例下面给出完整的可运行示例源自接口文档与仓库 examples/arch35/test_aclnn_hard_shrink.cpp 示例保持一致。示例输入特意构造了正值、负值、0、恰好等于 ±lambd 的值以及紧邻阈值的 0.49/0.51 等边界数据便于验证 HardShrink 的截断语义。具体编译与执行过程请参考仓库文档 编译与运行样例。#include iostream #include vector #include cstring #include acl/acl.h #include aclnn_hard_shrink.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; } 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); 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); 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); 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]; } *tensor aclCreateTensor( shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } 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; } int main() { // 1. ACL 初始化 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. 构造输入和输出 // 输入 self: shape[4, 4], dtypeFLOAT, 包含正值、负值和接近阈值的值 std::vectorint64_t selfShape {4, 4}; std::vectorfloat selfHostData { 1.0f, -1.0f, 0.3f, -0.3f, 0.5f, -0.5f, 0.0f, 2.0f, -2.0f, 0.1f, -0.1f, 10.0f, 0.49f, -0.49f, 0.51f, -0.51f }; aclTensor* self nullptr; void* selfDeviceAddr nullptr; ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 输出 out: 与 self 同 shape 同 dtype aclTensor* out nullptr; void* outDeviceAddr nullptr; std::vectorfloat outHostData(16, 0.0f); ret CreateAclTensor(outHostData, selfShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // lambd 标量参数double 类型 double lambd 0.5; // 3. 调用 aclnnHardShrink 第一段接口 uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; ret aclnnHardShrinkGetWorkspaceSize(self, lambd, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnHardShrinkGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 4. 申请 workspace void* workspaceAddr nullptr; if (workspaceSize static_castuint64_t(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); } // 5. 调用 aclnnHardShrink 第二段接口 ret aclnnHardShrink(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnHardShrink failed. ERROR: %d\n, ret); return ret); // 6. 同步等待计算完成 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 7. 获取输出的值将device侧内存上的结果拷贝至host侧 auto size GetShapeSize(selfShape); 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]); } // 8. 释放资源 aclDestroyTensor(self); aclDestroyTensor(out); aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize static_castuint64_t(0)) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }以 lambd 0.5 运行上述示例可预期的输出为1.0、-1.0、2.0、-2.0、10.0、0.51、-0.51等满足 |x| 0.5 的元素原样保留0.3、-0.3、0.5、-0.5、0.0、0.1、-0.1、0.49、-0.49等元素被置为 0。其中 0.5/-0.5恰等于阈值与 0.49/-0.49略小于阈值被置零直观验证了“绝对值小于等于阈值置零”的边界语义。从源码看 NPU 侧实现原理Host 侧算子定义与形状推导算子定义hard_shrink_def.cpp声明输入selfREQUIRED支持 FP16/FP32/BF16ND 格式AutoContiguous、输出out约束同 self并将 lambd 声明为Attr(lambd).AttrType(OPTIONAL).Float(0.5f)——即 lambd 以 Attr 形式经 TilingData 传递到 Kernel而非 Kernel 的运行时参数。同时开启动态编译、动态 rank、动态 shape 支持注册到ascend950架构。形状推导hard_shrink_infershape.cppHardShrink 是逐元素算子输出 shape 与 dtype 直接继承输入*outputShape *inputShape无需广播。Host 侧Tiling 切分策略hard_shrink_tiling.cpp 实现了运行时切分决策核心策略可归纳为多核切分blockFactor CeilDiv(totalNum, coreNum)按 AIV 核数GetCoreNumAiv把总元素均分到各核实际使用核数由CeilDiv(totalNum, blockFactor)得出。UB 切分根据 UB 内存大小GetCoreMemSize(UB)与缓冲模式计算ubFactor每轮循环处理的元素数并做 256B 对齐alignElems 256 / computeTypeSize。双缓冲阈值当totalNum 1024时BUFFER_MODE1双缓冲BUFFER_NUM2否则单缓冲用于隐藏搬数与计算延迟。bf16 特殊处理bf16 输入统一以 float 计算TilingKey 中IS_BF161Buffer 预算额外包含两个 float 临时缓冲区lambdBuf/negLambdBuf/tmpBuf/tmp2Buf。空 Tensor 处理HandleEmptyTensor将 blockDim 置 1、TilingData 清零、workspace 置 0并仍通过ASCENDC_TPL_SEL_PARAM完成模板参数选择保证编译期路径完备。最终 Tiling 结果totalNum/blockFactor/ubFactor/lambd写入 hard_shrink_tiling_data.h 定义的HardShrinkTilingData结构体随 Tiling 一并下发给 Kernel。Kernel 侧两次 Compare Select 的向量化实现Kernel 实现位于 op_kernel/hard_shrink.h入口 hard_shrink.cpp 通过REGISTER_TILING_DEFAULT读取 TilingData 后实例化模板类NsHardShrink::HardShrinkD_T, BUFFER_MODE, IS_BF16并执行。值得关注的是实现刻意采用“两次 Select”而非一次 Or 组合以避免 Or API 的兼容性风险计算流水为预填常量Duplicate将lambd与-lambd填充进两个 UB 常驻缓冲区lambdBuf/negLambdBuf。Compare(x lambd)生成 mask1Select(mask1 ? x : 0)得到中间结果 tmp保留大于 λ 的元素。Compare(x -lambd)生成 mask2Select(mask2 ? x : tmp)得到最终输出保留小于 -λ 的元素其余沿用 tmp 中的 0 或已保留值。对于 BF16 路径数据需先Cast bf16 → float参与比较与选择计算完成后再Cast float → bf16输出舍入模式CAST_RINTFP16/FP32 路径则直接以原生类型计算。Compare/Select 对 count 有 256B 对齐要求因此 Kernel 内按alignElems 256 / sizeof(COMPUTE_T)将currentNum向上对齐后调用向量指令。模板参数组合FP16/FP32/BF16 × 单/双缓冲由 hard_shrink_tiling_key.h 中的ASCENDC_TPL_ARGS_DECL/ASCENDC_TPL_SEL静态展开共 6 个编译实例。示例与测试配套仓库为该算子提供了多种验证入口可用于快速回归与精度确认test_aclnn_hard_shrink.cppaaclnn 两段式调用示例采用std::unique_ptr 自定义 deleterRAII管理 ACL 资源任何路径 return 都能保证 tensor、device 内存、stream、device、acl 正确释放是学习资源管理的良好范本。test_aclnn_hard_shrink_fp16.cpp / test_aclnn_hard_shrink_bf16.cpp分别验证 FP16、BF16 数据类型。test_aclnn_hard_shrink_large.cpp大 Tensor 场景覆盖多核切分与 UB 分块路径。总结aclnnHardShrink是 CANN ops-nn 中实现 HardShrink 硬阈值激活的标准入口遵循“GetWorkspaceSize 执行”的两段式调用范式支持 FLOAT/FLOAT16/BFLOAT16 与 0-8 维 ND 张量含空 Tensor 与非连续 Tensor默认确定性实现。从源码链路看其 NPU 实现依次经历算子定义Attr 传递 lambd、形状推导输出继承输入、Tiling多核 UB 双缓冲切分与 Kernel两次 CompareSelect 的向量化逐元素运算四层BF16 路径额外引入 float 中间计算以保证精度。开发者可直接复用本文示例完成单算子调用并结合 接口文档 与 模块 README 深入定制。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐CANN ops-nn 算子实战aclnnHardshrink 两段式接口的完整开发与调用指南CANN ops nn 算子实战aclnnHardshrink 两段式接口的完整开发与调用指南 导读 本文围绕 CANN 神经网络算子库 ops nn 中的人工智能算子库深度学习CANNAscendCANN ops-nn 算子开发指南aclnnSwiGlu 两段式接口详解与 SwiGlu 激活算子 NPU 实现剖析CANN ops nn 算子开发指南aclnnSwiGlu 两段式接口详解与 SwiGlu 激活算子 NPU 实现剖析 SwiGluSwish Gated人工智能算子库深度学习CANNAscendCANN ops-nn 算子开发指南aclnnCelu 与 aclnnInplaceCelu 两段式接口详解与 NPU 实战调用CANN ops nn 算子开发指南aclnnCelu 与 aclnnInplaceCelu 两段式接口详解与 NPU 实战调用 本篇技术指南以 activa人工智能算子库深度学习CANNAscend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考