CANN ops-nn 算子实战aclnnSoftshrink 两段式接口详解与 Softshrink 激活实现原理【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn本文以 CANN ops-nn 开源算子库中的 aclnnSoftshrink 接口文档 为核心系统讲解 Softshrink 软阈值激活算子的功能定义、产品支持矩阵、两段式 API 调用范式、参数约束与错误码语义并结合activation/softshrink目录下的 op_api、op_host、op_kernel 源码与单元测试深入剖析该算子在 NPU 上的计算流程、升精策略与 Tiling 调度原理。读完本文读者将能够独立完成 Softshrink 算子的 host 侧编程调用、参数校验与样例工程移植。1. 算子概述以元素为单位强制收缩 λ 范围内的输入Softshrink软阈值收缩是激活函数家族中的一员常用于信号去噪、稀疏化等场景。其核心语义是对输入张量中的每个元素执行按阈值收缩当元素值大于阈值 λ 时将其减去 λ当元素值小于 −λ 时将其加上 λ其余位于 [−λ, λ] 区间内的元素一律置为 0。该算子在 CANN ops-nn 仓库中位于 activation/softshrink对外通过两级 aclnn 接口aclnnSoftshrinkGetWorkspaceSize与aclnnSoftshrink暴露给上层应用是典型的元素级element-wiseNPU 加速算子。1.1 计算公式$$ Softshrink(x)\begin{cases} x-λ, if\ x λ \ xλ, if\ x -λ \ 0, otherwise \end{cases} $$其中 λ即接口参数lambd为阈值标量要求取值大于等于 0。1.2 产品支持情况根据 aclnnSoftshrink.md 的声明该接口在不同产品系列上的支持情况如下产品系列是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品支持Atlas 训练系列产品支持从源码看该算子的 AiCore Kernel 专门针对 arch35Ascend950架构实现见 op_kernel/softshrink.cpp 与 op_host/softshrink_def.cpp 中AICore().AddConfig(ascend950, ...)的注册同时通过通用两段式接口框架兼容上述训练/推理系列产品。2. 两段式接口调用范式aclnnSoftshrink 遵循 CANN 算子库统一的两段式接口规范必须先调用第一段接口aclnnSoftshrinkGetWorkspaceSize获取计算所需的 workspace 大小以及封装了算子计算流程的执行器executor再调用第二段接口aclnnSoftshrink真正执行计算。这种设计将资源规划与任务下发解耦便于框架层复用执行器、精确管理 Device 侧内存。2.1 函数原型aclnnStatus aclnnSoftshrinkGetWorkspaceSize( const aclTensor* self, const aclScalar* lambd, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnSoftshrink( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)两个接口均返回aclnnStatus状态码具体取值语义参见 aclnn 返回码说明。3. 第一段接口aclnnSoftshrinkGetWorkspaceSize 参数详解第一段接口完成入参校验、算子计算图构建与 workspace 大小估算参数说明如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入输入的张量公式中的 x。支持空 Tensor。FLOAT、FLOAT16、BFLOAT16ND0-8√lambdaclScalar*输入输入的标量公式中的输入 λ。数值要求大于等于 0。FLOATND0-8√outaclTensor*输出Softshrink 计算的出参。支持空 Tensorshape 需要与 self 一致。FLOAT、FLOAT16、BFLOAT16ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----平台差异提示对于 Atlas 推理系列产品与 Atlas 训练系列产品self与out的数据类型仅支持 FLOAT、FLOAT16即 BFLOAT16 不在支持范围内。3.1 入参校验规则与返回码第一段接口会执行严格的入参校验出现下列场景时返回对应错误码返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self、lambd 或 out 是空指针。ACLNN_ERR_PARAM_INVALID161002self、lambd 或 out 的数据类型不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002self 和 out 的 shape 不一致。ACLNN_ERR_PARAM_INVALID161002self 或 out 的维数大于 8。ACLNN_ERR_PARAM_INVALID161002lambd 0。这些校验逻辑在 op_api/aclnn_softshrink.cpp 中有完整对应实现CheckNotNull空指针检查、CheckDtypeValid数据类型检查且会通过CheckSocVersionIsSupportBf16判定当前 SoC 版本是否支持 BF16、CheckShapeshape 一致性 最大 8 维检查、CheckLambdValueλ ≥ 0 检查以及CheckFormatregbase 架构下禁止私有 format。五个检查依次执行任一失败即返回ACLNN_ERR_PARAM_NULLPTR或ACLNN_ERR_PARAM_INVALID与上表一一对应。4. 第二段接口aclnnSoftshrink 参数详解第二段接口负责将已构建好的执行器下发到指定 Stream 上执行参数说明如下参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnSoftshrinkGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。该接口内部为固定写法直接调用框架的CommonOpExecutorRun完成计算见 aclnn_softshrink.cpp返回值同样为aclnnStatus状态码。5. 约束说明确定性计算aclnnSoftshrink 默认为确定性实现deterministic即相同输入与相同硬件配置下多次执行的结果保持一致便于精度调试与结果复现。6. 完整调用示例可编译运行以下示例演示了从环境初始化、构造输入输出 Tensor、两段式调用、结果回拷到资源释放的完整流程。示例代码与仓库中的 examples/test_aclnn_softshrink.cpp 一致具体编译与执行过程请参考编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_softshrink.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 {4, 2}; std::vectorint64_t outShape {4, 2}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclScalar* lambd nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}; std::vectorfloat outHostData {0, 0, 0, 0, 0, 0, 0, 0}; float lambdValue 0.5f; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建lambd aclTensor lambd aclCreateScalar(lambdValue, aclDataType::ACL_FLOAT); CHECK_RET(lambd ! 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; // 调用aclnnSoftshrink第一段接口 ret aclnnSoftshrinkGetWorkspaceSize(self, lambd, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSoftshrinkGetWorkspaceSize 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); } // 调用aclnnSoftshrink第二段接口 ret aclnnSoftshrink(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSoftshrink 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(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需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyScalar(lambd); aclDestroyTensor(out); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }以示例数据λ0.5为例输入{0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}中前 5 个元素均落在 [−0.5, 0.5] 区间内输出 00.6 → 0.1、0.7 → 0.2、0.8 → 0.3最终输出为{0, 0, 0, 0, 0, 0.1, 0.2, 0.3}可据此快速验证调用正确性。7. 源码级实现原理剖析7.1 第一段接口内部的算子计算图从 op_api/aclnn_softshrink.h 的接口注释可以清晰看到aclnnSoftshrinkGetWorkspaceSize内部构建的计算流程为各环节作用如下对应 op_api/aclnn_softshrink.cpp 的实现Contiguous将self转为连续布局保证后续向量化计算的高效访存l0op::SoftShrink调用 L0 层算子完成核心 Softshrink 计算λ 以标量形式传入Cast将中间结果转换为out声明的数据类型用于 fp16/bf16 的精度处理与输出对齐ViewCopy将计算结果写回out——即使out是非连续 Tensor也能正确落位最后通过uniqueExecutor-GetWorkspaceSize()汇总整条链路的 workspace 需求并释放给调用方。此外接口对空 Tensor做了快速路径处理当self-IsEmpty()为真时直接返回workspaceSize 0并成功结束无需下发计算。7.2 AiCore Kernel 的升精计算策略Softshrink 的 AiCore 计算核心在 op_kernel/arch35/softshrink.h 中实现采用模板类SoftshrinkT, BUFFER_MODE, NEED_UPCAST其中T为 IO 数据类型half/float/bfloat16_tBUFFER_MODE固定为 1即启用双缓冲BUFFER_NUM 2隐藏访存延迟NEED_UPCAST决定是否升精到 fp32 计算。Kernel 入口 op_kernel/softshrink.cpp 按 dtype 分为三种模板实例schMode数据类型模板实例计算说明0FP32Softshrinkfloat, 1, 0fp32 直通计算1FP16Softshrinkhalf, 1, 1half → fp32 计算 → 回写 half2BF16Softshrinkbfloat16_t, 1, 1bf16 → fp32 计算 → 回写 bf16升精的原因fp16 无法精确表示 λ ∈ {0.1, 0.3, …} 这类小数若全程 fp16 计算边界判断与减法误差会累积导致与 golden 结果偏差。v2 版本将 fp16/bf16 统一升精到 fp32 计算与 PyTorch CPU 的vector_func升精路径aten/src/ATen/native/cpu/Activation.cpp行为对齐。计算内核采用两次 Compare Select 的组合实现Compare(x lambd)生成 mask1Select(mask1 ? x - lambd : 0)得到正侧中间结果Compare(x -lambd)生成 mask2Select(mask2 ? x lambd : tmp)得到最终输出。内核文件中同时标注了一个已知行为差异AscendC 上NaN lambd与NaN -lambd均为 false因此 NaN 会被静默置 0与 PyTorch CUDA 显式透传 NaN 的行为不一致精度对齐时需注意。7.3 算子定义、Shape 推导与 Tiling算子注册op_host/softshrink_def.cpp 定义算子Softshrink输入x与输出y均支持FLOAT16/FLOAT/BF16与ND格式并开启AutoContiguouslambd以可选 Attr形式注册默认值为 0.5与 PyTorch 语义一致同时声明DynamicShapeSupportFlag(true)、DynamicRankSupportFlag(true)、PrecisionReduceFlag(true)等动态能力标志。Shape 推导op_host/softshrink_infershape.cpp 直接复用InferShape4Elewise工具输出 shape 与输入完全一致这也是接口文档要求out.shape self.shape的底层依据。Tiling 调度op_host/arch35/softshrink_tiling.cpp 获取 AIV 核数与 UB 大小按元素总数切分为blockFactor核间切分与ubFactorUB 内切分UB 预算方面fp32 路径每元素 36 字节9 个 fp32 缓冲fp16/bf16 升精路径每元素 32 字节4×IO 元素 6×fp32 缓冲固定双缓冲并预留 1024 字节 UB 余量。7.4 单测覆盖与参数校验验证仓库在 tests/ut/op_host/op_api/test_aclnn_softshrink.cpp 中提供了针对性的 host 侧单元测试覆盖了正常两段式调用返回ACL_SUCCESS非法的lambd负值返回ACLNN_ERR_PARAM_INVALID空指针self/lambd/out 为 nullptr返回ACLNN_ERR_PARAM_NULLPTRshape 不一致、维度超限等非法输入场景的错误码断言。这些用例与本文第 3.1 节列出的返回码表格一一对应可作为自研算子 host 侧校验逻辑的参考模板。8. 总结与延伸阅读aclnnSoftshrink 是 CANN ops-nn 中结构完整、注释详尽的一个代表性元素级算子对外采用统一的两段式接口对内则完整呈现了算子定义 → Shape 推导 → Tiling 切分 → AiCore Kernel 计算的标准研发链路其 fp16/bf16 升精策略与 PyTorch 精度对齐的取舍也极具参考价值。延伸阅读算子两段式接口说明aclnn 返回码说明样例编译与运行指南Softshrink 算子 README含 aclnn 调用方式入口【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-nn 算子实战：aclnnSoftshrink 两段式接口详解与 Softshrink 激活实现原理
CANN ops-nn 算子实战aclnnSoftshrink 两段式接口详解与 Softshrink 激活实现原理【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn本文以 CANN ops-nn 开源算子库中的 aclnnSoftshrink 接口文档 为核心系统讲解 Softshrink 软阈值激活算子的功能定义、产品支持矩阵、两段式 API 调用范式、参数约束与错误码语义并结合activation/softshrink目录下的 op_api、op_host、op_kernel 源码与单元测试深入剖析该算子在 NPU 上的计算流程、升精策略与 Tiling 调度原理。读完本文读者将能够独立完成 Softshrink 算子的 host 侧编程调用、参数校验与样例工程移植。1. 算子概述以元素为单位强制收缩 λ 范围内的输入Softshrink软阈值收缩是激活函数家族中的一员常用于信号去噪、稀疏化等场景。其核心语义是对输入张量中的每个元素执行按阈值收缩当元素值大于阈值 λ 时将其减去 λ当元素值小于 −λ 时将其加上 λ其余位于 [−λ, λ] 区间内的元素一律置为 0。该算子在 CANN ops-nn 仓库中位于 activation/softshrink对外通过两级 aclnn 接口aclnnSoftshrinkGetWorkspaceSize与aclnnSoftshrink暴露给上层应用是典型的元素级element-wiseNPU 加速算子。1.1 计算公式$$ Softshrink(x)\begin{cases} x-λ, if\ x λ \ xλ, if\ x -λ \ 0, otherwise \end{cases} $$其中 λ即接口参数lambd为阈值标量要求取值大于等于 0。1.2 产品支持情况根据 aclnnSoftshrink.md 的声明该接口在不同产品系列上的支持情况如下产品系列是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品支持Atlas 训练系列产品支持从源码看该算子的 AiCore Kernel 专门针对 arch35Ascend950架构实现见 op_kernel/softshrink.cpp 与 op_host/softshrink_def.cpp 中AICore().AddConfig(ascend950, ...)的注册同时通过通用两段式接口框架兼容上述训练/推理系列产品。2. 两段式接口调用范式aclnnSoftshrink 遵循 CANN 算子库统一的两段式接口规范必须先调用第一段接口aclnnSoftshrinkGetWorkspaceSize获取计算所需的 workspace 大小以及封装了算子计算流程的执行器executor再调用第二段接口aclnnSoftshrink真正执行计算。这种设计将资源规划与任务下发解耦便于框架层复用执行器、精确管理 Device 侧内存。2.1 函数原型aclnnStatus aclnnSoftshrinkGetWorkspaceSize( const aclTensor* self, const aclScalar* lambd, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnSoftshrink( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)两个接口均返回aclnnStatus状态码具体取值语义参见 aclnn 返回码说明。3. 第一段接口aclnnSoftshrinkGetWorkspaceSize 参数详解第一段接口完成入参校验、算子计算图构建与 workspace 大小估算参数说明如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入输入的张量公式中的 x。支持空 Tensor。FLOAT、FLOAT16、BFLOAT16ND0-8√lambdaclScalar*输入输入的标量公式中的输入 λ。数值要求大于等于 0。FLOATND0-8√outaclTensor*输出Softshrink 计算的出参。支持空 Tensorshape 需要与 self 一致。FLOAT、FLOAT16、BFLOAT16ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----平台差异提示对于 Atlas 推理系列产品与 Atlas 训练系列产品self与out的数据类型仅支持 FLOAT、FLOAT16即 BFLOAT16 不在支持范围内。3.1 入参校验规则与返回码第一段接口会执行严格的入参校验出现下列场景时返回对应错误码返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self、lambd 或 out 是空指针。ACLNN_ERR_PARAM_INVALID161002self、lambd 或 out 的数据类型不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002self 和 out 的 shape 不一致。ACLNN_ERR_PARAM_INVALID161002self 或 out 的维数大于 8。ACLNN_ERR_PARAM_INVALID161002lambd 0。这些校验逻辑在 op_api/aclnn_softshrink.cpp 中有完整对应实现CheckNotNull空指针检查、CheckDtypeValid数据类型检查且会通过CheckSocVersionIsSupportBf16判定当前 SoC 版本是否支持 BF16、CheckShapeshape 一致性 最大 8 维检查、CheckLambdValueλ ≥ 0 检查以及CheckFormatregbase 架构下禁止私有 format。五个检查依次执行任一失败即返回ACLNN_ERR_PARAM_NULLPTR或ACLNN_ERR_PARAM_INVALID与上表一一对应。4. 第二段接口aclnnSoftshrink 参数详解第二段接口负责将已构建好的执行器下发到指定 Stream 上执行参数说明如下参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnSoftshrinkGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。该接口内部为固定写法直接调用框架的CommonOpExecutorRun完成计算见 aclnn_softshrink.cpp返回值同样为aclnnStatus状态码。5. 约束说明确定性计算aclnnSoftshrink 默认为确定性实现deterministic即相同输入与相同硬件配置下多次执行的结果保持一致便于精度调试与结果复现。6. 完整调用示例可编译运行以下示例演示了从环境初始化、构造输入输出 Tensor、两段式调用、结果回拷到资源释放的完整流程。示例代码与仓库中的 examples/test_aclnn_softshrink.cpp 一致具体编译与执行过程请参考编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_softshrink.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 {4, 2}; std::vectorint64_t outShape {4, 2}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclScalar* lambd nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}; std::vectorfloat outHostData {0, 0, 0, 0, 0, 0, 0, 0}; float lambdValue 0.5f; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建lambd aclTensor lambd aclCreateScalar(lambdValue, aclDataType::ACL_FLOAT); CHECK_RET(lambd ! 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; // 调用aclnnSoftshrink第一段接口 ret aclnnSoftshrinkGetWorkspaceSize(self, lambd, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSoftshrinkGetWorkspaceSize 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); } // 调用aclnnSoftshrink第二段接口 ret aclnnSoftshrink(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSoftshrink 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(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需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyScalar(lambd); aclDestroyTensor(out); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }以示例数据λ0.5为例输入{0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}中前 5 个元素均落在 [−0.5, 0.5] 区间内输出 00.6 → 0.1、0.7 → 0.2、0.8 → 0.3最终输出为{0, 0, 0, 0, 0, 0.1, 0.2, 0.3}可据此快速验证调用正确性。7. 源码级实现原理剖析7.1 第一段接口内部的算子计算图从 op_api/aclnn_softshrink.h 的接口注释可以清晰看到aclnnSoftshrinkGetWorkspaceSize内部构建的计算流程为各环节作用如下对应 op_api/aclnn_softshrink.cpp 的实现Contiguous将self转为连续布局保证后续向量化计算的高效访存l0op::SoftShrink调用 L0 层算子完成核心 Softshrink 计算λ 以标量形式传入Cast将中间结果转换为out声明的数据类型用于 fp16/bf16 的精度处理与输出对齐ViewCopy将计算结果写回out——即使out是非连续 Tensor也能正确落位最后通过uniqueExecutor-GetWorkspaceSize()汇总整条链路的 workspace 需求并释放给调用方。此外接口对空 Tensor做了快速路径处理当self-IsEmpty()为真时直接返回workspaceSize 0并成功结束无需下发计算。7.2 AiCore Kernel 的升精计算策略Softshrink 的 AiCore 计算核心在 op_kernel/arch35/softshrink.h 中实现采用模板类SoftshrinkT, BUFFER_MODE, NEED_UPCAST其中T为 IO 数据类型half/float/bfloat16_tBUFFER_MODE固定为 1即启用双缓冲BUFFER_NUM 2隐藏访存延迟NEED_UPCAST决定是否升精到 fp32 计算。Kernel 入口 op_kernel/softshrink.cpp 按 dtype 分为三种模板实例schMode数据类型模板实例计算说明0FP32Softshrinkfloat, 1, 0fp32 直通计算1FP16Softshrinkhalf, 1, 1half → fp32 计算 → 回写 half2BF16Softshrinkbfloat16_t, 1, 1bf16 → fp32 计算 → 回写 bf16升精的原因fp16 无法精确表示 λ ∈ {0.1, 0.3, …} 这类小数若全程 fp16 计算边界判断与减法误差会累积导致与 golden 结果偏差。v2 版本将 fp16/bf16 统一升精到 fp32 计算与 PyTorch CPU 的vector_func升精路径aten/src/ATen/native/cpu/Activation.cpp行为对齐。计算内核采用两次 Compare Select 的组合实现Compare(x lambd)生成 mask1Select(mask1 ? x - lambd : 0)得到正侧中间结果Compare(x -lambd)生成 mask2Select(mask2 ? x lambd : tmp)得到最终输出。内核文件中同时标注了一个已知行为差异AscendC 上NaN lambd与NaN -lambd均为 false因此 NaN 会被静默置 0与 PyTorch CUDA 显式透传 NaN 的行为不一致精度对齐时需注意。7.3 算子定义、Shape 推导与 Tiling算子注册op_host/softshrink_def.cpp 定义算子Softshrink输入x与输出y均支持FLOAT16/FLOAT/BF16与ND格式并开启AutoContiguouslambd以可选 Attr形式注册默认值为 0.5与 PyTorch 语义一致同时声明DynamicShapeSupportFlag(true)、DynamicRankSupportFlag(true)、PrecisionReduceFlag(true)等动态能力标志。Shape 推导op_host/softshrink_infershape.cpp 直接复用InferShape4Elewise工具输出 shape 与输入完全一致这也是接口文档要求out.shape self.shape的底层依据。Tiling 调度op_host/arch35/softshrink_tiling.cpp 获取 AIV 核数与 UB 大小按元素总数切分为blockFactor核间切分与ubFactorUB 内切分UB 预算方面fp32 路径每元素 36 字节9 个 fp32 缓冲fp16/bf16 升精路径每元素 32 字节4×IO 元素 6×fp32 缓冲固定双缓冲并预留 1024 字节 UB 余量。7.4 单测覆盖与参数校验验证仓库在 tests/ut/op_host/op_api/test_aclnn_softshrink.cpp 中提供了针对性的 host 侧单元测试覆盖了正常两段式调用返回ACL_SUCCESS非法的lambd负值返回ACLNN_ERR_PARAM_INVALID空指针self/lambd/out 为 nullptr返回ACLNN_ERR_PARAM_NULLPTRshape 不一致、维度超限等非法输入场景的错误码断言。这些用例与本文第 3.1 节列出的返回码表格一一对应可作为自研算子 host 侧校验逻辑的参考模板。8. 总结与延伸阅读aclnnSoftshrink 是 CANN ops-nn 中结构完整、注释详尽的一个代表性元素级算子对外采用统一的两段式接口对内则完整呈现了算子定义 → Shape 推导 → Tiling 切分 → AiCore Kernel 计算的标准研发链路其 fp16/bf16 升精策略与 PyTorch 精度对齐的取舍也极具参考价值。延伸阅读算子两段式接口说明aclnn 返回码说明样例编译与运行指南Softshrink 算子 README含 aclnn 调用方式入口【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考