CANN ops-math 算子开发指南aclnnRsqrt 与 aclnnInplaceRsqrt 接口详解与调用实践【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读本文围绕 CANN ops-math 仓库中 Rsqrt平方根倒数算子的两个 aclnn 单算子 API——aclnnRsqrt与aclnnInplaceRsqrt系统讲解其功能定义、两段式接口调用流程、参数约束、返回码语义、确定性保证并结合仓库内源码op_api / op_kernel / op_host / op_kernel_aicpu与实际调用示例给出可直接编译运行的完整 C 调用范式。读完本文你将掌握在 NPU 上以先查 workspace、后执行计算的两段式方式调用 Rsqrt 算子并能根据实际场景在新建输出张量与原地计算两种模式之间做出正确选择。一、算子功能与适用产品1.1 功能定义Rsqrt 算子对输入 Tensor 的每一个元素求平方根的倒数计算公式为$$ out \frac{1}{\sqrt{input}} $$其中input是输入张量selfout是输出张量。该算子在神经网络中被广泛用于 LayerNorm、RMSNorm、注意力分数归一化等场景的数值缩放即rsqrt。1.2 产品支持情况aclnnRsqrt与aclnnInplaceRsqrt在以下产品上均受支持对应算子文档 math/rsqrt/docs/aclnnRsqrtaclnnInplaceRsqrt.md 及 math/rsqrt/README.mdAscend 950PR / Ascend 950DTAtlas A3 训练系列产品 / Atlas A3 推理系列产品Atlas A2 训练系列产品 / Atlas A2 推理系列产品Atlas 200I/500 A2 推理产品Atlas 推理系列产品Atlas 训练系列产品说明不同产品线在数据类型支持范围上存在差异详见 四、参数与约束说明其中 BFLOAT16 的差异由源码级平台判断逻辑决定见 二、2.4 节。二、两段式接口与函数原型2.1 普通模式与原地模式的选择aclnnRsqrt和aclnnInplaceRsqrt实现完全相同的数学功能区别仅在于输出结果的存放方式接口输出方式适用场景aclnnRsqrt需新建一个输出张量对象out计算结果写入新张量内存需要保留原始输入数据的场景aclnnInplaceRsqrt无需新建输出张量直接在输入张量selfRef的内存中覆盖写回输入不再被后续逻辑使用的场景可节省一块输出内存原地模式Inplace在显存敏感或图融合场景中具有明显的内存收益但使用时必须确认原始输入不再需要否则会破坏上游数据。2.2 两段式接口机制两个算子均遵循 CANN 单算子 API 的两段式调用规范详见 docs/zh/context/two_phase_api.md第一段接口aclnnXxxGetWorkspaceSize完成入参校验并根据算子计算流程推导出在 Device 侧执行所需的临时内存workspace大小workspaceSize同时创建并返回算子执行器aclOpExecutor*。第二段接口aclnnXxx传入第一段接口申请的 workspace 与 executor真正在指定 Stream 上执行计算。其中workspace是指除输入/输出之外算子在 NPU 上完成计算所需的临时内存workspaceSize为其大小。需要注意第二段接口不能重复调用同一 executor 只能执行一次计算。2.3 函数原型普通模式见 math/rsqrt/op_api/aclnn_rsqrt.haclnnStatus 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)2.4 接口的底层实现流程从源码 math/rsqrt/op_api/aclnn_rsqrt.cpp 可以看出aclnnInplaceRsqrtGetWorkspaceSize在完成原地模式专属的入参检查后实际复用ExecRsqrtGetWorkspaceSize的实现将selfRef同时作为out传入因此两种模式的计算图构建流程完全一致参数空指针检查CheckNotNull2Tensor数据类型支持范围校验CheckDtypeValid/CheckInplaceDtypeValid检查self与out的 shape 是否一致CheckSameShape1In1Out空 Tensor 场景直接返回workspaceSize 0源码注释明确rsqrt算子的空tensor在kernel中支持将输入self转换为连续 Tensorl0op::Contiguous若输入类型不在输出支持列表内如整型、BOOL先用l0op::Cast将输入转为out的数据类型如 FLOAT调用l0op::Rsqrt构建核心计算节点将计算结果Cast回out的数据类型通过l0op::ViewCopy将结果拷贝到out上兼容非连续 Tensor汇总整个计算流程所需的workspaceSize并返回 executor。第二段接口aclnnRsqrt/aclnnInplaceRsqrt则统一通过CommonOpExecutorRun(workspace, workspaceSize, executor, stream)驱动执行器在指定 Stream 上完成计算。平台相关的 BF16 支持判断源码中的CheckSocVersionIsSupportBf16()通过GetCurrentPlatformInfo().GetCurNpuArch()判断当前 NPU 架构仅当架构为DAV_2201对应 Atlas A2 系列或IsRegBase(curArch)对应 Atlas A3 及 Ascend 950 系列时才允许输入使用 BFLOAT16。这正是文档中Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品不支持 BFLOAT16这一约束的底层依据op_api/aclnn_rsqrt.cpp。三、aclnnRsqrtGetWorkspaceSize 参数与返回值3.1 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入公式中的 input支持空 TensorFLOAT、FLOAT16、DOUBLE、COMPLEX64、COMPLEX128、UINT8、INT8、INT16、INT32、INT64、BOOL、BFLOAT16ND0-8√outaclTensor*输出公式中的 out支持空 Tensor。shape 需要与 self 一致FLOAT、FLOAT16、DOUBLE、COMPLEX64、COMPLEX128、BFLOAT16ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----平台差异Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品不支持 BFLOAT16数据类型此限制由上文 2.4 节的平台判断逻辑实现。注意self的输入数据类型范围含整型与 BOOL与out的输出数据类型范围仅浮点/复数并不相同。整型输入会在计算前通过 Cast 转换为浮点参与rsqrt运算这也是功能层面对求平方根倒数语义的合理扩展。3.2 返回值aclnnStatus返回状态码通用返回码定义见 docs/zh/context/aclnn_return_code.md其中与本文强相关的有状态码名称错误码说明ACLNN_SUCCESS0成功ACLNN_ERR_PARAM_NULLPTR161001参数中存在非法的 nullptrACLNN_ERR_PARAM_INVALID161002参数校验错误数据类型/格式不支持、shape 不匹配等第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self 或 out 是空指针ACLNN_ERR_PARAM_INVALID161002self 和 out 的数据类型和数据格式不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self 和 out 的 shape 不匹配这些校验与 op_api/aclnn_rsqrt.cpp 中CheckParamsRsqrt的检查顺序一一对应先查空指针161001再查 dtype/format161002最后查 shape 一致性161002。四、参数与约束说明4.1 aclnnRsqrt第二段接口参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnRsqrtGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值aclnnStatus具体含义同 3.2 节。4.2 aclnnInplaceRsqrtGetWorkspaceSize参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfRefaclTensor*输入/输出公式中的 input支持空 TensorFLOAT、FLOAT16、DOUBLE、COMPLEX64、COMPLEX128、BFLOAT16ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----平台差异Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品不支持 BFLOAT16。注意原地模式中selfRef的输入/输出角色表明其内存会被覆盖写且该接口的输入数据类型范围不包含整型与 BOOL仅浮点与复数与普通模式self的宽泛支持范围不同——原地模式要求能算、也能原位存回。错误场景返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 selfRef 是空指针ACLNN_ERR_PARAM_INVALID161002selfRef 的数据类型和数据格式不在支持的范围之内4.3 aclnnInplaceRsqrt第二段接口参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnInplaceRsqrtGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值aclnnStatus。4.4 确定性计算约束aclnnRsqrt与aclnnInplaceRsqrt默认为确定性实现deterministic即相同输入与配置下多次执行结果保持一致。这有利于调试与结果复现相关背景可参考 docs/zh/context/determinism_compute.md。4.5 数据类型与格式补充说明数据格式仅支持ND相关概念见 docs/zh/context/data_format.md输入输出均支持非连续 Tensor相关概念见 docs/zh/context/non_contiguous_tensor.md非连续输入会先被Contiguous规整为连续内存再参与计算输出侧则由ViewCopy负责写回非连续布局shape 维度支持 0-8 维从算子注册代码 math/rsqrt/op_host/rsqrt_def.cpp 看Rsqrt 的 AICore 配置开启了DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)即支持动态 shape / 动态 rank同时开启了PrecisionReduceFlag(true)允许在精度允许范围内使用低精度指令加速计算。五、调用示例完整可运行 C 代码以下示例完整演示了普通模式与原地模式两种调用方式对应仓库示例 math/rsqrt/examples/test_aclnn_rsqrt.cpp 及算子文档中的调用示例。具体编译和执行过程请参考 docs/zh/context/compile_and_run_sample.md。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_rsqrt.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_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2.构造输入与输出需要根据API的接口自定义构造 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 {1, 2, 3, 4}; std::vectorfloat outHostData {0, 0, 0, 0}; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, 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; // 调用aclnnRsqrt第一段接口 ret aclnnRsqrtGetWorkspaceSize(self, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnRsqrtGetWorkspaceSize 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); } // 调用aclnnRsqrt第二段接口 ret aclnnRsqrt(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnRsqrt 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]); } // aclnnInplaceXxx // 调用aclnnInplaceRsqrt第一段接口 ret aclnnInplaceRsqrtGetWorkspaceSize(self, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnInplaceRsqrtGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 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); } // 调用aclnnInplaceRsqrt第二段接口 ret aclnnInplaceRsqrt(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnInplaceRsqrt 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的接口定义修改 ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), selfDeviceAddr, 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和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(out); // 7.释放device资源需要根据具体API的接口定义参数 aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }5.1 示例关键流程拆解资源初始化aclInit→aclrtSetDevice→aclrtCreateStream建立 Host 与 NPU 的执行环境构造张量通过aclrtMalloc申请 Device 内存、aclrtMemcpy搬运 Host 数据、aclCreateTensor创建aclTensor注意连续 Tensor 的 strides 需要按 shape 从后往前推算两段式调用先调aclnnRsqrtGetWorkspaceSize拿到workspaceSize与executor若workspaceSize 0则用aclrtMalloc申请 workspace再调aclnnRsqrt(workspaceAddr, workspaceSize, executor, stream)执行同步等待aclrtSynchronizeStream确保计算在 Stream 上完成取回结果aclrtMemcpy将 Device 结果拷回 Host普通模式从outDeviceAddr读取原地模式从selfDeviceAddr即被覆盖的输入内存读取资源释放aclDestroyTensor→aclrtFree→aclrtDestroyStream→aclrtResetDevice→aclFinalize与初始化顺序相反。注意示例在同一个main中连续执行了普通与原地两次调用其中workspaceSize、executor、workspaceAddr被复用。实际工程中每对第一/第二段接口应严格配对且第二段接口不可重复调用。六、仓库源码佐证kernel 与测试6.1 NPU Kernel 实现AICore 侧 kernel 位于 math/rsqrt/op_kernel/rsqrt_apt.cpp通过TILING_KEY区分不同数据类型的计算路径TILING_KEY_IS(101UL)使用RsqrtDag::RsqrtOphalf::OpDag处理 FLOAT16TILING_KEY_IS(102UL)使用RsqrtDag::RsqrtOpbfloat16_t::OpDag处理 BFLOAT16TILING_KEY_IS(103UL)使用RsqrtDag::RsqrtOpfloat::OpDag处理 FLOAT。kernel 通过ElementwiseSch16B模板化的逐元素调度框架完成向量化计算AIV 核对应 arch35 的实现头文件为 math/rsqrt/op_kernel/arch35/rsqrt.h。6.2 AICPU 实现复数与多核并行对于 DOUBLE、COMPLEX64、COMPLEX128 等数据类型以及整型、BOOL 输入op_api 层会通过 Cast 统一转浮点后执行。AICPU 通路实现在 math/rsqrt/op_kernel_aicpu/rsqrt_aicpu.cpp实数路径outputy[i] static_castT(1) / sqrt(inputx[i])复数路径outputy[i] sqrt(conj(inputx[i])) / sqrt(inputx[i].real()^2 inputx[i].imag()^2)即先求模再按共轭方向归一保证复数域的平方根倒数语义并行策略当数据量超过阈值实数8*1024、复数4*1024时通过CpuKernelUtils::ParallelFor按dataNum / max_core_num切分任务充分利用多核 CPU。6.3 测试与精度验证仓库在 math/rsqrt/tests/assets/golden.py 中提供了基于 PyTorchtorch.rsqrt的 golden 对比实现rsqrt_goldenkernel / GEIR 通路输入为 numpy 数组aclnn_rsqrt_goldenaclnnRsqrt通路aclnn_inplace_rsqrt_goldenaclnnInplaceRsqrt通路对 FLOAT16 / BFLOAT16 采用先转 FP32 计算、再转回原类型的方式匹配 kernel 行为_rsqrt_torch并在_KERNEL_TOLERANCE中将三种类型的校验标准设为 L1 级交叉校验。ST 测试用例定义在 math/rsqrt/tests/st/aclnnRsqrt/atk_aclnnRsqrt.jsonUT 用例位于 math/rsqrt/tests/ut 下的 op_api / op_host / op_kernel_aicpu 目录可据此验证接口各数据类型的正确性。七、常见问题与排查建议返回161001ACLNN_ERR_PARAM_NULLPTR检查self、out/selfRef、workspaceSize、executor指针是否有效尤其out是否已通过aclCreateTensor创建。返回161002ACLNN_ERR_PARAM_INVALID确认数据类型在支持列表内普通模式self支持整型与 BOOL但out与原地模式selfRef仅支持浮点/复数确认数据格式为 ND确认self与out的 shape 完全一致。BFLOAT16 报错核对目标产品是否支持 BFLOAT16Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品不支持。aclnnRsqrt第二次执行异常第二段接口不可重复调用每次执行需重新走完整的第一段→第二段流程。结果不符合预期原地模式确认输入张量内存确实允许被覆盖若后续仍需要使用原始数据请改用普通模式aclnnRsqrt。结语aclnnRsqrt与aclnnInplaceRsqrt是 CANN ops-math 数学算子库中结构清晰、覆盖多产品线的逐元素算子接口普通模式保留输入、原地模式节省内存二者共享同一套两段式调用框架与底层 kernel。结合仓库中的 op_api 源码类型转换与平台判断、AICore/AICPU kernel实数与复数计算以及 golden 测试与torch.rsqrt交叉验证开发者可以快速将其集成进基于 aclnn 单算子 API 的推理或训练链路并以此为模板迁移到其他数学算子的调用开发中。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/21 15:05:32 · 阅读 1.2k CANN ops-math 算子开发指南：aclnnRsqrt 与 aclnnInplaceRsqrt 接口详解与调用实践
CANN ops-math 算子开发指南aclnnRsqrt 与 aclnnInplaceRsqrt 接口详解与调用实践【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读本文围绕 CANN ops-math 仓库中 Rsqrt平方根倒数算子的两个 aclnn 单算子 API——aclnnRsqrt与aclnnInplaceRsqrt系统讲解其功能定义、两段式接口调用流程、参数约束、返回码语义、确定性保证并结合仓库内源码op_api / op_kernel / op_host / op_kernel_aicpu与实际调用示例给出可直接编译运行的完整 C 调用范式。读完本文你将掌握在 NPU 上以先查 workspace、后执行计算的两段式方式调用 Rsqrt 算子并能根据实际场景在新建输出张量与原地计算两种模式之间做出正确选择。一、算子功能与适用产品1.1 功能定义Rsqrt 算子对输入 Tensor 的每一个元素求平方根的倒数计算公式为$$ out \frac{1}{\sqrt{input}} $$其中input是输入张量selfout是输出张量。该算子在神经网络中被广泛用于 LayerNorm、RMSNorm、注意力分数归一化等场景的数值缩放即rsqrt。1.2 产品支持情况aclnnRsqrt与aclnnInplaceRsqrt在以下产品上均受支持对应算子文档 math/rsqrt/docs/aclnnRsqrtaclnnInplaceRsqrt.md 及 math/rsqrt/README.mdAscend 950PR / Ascend 950DTAtlas A3 训练系列产品 / Atlas A3 推理系列产品Atlas A2 训练系列产品 / Atlas A2 推理系列产品Atlas 200I/500 A2 推理产品Atlas 推理系列产品Atlas 训练系列产品说明不同产品线在数据类型支持范围上存在差异详见 四、参数与约束说明其中 BFLOAT16 的差异由源码级平台判断逻辑决定见 二、2.4 节。二、两段式接口与函数原型2.1 普通模式与原地模式的选择aclnnRsqrt和aclnnInplaceRsqrt实现完全相同的数学功能区别仅在于输出结果的存放方式接口输出方式适用场景aclnnRsqrt需新建一个输出张量对象out计算结果写入新张量内存需要保留原始输入数据的场景aclnnInplaceRsqrt无需新建输出张量直接在输入张量selfRef的内存中覆盖写回输入不再被后续逻辑使用的场景可节省一块输出内存原地模式Inplace在显存敏感或图融合场景中具有明显的内存收益但使用时必须确认原始输入不再需要否则会破坏上游数据。2.2 两段式接口机制两个算子均遵循 CANN 单算子 API 的两段式调用规范详见 docs/zh/context/two_phase_api.md第一段接口aclnnXxxGetWorkspaceSize完成入参校验并根据算子计算流程推导出在 Device 侧执行所需的临时内存workspace大小workspaceSize同时创建并返回算子执行器aclOpExecutor*。第二段接口aclnnXxx传入第一段接口申请的 workspace 与 executor真正在指定 Stream 上执行计算。其中workspace是指除输入/输出之外算子在 NPU 上完成计算所需的临时内存workspaceSize为其大小。需要注意第二段接口不能重复调用同一 executor 只能执行一次计算。2.3 函数原型普通模式见 math/rsqrt/op_api/aclnn_rsqrt.haclnnStatus 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)2.4 接口的底层实现流程从源码 math/rsqrt/op_api/aclnn_rsqrt.cpp 可以看出aclnnInplaceRsqrtGetWorkspaceSize在完成原地模式专属的入参检查后实际复用ExecRsqrtGetWorkspaceSize的实现将selfRef同时作为out传入因此两种模式的计算图构建流程完全一致参数空指针检查CheckNotNull2Tensor数据类型支持范围校验CheckDtypeValid/CheckInplaceDtypeValid检查self与out的 shape 是否一致CheckSameShape1In1Out空 Tensor 场景直接返回workspaceSize 0源码注释明确rsqrt算子的空tensor在kernel中支持将输入self转换为连续 Tensorl0op::Contiguous若输入类型不在输出支持列表内如整型、BOOL先用l0op::Cast将输入转为out的数据类型如 FLOAT调用l0op::Rsqrt构建核心计算节点将计算结果Cast回out的数据类型通过l0op::ViewCopy将结果拷贝到out上兼容非连续 Tensor汇总整个计算流程所需的workspaceSize并返回 executor。第二段接口aclnnRsqrt/aclnnInplaceRsqrt则统一通过CommonOpExecutorRun(workspace, workspaceSize, executor, stream)驱动执行器在指定 Stream 上完成计算。平台相关的 BF16 支持判断源码中的CheckSocVersionIsSupportBf16()通过GetCurrentPlatformInfo().GetCurNpuArch()判断当前 NPU 架构仅当架构为DAV_2201对应 Atlas A2 系列或IsRegBase(curArch)对应 Atlas A3 及 Ascend 950 系列时才允许输入使用 BFLOAT16。这正是文档中Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品不支持 BFLOAT16这一约束的底层依据op_api/aclnn_rsqrt.cpp。三、aclnnRsqrtGetWorkspaceSize 参数与返回值3.1 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入公式中的 input支持空 TensorFLOAT、FLOAT16、DOUBLE、COMPLEX64、COMPLEX128、UINT8、INT8、INT16、INT32、INT64、BOOL、BFLOAT16ND0-8√outaclTensor*输出公式中的 out支持空 Tensor。shape 需要与 self 一致FLOAT、FLOAT16、DOUBLE、COMPLEX64、COMPLEX128、BFLOAT16ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----平台差异Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品不支持 BFLOAT16数据类型此限制由上文 2.4 节的平台判断逻辑实现。注意self的输入数据类型范围含整型与 BOOL与out的输出数据类型范围仅浮点/复数并不相同。整型输入会在计算前通过 Cast 转换为浮点参与rsqrt运算这也是功能层面对求平方根倒数语义的合理扩展。3.2 返回值aclnnStatus返回状态码通用返回码定义见 docs/zh/context/aclnn_return_code.md其中与本文强相关的有状态码名称错误码说明ACLNN_SUCCESS0成功ACLNN_ERR_PARAM_NULLPTR161001参数中存在非法的 nullptrACLNN_ERR_PARAM_INVALID161002参数校验错误数据类型/格式不支持、shape 不匹配等第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self 或 out 是空指针ACLNN_ERR_PARAM_INVALID161002self 和 out 的数据类型和数据格式不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self 和 out 的 shape 不匹配这些校验与 op_api/aclnn_rsqrt.cpp 中CheckParamsRsqrt的检查顺序一一对应先查空指针161001再查 dtype/format161002最后查 shape 一致性161002。四、参数与约束说明4.1 aclnnRsqrt第二段接口参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnRsqrtGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值aclnnStatus具体含义同 3.2 节。4.2 aclnnInplaceRsqrtGetWorkspaceSize参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfRefaclTensor*输入/输出公式中的 input支持空 TensorFLOAT、FLOAT16、DOUBLE、COMPLEX64、COMPLEX128、BFLOAT16ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----平台差异Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品不支持 BFLOAT16。注意原地模式中selfRef的输入/输出角色表明其内存会被覆盖写且该接口的输入数据类型范围不包含整型与 BOOL仅浮点与复数与普通模式self的宽泛支持范围不同——原地模式要求能算、也能原位存回。错误场景返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 selfRef 是空指针ACLNN_ERR_PARAM_INVALID161002selfRef 的数据类型和数据格式不在支持的范围之内4.3 aclnnInplaceRsqrt第二段接口参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnInplaceRsqrtGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值aclnnStatus。4.4 确定性计算约束aclnnRsqrt与aclnnInplaceRsqrt默认为确定性实现deterministic即相同输入与配置下多次执行结果保持一致。这有利于调试与结果复现相关背景可参考 docs/zh/context/determinism_compute.md。4.5 数据类型与格式补充说明数据格式仅支持ND相关概念见 docs/zh/context/data_format.md输入输出均支持非连续 Tensor相关概念见 docs/zh/context/non_contiguous_tensor.md非连续输入会先被Contiguous规整为连续内存再参与计算输出侧则由ViewCopy负责写回非连续布局shape 维度支持 0-8 维从算子注册代码 math/rsqrt/op_host/rsqrt_def.cpp 看Rsqrt 的 AICore 配置开启了DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)即支持动态 shape / 动态 rank同时开启了PrecisionReduceFlag(true)允许在精度允许范围内使用低精度指令加速计算。五、调用示例完整可运行 C 代码以下示例完整演示了普通模式与原地模式两种调用方式对应仓库示例 math/rsqrt/examples/test_aclnn_rsqrt.cpp 及算子文档中的调用示例。具体编译和执行过程请参考 docs/zh/context/compile_and_run_sample.md。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_rsqrt.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_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2.构造输入与输出需要根据API的接口自定义构造 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 {1, 2, 3, 4}; std::vectorfloat outHostData {0, 0, 0, 0}; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, 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; // 调用aclnnRsqrt第一段接口 ret aclnnRsqrtGetWorkspaceSize(self, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnRsqrtGetWorkspaceSize 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); } // 调用aclnnRsqrt第二段接口 ret aclnnRsqrt(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnRsqrt 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]); } // aclnnInplaceXxx // 调用aclnnInplaceRsqrt第一段接口 ret aclnnInplaceRsqrtGetWorkspaceSize(self, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnInplaceRsqrtGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 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); } // 调用aclnnInplaceRsqrt第二段接口 ret aclnnInplaceRsqrt(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnInplaceRsqrt 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的接口定义修改 ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), selfDeviceAddr, 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和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(out); // 7.释放device资源需要根据具体API的接口定义参数 aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }5.1 示例关键流程拆解资源初始化aclInit→aclrtSetDevice→aclrtCreateStream建立 Host 与 NPU 的执行环境构造张量通过aclrtMalloc申请 Device 内存、aclrtMemcpy搬运 Host 数据、aclCreateTensor创建aclTensor注意连续 Tensor 的 strides 需要按 shape 从后往前推算两段式调用先调aclnnRsqrtGetWorkspaceSize拿到workspaceSize与executor若workspaceSize 0则用aclrtMalloc申请 workspace再调aclnnRsqrt(workspaceAddr, workspaceSize, executor, stream)执行同步等待aclrtSynchronizeStream确保计算在 Stream 上完成取回结果aclrtMemcpy将 Device 结果拷回 Host普通模式从outDeviceAddr读取原地模式从selfDeviceAddr即被覆盖的输入内存读取资源释放aclDestroyTensor→aclrtFree→aclrtDestroyStream→aclrtResetDevice→aclFinalize与初始化顺序相反。注意示例在同一个main中连续执行了普通与原地两次调用其中workspaceSize、executor、workspaceAddr被复用。实际工程中每对第一/第二段接口应严格配对且第二段接口不可重复调用。六、仓库源码佐证kernel 与测试6.1 NPU Kernel 实现AICore 侧 kernel 位于 math/rsqrt/op_kernel/rsqrt_apt.cpp通过TILING_KEY区分不同数据类型的计算路径TILING_KEY_IS(101UL)使用RsqrtDag::RsqrtOphalf::OpDag处理 FLOAT16TILING_KEY_IS(102UL)使用RsqrtDag::RsqrtOpbfloat16_t::OpDag处理 BFLOAT16TILING_KEY_IS(103UL)使用RsqrtDag::RsqrtOpfloat::OpDag处理 FLOAT。kernel 通过ElementwiseSch16B模板化的逐元素调度框架完成向量化计算AIV 核对应 arch35 的实现头文件为 math/rsqrt/op_kernel/arch35/rsqrt.h。6.2 AICPU 实现复数与多核并行对于 DOUBLE、COMPLEX64、COMPLEX128 等数据类型以及整型、BOOL 输入op_api 层会通过 Cast 统一转浮点后执行。AICPU 通路实现在 math/rsqrt/op_kernel_aicpu/rsqrt_aicpu.cpp实数路径outputy[i] static_castT(1) / sqrt(inputx[i])复数路径outputy[i] sqrt(conj(inputx[i])) / sqrt(inputx[i].real()^2 inputx[i].imag()^2)即先求模再按共轭方向归一保证复数域的平方根倒数语义并行策略当数据量超过阈值实数8*1024、复数4*1024时通过CpuKernelUtils::ParallelFor按dataNum / max_core_num切分任务充分利用多核 CPU。6.3 测试与精度验证仓库在 math/rsqrt/tests/assets/golden.py 中提供了基于 PyTorchtorch.rsqrt的 golden 对比实现rsqrt_goldenkernel / GEIR 通路输入为 numpy 数组aclnn_rsqrt_goldenaclnnRsqrt通路aclnn_inplace_rsqrt_goldenaclnnInplaceRsqrt通路对 FLOAT16 / BFLOAT16 采用先转 FP32 计算、再转回原类型的方式匹配 kernel 行为_rsqrt_torch并在_KERNEL_TOLERANCE中将三种类型的校验标准设为 L1 级交叉校验。ST 测试用例定义在 math/rsqrt/tests/st/aclnnRsqrt/atk_aclnnRsqrt.jsonUT 用例位于 math/rsqrt/tests/ut 下的 op_api / op_host / op_kernel_aicpu 目录可据此验证接口各数据类型的正确性。七、常见问题与排查建议返回161001ACLNN_ERR_PARAM_NULLPTR检查self、out/selfRef、workspaceSize、executor指针是否有效尤其out是否已通过aclCreateTensor创建。返回161002ACLNN_ERR_PARAM_INVALID确认数据类型在支持列表内普通模式self支持整型与 BOOL但out与原地模式selfRef仅支持浮点/复数确认数据格式为 ND确认self与out的 shape 完全一致。BFLOAT16 报错核对目标产品是否支持 BFLOAT16Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品不支持。aclnnRsqrt第二次执行异常第二段接口不可重复调用每次执行需重新走完整的第一段→第二段流程。结果不符合预期原地模式确认输入张量内存确实允许被覆盖若后续仍需要使用原始数据请改用普通模式aclnnRsqrt。结语aclnnRsqrt与aclnnInplaceRsqrt是 CANN ops-math 数学算子库中结构清晰、覆盖多产品线的逐元素算子接口普通模式保留输入、原地模式节省内存二者共享同一套两段式调用框架与底层 kernel。结合仓库中的 op_api 源码类型转换与平台判断、AICore/AICPU kernel实数与复数计算以及 golden 测试与torch.rsqrt交叉验证开发者可以快速将其集成进基于 aclnn 单算子 API 的推理或训练链路并以此为模板迁移到其他数学算子的调用开发中。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考