CANN ops-nn Softshrink 算子实战解析数学原理、参数约束与 aclnn 两段式调用【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nnSoftshrink软收缩是一种带阈值的逐元素激活函数在神经网络中被用于对张量进行平滑稀疏化与噪声抑制。本文以 CANN ops-nn 仓库中 activation/softshrink 模块为对象完整讲解其数学定义、产品支持情况、参数规范、aclnn 两段式调用方式并结合op_api、op_host、op_kernel源码与单元测试深入剖析其实现原理。读完本文你将能够理解 Softshrink 算子在全流程中的定位掌握通过aclnnSoftshrink接口在 NPU 上正确调用该算子并校验结果的完整方法。一、算子功能与数学原理Softshrink 算子对输入张量x逐元素施加“软收缩”变换给定阈值λ将落在[-λ, λ]区间内的元素收缩为 0对区间外的元素则向 0 方向收缩一个λ的偏移量。计算公式如下见 activation/softshrink/README.md 与 aclnnSoftshrink 接口文档$$ Softshrink(x)\begin{cases} x-λ, \text{if } x λ \ xλ, \text{if } x -λ \ 0, \text{otherwise} \end{cases} $$从直观上看Softshrink 与 Hard Shrink硬收缩x -λ取xλx λ取x-λ否则为 0的区别在于Hard Shrink 在阈值处是跳跃不连续的而 Softshrink 在阈值处是连续过渡的因此梯度行为更平滑。该算子常用于去噪、特征稀疏化以及需要对小幅度信号进行抑制的场景。二、产品支持情况根据 activation/softshrink/README.md 中的产品支持矩阵Softshrink 算子在各系列产品上的支持情况如下产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品√Atlas 推理系列产品√Atlas 训练系列产品√说明算子 README 的产品支持矩阵为整体能力声明具体到aclnnSoftshrink接口层面aclnnSoftshrink.md 中列出了 Atlas 200I/500 A2 推理产品“不支持”的接口级说明并在 Atlas 推理/训练系列产品上仅支持 FLOAT、FLOAT16 两种数据类型。实际开发时应以目标版本与具体接口文档为准。三、参数说明Softshrink 算子共包含两个输入和一个输出参数定义如下摘自 activation/softshrink/README.md参数名输入/输出/属性描述数据类型数据格式x输入待进行 Softshrink 计算的输入张量公式中的 xfp16、fp32、bf16NDlambd输入Softshrink 的阈值参数默认值为 0.5fp321y输出Softshrink 计算的输出张量fp16、fp32、bf16ND从算子定义源码 op_host/softshrink_def.cpp 可以印证以上约束输入x与输出y的数据类型均限定为DT_FLOAT16、DT_FLOAT、DT_BF16格式为ND且开启AutoContiguous()处理lambd以算子属性Attr而非输入张量的形式注册类型为浮点默认值 0.5与 PyTorch 语义保持一致算子配置为动态 shape、动态 rank 支持DynamicShapeSupportFlag(true)、DynamicRankSupportFlag(true)能够处理编译期未知的张量形状。在 aclnn 接口层面见下文lambd作为aclScalar*传入数据类型为 FLOAT数值要求大于等于 0。四、aclnn 两段式接口详解CANN 的 aclnn 算子 API 采用两段式接口设计必须先调用aclnnSoftshrinkGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器再调用aclnnSoftshrink执行计算。两个接口的函数原型如下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)4.1 第一段接口aclnnSoftshrinkGetWorkspaceSize该接口完成入参校验、构建算子执行器并返回执行计算所需的 workspace 大小。参数说明如下参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 TensorselfaclTensor*输入输入的张量公式中的 x支持空 TensorFLOAT、FLOAT16、BFLOAT16ND0-8√lambdaclScalar*输入输入的标量公式中的输入 λ数值要求大于等于 0FLOATND0-8√outaclTensor*输出Softshrink 计算的出参支持空 Tensorshape 需要与 self 一致FLOAT、FLOAT16、BFLOAT16ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----需要说明的是在 Atlas 推理系列产品、Atlas 训练系列产品上数据类型仅支持 FLOAT、FLOAT16见 aclnnSoftshrink.md。返回值与错误码两个接口的返回值均为aclnnStatus状态码具体参见 aclnn 返回码说明。第一段接口会完成入参校验以下场景会报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self、lambd 或 out 是空指针ACLNN_ERR_PARAM_INVALID161002self、lambd 或 out 的数据类型不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self 和 out 的 shape 不一致ACLNN_ERR_PARAM_INVALID161002self 或 out 的维数大于 8ACLNN_ERR_PARAM_INVALID161002lambd 0以上校验逻辑在源码 op_api/aclnn_softshrink.cpp 的CheckParams流程中一一对应实现CheckNotNull检查self、lambd、out是否为空指针失败返回ACLNN_ERR_PARAM_NULLPTRCheckDtypeValid检查输入输出数据类型是否在{DT_FLOAT, DT_FLOAT16, DT_BF16}支持列表内并根据当前 SoC 版本判断是否支持 BF16仅ASCEND910B~ASCEND910E区间内的 SoC 支持否则返回ACLNN_ERR_PARAM_INVALIDCheckShape检查self与out的 shape 一致且维数不超过MAX_SUPPORT_DIMS_NUMS8 维CheckLambdValue检查lambd-ToFloat() 0CheckFormat在 Regbase 架构下拒绝私有格式Private format的输入输出。当self为空 Tensor 时第一段接口直接返回workspaceSize 0并释放执行器属于合法路径。4.2 第二段接口aclnnSoftshrink该接口在指定 Stream 上执行算子计算。参数说明如下参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnSoftshrinkGetWorkspaceSize获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream从源码看第二段接口内部调用CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成真正的算子下发执行属于框架固定写法。五、调用示例与完整调用流程以下示例代码摘自 examples/test_aclnn_softshrink.cpp与接口文档 aclnnSoftshrink.md 中的示例一致演示了在 NPU 上调用 Softshrink 算子的完整流程#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根据自己的需要处理 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; }该示例按固定套路分为七个阶段可归纳为“三步核心调用 四步固定收尾”资源初始化aclInit→aclrtSetDevice→aclrtCreateStream完成 ACL 运行环境与 Stream 创建构造输入输出使用aclrtMalloc申请 Device 侧内存aclrtMemcpy将 Host 数据拷贝至 Device再通过aclCreateTensor/aclCreateScalar创建aclTensor与aclScalar。注意示例中lambdValue 0.5f与默认阈值一致输入数据{0.1, ..., 0.8}均落在[-0.5, 0.5]内因此预期输出全为 0两段式调用先调用aclnnSoftshrinkGetWorkspaceSize获取workspaceSize与executor按需用aclrtMalloc申请 workspace再调用aclnnSoftshrink执行同步等待aclrtSynchronizeStream等待任务执行结束结果回拷aclrtMemcpy将 Device 侧结果拷贝回 Host 并打印释放张量对象aclDestroyTensor、aclDestroyScalar释放资源aclrtFree释放 Device 内存含 workspace、aclrtDestroyStream、aclrtResetDevice、aclFinalize。关于示例的具体编译与执行过程请参考仓库中的 编译与运行样例 文档在arch35目录下还提供了面向 Ascend 950 的 test_aclnn_softshrink.cpp 版本。六、源码级实现原理Softshrink 算子完整覆盖了 CANN 算子开发的“API 层 → L0 算子层 → 算子定义层 → Kernel 层”全链路以下按调用链自顶向下分析。6.1 API 层入参校验与算子图构建op_api/aclnn_softshrink.cpp 实现了两段式接口的对外封装。第一段接口内部除了前文所述的 5 步参数校验外还完成了算子计算图的构建对输入执行l0op::Contiguous将非连续 Tensor 规整为连续内存布局调用 L0 算子l0op::SoftShrink(selfContiguous, lambd-ToFloat(), executor)对结果执行l0op::Cast转换到输出要求的数据类型通过l0op::ViewCopy将结果拷贝到输出out上从而支持非连续输出 Tensor汇总得到workspaceSize并把uniqueExecutor的所有权转移给调用方。这也解释了接口文档中“支持非连续 Tensor√”的能力来源。6.2 L0 算子层AICore 路径分发op_api/softshrink.cpp 定义了 L0 算子SoftShrink通过OP_TYPE_REGISTER(SoftShrink)注册算子类型IsAiCoreSupport依据数据类型{DT_FLOAT, DT_FLOAT16, DT_BF16}判断是否走 AICore 路径支持时通过ADD_TO_LAUNCHER_LIST_AICORE(SoftShrink, ...)宏将算子加入任务队列lambd作为算子属性OP_ATTR随任务下发输出张量由executor-AllocTensor(input-GetViewShape(), input-GetDataType())自动分配shape 与数据类型与输入保持一致。6.3 算子定义与形状推导op_host/softshrink_def.cpp 使用OpDef框架注册算子声明输入x、输出yfp16/fp32/bf16ND 格式AutoContiguous以及可选属性lambd默认 0.5f并配置ascend950的 AICore 配置动态 shape/rank 支持、精度提升标志PrecisionReduceFlag(true)。形状推导分为两个层面op_host/softshrink_infershape.cpp输出 shape 与输入 shape 完全一致逐元素算子复用InferShape4Elewise通用实现op_graph/softshrink_graph_infer.cpp图模式下输出 dtype 与输入 dtype 保持一致InferDataType直接透传输入数据类型。6.4 Kernel 层schMode 模板分发与 fp16/bf16 升精op_kernel/softshrink.cpp 是 AICore Kernel 的入口采用单 schMode 模板按数据类型分发schMode数据类型模板实例化说明0FP32Softshrinkfloat, 1, 0直通计算NEED_UPCAST01FP16Softshrinkhalf, 1, 1升精计算NEED_UPCAST12BF16Softshrinkbfloat16_t, 1, 1升精计算NEED_UPCAST1其中BUFFER_MODE固定为 1双缓冲。源码注释指出一个重要实现细节fp16/bf16 路径均先 Cast 到 fp32 计算再 Cast 回原精度目的是对齐 PyTorch CPU 的 fp16/bf16 路径避免在λ ∈ {0.1, 0.3 ...}等无法被 fp16 精确表示的场景下因边界判断与减法误差累积导致与 golden 数据偏差。与之配套的 Tiling 实现 op_host/arch35/softshrink_tiling.cpp 按芯片核数AIV Core 数与 UB 内存大小切分任务并根据数据类型计算每元素所需的 UB 缓冲区占用fp32 路径为 36 字节/元素9 个 fp32 缓冲区fp16/bf16 升精路径为 32 字节/元素4 * sizeof(half/bf16) 6 * sizeof(float)。lambd属性在 Tiling 阶段从context-GetAttrs()-GetFloat(0)读取默认值为 0.5f。6.5 单元测试验证activation/softshrink/tests/ut/op_host/op_api/test_aclnn_softshrink.cpp 覆盖了多种数据类型组合的调用场景{5, 6, 7}形状的 FLOAT16 输入 FLOAT16 输出TestGetWorkspaceSize返回ACL_SUCCESS同样形状的 FLOAT 输入 FLOAT 输出返回ACL_SUCCESSBF16 输入 BF16 输出调用成功预期在支持 BF16 的平台通过DOUBLE 输入预期返回ACLNN_ERR_PARAM_INVALID验证了 dtype 校验逻辑FLOAT16 输入 FLOAT 输出输入输出 dtype 不一致返回ACL_SUCCESS印证了内部Cast环节对输出数据类型的适配能力。测试通过OP_API_UT(aclnnSoftshrink, INPUT(tensor_desc, lambd), OUTPUT(out_tensor_desc))框架封装了张量与标量描述含取值区间与精度阈值开发者可在接入环境后按需放开ut.TestPrecision()进行数值精度比对。七、约束与确定性计算约束说明算子 README 中明确指出 Softshrink 算子无额外约束activation/softshrink/README.md。实操中的限制均来自接口层面的入参校验数据类型需为 fp16/fp32/bf16、维数不超过 8、lambd ≥ 0、输出 shape 与输入一致等详见第四节错误码表。确定性计算aclnnSoftshrink默认采用确定性实现见 aclnnSoftshrink.md即相同输入在多次执行下结果可复现这对调试与结果比对十分友好。关于确定性计算的更完整约定可参考仓库文档 determinism_compute.md。八、总结Softshrink 是 CANN ops-nn 中结构完整、极具代表性的逐元素激活算子对外提供符合两段式约定的aclnnSoftshrink接口内部则贯通了参数校验、Contiguous/Cast/ViewCopy 算子图构建、AICore 模板分发、Tiling 切分与双缓冲 Kernel 的完整实现链路并通过 fp16/bf16 升精计算保证了与 PyTorch golden 的一致性。开发者无论是直接使用该算子还是将其作为学习 CANN 自定义算子开发范式的参考模板都可以从 activation/softshrink 目录下的 README、接口文档、示例与测试用例中获得完整的闭环指引。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考