CANN ops-math Rsqrt 算子深度指南aclnn 两段式接口调用、整型/浮点全类型支持与 AICore 核函数实现解析【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读本文围绕 CANN 开源数学算子库 ops-math 中新增的Rsqrt算子展开系统讲解其功能定义、产品与数据类型支持范围、aclnnAscend CANN Lite Neural Network两段式 API 的完整调用流程并结合仓库源码剖析算子从 Host 侧算子定义、shape 推导、tiling 计算到 Device 侧 AICore 核函数执行的完整实现链路。读完本文你将掌握在 Atlas A2/A3 系列产品上通过aclnnRsqrt与aclnnInplaceRsqrt编写可运行的 Rsqrt 推理程序并理解该算子内部如何针对 float/bfloat16/整型/bool 等不同数据类型做差异化计算。算子功能与计算公式RsqrtReciprocal Square Root平方根倒数算子对输入 Tensor 的每个元素执行先开平方、再取倒数的运算。在 README.md 中其功能被描述为将数据进行开方并取倒数运算对应的数学表达式为$$ out \frac{1}{\sqrt{input}} $$该算子在深度学习中常用于归一化类网络结构如 LayerNorm、BatchNorm 的反向传播中计算标准差倒数属于逐元素element-wise的基础数学算子。在 docs/aclnnRsqrtaclnnInplaceRsqrt.md 中给出的计算公式与上述一致且明确out与input的 shape 保持一一对应。支持的产品与数据类型产品支持情况根据 README.md 与 API 文档 docs/aclnnRsqrtaclnnInplaceRsqrt.mdRsqrt算子支持以下产品产品是否支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√这一产品支持范围与算子源码中的平台配置相互印证在 op_host/rsqrt_def.cpp 中算子通过this-AICore().AddConfig(ascend910b)与this-AICore().AddConfig(ascend910_93)注册了 AICore 平台配置分别对应 Atlas A2 系列与 Atlas A3 系列训练/推理产品。数据类型与数据格式README.md 中的原型信息表列出输入 xtensor数据类型支持float32, float16, bfloat16, int8, int16, int32, uint8, bool数据格式仅支持ND输出 ytensor数据类型支持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/aclnnRsqrtaclnnInplaceRsqrt.mdaclnnRsqrt需要用户新建一个输出张量对象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也不例外第一段GetWorkspaceSize 接口完成入参校验并根据当前输入 shape、数据类型等计算流程所需的 workspace 大小返回workspaceSize与封装了算子计算流程的executor第二段执行接口用户依据第一段返回的workspaceSize在 Device 侧申请 workspace 内存然后调用执行接口真正发起计算。这两段接口的声明位于 op_api/aclnn_rsqrt.h其中aclnnRsqrtGetWorkspaceSize的注释明确标注了根据具体的计算流程计算 workspace 大小且domain为aclnn_math。aclnnRsqrtGetWorkspaceSize 参数说明该接口的参数见 docs/aclnnRsqrtaclnnInplaceRsqrt.md如下参数类型说明selfconst aclTensor*计算输入公式中的 inputDevice 侧 aclTensor支持非连续的 Tensor数据格式支持 NDshape 需与 out 一致A2/A3 上支持 FLOAT、FLOAT16、UINT8、INT8、INT16、INT32、BOOL、BFLOAT16outaclTensor*计算输出公式中的 outDevice 侧 aclTensor支持非连续 Tensor格式支持 NDshape 需与 self 一致数据类型支持范围与 self 相同workspaceSizeuint64_t*出参返回需要在 Device 侧申请的 workspace 大小executoraclOpExecutor**出参返回 op 执行器包含算子计算流程返回值为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)参数类型说明workspacevoid*入参在 Device 侧申请的 workspace 内存地址workspaceSizeuint64_t入参在 Device 侧申请的 workspace 大小由第一段接口aclnnRsqrtGetWorkspaceSize获取executoraclOpExecutor*入参op 执行器包含算子计算流程streamaclrtStream入参指定执行任务的 Stream返回值为aclnnStatus状态码。aclnnInplaceRsqrt 系列接口aclnnInplaceRsqrtGetWorkspaceSize唯一计算入参为selfRefaclTensor*公式中的 input支持非连续 Tensor、ND 格式shape 维度不大于 8A2/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默认即为确定性实现详见 确定性计算即相同输入在多次运行中产生可复现的相同结果这对调试与精度回归验证非常重要数据格式仅支持 NDshape 约束非 inplace 版本要求self与outshape 一致inplace 版本要求selfRefshape 维度不大于 8inplace 语义inplace 版本会覆盖输入内存调用前需确认原输入数据不再需要。完整调用示例单测工程中的 test_aclnn_rsqrt.cpp仓库在 examples/test_aclnn_rsqrt.cpp 提供了完整可参考的调用工程同时 README.md 的运行验证章节也指向该文件说明其是官方认可的调用方式通过 aclnn 调用的方式调用 Rsqrt 算子。其核心流程可分为 7 步device/stream 初始化依次调用aclInit(nullptr)、aclrtSetDevice(deviceId)、aclrtCreateStream(stream)构造输入输出 aclTensor通过aclrtMalloc申请 Device 侧内存aclrtMemcpy将 Host 数据拷入计算连续 tensor 的 strides 后调用aclCreateTensor(shape, dims, dataType, strides, 0, ACL_FORMAT_ND, shape, dims, deviceAddr)创建 aclTensor调用第一段接口aclnnRsqrtGetWorkspaceSize(self, out, workspaceSize, executor)若workspaceSize 0则用aclrtMalloc申请 workspace调用第二段接口aclnnRsqrt(workspaceAddr, workspaceSize, executor, stream)执行计算同步等待aclrtSynchronizeStream(stream)等待任务执行结束回拷结果aclrtMemcpy将 Device 侧结果拷贝回 Host 并打印资源释放aclDestroyTensor销毁 tensoraclrtFree释放 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); // aclnnInplaceRsqrtinplace 版本结果写回 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_LISTFLOAT、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_INT322、DATA_NUM_FP323、DATA_NUM_BF166、DATA_NUM_INT86、DATA_NUM_UINT84、DATA_NUM_BOOL2、DATA_NUM_INT162等这是因为 bfloat16/int8 等类型在计算时需要额外的中间精度转换缓冲区如 bfloat16 需转 float 计算后再转回buffer 数量动态决策DetermineBufferNum根据单缓冲区需求与coreNum * ubSize的关系决定采用单缓冲SINGLE_BUFFER_NUM1还是双缓冲DOUBLE_BUFFER_NUM2并在SetTilingKey中写入对应的 tiling key0 对应双缓冲、1 对应单缓冲核函数据此模板参数选择流水模式核间负载均衡CalculateCoreBlockNums按块粒度将输入均分到各核区分 big core 与 small core前tailBlockNum个核多分一块并把每个核的数据进一步切成多个 tile 与一个 tail写入RsqrtTilingDataworkspace 申请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_NUM2或单缓冲SINGLE_BUFFER_NUM1实例并注册KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY)。真正逐元素计算的逻辑在 op_kernel/rsqrt.h 的KernelRsqrt模板类中其针对不同数据类型采用差异化的计算路径half / floatSqrt求平方根后用预填充的oneLocalDuplicate 填充 1.0执行Div(1/x)bfloat16先Cast到 float 计算避免 bf16 精度不足SqrtDiv后用CAST_RINT舍入回写 bf16uint8 / int8 / int16 / int32先 Cast 到 half 或 float 中间精度调用向量指令Rsqrt部分整型还经过ComputeIntCorrection的修正p1 min(p1,2) - 3*max(0, min(p1,2)-1)将结果约束到 [-1, 2] 区间最后CAST_RINT转回整型bool直接通过Duplicateint16_t(p1, 0x0101, ...)将结果填充为 1因为rsqrt(true)rsqrt(1)1并在 outQueueY 缓冲上按 int16 对齐处理奇偶数量避免越界写。在内存流水上KernelRsqrt采用经典的CopyIn - Compute - CopyOut三段式通过TQueQuePosition::VECIN/ VECOUT, BUFFER_NUM队列实现输入/输出缓冲的乒乓切换Init阶段依据 bool 与非 bool 类型差异化申请 out 队列缓冲bool 按((tileDataNum1)1) * sizeof(int16_t)申请这也是 tiling 阶段按类型区分DATA_NUM_*的原因。测试覆盖逐数据类型 UT算子配套了完整的单测工程tests/ut/op_kernel/按数据类型拆分为test_rsqrt.cppfloat 主用例、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 接口形成了完整的验证闭环。贡献记录贡献者贡献方贡献算子贡献时间贡献内容skywang2个人开发者Rsqrt2026/03/30新增 Rsqrt 算子Nice_try个人开发者Rsqrt2026/07/06Rsqrt 新增支持整型数据支持A2/A3小结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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
CANN ops-math Rsqrt 算子深度指南：aclnn 两段式接口调用、整型/浮点全类型支持与 AICore 核函数实现解析
CANN ops-math Rsqrt 算子深度指南aclnn 两段式接口调用、整型/浮点全类型支持与 AICore 核函数实现解析【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读本文围绕 CANN 开源数学算子库 ops-math 中新增的Rsqrt算子展开系统讲解其功能定义、产品与数据类型支持范围、aclnnAscend CANN Lite Neural Network两段式 API 的完整调用流程并结合仓库源码剖析算子从 Host 侧算子定义、shape 推导、tiling 计算到 Device 侧 AICore 核函数执行的完整实现链路。读完本文你将掌握在 Atlas A2/A3 系列产品上通过aclnnRsqrt与aclnnInplaceRsqrt编写可运行的 Rsqrt 推理程序并理解该算子内部如何针对 float/bfloat16/整型/bool 等不同数据类型做差异化计算。算子功能与计算公式RsqrtReciprocal Square Root平方根倒数算子对输入 Tensor 的每个元素执行先开平方、再取倒数的运算。在 README.md 中其功能被描述为将数据进行开方并取倒数运算对应的数学表达式为$$ out \frac{1}{\sqrt{input}} $$该算子在深度学习中常用于归一化类网络结构如 LayerNorm、BatchNorm 的反向传播中计算标准差倒数属于逐元素element-wise的基础数学算子。在 docs/aclnnRsqrtaclnnInplaceRsqrt.md 中给出的计算公式与上述一致且明确out与input的 shape 保持一一对应。支持的产品与数据类型产品支持情况根据 README.md 与 API 文档 docs/aclnnRsqrtaclnnInplaceRsqrt.mdRsqrt算子支持以下产品产品是否支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√这一产品支持范围与算子源码中的平台配置相互印证在 op_host/rsqrt_def.cpp 中算子通过this-AICore().AddConfig(ascend910b)与this-AICore().AddConfig(ascend910_93)注册了 AICore 平台配置分别对应 Atlas A2 系列与 Atlas A3 系列训练/推理产品。数据类型与数据格式README.md 中的原型信息表列出输入 xtensor数据类型支持float32, float16, bfloat16, int8, int16, int32, uint8, bool数据格式仅支持ND输出 ytensor数据类型支持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/aclnnRsqrtaclnnInplaceRsqrt.mdaclnnRsqrt需要用户新建一个输出张量对象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也不例外第一段GetWorkspaceSize 接口完成入参校验并根据当前输入 shape、数据类型等计算流程所需的 workspace 大小返回workspaceSize与封装了算子计算流程的executor第二段执行接口用户依据第一段返回的workspaceSize在 Device 侧申请 workspace 内存然后调用执行接口真正发起计算。这两段接口的声明位于 op_api/aclnn_rsqrt.h其中aclnnRsqrtGetWorkspaceSize的注释明确标注了根据具体的计算流程计算 workspace 大小且domain为aclnn_math。aclnnRsqrtGetWorkspaceSize 参数说明该接口的参数见 docs/aclnnRsqrtaclnnInplaceRsqrt.md如下参数类型说明selfconst aclTensor*计算输入公式中的 inputDevice 侧 aclTensor支持非连续的 Tensor数据格式支持 NDshape 需与 out 一致A2/A3 上支持 FLOAT、FLOAT16、UINT8、INT8、INT16、INT32、BOOL、BFLOAT16outaclTensor*计算输出公式中的 outDevice 侧 aclTensor支持非连续 Tensor格式支持 NDshape 需与 self 一致数据类型支持范围与 self 相同workspaceSizeuint64_t*出参返回需要在 Device 侧申请的 workspace 大小executoraclOpExecutor**出参返回 op 执行器包含算子计算流程返回值为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)参数类型说明workspacevoid*入参在 Device 侧申请的 workspace 内存地址workspaceSizeuint64_t入参在 Device 侧申请的 workspace 大小由第一段接口aclnnRsqrtGetWorkspaceSize获取executoraclOpExecutor*入参op 执行器包含算子计算流程streamaclrtStream入参指定执行任务的 Stream返回值为aclnnStatus状态码。aclnnInplaceRsqrt 系列接口aclnnInplaceRsqrtGetWorkspaceSize唯一计算入参为selfRefaclTensor*公式中的 input支持非连续 Tensor、ND 格式shape 维度不大于 8A2/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默认即为确定性实现详见 确定性计算即相同输入在多次运行中产生可复现的相同结果这对调试与精度回归验证非常重要数据格式仅支持 NDshape 约束非 inplace 版本要求self与outshape 一致inplace 版本要求selfRefshape 维度不大于 8inplace 语义inplace 版本会覆盖输入内存调用前需确认原输入数据不再需要。完整调用示例单测工程中的 test_aclnn_rsqrt.cpp仓库在 examples/test_aclnn_rsqrt.cpp 提供了完整可参考的调用工程同时 README.md 的运行验证章节也指向该文件说明其是官方认可的调用方式通过 aclnn 调用的方式调用 Rsqrt 算子。其核心流程可分为 7 步device/stream 初始化依次调用aclInit(nullptr)、aclrtSetDevice(deviceId)、aclrtCreateStream(stream)构造输入输出 aclTensor通过aclrtMalloc申请 Device 侧内存aclrtMemcpy将 Host 数据拷入计算连续 tensor 的 strides 后调用aclCreateTensor(shape, dims, dataType, strides, 0, ACL_FORMAT_ND, shape, dims, deviceAddr)创建 aclTensor调用第一段接口aclnnRsqrtGetWorkspaceSize(self, out, workspaceSize, executor)若workspaceSize 0则用aclrtMalloc申请 workspace调用第二段接口aclnnRsqrt(workspaceAddr, workspaceSize, executor, stream)执行计算同步等待aclrtSynchronizeStream(stream)等待任务执行结束回拷结果aclrtMemcpy将 Device 侧结果拷贝回 Host 并打印资源释放aclDestroyTensor销毁 tensoraclrtFree释放 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); // aclnnInplaceRsqrtinplace 版本结果写回 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_LISTFLOAT、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_INT322、DATA_NUM_FP323、DATA_NUM_BF166、DATA_NUM_INT86、DATA_NUM_UINT84、DATA_NUM_BOOL2、DATA_NUM_INT162等这是因为 bfloat16/int8 等类型在计算时需要额外的中间精度转换缓冲区如 bfloat16 需转 float 计算后再转回buffer 数量动态决策DetermineBufferNum根据单缓冲区需求与coreNum * ubSize的关系决定采用单缓冲SINGLE_BUFFER_NUM1还是双缓冲DOUBLE_BUFFER_NUM2并在SetTilingKey中写入对应的 tiling key0 对应双缓冲、1 对应单缓冲核函数据此模板参数选择流水模式核间负载均衡CalculateCoreBlockNums按块粒度将输入均分到各核区分 big core 与 small core前tailBlockNum个核多分一块并把每个核的数据进一步切成多个 tile 与一个 tail写入RsqrtTilingDataworkspace 申请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_NUM2或单缓冲SINGLE_BUFFER_NUM1实例并注册KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY)。真正逐元素计算的逻辑在 op_kernel/rsqrt.h 的KernelRsqrt模板类中其针对不同数据类型采用差异化的计算路径half / floatSqrt求平方根后用预填充的oneLocalDuplicate 填充 1.0执行Div(1/x)bfloat16先Cast到 float 计算避免 bf16 精度不足SqrtDiv后用CAST_RINT舍入回写 bf16uint8 / int8 / int16 / int32先 Cast 到 half 或 float 中间精度调用向量指令Rsqrt部分整型还经过ComputeIntCorrection的修正p1 min(p1,2) - 3*max(0, min(p1,2)-1)将结果约束到 [-1, 2] 区间最后CAST_RINT转回整型bool直接通过Duplicateint16_t(p1, 0x0101, ...)将结果填充为 1因为rsqrt(true)rsqrt(1)1并在 outQueueY 缓冲上按 int16 对齐处理奇偶数量避免越界写。在内存流水上KernelRsqrt采用经典的CopyIn - Compute - CopyOut三段式通过TQueQuePosition::VECIN/ VECOUT, BUFFER_NUM队列实现输入/输出缓冲的乒乓切换Init阶段依据 bool 与非 bool 类型差异化申请 out 队列缓冲bool 按((tileDataNum1)1) * sizeof(int16_t)申请这也是 tiling 阶段按类型区分DATA_NUM_*的原因。测试覆盖逐数据类型 UT算子配套了完整的单测工程tests/ut/op_kernel/按数据类型拆分为test_rsqrt.cppfloat 主用例、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 接口形成了完整的验证闭环。贡献记录贡献者贡献方贡献算子贡献时间贡献内容skywang2个人开发者Rsqrt2026/03/30新增 Rsqrt 算子Nice_try个人开发者Rsqrt2026/07/06Rsqrt 新增支持整型数据支持A2/A3小结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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
