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

CANN ops-math AxpyV2 算子 aclnnAxpyV2 接口深度指南从两段式调用到源码原理【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathAxpyV2 是 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产品是否支持Atlas A2 训练系列产品 / Atlas 800I A2 推理产品 / A200I A2 Box 异构组件√在算子定义层面这一支持范围体现在 axpy_v2_def.cpp 中的this-AICore().AddConfig(ascend910b)即该算子以 AICoreVector 核方式注册在 ascend910b 芯片配置上。二、功能说明与计算公式AxpyV2 完成的是经典的 AXPYa times X plus Y运算算子功能源操作数 2src2Tensor中每个元素与标量alphaScalar对应元素求积后与源操作数 1src1Tensor中的对应元素相加。计算公式$$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 也不例外调用顺序为先调用aclnnAxpyV2GetWorkspaceSize获取计算所需 workspace 大小同时获得包含算子计算流程的执行器executor再调用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、alphaScalarout对应 dstTensor。头文件为aclnn_axpy_v2.h。五、aclnnAxpyV2GetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorself输入待进行 axpy_v2 计算的入参公式中的 src1Tensor无FLOAT、FLOAT16、BFLOAT16、INT32ND0-8√other输入待进行 axpy_v2 计算的入参公式中的 src2Tensorshape 与 x1 相同FLOAT、FLOAT16、BFLOAT16、INT32ND0-8√alpha输入待进行 axpy_v2 计算的入参公式中的 alphaScalarshape 为 []FLOAT、FLOAT16、BFLOAT16、INT32ND0-8√out输出待进行 axpy_v2 计算的出参公式中的 dstTensorshape 与 x1 相同FLOAT、FLOAT16、BFLOAT16、INT32ND0-8√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小-----executor输出返回 op 执行器包含了算子计算流程-----关键点解读数据类型四个 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。数据格式统一为 NDOpDef 中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。第一段接口会完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 tensor 是空指针ACLNN_ERR_PARAM_INVALID161002x1 的数据类型和数据格式不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002x1 的数据维度超过了 8 维ACLNN_ERR_PARAM_INVALID161002x1 和 out 的数据形状不一致此外从 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 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnAxpyV2GetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream第二段接口返回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::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 otherShape {2, 2}; std::vectorint64_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::vectorfloat selfHostData {0, 1, 2, 3}; std::vectorfloat otherHostData {0, 1, 2, 3}; 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); // 创建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::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(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}alpha1.2result[i] self[i] 1.2 * other[i]即 0、2.2、4.4、6.6。九、源码级原理剖析从接口到 Kernel 的执行链路1. 算子定义注册OpDefaxpy_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 推导InferShapeaxpy_v2_infershape.cpp 中的InferShapeAxpyV2直接将输出y的 shape 逐维拷贝为输入x1的 shapeyShape-SetDim(i, dim)从实现上印证了“out 的 shape 与 x1 相同”这一使用说明。3. Tiling 切分Host 侧axpy_v2_tiling.cpp 是性能调优的关键核心策略包括平台信息获取通过PlatformAscendC读取 UB 大小与 AIV 核数GetCoreMemSize(UB)、GetCoreNumAiv32B 对齐inputLengthAlgin32将输入数据按 256 字节BLOCK_SIZE向上取整对齐单/双 Buffer 动态决策以DATA_NUM_32B332 位数据每 256B 可放 3 份和DATA_NUM_16B616 位数据每 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 中的alphaLocalFphalf/bf16/int32 均先转换到 float 再参与乘加类型提升计算FLOAT16/BFLOAT16 输入先Cast到 float再执行MulAddDst(x1LocalFp, x2LocalFp, alphaLocalFp, processDataNum)一次性完成乘加最后以CAST_RINT半精度输出或CAST_TRUNCINT32 输出写回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.cppKernel 侧算子单测位于 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考