CANN ops-nn Sigmoid 算子深度指南aclnn 两段式接口、AI Core 内核实现与图模式调用【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn本文以 CANN ops-nn 仓库中的 Sigmoid 算子的官方文档activation/sigmoid/README.md 与 aclnnSigmoidaclnnInplaceSigmoid.md为主线结合该算子模块的源码、样例与测试系统讲解 Sigmoid 在 NPU 上的功能定义、产品支持范围、aclnnSigmoid/aclnnInplaceSigmoid两段式调用方法、底层 AI Core 内核实现以及算子 IR 构图方式。读完本文你将能够独立完成 Sigmoid 算子在 Ascend 设备上的算子级开发与验证并理解从 aclnn 接口到 Tiling、再到 Kernel 的完整执行链路。一、功能说明与数学定义Sigmoid 是神经网络中最常用的激活函数之一将任意实数输入映射到 (0, 1) 开区间内常用于二分类输出层、LSTM/GRU 等循环神经网络的门控单元以及注意力机制的归一化环节。在本仓库中该算子位于activation/sigmoid目录其功能为对输入 Tensor 逐元素element-wise完成 Sigmoid 运算。计算公式为$$ out {\frac{1} {1{e}^{-input}}} $$该函数具有如下典型性质输出值域为 (0, 1)单调递增当输入趋近正无穷时输出趋近 1趋近负无穷时输出趋近 0关于原点中心对称sigmoid(-x) 1 - sigmoid(x)导数可用自身表示sigmoid(x) sigmoid(x) * (1 - sigmoid(x))这一性质使其在反向传播中计算高效也是 activation/sigmoid_grad 等梯度算子实现的基础。从算子原型看Sigmoid 的输入输出形状完全一致逐元素运算由 op_graph/sigmoid_proto.h 中的 IR 定义可见其输入x与输出y均使用TensorType::UnaryDataType()且该 IR 明确标注“Compatible with the TensorFlow operator Sigmoid”即与 TensorFlow 框架的 Sigmoid 算子语义对齐。二、产品支持情况根据官方文档Sigmoid 算子在各产品上的支持情况如下产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品√注以上为算子级支持情况aclnnSigmoid/aclnnInplaceSigmoid接口级支持情况与上表一致详见接口文档 aclnnSigmoidaclnnInplaceSigmoid.md。从源码可印证该支持范围在 op_host/sigmoid_def.cpp 中算子通过this-AICore().AddConfig(ascend950, aicoreConfig)与AddConfig(ascend350, aicoreConfig)注册了 ascend950、ascend350 两个平台的 AI Core 配置对应文档中的 Ascend 950 系列与 Atlas A2/A3 系列等而 Atlas 200I/500 A2 推理产品不在注册列表中因此标记为不支持。三、参数说明Sigmoid 算子仅有输入x和输出y两个 Tensor 参数具体规格如下参数名输入/输出描述数据类型数据格式x输入待进行 Sigmoid 计算的入参公式中的 inputFLOAT、FLOAT16、BFLOAT16NDy输出计算的出参FLOAT、FLOAT16、BFLOAT16ND需要说明的是上表为算子Kernel级支持的数据类型。在 aclnn 接口层通过 op_api/aclnn_sigmoid.cpp 中的DTYPE_SUPPORT_LIST可以看到接口输入self还额外支持DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOL、COMPLEX64、COMPLEX128等更宽泛的类型——这些类型会在接口层通过Cast算子统一转换后再送入 AI Core Kernel 计算这也是“算子层只支持三种浮点类型、而接口层支持更多类型”的原因所在。约束说明官方文档明确标注“无”额外约束。在实现层面唯一的形状约束是输入self与输出out的 shape 必须一致见下文参数校验逻辑。四、调用方式总览Sigmoid 算子在 CANN 中支持两种调用方式调用方式调用样例说明aclnn 调用test_aclnn_sigmoid.cpp通过aclnnSigmoid接口方式调用 Sigmoid 算子aclnn 调用test_aclnn_inplace_sigmoid.cpp通过aclnnInplaceSigmoid接口方式调用 Sigmoid 算子图模式-通过算子 IR 构图方式调用 Sigmoid 算子其中图模式对应的 IR 定义位于 op_graph/sigmoid_proto.h通过REG_OP(Sigmoid)注册算子原型Host 侧的形状推导在 op_host/sigmoid_infershape.cpp 中实现直接复用Ops::Base::InferShape4Elewise逐元素通用推断逻辑保证输出 shape 与输入一致。五、两段式 aclnn 接口详解acclnSigmoid与acclnInplaceSigmoid实现相同的计算功能区别仅在于结果存储方式aclnnSigmoid需新建一个输出张量对象out存储计算结果输入输出分离aclnnInplaceSigmoid无需新建输出张量对象直接在输入张量的内存中覆写计算结果in-place 操作节省一倍显存。与 CANN 其他算子一样这两个接口均采用两段式调用必须先调用第一段xxxGetWorkspaceSize接口获取 workspace 大小并完成入参校验、构建aclOpExecutor再调用第二段xxx接口在指定 Stream 上真正下发执行。两段式接口的通用原理可参考 docs/zh/context/basic_concept.md 中的 aclnn 两段式 API 说明。5.1 函数原型aclnnStatus aclnnSigmoidGetWorkspaceSize( const aclTensor *self, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor) aclnnStatus aclnnSigmoid( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream) aclnnStatus aclnnInplaceSigmoidGetWorkspaceSize( aclTensor* selfRef, uint64_t* workspaceSize, aclOpExecutor** executor) aclnnStatus aclnnInplaceSigmoid( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)5.2 aclnnSigmoidGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入待进行 Sigmoid 计算的入参公式中的 input支持空 Tensorshape 需要与 out 一致FLOAT、FLOAT16、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOL、COMPLEX64、COMPLEX128、BFLOAT16ND0-8√outaclTensor*输出计算的出参shape 需要与 self 一致FLOAT、FLOAT16、DOUBLE、COMPLEX64、COMPLEX128、BFLOAT16ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----5.3 aclnnInplaceSigmoidGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfRefaclTensor*输入/输出计算的入参结果原位覆写-FLOAT、FLOAT16、DOUBLE、COMPLEX64、COMPLEX128、BFLOAT16ND1-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----注意in-place 接口在部分老产品Atlas 推理系列产品、Atlas 训练系列产品上仅支持FLOAT、FLOAT16、DOUBLE、COMPLEX64、COMPLEX128不支持BFLOAT16。这一点在源码 op_api/aclnn_sigmoid.cpp 中有对应逻辑GetSelfRefDtypeList()会根据当前 NPU 架构DAV_2201或 RegBase 架构动态返回不同的 in-place 类型支持列表BF16 仅在 Ascend 950 及 Atlas A2/A3 系列等新架构上开放。5.4 两段接口aclnnSigmoid / aclnnInplaceSigmoid公共参数参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream5.5 返回值与错误码两段接口均返回aclnnStatus状态码。第一段接口会完成入参校验出现以下场景时报错aclnnSigmoidGetWorkspaceSize 校验错误返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self 或 out 是空指针ACLNN_ERR_PARAM_INVALID161002self 和 out 的数据类型和数据格式不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self 和 out 的 shape 不匹配aclnnInplaceSigmoidGetWorkspaceSize 校验错误返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 selfRef 是空指针ACLNN_ERR_PARAM_INVALID161002selfRef 的数据类型不在支持的范围之内这些校验在源码中的落点是 op_api/aclnn_sigmoid.cpp 的CheckParams/CheckInplaceParams函数依次检查空指针CheckNotNull、数据类型与 BF16 平台兼容性CheckDtypeValid其中OP_CHECK_DTYPE_NOT_SUPPORT映射到ACLNN_ERR_PARAM_INVALID161002、shape 一致性CheckShape要求 self 与 out 的 shape 完全一致且维度不超过 8。5.6 约束说明确定性计算aclnnSigmoid与aclnnInplaceSigmoid默认即为确定性实现同一输入多次执行结果一致无需额外配置。六、完整调用示例可直接运行下面以aclnnSigmoid为例给出完整的调用流程该示例可在 examples/test_aclnn_sigmoid.cpp 中找到完整可编译源码aclnnInplaceSigmoid的完整示例见 examples/test_aclnn_inplace_sigmoid.cpp。编译与运行环境准备可参考仓库文档中的样例编译说明。6.1 公共工具代码#include cinttypes #include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_sigmoid.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接口创建aclTensorND格式storageOffset为0 *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; }6.2 aclnnSigmoid 主流程非原地版int main() { // 1. 固定写法device/stream初始化参考acl APIdeviceId根据实际设备填写 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. 构造输入与输出 std::vectorint64_t selfShape {2, 2}; std::vectorint64_t outShape {2, 2}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 1, 2, 3}; std::vectorfloat outHostData {0, 0, 0, 0}; ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API两段式 uint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段获取workspace大小并完成入参校验 ret aclnnSigmoidGetWorkspaceSize(self, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSigmoidGetWorkspaceSize 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); } // 第二段在指定stream上执行计算 ret aclnnSigmoid(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSigmoid 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并打印 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[% PRId64 ] is: %f\n, i, resultData[i]); } // 6. 释放资源 aclDestroyTensor(self); aclDestroyTensor(out); aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }以上示例输入{0, 1, 2, 3}理论输出分别为{0.5, 0.731059, 0.880797, 0.952574}保留 6 位小数可用于快速校验算子结果正确性。6.3 aclnnInplaceSigmoid 与普通版的差异点in-place 版本只需申请一个输入 Tensorself第一段调用acclnInplaceSigmoidGetWorkspaceSize(self, workspaceSize, executor)执行完成后直接从self对应的 Device 内存selfDeviceAddr拷回结果即可无需构造out资源释放时也无需释放输出 Tensor// 调用第一段接口 ret aclnnInplaceSigmoidGetWorkspaceSize(self, inplaceWorkspaceSize, inplaceExecutor); // 按需申请workspace并调用第二段接口 ret aclnnInplaceSigmoid(inplaceWorkspaceAddr, inplaceWorkspaceSize, inplaceExecutor, stream); // 同步后直接从selfDeviceAddr读取结果 ret aclrtMemcpy(inplaceResultData.data(), ..., selfDeviceAddr, ..., ACL_MEMCPY_DEVICE_TO_HOST);七、源码级实现剖析从 aclnn 到 Kernel 的完整链路7.1 接口层op_api类型适配与图编排op_api/aclnn_sigmoid.cpp 是 aclnn 接口的核心实现其执行编排逻辑集中在ExecSigmoidGetWorkspaceSizeL126-L172主要步骤为创建 OpExecutorCREATE_EXECUTOR()固定写法参数校验CheckParams/CheckInplaceParams对应上文 5.5 节的错误码空 Tensor 短路当self-IsEmpty()时直接返回 0 workspace由 Kernel 侧处理空 Tensor文档标注“支持空 Tensor”连续性处理l0op::Contiguous(self, ...)将非连续输入转换为连续 Tensor类型适配若输入类型不在 Kernel 输出支持列表中如 INT8/BOOL 等则用l0op::Cast将输入转换为目标计算类型Kernel 调用l0op::Sigmoid(selfCast, ...)执行真正的 Sigmoid 计算结果回写l0op::Cast将结果转换回out的数据类型再通过l0op::ViewCopy拷贝到out兼容非连续输出 Tensor。由此可见接口层通过 Cast Contiguous ViewCopy 的组合让 aclnn 接口能覆盖远比 Kernel 本身更广的数据类型与内存布局。7.2 算子分派层op_api/sigmoid.cppAI Core 与 AI CPU 双路径op_api/sigmoid.cpp 中定义了底层算子l0op::Sigmoid其分派策略非常清晰AICORE_DTYPE_SUPPORT_LIST {DT_FLOAT, DT_FLOAT16, DT_BF16}当输入类型为 FLOAT/FLOAT16/BF16 时走AI Core 路径SigmoidAiCore通过ADD_TO_LAUNCHER_LIST_AICORE宏把算子加入 AI Core 任务队列其余类型DOUBLE、COMPLEX 等走AI CPU 路径SigmoidAiCpu通过ADD_TO_LAUNCHER_LIST_AICPU宏下发。输出 Tensor 由executor-AllocTensor(self-GetViewShape(), self-GetDataType(), FORMAT_ND)按输入 shape 与类型自动分配因此输出 shape 与输入天然一致。7.3 Host 侧op_host算子注册、shape 推导与 Tiling算子定义op_host/sigmoid_def.cpp 通过OpDef声明输入x、输出y均支持DT_FLOAT16/DT_BF16/DT_FLOAT与ND格式并注册 ascend950、ascend350 两个平台的 AI Core 配置同时开启DynamicShapeSupportFlag(true)动态 shape 支持、PrecisionReduceFlag(true)等特性指定 Kernel 文件为sigmoid_apt。shape 推导op_host/sigmoid_infershape.cpp 直接复用InferShape4Elewise逐元素通用推导保证y.shape x.shape。Tiling 计算op_host/arch35/sigmoid_tiling_arch35.cpp 负责在 Host 侧根据输入 dtype 生成 TilingKey 与分块参数dtype 组合F16/F16、BF16/BF16、F32/F32分别映射到OP_KEY_1/2/3L54-L70并通过GenerateTilingKey组合成最终 TilingKey同时通过PlatformAscendC获取 AI Vector Core 数量GetCoreNumAiv与 UB 内存大小GetCoreMemSize(CoreMemType::UB)据此计算每个核分块与 UB 内循环次数。7.4 Kernel 侧op_kernel基于 TilingKey 的多态分派op_kernel/sigmoid_apt.cpp 是 AI Core 侧的 Kernel 入口__global__ __aicore__ void sigmoid通过三个宏定义的 TilingKey 在设备侧进行多态分派#define SIGMOID_F16_TILING_KEY 100000000000100 // x是float16y是float16 #define SIGMOID_BF16_TILING_KEY 200000000000100 // x是bfloat16y是bfloat16 #define SIGMOID_F32_TILING_KEY 300000000000100 // x是float32y是float32Kernel 主体仅执行 AI Vector CoreKERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY)根据 TilingKey 分别实例化SigmoidF16、SigmoidBf16、SigmoidF32三种计算算子。三种 dtype 的具体实现位于 op_kernel/arch35/ 目录sigmoid_f16.h / sigmoid_bf16.h / sigmoid_f32.h。以 op_kernel/arch35/sigmoid_f32.h 为例可以看到典型的 AscendC 向量算子流水设计Init阶段设置 GlobalTensor 的全局内存地址SetGlobalBuffer并按elemNum初始化输入/输出双缓冲队列DOUBLE_BUFFER 2即乒乓缓冲隐藏搬运延迟Process阶段按块号GetBlockIdx判断当前是否为最后一个核从而选取ubLoopOfFormerBlock/ubLoopOfTailBlock与ubTailOfFormerBlock/ubTailOfTailBlock做循环次数与尾块处理每个循环依次执行CopyIn0DataCopyPad从 GM 拷贝到 UB→Compute1向量计算→CopyOut2结果写回 GM计算过程使用__VEC_SCOPE__向量作用域与寄存器级RegTensor/MaskReg优化实现 Sigmoid 的向量化计算。该设计使 Sigmoid 在 AI Core 上具备多核并行 双缓冲流水 尾块处理的完整性能优化结构同时通过 Host 侧 Tiling 与设备侧 TilingKey 的配合支持动态 shape。八、测试与验证仓库为 Sigmoid 提供了完整的单元测试与系统测试体系是验证算子正确性的最佳参照单元测试ut位于 activation/sigmoid/tests/utop_api/test_aclnn_sigmoid.cppaclnn 接口级测试op_host/test_sigmoid_infershape.cppshape 推导测试op_host/arch35/test_sigmoid_tiling.cppTiling 计算测试op_kernel/test_sigmoid_apt.cppKernel 层测试。系统测试st位于 activation/sigmoid/tests/staclnnSigmoid/atk_aclnnSigmoid.json与executor_aclnnSigmoid.pyATKAscend Test Kit驱动的 aclnn 接口系统测试arch35/ttk_aclnn_sigmoid_st.csv、ttk_aclnn_inplace_sigmoid_st.csv、ttk_kernel_sigmoid_st.csvTTK 驱动的接口与 Kernel 测试用例表assets/golden.pyGolden 数据生成脚本用于按公式1 / (1 exp(-x))生成预期结果。测试数据从公式层面验证了算子输出与数学定义的一致性同时覆盖了普通接口、in-place 接口与 Kernel 三条调用路径可作为二次开发时回归测试的基线。九、总结CANN ops-nn 中的 Sigmoid 算子activation/sigmoid是一个“小而全”的典型逐元素算子功能上严格遵循out 1/(1e^{-input})产品上覆盖 Ascend 950 系列与 Atlas A2/A3/推理/训练全系列Atlas 200I/500 A2 除外调用上同时支持aclnnSigmoid、acclnInplaceSigmoid两段式接口与图模式 IR 构图实现上形成了aclnn 接口层类型适配/连续性处理→ l0op 分派层AI Core / AI CPU 双路径→ Host Tiling多核分块→ AI Core KernelTilingKey 多态分派 双缓冲流水的完整链路。对于开发者而言理解本算子至少有三个层面的价值一是掌握 aclnn 两段式接口的标准调用模板可直接迁移到仓库内其他逐元素算子如 gelu、relu 等二是理解“接口层宽类型 Kernel 层窄类型 Cast 桥接”的分层设计思想三是借鉴其 TilingKey 分派与双缓冲 Kernel 写法用于开发新的高性能 AI Core 算子。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-nn Sigmoid 算子深度指南：aclnn 两段式接口、AI Core 内核实现与图模式调用
CANN ops-nn Sigmoid 算子深度指南aclnn 两段式接口、AI Core 内核实现与图模式调用【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn本文以 CANN ops-nn 仓库中的 Sigmoid 算子的官方文档activation/sigmoid/README.md 与 aclnnSigmoidaclnnInplaceSigmoid.md为主线结合该算子模块的源码、样例与测试系统讲解 Sigmoid 在 NPU 上的功能定义、产品支持范围、aclnnSigmoid/aclnnInplaceSigmoid两段式调用方法、底层 AI Core 内核实现以及算子 IR 构图方式。读完本文你将能够独立完成 Sigmoid 算子在 Ascend 设备上的算子级开发与验证并理解从 aclnn 接口到 Tiling、再到 Kernel 的完整执行链路。一、功能说明与数学定义Sigmoid 是神经网络中最常用的激活函数之一将任意实数输入映射到 (0, 1) 开区间内常用于二分类输出层、LSTM/GRU 等循环神经网络的门控单元以及注意力机制的归一化环节。在本仓库中该算子位于activation/sigmoid目录其功能为对输入 Tensor 逐元素element-wise完成 Sigmoid 运算。计算公式为$$ out {\frac{1} {1{e}^{-input}}} $$该函数具有如下典型性质输出值域为 (0, 1)单调递增当输入趋近正无穷时输出趋近 1趋近负无穷时输出趋近 0关于原点中心对称sigmoid(-x) 1 - sigmoid(x)导数可用自身表示sigmoid(x) sigmoid(x) * (1 - sigmoid(x))这一性质使其在反向传播中计算高效也是 activation/sigmoid_grad 等梯度算子实现的基础。从算子原型看Sigmoid 的输入输出形状完全一致逐元素运算由 op_graph/sigmoid_proto.h 中的 IR 定义可见其输入x与输出y均使用TensorType::UnaryDataType()且该 IR 明确标注“Compatible with the TensorFlow operator Sigmoid”即与 TensorFlow 框架的 Sigmoid 算子语义对齐。二、产品支持情况根据官方文档Sigmoid 算子在各产品上的支持情况如下产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品√注以上为算子级支持情况aclnnSigmoid/aclnnInplaceSigmoid接口级支持情况与上表一致详见接口文档 aclnnSigmoidaclnnInplaceSigmoid.md。从源码可印证该支持范围在 op_host/sigmoid_def.cpp 中算子通过this-AICore().AddConfig(ascend950, aicoreConfig)与AddConfig(ascend350, aicoreConfig)注册了 ascend950、ascend350 两个平台的 AI Core 配置对应文档中的 Ascend 950 系列与 Atlas A2/A3 系列等而 Atlas 200I/500 A2 推理产品不在注册列表中因此标记为不支持。三、参数说明Sigmoid 算子仅有输入x和输出y两个 Tensor 参数具体规格如下参数名输入/输出描述数据类型数据格式x输入待进行 Sigmoid 计算的入参公式中的 inputFLOAT、FLOAT16、BFLOAT16NDy输出计算的出参FLOAT、FLOAT16、BFLOAT16ND需要说明的是上表为算子Kernel级支持的数据类型。在 aclnn 接口层通过 op_api/aclnn_sigmoid.cpp 中的DTYPE_SUPPORT_LIST可以看到接口输入self还额外支持DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOL、COMPLEX64、COMPLEX128等更宽泛的类型——这些类型会在接口层通过Cast算子统一转换后再送入 AI Core Kernel 计算这也是“算子层只支持三种浮点类型、而接口层支持更多类型”的原因所在。约束说明官方文档明确标注“无”额外约束。在实现层面唯一的形状约束是输入self与输出out的 shape 必须一致见下文参数校验逻辑。四、调用方式总览Sigmoid 算子在 CANN 中支持两种调用方式调用方式调用样例说明aclnn 调用test_aclnn_sigmoid.cpp通过aclnnSigmoid接口方式调用 Sigmoid 算子aclnn 调用test_aclnn_inplace_sigmoid.cpp通过aclnnInplaceSigmoid接口方式调用 Sigmoid 算子图模式-通过算子 IR 构图方式调用 Sigmoid 算子其中图模式对应的 IR 定义位于 op_graph/sigmoid_proto.h通过REG_OP(Sigmoid)注册算子原型Host 侧的形状推导在 op_host/sigmoid_infershape.cpp 中实现直接复用Ops::Base::InferShape4Elewise逐元素通用推断逻辑保证输出 shape 与输入一致。五、两段式 aclnn 接口详解acclnSigmoid与acclnInplaceSigmoid实现相同的计算功能区别仅在于结果存储方式aclnnSigmoid需新建一个输出张量对象out存储计算结果输入输出分离aclnnInplaceSigmoid无需新建输出张量对象直接在输入张量的内存中覆写计算结果in-place 操作节省一倍显存。与 CANN 其他算子一样这两个接口均采用两段式调用必须先调用第一段xxxGetWorkspaceSize接口获取 workspace 大小并完成入参校验、构建aclOpExecutor再调用第二段xxx接口在指定 Stream 上真正下发执行。两段式接口的通用原理可参考 docs/zh/context/basic_concept.md 中的 aclnn 两段式 API 说明。5.1 函数原型aclnnStatus aclnnSigmoidGetWorkspaceSize( const aclTensor *self, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor) aclnnStatus aclnnSigmoid( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream) aclnnStatus aclnnInplaceSigmoidGetWorkspaceSize( aclTensor* selfRef, uint64_t* workspaceSize, aclOpExecutor** executor) aclnnStatus aclnnInplaceSigmoid( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)5.2 aclnnSigmoidGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入待进行 Sigmoid 计算的入参公式中的 input支持空 Tensorshape 需要与 out 一致FLOAT、FLOAT16、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOL、COMPLEX64、COMPLEX128、BFLOAT16ND0-8√outaclTensor*输出计算的出参shape 需要与 self 一致FLOAT、FLOAT16、DOUBLE、COMPLEX64、COMPLEX128、BFLOAT16ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----5.3 aclnnInplaceSigmoidGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfRefaclTensor*输入/输出计算的入参结果原位覆写-FLOAT、FLOAT16、DOUBLE、COMPLEX64、COMPLEX128、BFLOAT16ND1-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----注意in-place 接口在部分老产品Atlas 推理系列产品、Atlas 训练系列产品上仅支持FLOAT、FLOAT16、DOUBLE、COMPLEX64、COMPLEX128不支持BFLOAT16。这一点在源码 op_api/aclnn_sigmoid.cpp 中有对应逻辑GetSelfRefDtypeList()会根据当前 NPU 架构DAV_2201或 RegBase 架构动态返回不同的 in-place 类型支持列表BF16 仅在 Ascend 950 及 Atlas A2/A3 系列等新架构上开放。5.4 两段接口aclnnSigmoid / aclnnInplaceSigmoid公共参数参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream5.5 返回值与错误码两段接口均返回aclnnStatus状态码。第一段接口会完成入参校验出现以下场景时报错aclnnSigmoidGetWorkspaceSize 校验错误返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self 或 out 是空指针ACLNN_ERR_PARAM_INVALID161002self 和 out 的数据类型和数据格式不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self 和 out 的 shape 不匹配aclnnInplaceSigmoidGetWorkspaceSize 校验错误返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 selfRef 是空指针ACLNN_ERR_PARAM_INVALID161002selfRef 的数据类型不在支持的范围之内这些校验在源码中的落点是 op_api/aclnn_sigmoid.cpp 的CheckParams/CheckInplaceParams函数依次检查空指针CheckNotNull、数据类型与 BF16 平台兼容性CheckDtypeValid其中OP_CHECK_DTYPE_NOT_SUPPORT映射到ACLNN_ERR_PARAM_INVALID161002、shape 一致性CheckShape要求 self 与 out 的 shape 完全一致且维度不超过 8。5.6 约束说明确定性计算aclnnSigmoid与aclnnInplaceSigmoid默认即为确定性实现同一输入多次执行结果一致无需额外配置。六、完整调用示例可直接运行下面以aclnnSigmoid为例给出完整的调用流程该示例可在 examples/test_aclnn_sigmoid.cpp 中找到完整可编译源码aclnnInplaceSigmoid的完整示例见 examples/test_aclnn_inplace_sigmoid.cpp。编译与运行环境准备可参考仓库文档中的样例编译说明。6.1 公共工具代码#include cinttypes #include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_sigmoid.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接口创建aclTensorND格式storageOffset为0 *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; }6.2 aclnnSigmoid 主流程非原地版int main() { // 1. 固定写法device/stream初始化参考acl APIdeviceId根据实际设备填写 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. 构造输入与输出 std::vectorint64_t selfShape {2, 2}; std::vectorint64_t outShape {2, 2}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 1, 2, 3}; std::vectorfloat outHostData {0, 0, 0, 0}; ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API两段式 uint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段获取workspace大小并完成入参校验 ret aclnnSigmoidGetWorkspaceSize(self, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSigmoidGetWorkspaceSize 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); } // 第二段在指定stream上执行计算 ret aclnnSigmoid(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSigmoid 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并打印 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[% PRId64 ] is: %f\n, i, resultData[i]); } // 6. 释放资源 aclDestroyTensor(self); aclDestroyTensor(out); aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }以上示例输入{0, 1, 2, 3}理论输出分别为{0.5, 0.731059, 0.880797, 0.952574}保留 6 位小数可用于快速校验算子结果正确性。6.3 aclnnInplaceSigmoid 与普通版的差异点in-place 版本只需申请一个输入 Tensorself第一段调用acclnInplaceSigmoidGetWorkspaceSize(self, workspaceSize, executor)执行完成后直接从self对应的 Device 内存selfDeviceAddr拷回结果即可无需构造out资源释放时也无需释放输出 Tensor// 调用第一段接口 ret aclnnInplaceSigmoidGetWorkspaceSize(self, inplaceWorkspaceSize, inplaceExecutor); // 按需申请workspace并调用第二段接口 ret aclnnInplaceSigmoid(inplaceWorkspaceAddr, inplaceWorkspaceSize, inplaceExecutor, stream); // 同步后直接从selfDeviceAddr读取结果 ret aclrtMemcpy(inplaceResultData.data(), ..., selfDeviceAddr, ..., ACL_MEMCPY_DEVICE_TO_HOST);七、源码级实现剖析从 aclnn 到 Kernel 的完整链路7.1 接口层op_api类型适配与图编排op_api/aclnn_sigmoid.cpp 是 aclnn 接口的核心实现其执行编排逻辑集中在ExecSigmoidGetWorkspaceSizeL126-L172主要步骤为创建 OpExecutorCREATE_EXECUTOR()固定写法参数校验CheckParams/CheckInplaceParams对应上文 5.5 节的错误码空 Tensor 短路当self-IsEmpty()时直接返回 0 workspace由 Kernel 侧处理空 Tensor文档标注“支持空 Tensor”连续性处理l0op::Contiguous(self, ...)将非连续输入转换为连续 Tensor类型适配若输入类型不在 Kernel 输出支持列表中如 INT8/BOOL 等则用l0op::Cast将输入转换为目标计算类型Kernel 调用l0op::Sigmoid(selfCast, ...)执行真正的 Sigmoid 计算结果回写l0op::Cast将结果转换回out的数据类型再通过l0op::ViewCopy拷贝到out兼容非连续输出 Tensor。由此可见接口层通过 Cast Contiguous ViewCopy 的组合让 aclnn 接口能覆盖远比 Kernel 本身更广的数据类型与内存布局。7.2 算子分派层op_api/sigmoid.cppAI Core 与 AI CPU 双路径op_api/sigmoid.cpp 中定义了底层算子l0op::Sigmoid其分派策略非常清晰AICORE_DTYPE_SUPPORT_LIST {DT_FLOAT, DT_FLOAT16, DT_BF16}当输入类型为 FLOAT/FLOAT16/BF16 时走AI Core 路径SigmoidAiCore通过ADD_TO_LAUNCHER_LIST_AICORE宏把算子加入 AI Core 任务队列其余类型DOUBLE、COMPLEX 等走AI CPU 路径SigmoidAiCpu通过ADD_TO_LAUNCHER_LIST_AICPU宏下发。输出 Tensor 由executor-AllocTensor(self-GetViewShape(), self-GetDataType(), FORMAT_ND)按输入 shape 与类型自动分配因此输出 shape 与输入天然一致。7.3 Host 侧op_host算子注册、shape 推导与 Tiling算子定义op_host/sigmoid_def.cpp 通过OpDef声明输入x、输出y均支持DT_FLOAT16/DT_BF16/DT_FLOAT与ND格式并注册 ascend950、ascend350 两个平台的 AI Core 配置同时开启DynamicShapeSupportFlag(true)动态 shape 支持、PrecisionReduceFlag(true)等特性指定 Kernel 文件为sigmoid_apt。shape 推导op_host/sigmoid_infershape.cpp 直接复用InferShape4Elewise逐元素通用推导保证y.shape x.shape。Tiling 计算op_host/arch35/sigmoid_tiling_arch35.cpp 负责在 Host 侧根据输入 dtype 生成 TilingKey 与分块参数dtype 组合F16/F16、BF16/BF16、F32/F32分别映射到OP_KEY_1/2/3L54-L70并通过GenerateTilingKey组合成最终 TilingKey同时通过PlatformAscendC获取 AI Vector Core 数量GetCoreNumAiv与 UB 内存大小GetCoreMemSize(CoreMemType::UB)据此计算每个核分块与 UB 内循环次数。7.4 Kernel 侧op_kernel基于 TilingKey 的多态分派op_kernel/sigmoid_apt.cpp 是 AI Core 侧的 Kernel 入口__global__ __aicore__ void sigmoid通过三个宏定义的 TilingKey 在设备侧进行多态分派#define SIGMOID_F16_TILING_KEY 100000000000100 // x是float16y是float16 #define SIGMOID_BF16_TILING_KEY 200000000000100 // x是bfloat16y是bfloat16 #define SIGMOID_F32_TILING_KEY 300000000000100 // x是float32y是float32Kernel 主体仅执行 AI Vector CoreKERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY)根据 TilingKey 分别实例化SigmoidF16、SigmoidBf16、SigmoidF32三种计算算子。三种 dtype 的具体实现位于 op_kernel/arch35/ 目录sigmoid_f16.h / sigmoid_bf16.h / sigmoid_f32.h。以 op_kernel/arch35/sigmoid_f32.h 为例可以看到典型的 AscendC 向量算子流水设计Init阶段设置 GlobalTensor 的全局内存地址SetGlobalBuffer并按elemNum初始化输入/输出双缓冲队列DOUBLE_BUFFER 2即乒乓缓冲隐藏搬运延迟Process阶段按块号GetBlockIdx判断当前是否为最后一个核从而选取ubLoopOfFormerBlock/ubLoopOfTailBlock与ubTailOfFormerBlock/ubTailOfTailBlock做循环次数与尾块处理每个循环依次执行CopyIn0DataCopyPad从 GM 拷贝到 UB→Compute1向量计算→CopyOut2结果写回 GM计算过程使用__VEC_SCOPE__向量作用域与寄存器级RegTensor/MaskReg优化实现 Sigmoid 的向量化计算。该设计使 Sigmoid 在 AI Core 上具备多核并行 双缓冲流水 尾块处理的完整性能优化结构同时通过 Host 侧 Tiling 与设备侧 TilingKey 的配合支持动态 shape。八、测试与验证仓库为 Sigmoid 提供了完整的单元测试与系统测试体系是验证算子正确性的最佳参照单元测试ut位于 activation/sigmoid/tests/utop_api/test_aclnn_sigmoid.cppaclnn 接口级测试op_host/test_sigmoid_infershape.cppshape 推导测试op_host/arch35/test_sigmoid_tiling.cppTiling 计算测试op_kernel/test_sigmoid_apt.cppKernel 层测试。系统测试st位于 activation/sigmoid/tests/staclnnSigmoid/atk_aclnnSigmoid.json与executor_aclnnSigmoid.pyATKAscend Test Kit驱动的 aclnn 接口系统测试arch35/ttk_aclnn_sigmoid_st.csv、ttk_aclnn_inplace_sigmoid_st.csv、ttk_kernel_sigmoid_st.csvTTK 驱动的接口与 Kernel 测试用例表assets/golden.pyGolden 数据生成脚本用于按公式1 / (1 exp(-x))生成预期结果。测试数据从公式层面验证了算子输出与数学定义的一致性同时覆盖了普通接口、in-place 接口与 Kernel 三条调用路径可作为二次开发时回归测试的基线。九、总结CANN ops-nn 中的 Sigmoid 算子activation/sigmoid是一个“小而全”的典型逐元素算子功能上严格遵循out 1/(1e^{-input})产品上覆盖 Ascend 950 系列与 Atlas A2/A3/推理/训练全系列Atlas 200I/500 A2 除外调用上同时支持aclnnSigmoid、acclnInplaceSigmoid两段式接口与图模式 IR 构图实现上形成了aclnn 接口层类型适配/连续性处理→ l0op 分派层AI Core / AI CPU 双路径→ Host Tiling多核分块→ AI Core KernelTilingKey 多态分派 双缓冲流水的完整链路。对于开发者而言理解本算子至少有三个层面的价值一是掌握 aclnn 两段式接口的标准调用模板可直接迁移到仓库内其他逐元素算子如 gelu、relu 等二是理解“接口层宽类型 Kernel 层窄类型 Cast 桥接”的分层设计思想三是借鉴其 TilingKey 分派与双缓冲 Kernel 写法用于开发新的高性能 AI Core 算子。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考