企业官网建设流程全解析

CANN ops-math Rsqrt 算子深度指南：aclnn 两段式接口调用、整型/浮点全类型支持与 AICore 核函数实现解析

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math

导读

本文围绕 CANN 开源数学算子库 ops-math 中新增的Rsqrt算子展开，系统讲解其功能定义、产品与数据类型支持范围、aclnn（Ascend CANN Lite Neural Network）两段式 API 的完整调用流程，并结合仓库源码剖析算子从 Host 侧算子定义、shape 推导、tiling 计算到 Device 侧 AICore 核函数执行的完整实现链路。读完本文，你将掌握在 Atlas A2/A3 系列产品上通过aclnnRsqrt与aclnnInplaceRsqrt编写可运行的 Rsqrt 推理程序，并理解该算子内部如何针对 float/bfloat16/整型/bool 等不同数据类型做差异化计算。

算子功能与计算公式

Rsqrt（Reciprocal Square Root，平方根倒数）算子对输入 Tensor 的每个元素执行"先开平方、再取倒数"的运算。在 README.md 中其功能被描述为"将数据进行开方并取倒数运算"，对应的数学表达式为：

$$ out = \frac{1}{\sqrt{input}} $$

该算子在深度学习中常用于归一化类网络结构（如 LayerNorm、BatchNorm 的反向传播中计算标准差倒数），属于逐元素（element-wise）的基础数学算子。在 docs/aclnnRsqrt&aclnnInplaceRsqrt.md 中给出的计算公式与上述一致，且明确out与input的 shape 保持一一对应。

支持的产品与数据类型

产品支持情况

根据 README.md 与 API 文档 docs/aclnnRsqrt&aclnnInplaceRsqrt.md，Rsqrt算子支持以下产品：

这一产品支持范围与算子源码中的平台配置相互印证：在 op_host/rsqrt_def.cpp 中，算子通过this->AICore().AddConfig("ascend910b")与this->AICore().AddConfig("ascend910_93")注册了 AICore 平台配置，分别对应 Atlas A2 系列与 Atlas A3 系列训练/推理产品。

数据类型与数据格式

README.md 中的原型信息表列出：

• 输入 x：tensor，数据类型支持float32, float16, bfloat16, int8, int16, int32, uint8, bool，数据格式仅支持ND；
• 输出 y：tensor，数据类型支持float32, float16, bfloat16, int8, int16, int32, uint8, bool，数据格式仅支持ND；
• 核函数名：rsqrt。

约束与限制为：x、y 的数据类型只支持上述列表，数据格式只支持 ND。在算子定义源码 op_host/rsqrt_def.cpp 中，Input("x")与Output("y")的DataType与Format列表与此完全一致，且UnknownShapeFormat同样限定为FORMAT_ND，说明该算子在静态与动态 shape 场景下均只走 ND 布局。API 文档进一步说明在 Atlas A2 与 Atlas A3 上，self与out均支持FLOAT、FLOAT16、UINT8、INT8、INT16、INT32、BOOL、BFLOAT16八种类型。

值得说明的是，从 贡献说明 可见，该算子的整型（int8/int16/int32/uint8/bool）支持是 2026/07/06 由贡献者 Nice_try 新增的，因此 README 中的产品支持型号最初仅列 Atlas A2，而 API 文档已同步更新为 A2/A3 双系列。

两段式 aclnn 接口：aclnnRsqrt 与 aclnnInplaceRsqrt

接口选型：何时用 inplace 版本

CANN 为Rsqrt提供了两套功能完全相同的接口，区别仅在于输出张量的管理方式（详见 docs/aclnnRsqrt&aclnnInplaceRsqrt.md）：

• aclnnRsqrt：需要用户新建一个输出张量对象out来存储计算结果，适合需要保留原始输入数据的场景；
• aclnnInplaceRsqrt：无需新建输出张量对象，直接在输入张量self的内存中覆盖写入计算结果，可省去额外的输出内存分配，适合输入数据不再被使用的场景（如就地归一化）。

对应的四个接口函数原型如下：

aclnnStatus aclnnRsqrtGetWorkspaceSize(const aclTensor *self, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor) aclnnStatus aclnnRsqrt(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream) aclnnStatus aclnnInplaceRsqrtGetWorkspaceSize(aclTensor* selfRef, uint64_t* workspaceSize, aclOpExecutor** executor) aclnnStatus aclnnInplaceRsqrt(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

两段式调用机制

每个 aclnn 算子都遵循 两段式接口 设计，Rsqrt也不例外：

1. 第一段（GetWorkspaceSize 接口）：完成入参校验，并根据当前输入 shape、数据类型等计算流程所需的 workspace 大小，返回workspaceSize与封装了算子计算流程的executor；
2. 第二段（执行接口）：用户依据第一段返回的workspaceSize在 Device 侧申请 workspace 内存，然后调用执行接口真正发起计算。

这两段接口的声明位于 op_api/aclnn_rsqrt.h，其中aclnnRsqrtGetWorkspaceSize的注释明确标注了"根据具体的计算流程，计算 workspace 大小"，且@domain为aclnn_math。

aclnnRsqrtGetWorkspaceSize 参数说明

该接口的参数（见 docs/aclnnRsqrt&aclnnInplaceRsqrt.md）如下：

返回值为aclnnStatus状态码，具体可参考 aclnn 返回码。

第一段接口的入参校验与报错场景（文档原文说明）：

161001(ACLNN_ERR_PARAM_NULLPTR)：1. 传入的 self 或 out 是空指针。 161002(ACLNN_ERR_PARAM_INVALID)：1. self 和 out 的数据类型和数据格式不在支持的范围之内。 2. self 和 out 的 shape 不匹配。 3. self 和 out 的数据类型不一致。

也就是说，非 inplace 版本要求self与out的 shape 一致、数据类型一致，且两者都必须在支持的数据类型与 ND 格式范围内，否则第一段接口就会直接返回错误，不会进入实际计算阶段。

aclnnRsqrt 执行接口参数说明

aclnnStatus aclnnRsqrt(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

返回值为aclnnStatus状态码。

aclnnInplaceRsqrt 系列接口

• aclnnInplaceRsqrtGetWorkspaceSize：唯一计算入参为selfRef（aclTensor*，公式中的 input），支持非连续 Tensor、ND 格式，shape 维度不大于 8，A2/A3 上数据类型支持 FLOAT、FLOAT16、UINT8、INT8、INT16、INT32、BOOL、BFLOAT16；其余workspaceSize、executor出参含义与非 inplace 版本相同。

  其校验报错场景为：

  161001(ACLNN_ERR_PARAM_NULLPTR)：1. 传入的 selfRef 是空指针。 161002(ACLNN_ERR_PARAM_INVALID)：1. selfRef 的数据类型和数据格式不在支持的范围之内。

• aclnnInplaceRsqrt：入参为workspace、workspaceSize、executor、stream，含义与aclnnRsqrt完全相同，计算结果直接写回selfRef对应的 Device 内存。

约束说明

• 确定性计算：aclnnRsqrt与aclnnInplaceRsqrt默认即为确定性实现（详见 确定性计算），即相同输入在多次运行中产生可复现的相同结果，这对调试与精度回归验证非常重要；
• 数据格式：仅支持 ND；
• shape 约束：非 inplace 版本要求self与outshape 一致；inplace 版本要求selfRefshape 维度不大于 8；
• inplace 语义：inplace 版本会覆盖输入内存，调用前需确认原输入数据不再需要。

完整调用示例：单测工程中的 test_aclnn_rsqrt.cpp

仓库在 examples/test_aclnn_rsqrt.cpp 提供了完整可参考的调用工程，同时 README.md 的运行验证章节也指向该文件，说明其是官方认可的调用方式（通过 aclnn 调用的方式调用 Rsqrt 算子）。其核心流程可分为 7 步：

1. device/stream 初始化：依次调用aclInit(nullptr)、aclrtSetDevice(deviceId)、aclrtCreateStream(&stream)；
2. 构造输入输出 aclTensor：通过aclrtMalloc申请 Device 侧内存，aclrtMemcpy将 Host 数据拷入，计算连续 tensor 的 strides 后调用aclCreateTensor(shape, dims, dataType, strides, 0, ACL_FORMAT_ND, shape, dims, deviceAddr)创建 aclTensor；
3. 调用第一段接口aclnnRsqrtGetWorkspaceSize(self, out, &workspaceSize, &executor)，若workspaceSize > 0则用aclrtMalloc申请 workspace；
4. 调用第二段接口aclnnRsqrt(workspaceAddr, workspaceSize, executor, stream)执行计算；
5. 同步等待：aclrtSynchronizeStream(stream)等待任务执行结束；
6. 回拷结果：aclrtMemcpy将 Device 侧结果拷贝回 Host 并打印；
7. 资源释放：aclDestroyTensor销毁 tensor，aclrtFree释放 device 内存与 workspace，最后aclrtDestroyStream、aclrtResetDevice、aclFinalize。

示例中的核心调用代码（含 inplace 版本，示例输入为{1, 2, 3, 4}的 2×2 FLOAT tensor）：

#include "acl/acl.h" #include "aclnnop/aclnn_rsqrt.h" uint64_t workspaceSize = 0; aclOpExecutor* executor; // aclnnRsqrt：非 inplace 版本 ret = aclnnRsqrtGetWorkspaceSize(self, out, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnRsqrtGetWorkspaceSize failed. ERROR: %d\n", ret); return ret); 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); } ret = aclnnRsqrt(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnRsqrt failed. ERROR: %d\n", ret); return ret); ret = aclrtSynchronizeStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret); // 回拷非 inplace 结果（out） aclrtMemcpy(resultData.data(), size * sizeof(float), outDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); // aclnnInplaceRsqrt：inplace 版本，结果写回 self ret = aclnnInplaceRsqrtGetWorkspaceSize(self, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnInplaceRsqrtGetWorkspaceSize failed. ERROR: %d\n", ret); return ret); 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); } ret = aclnnInplaceRsqrt(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnInplaceRsqrt failed. ERROR: %d\n", ret); return ret); ret = aclrtSynchronizeStream(stream); // 回拷 inplace 结果（此时数据在 self 的 device 内存中） aclrtMemcpy(resultData.data(), size * sizeof(float), selfDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST);

示例中创建连续 tensor 时 strides 的计算方式为：初始化全 1 后从倒数第二维向前累乘，即strides[i] = shape[i + 1] * strides[i + 1]，这与 CANN 对 ND 连续 tensor 的约定一致。工程对应的编译与运行方式可参考 编译与运行样例，算子的快速验证命令可参考 build.sh 调用方式。

源码级实现剖析

算子注册与 shape 推导

• 算子定义：op_host/rsqrt_def.cpp 通过class Rsqrt : public OpDef注册 OpType 为Rsqrt的算子，声明必选输入x与必选输出y，两者的 DataType 列表、ND 格式列表以及 UnknownShapeFormat 一一对应，并通过AddConfig("ascend910b")、AddConfig("ascend910_93")注册 AICore 平台；
• shape 推导：op_host/rsqrt_infershape.cpp 中的InferShapeRsqrt将输出 shape 直接赋值为输入 shape（*y_shape = *x1_shape），印证了 Rsqrt 是逐元素算子、输出与输入 shape 完全一致的语义，这也是第一段接口要求 self 与 out shape 一致的底层原因。

算子分发：AICore 与 AICPU 双路径

op_api/rsqrt.cpp 展示了算子在 aclnn 层的分发逻辑：

• 定义AICORE_DTYPE_SUPPORT_LIST（FLOAT、FLOAT16、BF16、INT8、INT16、INT32、UINT8、BOOL），与文档中八种支持类型完全对应；
• 当输入数据类型在 AICore 支持列表中时，走ADD_TO_LAUNCHER_LIST_AICORE(Rsqrt, ...)启动 AICore 核函数；否则（如其他未支持类型）走ADD_TO_LAUNCHER_LIST_AICPU(...)的 AICPU 兜底路径；
• 非 inplace 的Rsqrt函数通过executor->AllocTensor(self->GetViewShape(), self->GetDataType(), FORMAT_ND)自动分配输出张量，这也解释了为什么aclnnInplaceRsqrtGetWorkspaceSize只需传入selfRef一个计算入参。

同时 op_api/rsqrt.h 还声明了 level0 单算子接口l0op::Rsqrt(const aclTensor* self, aclOpExecutor* executor)，供更底层的调用场景使用。

Tiling 策略：UB 资源测算与负载均衡

op_host/rsqrt_tiling.cpp 实现了 Host 侧 tiling 计算，是理解该算子性能设计的关键：

• 平台信息获取：通过PlatformAscendC::GetCoreMemSize(UB, ubSize)获取片上 Unified Buffer 大小、GetCoreNumAiv()获取可用 AIV 核数；
• 按数据类型差异化测算：不同数据类型在 UB 中需要不同的缓冲区份数（DATA_NUM_INT32=2、DATA_NUM_FP32=3、DATA_NUM_BF16=6、DATA_NUM_INT8=6、DATA_NUM_UINT8=4、DATA_NUM_BOOL=2、DATA_NUM_INT16=2等），这是因为 bfloat16/int8 等类型在计算时需要额外的中间精度转换缓冲区（如 bfloat16 需转 float 计算后再转回）；
• buffer 数量动态决策：DetermineBufferNum根据单缓冲区需求与coreNum * ubSize的关系，决定采用单缓冲（SINGLE_BUFFER_NUM=1）还是双缓冲（DOUBLE_BUFFER_NUM=2），并在SetTilingKey中写入对应的 tiling key（0 对应双缓冲、1 对应单缓冲），核函数据此模板参数选择流水模式；
• 核间负载均衡：CalculateCoreBlockNums按块粒度将输入均分到各核，区分 big core 与 small core（前tailBlockNum个核多分一块），并把每个核的数据进一步切成多个 tile 与一个 tail，写入RsqrtTilingData；
• workspace 申请：GetWorkspaceSize通过GetLibApiWorkSpaceSize()获取系统库所需 workspace，写入context->GetWorkspaceSizes(1)（当前限制使用一块）。

tiling 数据结构的字段（op_kernel/rsqrt_tiling_data.h）包括smallCoreDataNum / bigCoreDataNum、finalBigTileNum / finalSmallTileNum、tileDataNum、smallTailDataNum / bigTailDataNum、tailBlockNum，核函数按核号与tailBlockNum比较选择大小核分支。tiling key 通过 op_kernel/rsqrt_tiling_key.h 中的ASCENDC_TPL_ARGS_DECL(Rsqrt, ...)声明了两种调度模式（ELEMENTWISE_TPL_SCH_MODE_0/1）。

核函数实现：按数据类型分派计算

op_kernel/rsqrt.cpp 是算子入口，声明了rsqrt(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)核函数，根据模板参数schMode选择KernelRsqrt的双缓冲（DOUBLE_BUFFER_NUM=2）或单缓冲（SINGLE_BUFFER_NUM=1）实例，并注册KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY)。

真正逐元素计算的逻辑在 op_kernel/rsqrt.h 的KernelRsqrt模板类中，其针对不同数据类型采用差异化的计算路径：

• half / float：Sqrt求平方根后，用预填充的oneLocal（Duplicate 填充 1.0）执行Div(1/x)；
• bfloat16：先Cast到 float 计算（避免 bf16 精度不足），Sqrt+Div后用CAST_RINT舍入回写 bf16；
• uint8 / int8 / int16 / int32：先 Cast 到 half 或 float 中间精度，调用向量指令Rsqrt，部分整型还经过ComputeIntCorrection的修正（p1 = min(p1,2) - 3*max(0, min(p1,2)-1)，将结果约束到 [-1, 2] 区间），最后CAST_RINT转回整型；
• bool：直接通过Duplicate<int16_t>(p1, 0x0101, ...)将结果填充为 1（因为rsqrt(true)=rsqrt(1)=1），并在 outQueueY 缓冲上按 int16 对齐处理奇偶数量，避免越界写。

在内存流水上，KernelRsqrt采用经典的CopyIn -> Compute -> CopyOut三段式，通过TQue<QuePosition::VECIN/ VECOUT, BUFFER_NUM>队列实现输入/输出缓冲的乒乓切换；Init阶段依据 bool 与非 bool 类型差异化申请 out 队列缓冲（bool 按((tileDataNum+1)>>1) * sizeof(int16_t)申请），这也是 tiling 阶段按类型区分DATA_NUM_*的原因。

测试覆盖：逐数据类型 UT

算子配套了完整的单测工程（tests/ut/op_kernel/），按数据类型拆分为test_rsqrt.cpp（float 主用例）、test_rsqrt_bf16.cpp、test_rsqrt_bool.cpp、test_rsqrt_fp16.cpp、test_rsqrt_int16.cpp、test_rsqrt_int32.cpp、test_rsqrt_int8.cpp、test_rsqrt_uint8.cpp，与文档宣称的八种数据类型逐一对应；rsqrt_test_base.h 提供共享的测试基类，并在 SetUp 时拷贝rsqrt_data目录下的预生成数据。此外，tests/ut/op_host/ 下还有test_rsqrt_infershape.cpp（验证 shape 推导）与test_rsqrt_tiling.cpp（验证 tiling 计算），tests/ut/op_api/ 下则有test_aclnn_rsqrt.cpp覆盖 aclnn 接口层，从算子定义、tiling、核函数到 aclnn 接口形成了完整的验证闭环。

贡献记录

小结

Rsqrt是 ops-math 中一个典型的逐元素数学算子：对外提供aclnnRsqrt/aclnnInplaceRsqrt两套两段式 aclnn 接口，覆盖 Atlas A2/A3 系列产品与八种数据类型；对内则通过"算子定义 → shape 推导 → tiling 计算 → AICore 核函数"的标准链路实现，其中按类型区分的 UB 缓冲测算、单/双缓冲动态切换、int8 修正逻辑等细节均可在 experimental/math/rsqrt 目录下的源码与测试中逐一找到依据，可作为理解 CANN 数学算子开发范式与贡献新算子的参考模板。

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math