企业官网建设流程全解析

CANN ops-math AxpyV2 算子 aclnnAxpyV2 接口深度指南：从两段式调用到源码原理

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math

AxpyV2 是 CANN ops-math 开源仓库中实现线性组合计算dst = src1 + alpha * src2的基础数学算子，本文以其官方接口文档 aclnnAxpyV2.md 为骨架，结合 axpy_v2 算子目录 下的 op_api、op_host、op_kernel 源码与测试用例，系统讲解 aclnnAxpyV2 的产品支持范围、两段式 API 调用流程、参数约束与返回码，并深入剖析其在 NPU 上的 shape 推导、tiling 切分与 Vector 计算实现。读完本文，你将能够独立编写、编译并运行调用 aclnnAxpyV2 的 C++ 样例程序，并理解该算子在昇腾 AI Core 上从接口到 kernel 的完整执行链路。

一、产品支持情况

AxpyV2 算子当前仅支持以下产品系列（来源于 aclnnAxpyV2.md 与 README.md）：

在算子定义层面，这一支持范围体现在 axpy_v2_def.cpp 中的this->AICore().AddConfig("ascend910b")，即该算子以 AICore（Vector 核）方式注册在 ascend910b 芯片配置上。

二、功能说明与计算公式

AxpyV2 完成的是经典的 AXPY（a times X plus Y）运算：

• 算子功能：源操作数 2（src2Tensor）中每个元素与标量（alphaScalar）对应元素求积后，与源操作数 1（src1Tensor）中的对应元素相加。
• 计算公式：

$$dstTensor_i = src1Tensor_i + alphaScalar * src2Tensor_i$$

关于精度提升，文档明确说明：对于数据类型为 FLOAT16 和 BFLOAT16 的情况，需要类型转换为 FLOAT32 进行计算；同时支持 alphaScalar 与 src1Tensor 数据类型不一致的情况（例如 alpha 为 FLOAT，而 x1/x2 为 FLOAT16）。这一约束在 axpy_v2.h 的 kernel 实现中得到了印证：当输入为 half 或 bfloat16_t 时，数据先经DataCopy搬运进 UB，再通过Cast(x1LocalFp, x1LocalLast, RoundMode::CAST_NONE, ...)提升为 float 计算，最后在写出前用Cast(yLocal, x1LocalFp, RoundMode::CAST_RINT, ...)转回原精度。

三、两段式接口（Two-Phase API）概述

CANN 算子库中每个算子都采用两段式接口设计，具体机制见 two_phase_api.md。AxpyV2 也不例外，调用顺序为：

1. 先调用aclnnAxpyV2GetWorkspaceSize获取计算所需 workspace 大小，同时获得包含算子计算流程的执行器（executor）；
2. 再调用aclnnAxpyV2执行实际计算。

这种设计的好处是：第一段接口在 Host 侧完成全部参数校验与 shape 推导（编译准备），第二段接口只负责异步下发任务，便于调用方统一管理 workspace 内存与执行器生命周期。

四、函数原型

aclnnStatus aclnnAxpyV2GetWorkspaceSize( const aclTensor *self, const aclTensor *other, const aclTensor *alpha, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)

aclnnStatus aclnnAxpyV2( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)

其中self、other、alpha分别对应公式中的 src1Tensor、src2Tensor、alphaScalar，out对应 dstTensor。头文件为aclnn_axpy_v2.h。

五、aclnnAxpyV2GetWorkspaceSize 参数说明

关键点解读：

• 数据类型：四个 tensor 均支持 FLOAT、FLOAT16、BFLOAT16、INT32 四种类型，但第一段接口内部要求self、other、alpha、out的 dtype 全部一致（见 axpy_v2.cpp 的 dtype 一致性检查）；对应地，OpDef 注册中的输入输出数据类型列表也为{DT_FLOAT, DT_INT32, DT_FLOAT16, DT_BF16}（见 axpy_v2_def.cpp）。
• 数据格式：统一为 ND，OpDef 中FormatList({ge::FORMAT_ND})与之对应。
• 维度：0-8 维，超过 8 维会报参数错误。
• 非连续 Tensor：均支持 √，且 OpDef 中对输入输出均声明了AutoContiguous()，即框架会自动完成非连续内存的连续化处理。
• alpha 的标量特性：alpha 的 shape 固定为 []，对应 OpDef 中 alpha 输入注册的.Scalar()声明（见 axpy_v2_def.cpp），因此在样例中以aclCreateScalar方式创建。

返回值与错误码

第一段接口返回aclnnStatus状态码，完整定义参见 aclnn_return_code.md。第一段接口会完成入参校验，出现以下场景时报错：

此外，从 axpy_v2.cpp 的实现还可以看到，即使参数校验通过，若self的数据类型不在 AICore 支持列表{DT_FLOAT, DT_INT32, DT_FLOAT16, DT_BF16, DT_UINT8, DT_INT8, DT_INT64, DT_BOOL}（注册基座环境）中，同样会以ACLNN_ERR_PARAM_INVALID报错拒绝执行。

六、aclnnAxpyV2 参数说明

第二段接口返回aclnnStatus状态码，同样参见 aclnn_return_code.md。注意workspace与workspaceSize必须来自第一段接口的产出，且仅当workspaceSize > 0时才需要实际申请 Device 内存。

七、约束说明

官方文档标注“无”额外约束。不过结合源码仍可归纳出以下隐含约束，供调用时参考：

• 输入输出维度最高 8 维；
• self、other、alpha、out数据类型需保持一致（dtype 校验在 axpy_v2.cpp）；
• 数据类型需落在 FLOAT、FLOAT16、BFLOAT16、INT32 范围内，格式为 ND；
• 支持非连续 Tensor（由AutoContiguous()保证内存自动连续化）。

八、完整调用示例

以下示例代码摘自 aclnnAxpyV2.md（仓库内另有完整可运行版本 test_aclnn_axpy_v2.cpp），以 shape 为 {2, 2} 的 FLOAT 数据演示完整流程，具体编译和执行过程请参考 compile_and_run_sample.md。

#include <iostream> #include <vector> #include "acl/acl.h" #include "aclnn_axpy_v2.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 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 = {2, 2}; std::vector<int64_t> otherShape = {2, 2}; std::vector<int64_t> outShape = {2, 2}; void* selfDeviceAddr = nullptr; void* otherDeviceAddr = nullptr; void* outDeviceAddr = nullptr; aclTensor* self = nullptr; aclTensor* other = nullptr; aclScalar* alpha = nullptr; aclTensor* out = nullptr; std::vector<float> selfHostData = {0, 1, 2, 3}; std::vector<float> otherHostData = {0, 1, 2, 3}; std::vector<float> outHostData = {0, 0, 0, 0}; // 创建self aclTensor ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建other aclTensor ret = CreateAclTensor(otherHostData, otherShape, &otherDeviceAddr, aclDataType::ACL_FLOAT, &other); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建alpha aclScalar float scalarValue = 1.2f; alpha = aclCreateScalar(&scalarValue, ACL_FLOAT); // 创建out aclTensor ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out); CHECK_RET(ret == ACL_SUCCESS, return ret); // 3. 调用CANN算子库API，需要修改为具体的API名称 // aclnnAxpyV2接口调用示例 uint64_t workspaceSize = 0; aclOpExecutor* executor; // 调用aclnnAxpyV2第一段接口 ret = aclnnAxpyV2GetWorkspaceSize(self, other, alpha, out, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAxpyV2GetWorkspaceSize 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); } // 调用aclnnAxpyV2第二段接口 ret = aclnnAxpyV2(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAxpyV2 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("aclnnAxpyV2 result[%ld] is: %f\n", i, resultData[i]); } // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(other); aclDestroyScalar(alpha); aclDestroyTensor(out); // 7. 释放device资源，需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(otherDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }

示例执行结果（self={0,1,2,3}，other={0,1,2,3}，alpha=1.2）：result[i] = self[i] + 1.2 * other[i]，即 0、2.2、4.4、6.6。

九、源码级原理剖析：从接口到 Kernel 的执行链路

1. 算子定义注册（OpDef）

axpy_v2_def.cpp 通过OP_ADD(AxpyV2)将算子注册进算子信息库，声明了三个必选输入x1、x2、alpha与一个输出y，其中alpha被标记为.Scalar()。所有输入输出均声明了数据类型列表{DT_FLOAT, DT_INT32, DT_FLOAT16, DT_BF16}、格式FORMAT_ND以及AutoContiguous()自动连续化。正是这份声明，决定了第一段接口参数校验的合法取值集合。

2. Shape 推导（InferShape）

axpy_v2_infershape.cpp 中的InferShapeAxpyV2直接将输出y的 shape 逐维拷贝为输入x1的 shape（yShape->SetDim(i, dim)），从实现上印证了“out 的 shape 与 x1 相同”这一使用说明。

3. Tiling 切分（Host 侧）

axpy_v2_tiling.cpp 是性能调优的关键，核心策略包括：

• 平台信息获取：通过PlatformAscendC读取 UB 大小与 AIV 核数（GetCoreMemSize(UB)、GetCoreNumAiv）；
• 32B 对齐：inputLengthAlgin32将输入数据按 256 字节（BLOCK_SIZE）向上取整对齐；
• 单/双 Buffer 动态决策：以DATA_NUM_32B=3（32 位数据每 256B 可放 3 份）和DATA_NUM_16B=6（16 位数据每 256B 可放 6 份）估算单 Buffer 所需 UB 大小。若singleBufferNeedSize <= coreNum * ubSize，采用单流水SINGLE_BUFFER_NUM（性能更优）；否则退化为双流水DOUBLE_BUFFER_NUM；
• 核间负载均衡：通过CalculateCoreBlockNums将数据划分为小核数据量/大核数据量，并把尾部数据（tailBlockNum）均匀分摊，避免少数核成为瓶颈；
• TilingKey 选择：单 Buffer 对应ELEMENTWISE_TPL_SCH_MODE_1，双 Buffer 对应ELEMENTWISE_TPL_SCH_MODE_0，通过context->SetTilingKey与SetBlockDim(coreNum)将切分结果下发给 Kernel。

切分结果通过 axpy_v2_tiling_data.h 中定义的AxpyV2TilingData结构体（含 smallCoreDataNum、bigCoreDataNum、tileDataNum、bufferNum 等字段）传递。

4. Kernel 实现（Device 侧）

axpy_v2.cpp 中axpy_v2内核根据 TilingKey 选择实例化NsAxpyV2::AxpyV2<..., 2>（双 Buffer）或<..., 1>（单 Buffer），随后执行Init → Process流水。

真正完成计算的 axpy_v2.h 实现要点：

• 标量广播：InitAlphaLocal将 GM 上的标量 alpha 读取后按 tile 大小Duplicate展开到 UB 中的alphaLocalFp（half/bf16/int32 均先转换到 float 再参与乘加）；
• 类型提升计算：FLOAT16/BFLOAT16 输入先Cast到 float，再执行MulAddDst(x1LocalFp, x2LocalFp, alphaLocalFp, processDataNum)一次性完成乘加，最后以CAST_RINT（半精度输出）或CAST_TRUNC（INT32 输出）写回；
• INT32 专用路径：当 x1、x2、alpha 全为 int32 时，直接以 int32 执行Mul与Add，输出 float 或 half/bf16 时再做类型转换，避免不必要的浮点开销；
• 事件同步：通过SetFlag/WaitFlag对MTE2→V、V→MTE3硬事件进行同步，保证数据搬运、向量计算、结果回写三级流水之间的正确性。

十、编译、运行与测试验证

• 编译运行：完整可运行样例见 test_aclnn_axpy_v2.cpp，编译与执行环境准备请遵循 compile_and_run_sample.md 的指引；
• UT 测试：Host 侧 tiling 单测位于 test_axpy_v2_tiling.cpp，Kernel 侧算子单测位于 test_axpy_v2.cpp，配套测试数据生成与比对脚本为 gen_data.py 与 compare_data.py，可用于验证不同 dtype、不同 shape 下的计算结果正确性；
• 仓库配套文档：如需了解算子 API 背后的数据结构（aclTensor/aclScalar 的创建与释放）、ND 数据格式、非连续 Tensor 支持与两段式接口设计，可继续阅读 basic_concept.md、data_structure.md、non_contiguous_tensor.md 与 two_phase_api.md。

综上，aclnnAxpyV2 作为一个典型的逐元素线性组合算子，其接口层约束清晰、调用流程规范，而其底层在 shape 推导、UB 切分、流水调度与类型提升上的实现，也为在 Atlas A2 系列产品上编写和优化同类 Elementwise 算子提供了可直接复用的参考范式。

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math