人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载本文基于 CANN ops-nn 仓库中experimental/activation/hard_shrink模块的算子说明文档系统讲解 HardShrink 算子的功能定义、aclnn 两段式调用方法、参数与约束并结合算子定义、形状推导、Tiling 与 Kernel 源码剖析其在 Ascend 950 系列产品上的实现原理。读完本文你将能够独立完成 HardShrink 算子的 aclnn 接口调用、数据验证并理解其多核切分与双缓冲调度机制。一、算子功能与数学定义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 中的元素$\lambda$ 为阈值参数 lambd默认值为 0.5。从数值行为上看该函数在 $(-\lambda, \lambda)$ 开区间内输出 0在 $x \ge \lambda$ 或 $x \le -\lambda$ 处原样输出 x是典型的稀疏化/去噪激活手段常用于需要抑制小幅度噪声特征、保留显著响应的网络结构。二、产品支持情况根据 hard_shrink/README.md 与接口文档 aclnnHardShrink.mdHardShrink 算子的产品支持矩阵如下产品是否支持Ascend 950PR/Ascend 950DT√从算子定义代码 hard_shrink_def.cpp 可以看到算子通过this-AICore().AddConfig(ascend950, aicoreConfig)注册了 ascend950 平台的 AICore 配置与文档中的产品支持情况一致。三、目录结构HardShrink 算子模块在仓库中位于experimental/activation/hard_shrink/采用 CANN 算子标准的 Host/Kernel 双层代码组织hard_shrink/ ├── op_host/ # Host 侧代码 │ ├── CMakeLists.txt # Host 侧构建配置 │ ├── hard_shrink_def.cpp # 算子定义 │ ├── hard_shrink_infershape.cpp # 形状推导 │ └── hard_shrink_tiling.cpp # Tiling 实现 ├── op_kernel/ # Kernel 侧代码 │ ├── hard_shrink_apt.cpp # Kernel 入口对应 hard_shrink.cpp │ ├── hard_shrink.h # Kernel 类定义 │ ├── hard_shrink_tiling_data.h # TilingData 结构体 │ └── hard_shrink_tiling_key.h # TilingKey 定义 ├── docs/ # 接口文档 │ └── aclnnHardShrink.md # aclnnHardShrink 接口文档 ├── examples/ # 调用示例 │ └── arch35/ # Ascend 950 架构示例 │ ├── test_aclnn_hard_shrink.cpp # aclnn 两段式调用示例 │ ├── test_aclnn_hard_shrink_fp16.cpp # FP16 数据类型验证 │ ├── test_aclnn_hard_shrink_bf16.cpp # BF16 数据类型验证 │ └── test_aclnn_hard_shrink_large.cpp # 大 Tensor 多核切分验证 ├── tests/ # 测试代码 ├── CMakeLists.txt # 构建配置 └── README.md # 说明文档其中 Host 侧负责算子注册、形状推导与 Tiling 策略计算Kernel 侧负责在 AICore 上完成实际的向量计算。四、参数说明HardShrink 算子共包含两个输入参数和一个输出参数参数名输入/输出/属性描述数据类型数据格式self输入输入张量对应公式中的 x。支持 0-8 维支持空 Tensor。FLOAT、FLOAT16、BFLOAT16NDlambd输入阈值参数对应公式中的 λfloat 类型标量默认值 0.5。FLOAT-out输出输出张量与 self 同 shape 同 dtype。FLOAT、FLOAT16、BFLOAT16ND各参数在源码中的注册情况如下self 与 out在 hard_shrink_def.cpp 中通过this-Input(self)与this-Output(out)注册数据类型限定为ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16格式为ge::FORMAT_ND并设置了AutoContiguous()支持动态 shapeDynamicShapeSupportFlag(true)与动态 rankDynamicRankSupportFlag(true)。lambd作为算子属性Attr而非张量参数注册this-Attr(lambd).AttrType(OPTIONAL).Float(0.5f)默认值 0.5f由 Host 侧 Tiling 阶段读取后写入 TilingData 传递给 Kernel。五、约束说明使用 HardShrink 算子需遵守以下约束self 与 out 的数据类型必须一致支持 FLOAT、FLOAT16、BFLOAT16。self 与 out 的 shape 必须一致不涉及广播。self 支持 0-8 维。self 支持空 Tensor0 元素此时 out 也为空 Tensor不执行计算。lambd 为 float 类型标量取值无限制但 nan/inf 时输出可能全为 0见下文 Tiling 分析。aclnnHardShrink 为默认确定性实现。上述约束在接口层由第一段接口aclnnHardShrinkGetWorkspaceSize完成入参校验在算子定义层则由数据类型/格式注册与形状推导逻辑双重保证。六、aclnn 调用说明HardShrink 算子通过 CANN 标准的aclnn 两段式接口调用完整调用样例见 test_aclnn_hard_shrink.cpp接口细节见 aclnnHardShrink.md。6.1 两段式接口原型第一段获取 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 大小以及包含算子计算流程的执行器再调用第二段执行计算。6.2 第一段接口参数说明参数名输入/输出描述使用说明数据类型数据格式维度(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 执行器包含了算子计算流程。-----值得注意的是接口层支持非连续 Tensorstrides 非连续的内存布局这得益于算子定义中的AutoContiguous()配置Host 侧会自动完成连续化处理。6.3 第一段接口返回值与错误码第一段接口返回aclnnStatus状态码具体返回码说明参见 aclnn返回码。常见错误场景如下返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001self、lambd、out 存在空指针。ACLNN_ERR_PARAM_INVALID161002self 的数据类型不在支持的范围之内out 的数据类型与 self 不一致out 的 shape 与 self 不一致。6.4 第二段接口参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnHardShrinkGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。6.5 完整调用示例以下代码演示完整的 aclnn 两段式调用流程包括 ACL 初始化、Tensor 构造、两段接口调用、结果回拷与资源释放完整版见 test_aclnn_hard_shrink.cpp#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) 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, 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]; } *tensor aclCreateTensor( shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. ACL 初始化 int32_t deviceId 0; aclrtStream stream; CHECK_RET(aclInit(nullptr) ACL_SUCCESS, return -1); CHECK_RET(aclrtSetDevice(deviceId) ACL_SUCCESS, return -1); CHECK_RET(aclrtCreateStream(stream) ACL_SUCCESS, return -1); // 2. 构造输入和输出shape[4,4], FLOAT含正值/负值/阈值边界值 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; CHECK_RET(CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self) 0, return -1); aclTensor* out nullptr; void* outDeviceAddr nullptr; std::vectorfloat outHostData(16, 0.0f); CHECK_RET(CreateAclTensor(outHostData, selfShape, outDeviceAddr, aclDataType::ACL_FLOAT, out) 0, return -1); // lambd 阈值对应接口原型中的 aclScalar示例中直接以标量传入 double lambd 0.5; // 3. 调用第一段接口获取 workspace 大小与执行器 uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; CHECK_RET(aclnnHardShrinkGetWorkspaceSize(self, lambd, out, workspaceSize, executor) ACL_SUCCESS, return -1); // 4. 申请 workspace void* workspaceAddr nullptr; if (workspaceSize 0) { CHECK_RET(aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST) ACL_SUCCESS, return -1); } // 5. 调用第二段接口执行计算 CHECK_RET(aclnnHardShrink(workspaceAddr, workspaceSize, executor, stream) ACL_SUCCESS, return -1); // 6. 同步等待计算完成 CHECK_RET(aclrtSynchronizeStream(stream) ACL_SUCCESS, return -1); // 7. 将结果从 Device 拷贝回 Host 并打印 auto size GetShapeSize(selfShape); std::vectorfloat resultData(size, 0); CHECK_RET(aclrtMemcpy(resultData.data(), resultData.size() * sizeof(float), outDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST) ACL_SUCCESS, return -1); for (int64_t i 0; i size; i) { printf(result[%ld] is: %f\n, i, resultData[i]); } // 8. 释放资源 aclDestroyTensor(self); aclDestroyTensor(out); aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }预期结果分析以 lambd0.5 为例输入数据中0.3、-0.3、0.5、-0.5、0.0、0.1、-0.1、0.49、-0.49的绝对值均小于等于 0.5输出为 0而1.0、-1.0、2.0、-2.0、10.0、0.51、-0.51绝对值大于 0.5保持原值输出。其中0.51恰好略大于阈值体现了边界处大于而非大于等于的判定语义。仓库示例 test_aclnn_hard_shrink.cpp 还展示了更健壮的RAII 资源管理写法使用std::unique_ptr配合自定义 deleter 封装aclTensor、Device 内存、Stream并用带析构的 guard 管理aclrtResetDevice/aclFinalize保证函数任意路径 return 时资源都能正确释放是编写生产级调用代码的推荐范式。6.6 更多验证示例examples 目录下还提供了三个针对性验证示例均为 arch35 / Ascend 950 架构test_aclnn_hard_shrink_fp16.cpp验证 FLOAT16 数据类型下的算子行为test_aclnn_hard_shrink_bf16.cpp验证 BFLOAT16 数据类型下的算子行为对应 Kernel 中 bf16 需要 cast 到 float 计算的特殊路径test_aclnn_hard_shrink_large.cpp构造 shape[1024, 1024] 共 1M 元素的大 Tensor覆盖多核切分、UB 分片循环与双缓冲路径并与 CPU 参考值比对统计 PASS/FAIL。七、实现原理从 Host Tiling 到 Kernel 计算7.1 形状推导逐元素算子的天然一致性hard_shrink_infershape.cpp 中的InferShape4HardShrink逻辑非常简洁取出输入 self 的 shape 后直接赋给输出 out*outputShape *inputShape。HardShrink 是逐元素算子输出 shape 恒等于输入 shape输出 dtype 恒等于输入 dtype不涉及广播与维度变换。7.2 Tiling 策略多核切分 UB 分片 双缓冲hard_shrink_tiling.cpp 实现了完整的 Tiling 计算核心策略如下多核切分blockFactor CeilDiv(totalNum, coreNum)将总元素数按 AICore 核数GetCoreNumAiv()获取均分每个核处理blockFactor个元素实际使用的核数usedCoreNum CeilDiv(totalNum, blockFactor)保证空余核不被分配任务。UB 分片根据 UB 容量GetCoreMemSize(CoreMemType::UB)计算ubFactor单次 UB 循环处理元素数并做256B 对齐alignElems 256 / computeTypeSize即按计算类型大小换算成元素对齐粒度预留UB_RESERVED_BYTE 1024字节安全余量。双缓冲阈值bufferMode (totalNum 1024) ? 1 : 0即总元素数超过 1024 时启用双缓冲BUFFER_MODE1BUFFER_NUM2隐藏 CopyIn/Compute/CopyOut 之间的访存延迟否则使用单缓冲。bf16 特殊处理isBf16 (dataType DT_BF16)bf16 输入在 Kernel 中需要 cast 到 float 计算Tiling 阶段通过ASCENDC_TPL_SEL_PARAM下发dType / bufferMode / isBf16三个模板参数完成 Kernel 实例的选择。lambd 传递GetLambdAttr从 Attr 中读取 lambd默认 0.5f若为 nan/inf 则打印告警输出可能全为 0随后写入tiling-lambd随 TilingData 下发。TilingData 结构体定义在 hard_shrink_tiling_data.hstruct HardShrinkTilingData { int64_t totalNum 0; // 总元素数量 int64_t blockFactor 0; // 每个核处理的元素数量 int64_t ubFactor 0; // 每次 UB 循环处理的元素数量已 256B 对齐 float lambd 0.5f; // 阈值参数 λ };另外空 Tensor 场景由HandleEmptyTensor处理设置SetBlockDim(1)、TilingData 清零、workspace 大小为 0Kernel 侧在blockLength_ 0时直接返回不执行任何计算见 hard_shrink.h。7.3 Kernel 实现两次 Compare SelectKernel 侧的计算核心在 hard_shrink.h 的Compute方法中采用两次 Compare 两次 Select的方案实现刻意避免使用 Or 位运算 API规避兼容性风险Compare(mask, input, lambd, GT)生成input lambd的 bit maskSelect(tmp, mask, input, 0)mask 为真时保留输入否则置 0处理 x λ 分支Compare(mask, input, negLambd, LT)生成input -lambd的 bit maskSelect(out, mask, input, tmp)mask 为真时保留输入否则取 tmp处理 x -λ 分支。两次 Select 叠加后落在 $[-\lambda, \lambda]$ 区间内的元素必然输出 0区间外元素保留原值与公式完全等价。针对 Compare/Select 硬件指令对 count 所占空间 256B 对齐的要求Kernel 将currentNum向上对齐到256 / sizeof(COMPUTE_T)个元素并限制不超过ubLength_保证 buffer 安全。bf16 路径IS_BF161略有不同由于 bf16 精度不足计算前先Castbf16→float以 float 完成两次 Compare/Select 后再CastCAST_RINT舍入模式回 bf16 输出。为此 Kernel 类内部通过std::conditional分别定义了 IO 类型IO_T与计算类型COMPUTE_T// IO 类型: IS_BF16 ? bfloat16_t : T using IO_T typename std::conditionalIS_BF16 1, bfloat16_t, T::type; // 计算类型: IS_BF16 ? float : Tbf16 需要 cast 到 float 进行计算 using COMPUTE_T typename std::conditionalIS_BF16 1, float, T::type;Init 阶段还会用AscendC::Duplicate一次性将lambd与-lambd常量填充到 UB 的 VECCALC buffer 中避免在循环内重复生成常量。7.4 Kernel 模板参数TilingKeyhard_shrink_tiling_key.h 通过ASCENDC_TPL_ARGS_DECL声明了三个模板参数并通过ASCENDC_TPL_SEL枚举出 6 个 Kernel 实例组合模板参数含义取值D_T计算数据类型FLOAT16、FLOAT、BF16bf16 场景实际以 float 计算BUFFER_MODE缓冲模式0单缓冲1双缓冲IS_BF16是否为 bf16 输入0否1是触发 Cast bf16↔float组合覆盖为fp16 单/双缓冲、fp32 单/双缓冲、bf16 单/双缓冲共 6 种。Kernel 入口 hard_shrink.cpp文档中称 hard_shrink_apt.cpp通过宏GET_TILING_DATA_WITH_STRUCT解析 TilingData实例化NsHardShrink::HardShrinkD_T, BUFFER_MODE, IS_BF16并依次调用Init与Process。7.5 流水线主循环Process方法hard_shrink.h按blockLength_ / ubLength_计算循环次数逐片执行CopyIn → Compute → CopyOutCopyIn用DataCopyPad将 GM 数据按currentNum * sizeof(IO_T)块长拷入输入队列支持非对齐边界Compute执行上述两次 Compare/Select 计算结果写入输出队列CopyOut用DataCopyPad将输出队列结果拷回 GM。当启用双缓冲BUFFER_NUM2时输入/输出队列各含两个 buffer流水线可在前一片数据计算的同时预取后一片数据有效提升大 Tensor 场景的吞吐。八、总结HardShrink 是 CANN ops-nn 中一个实现简洁但工程细节完整的激活算子功能上对标 PyTorchtorch.nn.functional.hardshrink支持 FLOAT/FLOAT16/BFLOAT16 与 0-8 维、空 Tensor、非连续 Tensor接口上采用标准 aclnn 两段式调用实现上以两次 Compare 两次 Select规避位运算兼容风险通过多核均分、UB 分片、256B 对齐、双缓冲切换与 bf16 cast-to-float 等策略覆盖从 1 个元素到百万级大 Tensor 的各类场景。开发者可直接参考 examples/arch35 下的四个示例快速上手调用与验证并结合本文对 Host/Kernel 源码的剖析深入理解其调度与计算原理。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐开源阅读鸿蒙版如何彻底解决你的数字阅读三大痛点开源阅读鸿蒙版如何彻底解决你的数字阅读三大痛点 你是否厌倦了广告满天飞的阅读应用是否希望拥有一个完全掌控阅读体验的工具开源阅读鸿蒙版Legado正是算子库人工智能CANNRuView ESP32 CSI 传感网格从节点固件到汇聚端的分布式无摄像感知落地指南RuView ESP32 CSI 传感网格从节点固件到汇聚端的分布式无摄像感知落地指南 本文基于仓库架构决策记录 ADR 012 https://link.g人工智能算子库深度学习CANNAscendCANN ops-nn 算子库 aclnnSwish 接口详解Swish 激活算子的两段式调用与源码实现CANN ops nn 算子库 aclnnSwish 接口详解Swish 激活算子的两段式调用与源码实现 本文以 CANN ops nn 开源仓库中的 acl人工智能算子库深度学习CANNAscend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考