CANN ops-math 算子详解Pow2 张量指数运算的 aclnn 接口与 AscendC 实现剖析【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本文是 CANN 开源数学算子库 ops-math 中 Pow2 算子的完整技术指南。Pow2 在experimental/math/pow2目录下实现提供张量元素级指数运算out_i pow(x1_i, x2_i)支持底数与指数两个输入张量之间的广播语义可用于网络模型中任意逐元素幂运算加速。读完本文你将掌握 Pow2 算子的功能与约束、参数语义、两段式 aclnn 编程接口aclnnPow2GetWorkspaceSize/aclnnPow2的完整调用流程、可直接运行的示例代码以及从算子注册、形状推导、tiling 切分到 AscendC Kernel 指令级实现的全链路源码原理。一、算子概述与产品支持情况Pow2 是 ops-math 仓库中实验性数学算子集合experimental/math的一员对应目录为 experimental/math/pow2其定位是在昇腾 NPU 上完成张量级的逐元素指数幂运算。目录中同时包含算子文档README、接口调用文档、调用示例、Host 侧实现算子定义/形状推导/tiling与 Device 侧 Kernel 实现是一个结构完整、可直接阅读和复用的算子样例。根据 experimental/math/pow2/README.md 中的产品支持情况说明当前 Pow2 算子支持以下产品形态产品是否支持Atlas A2 训练系列产品 / Atlas 800I A2 推理产品 / A200I A2 Box 异构组件√同时在算子注册代码中Pow2 通过this-AICore().AddConfig(ascend910b)见 op_host/pow2_def.cpp声明其 AICore 侧的配置从源码结构看该实现面向 910B 系列的 AICore 指令集。适用前提运行与编译该算子需要对应昇腾硬件环境及配套 CANN 工具链。二、功能说明与计算公式Pow2 算子实现张量的指数运算对输入张量 x1底数和 x2指数逐元素计算幂算子功能实现张量的指数运算功能对输入张量 x1底数和 x2指数进行元素级计算。计算公式$$ out_i \text{pow}(x1_i, x2_i) $$其中 x1 与 x2 均可为标量或多维张量两者形状不必完全一致通过广播规则匹配后输出 out。也就是说该算子同时覆盖了pow(tensor, tensor)、pow(tensor, scalar)、pow(scalar, tensor)三种常见调用形态示例代码中对此均有对应测试分支。从 Kernel 实现看op_kernel/pow2.h 中的PowCompute把幂运算拆解为对数-指数恒等变换与边界修正的组合先计算exp(x2 * ln(|x1|))得到主结果再通过Compare/Select指令修正三类特殊情况——底数为负时补符号位、指数为 0 时结果置 1、底数为 0 时按幂运算语义修正从而在向量指令层面完成通用幂函数计算。三、参数说明Pow2 算子的输入、输出与属性如下来源于 README.md 与 docs/test_aclnn_pow2.md 的参数表参数名输入/输出/属性描述数据类型数据格式x1输入底数张量可为标量或多维张量支持广播到输出张量。FLOAT16、FLOAT32、BFLOAT16、INT8、UINT8、INT16、INT32NDx2输入指数张量可为标量或多维张量支持广播到输出张量。FLOAT16、FLOAT32、BFLOAT16、INT8、UINT8、INT16、INT32NDout输出元素级计算结果张量输出数据类型与输入类型一致或通过 Cast 转换。FLOAT16、FLOAT32、BFLOAT16、INT8、UINT8、INT16、INT32ND在 aclnn 接口层面aclnnPow2GetWorkspaceSize对参数的进一步约束为维度shapex1、x2、out 均支持 0~8 维其中 0 维即标量。非连续 Tensor三个 Tensor 均支持非连续内存布局表格中标注为 √Host 侧通过AutoContiguous()完成内存自动连续化见下文源码解析。out 的 shape与广播后的输出形状一致接口文档中表述为“shape 与 self 相同”。混合类型从算子定义的数据类型组合表看x1、x2 可以为不同类型如 INT8 底数 UINT8 指数、INT8 底数 FP32 指数输出类型按组合表由框架推导可不同于任一输入类型通过 Cast 转换完成。四、约束说明使用 Pow2 算子时需注意以下约束见 README.md 约束说明章节暂不支持 int64/uint64 类型输入与输出数据类型仅限 FLOAT16、FLOAT32、BFLOAT16、INT8、UINT8、INT16、INT32。广播规则匹配输入张量维度可以不同但需要通过广播规则匹配不满足广播条件时形状推导InferShape阶段会直接报错返回GRAPH_FAILED。此外接口文档中 aclnn 第一段接口的入参校验还包含Tensor 为空指针、数据类型/数据格式超范围、维度超过 8 维、数据形状不一致等场景均会在调用aclnnPow2GetWorkspaceSize时被拦截并返回对应错误码。五、aclnn 两段式接口调用说明Pow2 算子通过aclnnAscend CANN Neural Network两段式接口对外提供调用能力即先调用aclnnPow2GetWorkspaceSize获取计算所需 workspace 大小与执行器再调用aclnnPow2真正执行计算。完整的接口约定参见 experimental/math/pow2/docs/test_aclnn_pow2.md。5.1 函数原型第一段接口获取 workspace 大小与执行器aclnnStatus aclnnPow2GetWorkspaceSize( const aclTensor *x1, const aclTensor *x2, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)第二段接口执行计算aclnnStatus aclnnPow2( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)5.2 aclnnPow2GetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorx1输入底数张量可为标量或多维张量支持广播到输出张量公式中的 x1。无FLOAT16、FLOAT32、BFLOAT16、INT8、UINT8、INT16、INT32ND0-8√x2输入指数张量可为标量或多维张量支持广播到输出张量公式中的 x2。无FLOAT16、FLOAT32、BFLOAT16、INT8、UINT8、INT16、INT32ND0-8√out输出元素级计算结果张量输出数据类型与输入类型一致或通过 Cast 转换公式中的 out。shape 与输出一致。FLOAT16、FLOAT32、BFLOAT16、INT8、UINT8、INT16、INT32ND0-8√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小。-----executor输出返回 op 执行器包含了算子计算流程。-----5.3 aclnnPow2GetWorkspaceSize 返回值与报错场景第一段接口返回aclnnStatus状态码。在入参校验阶段以下场景会直接报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 tensor 是空指针。ACLNN_ERR_PARAM_INVALID161002x1、x2 和 out 的数据类型和数据格式不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002x1、x2 和 out 的数据维度超过了 8 维。ACLNN_ERR_PARAM_INVALID161002x1、x2 和 out 的数据形状不一致。5.4 aclnnPow2 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnPow2GetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。六、完整调用示例以下示例来自 examples/test_aclnn_pow2.cpp演示 FP32 场景下从环境初始化、构造 Tensor、两段式调用到结果回读与资源释放的完整流程示例中的 shape 为{10240}前 4 个元素为预设的边界用例x1 {99, -1, -2, 0}、x2 {0, 1, 2, 0}分别覆盖指数为 0、底数为负、底数为 0 等特殊分支其余元素按x2 2填充。#include iostream #include vector #include acl/acl.h #include aclnn_pow2.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 testFp32() { LOG_PRINT(Test for Fp32\n); // 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 {10240}; std::vectorint64_t exponentShape {10240}; std::vectorint64_t outShape {10240}; void* selfDeviceAddr nullptr; void* exponentDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* exponent nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {99, -1, -2, 0}; std::vectorfloat exponentHostData {0, 1, 2, 0}; std::vectorfloat outHostData {0, 0, 0, 0}; for (int i 4; i 10240; i) { selfHostData.push_back(i); exponentHostData.push_back(2); outHostData.push_back(0); } // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建 exponent aclTensor ret CreateAclTensor(exponentHostData, exponentShape, exponentDeviceAddr, aclDataType::ACL_FLOAT, exponent); 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 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnPow2第一段接口 ret aclnnPow2GetWorkspaceSize(self, exponent, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnPow2GetWorkspaceSize 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); } // 调用aclnnPow2第二段接口 ret aclnnPow2(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnPow2 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 5; i) { LOG_PRINT(aclnnPow2 result[%ld] is: %f\n, i, resultData[i]); } // 6. 释放aclTensor和aclScalar aclDestroyTensor(self); aclDestroyTensor(exponent); aclDestroyTensor(out); // 7. 释放device资源 aclrtFree(selfDeviceAddr); aclrtFree(exponentDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例代码的执行链路可归纳为 7 步初始化环境 → 构造 Tensor → 两段式调用GetWorkspaceSize 执行→ 流同步 → 结果回读 → 释放 Tensor → 释放设备资源。其中CreateAclTensor模板按连续布局计算 strides 后调用aclCreateTensor创建ACL_FORMAT_ND格式的 aclTensor是可复用的通用辅助函数。除 FP32 外examples/test_aclnn_pow2.cpp 还提供了 INT8、UINT8、INT16、INT32、INT8UINT8 混合类型以及广播场景x1: {2,1,9}、x2: {1,4,1}→out: {2,4,9}等多个测试函数main中按需取消注释即可切换用例。编译与运行的完整流程可参考仓库 docs/zh 下 context 目录中关于样例编译运行的说明文档。七、源码级实现解析Pow2 算子的实现分布在 Host 侧与 Device 侧下面沿着“算子定义 → 形状推导 → tiling 切分 → Kernel 计算”的主线逐层剖析。7.1 算子定义OpDef 注册op_host/pow2_def.cpp 通过算子定义注册框架声明 Pow2输入 x1、x2 均为必选ParamType(REQUIRED)输入与输出通过长度相等的数据类型列表 格式列表按位对齐的方式声明所有合法组合例如 x1 的 49 种声明中覆盖 INT8/UINT8/INT16/INT32/FLOAT16/BF16/FLOATx2 的 49 种声明对应不同组合输出 y 的 49 种声明则表达了“输出类型随输入组合推导”的规则如 INT8 底数配 UINT8 指数时输出为 INT16INT8 底数配 FP32 指数时输出为 FP32所有格式均为FORMAT_ND并调用.AutoContiguous()实现非连续输入的内存自动连续化通过OP_ADD(Pow2)将算子信息注册进算子信息库。7.2 形状推导InferShape 与广播op_host/pow2_infershape.cpp 实现输出形状推导InferShapePow2分别读取 x1、x2 的输入 shape取两者维度的较大者作为输出维度maxDim输出各维初始化为 1从最后一个维度开始向前对齐按 NumPy 广播规则处理若d1 d2或d1 1或d2 1则输出该维取max(d1, d2)否则记录Dimension mismatch错误并返回GRAPH_FAILED推导成功后将结果写入输出 shape。这印证了 README 中“输入张量维度可不同但需要通过广播规则匹配”的约束——不满足广播条件的输入在构图阶段就会被拦截。7.3 tiling 切分多核调度与 workspaceop_host/pow2_tiling.cpp 负责在 Host 侧完成数据切分规划核心逻辑包括平台信息获取通过platform_ascendc::PlatformAscendC获取 UB 大小与可用核数coreNum见GetPlatformInfo数据量计算由输出 shape 总元素数与数据类型长度计算总字节数inputLength并按 256BBLOCK_SIZE 256对齐得到inputLengthAlign256单块 tile 数据量tileDataNum (tileBlockNum * 256) / inputBytesUB 容量自适应DEFAULT_UB_NUM 10当输入或输出涉及 INT8/UINT8 时改用INT8_UB_NUM 248 位类型在 UB 中可容纳更多元素并据此计算tileBlockNum (ubSize / BUFFER_NUM / 256) / ubDataNumber其中BUFFER_NUM 2对应双缓冲多核负载均衡CalculateCoreBlockNums按核数均分 256B 块分出smallCoreDataNum/bigCoreDataNum及对应的 tile 数、尾块数据量前tailBlockNum个核多处理一个块广播预计算在 Host 侧对输入形状做右对齐补齐alignedX1/alignedX2/alignedY预计算各维 stride并将广播维的有效 stride 置 0effStride同时标记isSameX1/isSameX2、is_input0_scalar/is_input1_scalar写入 tiling 数据供 Kernel 使用workspace 规划GetWorkspaceSize通过框架申请一块 workspace用户大小 系统库 API 所需空间最终通过context-SetBlockDim(coreNum)设置核数并以ELEMENTWISE_TPL_SCH_MODE_0作为模板调度键tilingKey见 op_kernel/pow2_tiling_key.h。切分结果通过结构体Pow2TilingDataop_kernel/pow2_tiling_data.h传递到 Kernel 侧字段包括各核数据量、tile 数、尾块数据量、广播标记、各输入/输出 stride 及标量标记等。7.4 Kernel 计算AscendC 实现op_kernel/pow2.cpp 是 Kernel 入口模板参数DTYPE_X1/DTYPE_X2/DTYPE_Y由编译期实例化NsPow2::Pow2类完成 Init → Process 的执行流程。核心实现在 op_kernel/pow2.h数据搬运CopyIn根据 tiling 标记分四种路径处理——双标量、单标量、无广播、部分/全广播无广播时走DataCopy直接 DMA 搬运广播时通过GetBroadcastIndexEff利用预计算的有效 stride 将线性索引换算为输入偏移后GetValue逐元素取数int8/uint8 的标量复制因Duplicate不支持而改用循环SetValue类型提升Compute阶段将各类输入统一 Cast 为 FP32 参与运算int8/uint8 经 half 中转FP32 直接ReinterpretCast输出再按目标类型 Cast 回写整型输出使用CAST_RINT舍入幂运算实现PowCompute用Abs Ln Mul Exp完成exp(x2 * ln(|x1|))主计算再以三条Compare/Select规则修正底数为负时翻转符号、指数为 0 时结果为 1、底数为 0 时按边界语义修正最终结果写入yLocal流水与双缓冲Process对 tile 循环执行 CopyIn/Compute/CopyOut通过TPipe与TQueBUFFER_NUM 2实现双缓冲流水最后一块使用tailDataNum处理尾数据。这种“Host 预计算 stride Device 端按有效 stride 寻址”的设计把广播开销从运行时计算转移到编译期/tiling 期是理解该算子性能表现的源码依据。八、构建与验证从仓库构建体系看Pow2 通过 experimental/math/pow2/CMakeLists.txt 中的add_all_modules_sources(OPTYPE pow2 ACLNNTYPE aclnn)接入整体构建该命令同时将 pow2 注册为算子类型OPTYPE并生成 aclnn 类型接口ACLNNTYPE。也就是说Pow2 的对外发布形态即aclnnPow2系列接口调用侧只需包含aclnn_pow2.h头文件并链接对应算子库即可。验证方面experimental/math/pow2/tests/ut目录预留了单测UT骨架开发者可按仓库 tests/ut 与 docs/zh 中提供的单测组织方式补充用例日常快速验证可直接运行 examples/test_aclnn_pow2.cpp 中的各类型测试函数。九、贡献信息根据 README.md 的贡献说明该算子由个人开发者 Shi xiangyang 于 2025/12/16 贡献贡献内容为“Pow 算子适配开源仓”。作为 ops-math 开源生态的一部分Pow2 的文档、示例、Host 实现与 Kernel 实现全部位于 experimental/math/pow2 目录下可作为学习 aclnn 算子接入流程与 AscendC 算子开发的完整参考样例。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考