企业官网建设流程全解析

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 算子支持以下产品形态：

同时在算子注册代码中，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 的参数表）：

在 aclnn 接口层面（aclnnPow2GetWorkspaceSize）对参数的进一步约束为：

• 维度（shape）：x1、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 算子通过aclnn（Ascend 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 参数说明

5.3 aclnnPow2GetWorkspaceSize 返回值与报错场景

第一段接口返回aclnnStatus状态码。在入参校验阶段，以下场景会直接报错：

5.4 aclnnPow2 参数说明

六、完整调用示例

以下示例来自 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::vector<int64_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::vector<T>& hostData, const std::vector<int64_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::vector<int64_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::vector<int64_t> selfShape = {10240}; std::vector<int64_t> exponentShape = {10240}; std::vector<int64_t> outShape = {10240}; void* selfDeviceAddr = nullptr; void* exponentDeviceAddr = nullptr; void* outDeviceAddr = nullptr; aclTensor* self = nullptr; aclTensor* exponent = nullptr; aclTensor* out = nullptr; std::vector<float> selfHostData = {99, -1, -2, 0}; std::vector<float> exponentHostData = {0, 1, 2, 0}; std::vector<float> 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::vector<float> 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、INT8+UINT8 混合类型以及广播场景（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/FLOAT，x2 的 49 种声明对应不同组合，输出 y 的 49 种声明则表达了“输出类型随输入组合推导”的规则（如 INT8 底数配 UINT8 指数时输出为 INT16，INT8 底数配 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 切分（多核调度与 workspace）

op_host/pow2_tiling.cpp 负责在 Host 侧完成数据切分规划，核心逻辑包括：

• 平台信息获取：通过platform_ascendc::PlatformAscendC获取 UB 大小与可用核数coreNum（见GetPlatformInfo）；
• 数据量计算：由输出 shape 总元素数与数据类型长度计算总字节数inputLength，并按 256B（BLOCK_SIZE = 256）对齐得到inputLengthAlign256；单块 tile 数据量tileDataNum = (tileBlockNum * 256) / inputBytes；
• UB 容量自适应：DEFAULT_UB_NUM = 10，当输入或输出涉及 INT8/UINT8 时改用INT8_UB_NUM = 24（8 位类型在 UB 中可容纳更多元素），并据此计算tileBlockNum = (ubSize / BUFFER_NUM / 256) / ubDataNumber，其中BUFFER_NUM = 2对应双缓冲；
• 多核负载均衡：CalculateCoreBlockNums按核数均分 256B 块，分出smallCoreDataNum/bigCoreDataNum及对应的 tile 数、尾块数据量，前tailBlockNum个核多处理一个块；
• 广播预计算：在 Host 侧对输入形状做右对齐补齐（alignedX1/alignedX2/alignedY），预计算各维 stride，并将广播维的有效 stride 置 0（effStride），同时标记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）。

切分结果通过结构体Pow2TilingData（op_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与TQue（BUFFER_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