CANN ops-math 开源仓 Sqrt 算子深度解析从 aclnnSqrt 接口调用到 Ascend C 内核实现【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本文以 experimental/math/sqrt/README.md 为主线系统讲解 CANN ops-math 开源仓库中 Sqrt 算子平方根算子的功能定义、支持的产品与数据类型、编译部署与算子调用方式并深入源码剖析其 Ascend C 内核实现、Tiling 分块策略与 aclnnSqrt 两段式接口的完整调用链路。读完本文你将掌握如何在 Atlas 系列 NPU 上编译、部署并调用 Sqrt 算子同时理解其从 Host 侧算子定义到 Device 侧向量计算的全流程设计。一、Sqrt 算子概览1.1 功能描述Sqrt算子是 CANN ops-math 提供的数学类基础算子之一功能是返回输入数据经过开方运算平方根计算后的结果。其计算公式为$$ y \sqrt{x} $$在图像处理、信号处理、物理模拟等工程领域平方根运算是一种基础且高频的操作。Sqrt 算子能够在 NPU 上高效批量处理数值的平方根计算支持浮点数输入。1.2 贡献信息该算子由社区贡献者合入开源仓库具体贡献信息如下表贡献者贡献方贡献算子贡献时间贡献内容Nice_try西北工业大学-智能感知交互实验室Sqrt2024/12/24新增Sqrt算子liuxiqiang个人开发者Sqrt2025/11/12Sqrt算子适配开源仓1.3 支持的产品型号根据 experimental/math/sqrt/README.md 与 experimental/math/sqrt/docs/aclnnSqrt.md 的说明Sqrt 算子支持的产品型号为Atlas A2 训练系列产品Atlas 200I/500 A2 推理产品这一产品适配信息在算子注册源码 op_host/sqrt_def.cpp 中也有直接对应即通过this-AICore().AddConfig(ascend910b).AddConfig(ascend310b)为 ascend910bAtlas A2 训练系列与 ascend310b 两个芯片平台注册 AICore 配置op_host/config/目录下也分别维护了 ascend910b/sqrt_binary.json 与 ascend310b 两套算子二进制描述文件。二、算子原型与约束限制2.1 原型信息Sqrt 算子的原型定义如下算子类型(OpType)Sqrt算子输入xtensor数据类型 float32/float16/bfloat16数据格式 ND算子输出ytensor数据类型 float32/float16/bfloat16数据格式 ND核函数名sqrt2.2 约束与限制x、y、out 的数据类型仅支持 float32、float16、bfloat16数据格式仅支持 ND普通稠密格式即按行连续存储的非排他格式可参考 docs/zh/context/data_format.md 了解 ND 格式定义张量维度支持 0~8 维且支持非连续 Tensornon-contiguous tensor相关概念可参考 docs/zh/context/non_contiguous_tensor.md输入 x 与输出 y 的 shape 必须一致。2.3 源码中的原型注册上述原型信息在算子注册源码 op_host/sqrt_def.cpp 中有完整的实现对应this-Input(x) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); this-Output(y) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16}) .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});同时该注册代码通过OpAICoreConfig声明了算子的动态能力支持动态编译DynamicCompileStaticFlag(true)、动态 shapeDynamicShapeSupportFlag(true)、动态 rankDynamicRankSupportFlag(true)并设置PrecisionReduceFlag(true)开启精度降低优化。此外算子二进制描述文件 ascend910b/sqrt_binary.json 中为 float32、float16、bfloat16 三种数据类型各生成了一份二进制配置Sqrt_a1532827238e1555db7b997c7bce2928/2929/2930shape 均声明为-2动态 shape印证了算子对动态 shape 的支持。三、环境准备、编译部署与算子调用3.1 开发运行环境部署使用该算子前需要先完成开发运行环境的部署请参考社区版CANN开发套件包安装文档。环境部署完成后通过以下命令获取开源仓代码并进入仓库目录git clone https://gitcode.com/cann/ops-math cd ${git_clone_path}/ops-math说明构建命令以当前仓库实际提供的构建脚本为准仓库根目录与 scripts 目录下提供了build_cmake.sh、build_lib.sh、build_example.sh、build_ut.sh等构建辅助脚本也可参考 docs/zh/install/compile.md 了解整体编译流程。3.2 编译算子包在仓库目录下执行编译只编译 Sqrt 算子对应的 experimental 包bash build.sh --pkg --experimental --socascend910b --opssqrt其中关键参数含义--pkg执行打包流程产出可安装的算子包--experimental指定构建 experimental 目录下的算子Sqrt 位于 experimental/math/sqrt--socascend910b指定目标 SoC 型号为 ascend910b对应 Atlas A2 训练系列产品--opssqrt指定本次仅编译 Sqrt 算子缩小构建范围、加快编译速度。3.3 部署算子包编译完成后在仓库根目录的build_out目录下会生成安装包执行部署./build_out/cann-ops-vendor_name-linux.arch.run其中vendor_name与arch为实际构建时指定的 vendor 名称与目标架构例如 x86_64、aarch64。3.4 执行算子调用样例算子包部署完成后可直接运行仓库内置的调用样例进行验证bash build.sh --run_example sqrt eager cust --vendor_namecustom该命令会运行 Sqrt 算子的调用示例eager 调用模式其中--vendor_namecustom与部署算子包时的 vendor 名称保持一致。调用示例对应的完整代码见 experimental/math/sqrt/examples/test_aclnn_sqrt.cpp。3.5 调用方式总览Sqrt 算子提供 aclnn 接口调用方式具体如下调用方式调用样例说明aclnn调用test_aclnn_sqrt.cpp通过 aclnnSqrt 接口方式调用Sqrt算子四、aclnnSqrt 两段式接口与调用示例4.1 两段式接口机制aclnnSqrt 遵循 CANN 算子库通用的两段式接口设计必须先调用aclnnSqrtGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器再调用aclnnSqrt执行实际计算。第一段接口获取 workspace 大小与执行器aclnnStatus aclnnSqrtGetWorkspaceSize( const aclTensor *self, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **opExecutor)第二段接口执行计算aclnnStatus aclnnSqrt( void *workspace, uint64_t workspaceSize, aclOpExecutor *opExecutor, const aclrtStream stream)4.2 aclnnSqrtGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorself输入待进行 sqrt 计算的入参公式中的 self无FLOAT、FLOAT16、BFLOAT16ND0-8√out输出待进行 sqrt 计算的出参公式中的 outshape 与 self 相同FLOAT、FLOAT16、BFLOAT16ND0-8√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小-----opExecutor输出返回 op 执行器包含了算子计算流程-----4.3 aclnnSqrt 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnSqrtGetWorkspaceSize 获取opExecutor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream4.4 返回值与错误码两个接口均返回aclnnStatus状态码具体含义参见 aclnn 返回码。第一段接口aclnnSqrtGetWorkspaceSize会完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 tensor 是空指针ACLNN_ERR_PARAM_INVALID161002self 的数据类型和数据格式不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self 的数据维度超过了 8 维ACLNN_ERR_PARAM_INVALID161002self 和 out 的数据形状不一致4.5 完整调用示例解读experimental/math/sqrt/examples/test_aclnn_sqrt.cpp 给出了一个完整的、可直接编译运行的调用示例其流程可拆解为七个步骤初始化资源调用aclInit、aclrtSetDevice、aclrtCreateStream完成 ACL 环境与 Stream 初始化示例中 deviceId 为 0构造输入输出 aclTensor示例输入为 shape{2, 2}的 float32 张量{0, 1, 2, 3}输出初始为{0, 0, 0, 0}。通过aclrtMalloc申请 Device 侧内存、aclrtMemcpy将 Host 数据拷贝到 Device、aclCreateTensor创建ACL_FORMAT_ND格式的 aclTensor调用第一段接口aclnnSqrtGetWorkspaceSize(self, out, workspaceSize, executor)获取 workspace 大小与执行器申请 workspace 并调用第二段接口若workspaceSize 0则aclrtMalloc申请 workspace 内存随后调用aclnnSqrt(workspaceAddr, workspaceSize, executor, stream)执行计算同步等待aclrtSynchronizeStream(stream)等待任务执行结束保证结果可读回拷结果通过aclrtMemcpy将 Device 侧结果拷回 Host 并逐元素打印对于输入{0, 1, 2, 3}预期输出约为{0, 1, 1.414, 1.732}释放资源依次释放 aclTensor、Device 内存、Stream并调用aclrtResetDevice与aclFinalize收尾。该示例中的CHECK_RET宏统一封装了错误检查与日志打印是编写 CANN 算子调用程序的推荐范式。五、实现原理Ascend C 内核源码剖析5.1 整体实现思路根据 experimental/math/sqrt/docs/aclnnSqrt.md 的实现原理说明Sqrt 算子直接调用 Ascend C 的API接口Sqrt实现平方根计算对于 16 位数据类型float16、bfloat16会先通过Cast接口转换为 32 位浮点数进行计算计算完成后再转回原类型输出。5.2 核函数入口与模板调度内核入口位于 op_kernel/sqrt.cppenum class SqrtTilingKey : uint32_t { TILING_KEY_EXAMPLE_FLOAT 0, TILING_KEY_EXAMPLE_OTHER 1, }; template uint32_t schMode __global__ __aicore__ void sqrt(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) { REGISTER_TILING_DEFAULT(SqrtTilingData); GET_TILING_DATA_WITH_STRUCT(SqrtTilingData, tilingData, tiling); MySqrt::KernelSqrtDTYPE_X, DTYPE_Y op; op.Init(...); // 算子kernel实例初始化 op.Process(); }核函数签名包含四个 GM 地址参数输入x、输出y、workspace 与 tiling 参数通过GET_TILING_DATA_WITH_STRUCT从 tiling 地址还原 Host 侧下发的分块参数然后实例化MySqrt::KernelSqrt模板类并依次执行Init与Process。模板参数DTYPE_X、DTYPE_Y由编译期根据算子二进制配置见 sqrt_binary.json实例化schMode则由 op_kernel/sqrt_tiling_key.h 中声明的 TilingKey 模板参数机制ASCENDC_TPL_ARGS_DECL/ASCENDC_TPL_SEL提供调度选择。5.3 流水线设计与双缓冲内核类KernelSqrt定义在 op_kernel/sqrt.h其核心设计包括TPipe 流水线使用AscendC::TPipe pipe管理 Vector 流水线双缓冲队列BUFFER_NUM 2输入队列inQueueX与输出队列outQueueY均为TQueVECIN/VECOUT, 2双缓冲队列实现搬入CopyIn与计算Compute的重叠隐藏数据搬运延迟临时缓冲区TBufVECCALC tmp1仅在对非 float32 类型计算时使用用于存放 Cast 到 float 后的中间结果。Process()sqrt.h按主循环 尾块的方式驱动流水线对前loopCount - 1个 tile 执行完整的CopyIn - Compute - CopyOut最后一个 tile 以tailDataNum尾块元素数作为处理长度处理剩余数据。5.4 Compute 计算逻辑Computesqrt.h实现了两种计算路径float32 路径直接调用向量指令AscendC::Sqrt(yLocal, xLocal, processDataNum)完成开方16 位类型路径先AscendC::Cast(p1, xLocal, RoundMode::CAST_NONE, ...)将 float16/bfloat16 提升为 float接着对 float 中间结果执行AscendC::Sqrt最后再以RoundMode::CAST_RINT四舍五入模式Cast回原 16 位类型输出。两条路径之间通过PipeBarrierPIPE_V保证向量指令的执行顺序。CopyIn/CopyOut则分别使用AscendC::DataCopy完成 GMGlobal Memory与 Local Tensor 之间的批量搬入、搬出搬移长度为processDataNum个元素。六、Tiling 分块与 Shape 推导6.1 TilingData 结构Host 侧 tiling 计算下发的参数结构定义在 op_kernel/sqrt_tiling_data.hstruct SqrtTilingData { uint64_t smallCoreDataNum; uint64_t bigCoreDataNum; uint64_t finalBigTileNum; uint64_t finalSmallTileNum; uint64_t tileDataNum; uint64_t smallTailDataNum; uint64_t bigTailDataNum; uint64_t tailBlockNum; };可以看到 tiling 参数将核心区分为大数据量核与小数据量核两类big/small 前缀并分别给出整 tile 数、每个核处理的数据量与尾块数据量用于在核间做负载均衡。6.2 Tiling 计算流程Tiling 主流程位于 op_host/sqrt_tiling.cpp整体分为四步获取平台信息GetPlatformInfo通过platform_ascendc::PlatformAscendC获取 UBUnified Buffer大小与可用核数coreNum获取输入 shape 信息GetShapeAttrsInfo读取输入存储 shape 的元素总数与数据类型长度得到输入总字节数以 32 字节BLOCK_SIZE 32U为对齐粒度计算对齐后长度并按数据类型确定 UB 内可容纳的分块数float32 对应UB_DATA_NUM_FLOAT 4U其他类型对应UB_DATA_NUM_OTHER 6U从而算出每个 tile 的元素数tileDataNum计算各核分块CalculateCoreBlockNums按每个核至少分配 32 字节数据的原则确定实际使用的核数再计算大小核各自的数据量、tile 数与尾块数保证负载均衡下发 tiling将计算结果写入SqrtTilingData通过GetWorkspaceSize计算 workspace 大小sysWorkspaceSize来自ascendcPlatform.GetLibApiWorkSpaceSize()最后context-SetTilingKey(0)并context-SetBlockDim(coreNum)设置核函数调度维度。该 Host 侧逻辑通过IMPL_OP_OPTILING(Sqrt).Tiling(SqrtTilingFunc)注册配套的单测见 tests/ut/op_host/test_sqrt_tiling.cpp。6.3 Shape 推导op_host/sqrt_infershape.cpp 中InferShapeSqrt实现了 shape 推导读取输入 x 的 shape 后直接将输出 y 的 shape 赋值为相同值*yShape *xShape并通过IMPL_OP_INFERSHAPE(Sqrt)完成注册。这对应了文档中out 的 shape 与 self 相同的约束。七、测试与验证仓库为 Sqrt 算子提供了 Host 侧与 Kernel 侧两级单元测试Kernel 侧测试tests/ut/op_kernel/test_sqrt.cpp 基于 gtest 与tikicpulib.h在 CPU 上模拟执行核函数通过SqrtCompileInfo如{64, 262144, false}分别对应核数、数据量等编译信息构造 tiling 上下文覆盖 float16、float32、bfloat16 等不同数据类型的用例测试数据由 tests/ut/op_kernel/sqrt_data/gen_data.py 生成并由 compare_data.py 与参考结果对比校验Host 侧测试tests/ut/op_host/test_sqrt_tiling.cpp 与 test_sqrt_infershape.cpp 分别验证 tiling 参数计算与 shape 推导的正确性。八、总结Sqrt 算子是 CANN ops-math 仓库中一个典型的小而完整的向量类算子样例对外通过aclnnSqrt两段式接口提供标准调用方式对内则由算子定义sqrt_def.cpp、shape 推导sqrt_infershape.cpp、tiling 计算sqrt_tiling.cpp与 Ascend C 内核sqrt.cpp / sqrt.h四部分协同完成。其16 位类型先 Cast 到 float 再计算的实现策略、基于 TPipe 双缓冲的流水线设计、以及按核数均衡切分数据块的 tiling 策略对于理解 CANN 数学类算子的通用开发范式具有很好的参考价值。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考