- 算子库
- 人工智能
- CANN
【免费下载链接】ops-math
本项目是CANN提供的数学类基础计算算子库,实现网络在NPU上加速计算。
本文是 CANN/ops-math 仓库中 Pow 数学算子的实战指南,围绕 math/pow/docs/aclnnPowTensorTensor&aclnnInplacePowTensorTensor.md 展开,介绍aclnnPowTensorTensor(元素级幂运算,底数与指数均为 Tensor)及原地版本aclnnInplacePowTensorTensor的接口原型、参数约束、错误码、源码级实现原理与完整可运行示例。读完本文,你将掌握如何在 NPU 上正确调用这两组接口完成 Pow 计算,理解两段式 API 的 workspace/executor 申请与执行机制,并能够对照仓库源码与单测用例排查参数校验失败问题。
功能说明与产品支持情况
Pow 算子计算exponent每个元素作为input(self)对应元素的幂:
$$ out_i = x_i^{exponent_i} $$
其中x为底数(self),exponent为指数,两者都是 Tensor,逐元素完成幂运算,并支持 broadcast 广播语义。该算子属于数学类基础算子,在 CANN/ops-math 仓库中位于 math/pow,同一算子目录下还提供了aclnnPowScalarTensor、aclnnPowTensorScalar、aclnnExp2等变体接口。
根据文档,该接口在以下产品上均获得支持:
- Ascend 950PR / Ascend 950DT
- Atlas A3 训练系列产品 / Atlas A3 推理系列产品
- Atlas A2 训练系列产品 / Atlas A2 推理系列产品
- Atlas 200I/500 A2 推理产品
- Atlas 推理系列产品
- Atlas 训练系列产品
注意:在 Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品上,本接口不支持 BFLOAT16 数据类型(详见下文数据类型小节)。
INT32 整型计算的性能约束
文档明确指出该算子存在一项重要的性能约束:INT32 整型计算在如下范围以外的场景会出现超时。表格含义为:当shape(Tensor 元素规模)不超过某量级时,exponent_value的取值应落在对应的指数范围内:
| shape(元素规模) | exponent_value(指数取值范围) |
|---|---|
| ≤ 100000(十万) | -200000000 ~ 200000000(两亿) |
| ≤ 1000000(百万) | -20000000 ~ 20000000(两千万) |
| ≤ 10000000(千万) | -2000000 ~ 2000000(两百万) |
| ≤ 100000000(亿) | -200000 ~ 200000(二十万) |
| ≤ 1000000000(十亿) | -20000 ~ 20000(两万) |
可以理解为「元素规模 × 指数绝对值」共同决定整型幂运算的计算量:元素越多,允许的指数绝对值就越小。例如 shape 达到一亿量级时,指数应控制在 -200000 ~ 200000 之间,否则可能出现超时。这与整型幂通过快速幂/逐位相乘实现、计算量与位宽和指数大小相关的底层机制一致。
两段式接口(Two-Phase API)与函数原型
Pow 算子按 CANN 算子库的规范采用两段式接口设计,相关通用机制详见 两段式接口说明:
- 第一段接口(GetWorkspaceSize):完成入参校验,根据计算流程推导出需要在 Device 侧申请的 workspace 大小,并创建包含完整算子计算流程的 op 执行器(
aclOpExecutor)。 - 第二段接口(执行接口):由调用方按第一段接口返回的
workspaceSize申请 Device 内存后,传入 workspace、executor 与 stream 真正触发计算。
四个函数原型如下:
// 第一段接口:获取 workspace 大小与执行器 aclnnStatus aclnnPowTensorTensorGetWorkspaceSize( const aclTensor* self, const aclTensor* exponent, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)// 第二段接口:执行计算 aclnnStatus aclnnPowTensorTensor( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)原地版本(计算结果直接写回 self,无需单独 out 参数):
aclnnStatus aclnnInplacePowTensorTensorGetWorkspaceSize( const aclTensor* self, const aclTensor* exponent, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnInplacePowTensorTensor( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)接口声明位于 math/pow/op_api/aclnn_pow_tensor_tensor.h,实现位于 math/pow/op_api/aclnn_pow_tensor_tensor.cpp。头文件中将接口归属域标记为aclnn_math,并使用extern "C"保证 C/C++ 混合调用时符号一致。
aclnnPowTensorTensorGetWorkspaceSize 参数说明与返回值
参数说明
| 参数名 | 输入/输出 | 描述 | 使用说明 | 数据类型 | 数据格式 | 维度(shape) | 非连续Tensor |
|---|---|---|---|---|---|---|---|
| self(aclTensor*) | 输入 | 公式中的 x,pow 运算的底数 | self 和 exponent 数据类型不支持同时为 BOOL,shape 需要与 exponent 满足 broadcast 关系 | FLOAT、FLOAT16、DOUBLE、INT16、BOOL、INT32、INT64、INT8、UINT8、COMPLEX64、COMPLEX128、BFLOAT16 | ND | 0-8 | √ |
| exponent(aclTensor*) | 输入 | 公式中的 exponent,pow 运算的指数 | self 和 exponent 数据类型不支持同时为 BOOL,shape 需要与 self 满足 broadcast 关系 | FLOAT、FLOAT16、DOUBLE、INT16、BOOL、INT32、INT64、INT8、UINT8、COMPLEX64、COMPLEX128、BFLOAT16 | ND | - | √ |
| out(aclTensor*) | 输出 | 公式中的 out,存储计算结果 | 数据类型需要是 self 与 exponent 推导之后可转换的数据类型(参见互推导关系),shape 需要是 self 与 exponent broadcast 之后的 shape | FLOAT、FLOAT16、DOUBLE、BOOL、INT16、INT32、INT64、INT8、UINT8、COMPLEX64、COMPLEX128、BFLOAT16、UINT16、UINT32、UINT64 | ND | - | √ |
| workspaceSize(uint64_t*) | 输出 | 返回需要在 Device 侧申请的 workspace 大小 | - | - | - | - | - |
| executor(aclOpExecutor**) | 输出 | 返回 op 执行器,包含了算子计算流程 | - | - | - | - | - |
需要特别说明的几点:
- 维度上限:
self支持 0-8 维;从实现看 aclnn_pow_tensor_tensor.cpp 中定义了constexpr size_t MAX_DIM_LEN = 8,CheckShape通过OP_CHECK_MAX_DIM对 self 和 exponent 同时做 8 维上限校验。 - 平台差异:Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品不支持 BFLOAT16 数据类型。源码中通过
CheckSocVersionIsSupportBf16()(判断 SoC 版本是否处于 ASCEND910B~ASCEND910E 区间或是否注册为 RegBase 平台)决定是否放行 BF16 输入。 - BOOL 限制:self 和 exponent 不允许同时为 BOOL(但允许其中一方为 BOOL,单测中有
float + bool通过、bool + bool返回ACLNN_ERR_PARAM_INVALID的用例)。 - 非连续 Tensor:两者均支持非连续 Tensor(√),接口内部会自动做 Contiguous 与 ViewCopy 处理。
返回值(第一段接口错误码)
第一段接口完成入参校验,出现以下场景时报错。返回码含义可参见 aclnn 返回码说明:
| 返回值 | 错误码 | 描述 |
|---|---|---|
| ACLNN_ERR_PARAM_NULLPTR | 161001 | 传入的 self、exponent、out 是空指针 |
| ACLNN_ERR_PARAM_INVALID | 161002 | self 和 exponent 的数据类型不在支持的范围之内 |
| ACLNN_ERR_PARAM_INVALID | 161002 | self 和 exponent 无法做数据类型推导 |
| ACLNN_ERR_PARAM_INVALID | 161002 | 推导出的数据类型无法转换为指定输出 out 的类型 |
| ACLNN_ERR_PARAM_INVALID | 161002 | self 和 exponent 的 shape 无法做 broadcast |
| ACLNN_ERR_PARAM_INVALID | 161002 | self 和 exponent 同时为 bool 数据类型 |
对照 aclnn_pow_tensor_tensor.cpp 中的CheckParams实现,可以确认校验顺序依次为:空指针检查(CheckNotNull)→ 数据类型合法性检查(CheckDtypeValid)→ 类型推导/可转换检查(CheckPromoteType)→ broadcast 与输出 shape 检查(CheckShape)。任一步失败即返回对应的ACLNN_ERR_PARAM_NULLPTR或ACLNN_ERR_PARAM_INVALID。
aclnnPowTensorTensor 参数说明
| 参数名 | 输入/输出 | 描述 |
|---|---|---|
| workspace | 输入 | 在 Device 侧申请的 workspace 内存地址 |
| workspaceSize | 输入 | 在 Device 侧申请的 workspace 大小,由第一段接口 aclnnPowTensorTensorGetWorkspaceSize 获取 |
| executor | 输入 | op 执行器,包含了算子计算流程 |
| stream | 输入 | 指定执行任务的 Stream |
返回值仍为aclnnStatus。第二段接口内部通过CommonOpExecutorRun(workspace, workspaceSize, executor, stream)驱动框架完成计算(见 aclnn_pow_tensor_tensor.cpp)。
aclnnInplacePowTensorTensorGetWorkspaceSize 参数说明与返回值
原地版本将计算结果直接写回self,因此不再需要独立的out参数:
| 参数名 | 输入/输出 | 描述 | 使用说明 | 数据类型 | 数据格式 | 维度(shape) | 非连续Tensor |
|---|---|---|---|---|---|---|---|
| self(aclTensor*) | 输入|输出 | 公式中的 x 和 out,pow 运算的底数,计算结果写回此 tensor | 数据类型是 self 与 exponent 推导之后可转换的数据类型,shape 需要与 exponent 满足 broadcast 关系 | FLOAT、FLOAT16、DOUBLE、INT32、INT64、INT8、UINT8、COMPLEX64、COMPLEX128、INT16、BFLOAT16 | ND | - | √ |
| exponent(aclTensor*) | 输入 | 公式中的 exponent,pow 运算的指数 | 数据类型是 self 与 exponent 推导之后可转换的数据类型(参见互推导关系),shape 需要与 self 满足 broadcast 关系 | FLOAT、FLOAT16、DOUBLE、INT32、INT64、INT8、UINT8、COMPLEX64、COMPLEX128、INT16、BFLOAT16 | ND | - | √ |
| workspaceSize(uint64_t*) | 输出 | 返回需要在 Device 侧申请的 workspace 大小 | - | - | - | - | - |
| executor(aclOpExecutor**) | 输出 | 返回 op 执行器,包含了算子计算流程 | - | - | - | - | - |
需要特别留意:
- 原地语义要求:self 和 exponent 进行 broadcast 后的 shape 必须等于 self 的 shape(即广播结果不能比 self 更大),否则无法写回,接口会报错。
- BFLOAT16 平台限制:Atlas 训练系列产品、Atlas 200I/500 A2 推理产品、Atlas 推理系列产品不支持 BFLOAT16。
- 原地版本不支持 BOOL作为输入数据类型(支持列表中不含 BOOL)。
从源码实现看,aclnnInplacePowTensorTensorGetWorkspaceSize的实现非常简洁,它直接将out指向self后复用普通版本的第一段接口逻辑(aclnn_pow_tensor_tensor.cpp):
aclnnStatus aclnnInplacePowTensorTensorGetWorkspaceSize(const aclTensor* self, const aclTensor* exponent, uint64_t* workspaceSize, aclOpExecutor** executor) { auto out = const_cast<aclTensor*>(self); return aclnnPowTensorTensorGetWorkspaceSize(self, exponent, out, workspaceSize, executor); }返回值(第一段接口错误码)
| 返回码 | 错误码 | 描述 |
|---|---|---|
| ACLNN_ERR_PARAM_NULLPTR | 161001 | 传入的 self、exponent 空指针 |
| ACLNN_ERR_PARAM_INVALID | 161002 | self 和 exponent 的数据类型不在支持的范围之内 |
| ACLNN_ERR_PARAM_INVALID | 161002 | self 和 exponent 的 shape 大于 8 维 |
| ACLNN_ERR_PARAM_INVALID | 161002 | self 和 exponent 无法满足数据类型推导规则 |
| ACLNN_ERR_PARAM_INVALID | 161002 | self 和 exponent 的 shape 无法做 broadcast |
| ACLNN_ERR_PARAM_INVALID | 161002 | self 和 exponent 进行 broadcast 后的 shape 不等于 self 的 shape |
| ACLNN_ERR_PARAM_INVALID | 161002 | self 和 exponent 同时为 bool 数据类型 |
aclnnInplacePowTensorTensor 参数说明
| 参数名 | 输入/输出 | 描述 |
|---|---|---|
| workspace | 输入 | 在 Device 侧申请的 workspace 内存地址 |
| workspaceSize | 输入 | 在 Device 侧申请的 workspace 大小,由第一段接口 aclnnInplacePowTensorTensorGetWorkspaceSize 获取 |
| executor | 输入 | op 执行器,包含了算子计算流程 |
| stream | 输入 | 指定执行任务的 Stream |
约束说明
- 确定性计算:
aclnnPowTensorTensor与aclnnInplacePowTensorTensor默认采用确定性实现,同一输入多次执行结果可复现。确定性计算的通用机制可参见 确定性计算说明。 - 溢出边界行为:在 Atlas 训练系列产品、Atlas 推理系列产品上,如果计算结果取值超过了设定的数据类型取值范围,则会以该数据类型的边界值(饱和值)作为结果返回,即「溢出饱和」而非回绕。
- 数据格式仅支持 ND;从实现看,
CheckFormat会对FORMAT_FRACTAL_NZ格式给出告警日志(不阻断执行)。
源码级实现原理:从 aclnn 接口到 Kernel 的分层链路
理解底层调用链有助于更好地使用接口。aclnnPowTensorTensorGetWorkspaceSize的第一段实现中,注释清晰画出了完整的算子计算流程(见 aclnn_pow_tensor_tensor.cpp):
self exponent | | \ / Contiguous(workspace_0) Contiguous(workspace_2) \ / Cast(workspace_1) Cast(workspace_3) \ / Pow(workspace_4) | Cast(workspace_5) | ViewCopy | result即:先将两个输入分别转为连续 Tensor(Contiguous)→ 按类型推导结果统一数据类型(Cast)→ 执行 Pow 核心计算 → 将结果 Cast 回 out 的数据类型 → 通过 ViewCopy 写入(可能非连续的)输出 Tensor。每一个中间步骤都会在 executor 中登记,最终通过uniqueExecutor->GetWorkspaceSize()汇总返回需要的 workspace 总大小。
两个值得注意的实现细节:
- 空 Tensor 支持:如果 self 或 exponent 为空 Tensor(任一维度为 0),接口直接返回
workspaceSize = 0并成功退出,无需真正调度 Kernel。 - 类型推导:
op::PromoteType(self->GetDataType(), exponent->GetDataType())决定中间计算类型,若推导结果非法则返回ACLNN_ERR_PARAM_INVALID;最终输出再 Cast 回out指定的类型。
核心 Kernel 入口位于 math/pow/op_kernel/pow_apt.cpp,函数签名为pow(GM_ADDR base, GM_ADDR exponent, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling),内部通过 tiling key 分发到不同实现:
- 浮点路径(FP16/BF16/FP32)走
PowTensorScalarFloat<T, T>; - 整型路径(U8/S8/S16/S32)走
PowTensorScalarInteger<T, T>; - Tensor-Tensor 场景走专门的 NDDMA 实现(
PowF16NddmaWith/WithoutLoops、PowBf16Nddma...、PowF32Nddma...、PowU8/S8/S16/S32Nddma...),其中 with/without loops 的区别在于数据是否能够一次性装入 UB(UB 内可容纳 5 维与 8 维两种形态)。
对应的 tiling 生成逻辑在 math/pow/op_host/arch35/pow_tensor_tensor_tiling_arch35.cpp,其中SetOpKey()将 (输入 X dtype, 输入 Y dtype, 输出 Z dtype) 三元组映射到内部算子 key,再结合 broadcast 参数生成最终 tiling key。这解释了文档中"数据格式需与对方一致、shape 需满足 broadcast"等约束的底层来源。
完整调用示例(可编译运行)
以下示例代码与仓库 math/pow/examples/test_aclnn_pow_tensor_tensor.cpp 一致,编译与运行的整体流程请参考 编译与运行样例。
aclnnPowTensorTensor 示例代码
#include <iostream> #include <vector> #include "acl/acl.h" #include "aclnnop/aclnn_pow_tensor_tensor.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); aclFinalize(); return ret); ret = aclrtCreateStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); aclrtResetDevice(deviceId); aclFinalize(); 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 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::vector<int64_t> selfShape = {4, 2}; std::vector<int64_t> expShape = {4, 2}; std::vector<int64_t> outShape = {4, 2}; void* selfDeviceAddr = nullptr; void* expDeviceAddr = nullptr; void* outDeviceAddr = nullptr; aclTensor* self = nullptr; aclTensor* exp = nullptr; aclTensor* out = nullptr; std::vector<float> selfHostData = {0, 1, 2, 3, 4, 5, 6, 7}; std::vector<float> expHostData = {1, 1, 1, 2, 2, 2, 3, 3}; std::vector<float> outHostData = {0, 0, 0, 0, 0, 0, 0, 0}; // 创建self aclTensor ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建exp aclTensor ret = CreateAclTensor(expHostData, expShape, &expDeviceAddr, aclDataType::ACL_FLOAT, &exp); 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; // 调用aclnnPowTensorTensor第一段接口 ret = aclnnPowTensorTensorGetWorkspaceSize(self, exp, out, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnPowTensorTensorGetWorkspaceSize 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); } // 调用aclnnPowTensorTensor第二段接口 ret = aclnnPowTensorTensor(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnPowTensorTensor 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::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 < size; i++) { LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]); } // 6.释放aclTensor和aclScalar,需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(exp); aclDestroyTensor(out); // 7.释放device资源 aclrtFree(selfDeviceAddr); aclrtFree(expDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }以上示例中,self = {0,1,2,3,4,5,6,7}、exponent = {1,1,1,2,2,2,3,3},输出依次为0^1, 1^1, 2^1, 3^2, 4^2, 5^2, 6^3, 7^3。整体调用流程可概括为七个步骤:初始化 device/stream → 构造输入输出 aclTensor → 两段式调用算子 API(先 GetWorkspaceSize,再按需申请 workspace 后执行)→ 同步等待 → 拷贝结果到 host → 释放 tensor → 释放 device 资源。
aclnnInplacePowTensorTensor 示例代码
#include <iostream> #include <vector> #include "acl/acl.h" #include "aclnnop/aclnn_pow_tensor_tensor.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); aclFinalize(); return ret); ret = aclrtCreateStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); aclrtResetDevice(deviceId); aclFinalize(); 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 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::vector<int64_t> selfRefShape = {4, 2}; std::vector<int64_t> expShape = {4, 2}; void* selfRefDeviceAddr = nullptr; void* expDeviceAddr = nullptr; aclTensor* selfRef = nullptr; aclTensor* exp = nullptr; std::vector<float> selfRefHostData = {0, 1, 2, 3, 4, 5, 6, 7}; std::vector<float> expHostData = {1, 1, 1, 2, 2, 2, 3, 3}; // 创建selfRef aclTensor ret = CreateAclTensor(selfRefHostData, selfRefShape, &selfRefDeviceAddr, aclDataType::ACL_FLOAT, &selfRef); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建exp aclTensor ret = CreateAclTensor(expHostData, expShape, &expDeviceAddr, aclDataType::ACL_FLOAT, &exp); CHECK_RET(ret == ACL_SUCCESS, return ret); // 3.调用CANN算子库API,需要修改为具体的API名称 uint64_t workspaceSize = 0; aclOpExecutor* executor; // 调用aclnnInplacePowTensorTensor第一段接口 ret = aclnnInplacePowTensorTensorGetWorkspaceSize(selfRef, exp, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnInplacePowTensorTensorGetWorkspaceSize 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); } // 调用aclnnInplacePowTensorTensor第二段接口 ret = aclnnInplacePowTensorTensor(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnInplacePowTensorTensor 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(selfRefShape); std::vector<float> resultData(size, 0); ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), selfRefDeviceAddr, 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,需要根据具体API的接口定义修改 aclDestroyTensor(selfRef); aclDestroyTensor(exp); // 7.释放device资源 aclrtFree(selfRefDeviceAddr); aclrtFree(expDeviceAddr); if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }与普通版本相比,原地版本示例的关键差异在于:不单独创建 out tensor,计算完成后直接读取selfRef所在 device 内存即为结果;相应地,第一段接口的入参也从 5 个变为 4 个(少了out)。
单测与验证:如何确认接口行为符合预期
仓库在 math/pow/tests/ut/op_api/test_aclnn_pow_tensor_tensor.cpp 提供了大量基于 gtest 与 OP_API_UT 框架的用例,可作为使用与排障的参考:
- broadcast 场景:如
self {1,11}+exponent {23,11}→out {23,11};self {1,11,21,16}+exponent {12,1,1,1}→out {12,11,21,16}(跨 dtype:FLOAT32 与 FLOAT16 混用)。 - 空 Tensor 场景:
self {0,2}+exponent {1,2}→out {0,2},GetWorkspaceSize 返回ACLNN_SUCCESS。 - 非连续 Tensor 场景:通过 viewDim/storageDim/offset 构造非连续输入,验证返回成功。
- 溢出边界场景:FLOAT16 使用
65504.0 / -65504.0边界值;INT8 使用ValueRange(257, 300)验证溢出饱和行为;INT32 常规取值验证。 - 异常入参场景:空指针(返回
ACLNN_ERR_PARAM_NULLPTR)、输出 shape 与 broadcast 结果不一致({11,2}vs{8,2},返回ACLNN_ERR_PARAM_INVALID)、输入超过 8 维(9 维 shape,返回ACLNN_ERR_PARAM_INVALID)、不支持的 dtype(UINT64,返回ACLNN_ERR_PARAM_INVALID)、无法转换的输出类型(FLOAT 输入推导结果转 INT32 输出,返回ACLNN_ERR_PARAM_INVALID)、bool+bool 组合(返回ACLNN_ERR_PARAM_INVALID)。
这些用例恰好与文档「返回值」一节中的错误码一一对应,是理解接口校验规则的直接证据。同时 math/pow/README.md 还提供了其余调用入口:通过aclnnExp2接口调用 Pow(test_aclnn_exp2.cpp)以及通过算子 IR 构图方式调用(pow_proto.h 与 test_geir_pow.cpp),可按需选择调用方式。
小结
aclnnPowTensorTensor/aclnnInplacePowTensorTensor是 CANN/ops-math 中 Pow 数学算子的标准两段式 aclnn 接口。使用时的核心要点可归纳为:先调用 GetWorkspaceSize 完成校验并获取 workspace 大小与 executor,再按返回值申请 Device 内存并调用执行接口;注意 self 与 exponent 需满足 broadcast 关系且不能同时为 BOOL、输出 shape 必须等于 broadcast 后的 shape(原地版本还要求 broadcast 结果不得大于 self)、输入不超过 8 维、特定平台不支持 BFLOAT16、INT32 计算需遵循元素规模与指数取值范围的对照表。结合仓库中 op_api 实现 与 单测用例 可以快速定位参数问题,并进一步深入 Kernel 层(pow_apt.cpp)了解 tiling 与 NDDMA 分派逻辑。
- 算子库
- 人工智能
- CANN
【免费下载链接】ops-math
本项目是CANN提供的数学类基础计算算子库,实现网络在NPU上加速计算。
相关推荐
CANN ops-math 算子 aclnnAcos 与 aclnnInplaceAcos 两段式接口完全指南
CANN ops math 算子 aclnnAcos 与 aclnnInplaceAcos 两段式接口完全指南 导读 本文以 CANN 数学算子库 ops ma
算子库人工智能CANNCANN ops-math 算子 aclnnInplaceMaskedFillTensor 两段式接口开发指南
CANN ops math 算子 aclnnInplaceMaskedFillTensor 两段式接口开发指南 aclnnInplaceMaskedFillTe
算子库人工智能CANNCANN ops-math 算子指南:aclnnAddcmul 与 aclnnInplaceAddcmul 两段式接口详解
CANN ops math 算子指南:aclnnAddcmul 与 aclnnInplaceAddcmul 两段式接口详解 aclnnAddcmul / acl
算子库人工智能CANN
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考