算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载本文面向在华为昇腾 NPU 上开发算子的工程师系统讲解 CANN ops-math 数学算子库中 Addr 算子的功能语义、产品支持情况、参数约束、两段式 aclnn 调用流程与源码级实现原理。通过阅读本文你将掌握如何在 Ascend 950PR/Ascend 950DT 等平台上使用aclnnAddr/aclnnInplaceAddr接口完成 外积 加权求和 的融合计算并理解其背后的算子注册、Shape 推导、Tiling 与 Kernel 实现。Addr 算子功能与数学定义Addr 算子完成的是 PyTorch 中torch.addr的等价语义对一维向量vec1与vec2求外积Outer Product得到一个二维矩阵再将外积结果乘上系数α与自身self乘系数β之后相加输出。其计算公式为$$ \text{out} \beta \cdot \text{self} \alpha \cdot (\text{vec1} \otimes \text{vec2}) $$其中vec1 ⊗ vec2表示向量外积若vec1长度为m、vec2长度为n则外积结果为m × n的矩阵矩阵元素为out[i][j] vec1[i] * vec2[j]。该功能在 math/addr/README.md 中有明确描述。Addr 算子位于 CANN ops-math 仓库的 math/addr 目录目录下完整包含算子定义、Host 侧 Shape 推导与 Tiling、Kernel 实现、op_api 接口以及单元测试等全部交付件是一个标准的单算子交付目录结构。产品支持情况README 中给出了 Addr 算子的产品支持矩阵核心结论是当前仓库中 Addr 算子含 aclnn 接口仅声明支持 Ascend 950 系列产品产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品×Atlas A2 训练系列产品/Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×这一结论可以从算子注册代码中得到印证math/addr/op_host/addr_def.cpp 中Addr的 OpDef 仅通过this-AICore().AddConfig(ascend950, aicoreConfig)注册了ascend950一种 AICore 配置其中开启了DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)与PrecisionReduceFlag(true)而 math/addr/op_kernel 下的 kernel 实现也仅存在于arch35/对应 Ascend 950 架构目录。需要说明的是接口文档 math/addr/docs/aclnnAddraclnnInplaceAddr.md 中标注 Atlas 训练系列、Atlas A2/A3 系列也支持这是 aclnn 接口层的能力——当平台不支持独立 Addr 内核时op_api 层会通过组合 Unsqueeze/Mul/Add 等基础算子来模拟等价计算下文源码级实现原理一节会详述。若需要调用独立内核算子以 README 的产品支持矩阵为准。参数说明Addr 算子共 6 个参数5 个输入x1、x2、x3、beta、alpha与 1 个输出y全部为 ND 数据格式。在 aclnn 接口层x1/x2/x3/y对应aclTensorbeta/alpha对应aclScalar。参数名输入/输出/属性描述数据类型数据格式x1输入待进行 addr 计算的入参公式中的 selfFLOAT、FLOAT16、BFLOAT16、INT8、UINT8、BOOLNDx2输入待进行 addr 计算的入参公式中的 vec1FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、BOOLNDx3输入待进行 addr 计算的入参公式中的 vec2FLOAT、FLOAT16、BFLOAT16、INT8、UINT8、BOOLNDbeta输入待进行 addr 计算的入参公式中的 βFLOAT、FLOAT16、BFLOAT16、INT8、UINT8、BOOLNDalpha输入待进行 addr 计算的入参公式中的 αFLOAT、FLOAT16、BFLOAT16、INT8、UINT8、BOOLNDy输出待进行 addr 计算的出参公式中的 outFLOAT、FLOAT16、BFLOAT16、INT8、UINT8、BOOLND平台差异与数据类型支持README 补充说明不同平台的数据类型支持范围存在差异Atlas 训练系列产品支持 FLOAT、FLOAT16、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOL不支持 BFLOAT16。Atlas A2 训练系列产品/Atlas A2 推理系列产品、Atlas A3 训练系列产品/Atlas A3 推理系列产品、Ascend 950PR/Ascend 950DT支持 FLOAT、FLOAT16、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOL、BFLOAT16。这一差异与 op_api 层源码中的数据类型支持列表一一对应见 math/addr/op_api/aclnn_addr.cppASCEND910_DTYPE_SUPPORT_LISTAtlas 训练系列DOUBLE、FLOAT、FLOAT16、INT64、INT32、INT16、INT8、UINT8、BOOL不含 BFLOAT16ASCEND910B_DTYPE_SUPPORT_LISTAtlas A2 系列在上述基础上增加 BFLOAT16ASCEND950_ADDR_DTYPE_SUPPORT_LISTAscend 950 系列仅 FLOAT、FLOAT16、BF16、INT8、UINT8、BOOL——这与 README 参数表格中的数据类型完全一致也与addr_def.cpp中xDType {ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_INT8, ge::DT_UINT8, ge::DT_BOOL}一致。Shape 与约束说明README 的约束说明为无但实际 Shape 校验有明确规则见接口文档与 math/addr/op_host/addr_infershape.cpp 中的InferShape4Addrvec1x2与vec2x3必须是一维向量selfx1维度不能超过 2可以是 1 维或 2 维self的 shape 需要能与vec1与vec2的外积结果[len(vec1), len(vec2)]满足 broadcast 关系输出y的 shape 恒为[len(vec1), len(vec2)]即InferShape4Addr中out_shape-SetDim(0, x2_shape-GetDim(0)); out_shape-SetDim(1, x3_shape-GetDim(0))输入为未知 rankshape 为[-2]时输出也置为[-2]交由动态 shape 流程处理。beta与alpha另有两条组合约束op_api 层CheckBetaAndAlphaDtyeValid实现若beta/alpha为 bool 类型则self/vec1/vec2的数据类型只能是 bool若self/vec1/vec2均为整型或整型 bool则beta/alpha不能为浮点型。两段式调用接口aclnnAddr 与 aclnnInplaceAddrAddr 算子的 aclnn 接口遵循 CANN 标准的两段式调用模式第一段GetWorkspaceSize接口完成入参校验与 workspace 大小计算第二段接口在指定 stream 上真正执行计算。aclnnAddr 与 aclnnInplaceAddr 功能等价区别在于aclnnAddr需新建一个输出张量对象存储计算结果对应yaclnnInplaceAddr无需新建输出张量对象直接在输入张量的内存中存储计算结果selfRef同时作为输入与输出。从实现上看aclnnInplaceAddrGetWorkspaceSize就是将out参数替换为selfRef后直接复用aclnnAddrGetWorkspaceSize的封装见 math/addr/op_api/aclnn_addr.cppaclnnStatus aclnnInplaceAddrGetWorkspaceSize(aclTensor* selfRef, const aclTensor* vec1, const aclTensor* vec2, const aclScalar* betaOptional, const aclScalar* alphaOptional, uint64_t* workspaceSize, aclOpExecutor** executor) { return aclnnAddrGetWorkspaceSize(selfRef, vec1, vec2, betaOptional, alphaOptional, selfRef, workspaceSize, executor); }第一段接口原型aclnnStatus aclnnAddrGetWorkspaceSize( const aclTensor* self, // 外积扩展矩阵公式中的 self const aclTensor* vec1, // 外积入参第一向量一维 const aclTensor* vec2, // 外积入参第二向量一维 const aclScalar* betaOptional,// 扩展矩阵比例因子 β const aclScalar* alphaOptional,// 外积比例因子 α aclTensor* out, // 输出张量 uint64_t* workspaceSize, // 返回需要在 Device 侧申请的 workspace 大小 aclOpExecutor** executor) // 返回 op 执行器aclnnStatus aclnnInplaceAddrGetWorkspaceSize( aclTensor* selfRef, // 输入/输出外积扩展矩阵及输出矩阵shape 必须为 2 维 const aclTensor* vec1, const aclTensor* vec2, const aclScalar* betaOptional, const aclScalar* alphaOptional, uint64_t* workspaceSize, aclOpExecutor** executor)第二段接口原型aclnnStatus aclnnAddr(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream); aclnnStatus aclnnInplaceAddr(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream);关键参数说明self/selfRef外积扩展矩阵。shape 维度不超过 2inplace 版本要求为 2 维需与vec1、vec2的外积结果满足 broadcast 关系支持非连续 Tensor。vec1/vec2一维向量shape 需与self满足 broadcast 关系支持非连续 Tensor。betaOptional/alphaOptionalaclScalar类型比例因子可以为空指针。为空时按默认值1处理见 math/addr/op_api/aclnn_addr.cpp 中addrProc对空指针分配1标量的逻辑。workspaceSize第一段接口计算出的 workspace 大小第二段接口调用前需据此在 Device 侧aclrtMalloc申请内存。executorop 执行器封装了算子计算流程由第一段接口产出、第二段接口消费。stream指定执行任务的 Stream。返回值与错误码两段接口均返回aclnnStatus状态码。第一段接口完成入参校验以下场景会报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 tensor 或 out 是空指针ACLNN_ERR_PARAM_INVALID161002self/vec1/vec2 的数据类型和数据格式不在支持范围之内ACLNN_ERR_PARAM_INVALID161002vec1 和 vec2 维度不为 1或 self 维度超过 2inplace 版本要求 selfRef 维度为 2ACLNN_ERR_PARAM_INVALID161002self 不能扩展成为 vec1 和 vec2 的外积结果形状ACLNN_ERR_PARAM_INVALID161002beta 或 alpha 为 bool 类型时self/vec1/vec2 数据类型非 boolACLNN_ERR_PARAM_INVALID161002self/vec1/vec2 均为整型或 bool 时beta 或 alpha 为浮点型这些校验逻辑对应 math/addr/op_api/aclnn_addr.cpp 中的CheckParams依次调用CheckNotNull、CheckDtypeValid、CheckShape其中CheckShape还会校验 self 与外积结果 broadcast 后的 shape 必须完全一致。完整调用样例仓库提供了可直接运行的完整示例 math/addr/examples/test_aclnn_addr.cpp同一程序中依次演示了aclnnAddr非原地与aclnnInplaceAddr原地两种调用方式。核心流程如下初始化aclInit→aclrtSetDevice→aclrtCreateStream构造张量aclrtMalloc申请 Device 内存、aclrtMemcpy拷贝 Host 数据、按 shape 计算 strides、aclCreateTensor创建aclTensor创建标量aclCreateScalar创建beta、alpha两段式调用先调GetWorkspaceSize再按返回的workspaceSize申请内存并调用第二段接口同步与取数aclrtSynchronizeStream同步后aclrtMemcpy将结果从 Device 拷回 Host 并打印资源释放aclDestroyTensor、aclrtFree、aclrtDestroyStream、aclrtResetDevice、aclFinalize。示例中的关键调用片段// 构造输入与输出以 FLOAT 为例 std::vectorint64_t inputShape {3, 2}; // self std::vectorint64_t vec1Shape {3}; // vec1 std::vectorint64_t vec2Shape {2}; // vec2 std::vectorint64_t outShape {3, 2}; // out float betaValue 1.5f; float alphaValue 1.5f; // 创建 beta 和 alpha scalar 值 beta aclCreateScalar(betaValue, aclDataType::ACL_FLOAT); alpha aclCreateScalar(alphaValue, aclDataType::ACL_FLOAT); uint64_t workspaceSize 0; aclOpExecutor* executor; // aclnnAddr 接口调用 // 第一段接口校验入参并计算 workspace 大小 ret aclnnAddrGetWorkspaceSize(input, vec1, vec2, beta, alpha, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAddrGetWorkspaceSize 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); } // 第二段接口执行计算 ret aclnnAddr(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAddr failed. ERROR: %d\n, ret); return ret); // 同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 将 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); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %f\n, i, resultData[i]); } // aclnnInplaceAddr 接口调用原地计算 // selfRef 为输入输出张量此处传入 out ret aclnnInplaceAddrGetWorkspaceSize(out, vec1, vec2, beta, alpha, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnInplaceAddrGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); 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); } ret aclnnInplaceAddr(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnInplaceAddr failed. ERROR: %d\n, ret); return ret); ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret);注意示例中inputShape {3, 2}、vec1Shape {3}、vec2Shape {2}即self的 shape 与vec1 ⊗ vec2的外积 shape3×2完全一致无需 broadcast 扩展是最典型的调用形态。若self为{1, 2}或{3, 1}等可广播 shape同样合法。源码级实现原理op_api 层的组合实现与 950 内核实现math/addr/op_api/aclnn_addr.cpp 是 Addr 算子接口的核心实现。addrProc的完整计算流程如下类型提升通过op::PromoteType计算self、vec1、vec2三者中的最高数据类型hightDtype后续统一按该类型计算连续性转换对三个输入分别调用l0op::Contiguous从而支持非连续 Tensor 输入类型转换对三个输入调用l0op::Cast转换到hightDtype标量处理alpha/beta为空指针时按默认值 1 处理否则通过executor-ConvertToTensor转换为张量参与计算分支调度IsSupportAddr(hightDtype)判断当前平台注册态/950是否支持独立 Addr 内核——支持则调用l0op::Addr直达内核否则走addrStub组合路径。addrStub组合路径揭示了 Addr 的数学本质它被拆解为三步将vec1经vecUnsqueezeWithDim(vec1, 1)变为m×1的列向量将vec2经vecUnsqueezeWithDim(vec2, 0)变为1×n的行向量两者Mul即为外积对vec1 × alpha与vec2做乘法得到α·(vec1⊗vec2)对self × beta与上述结果做加法得到β·self α·(vec1⊗vec2)。针对 bool 数据类型addrStubMul/addrStubAdd会特殊处理为LogicalAnd/LogicalOr保证 bool 语义正确。此外还有若干细节优化beta为 0 时跳过 self 的扩展与缩放IgnoreInput模式beta为 1 时跳过缩放IgnoreInputScaling模式空 Tensor 处理上aclnnAddrGetWorkspaceSize检测到vec1或vec2为空时直接返回workspaceSize 0与成功状态。Kernel 侧的 Tiling 与执行在 Ascend 950arch35上Addr 有独立的 Kernel 实现。Tiling 逻辑见 math/addr/op_host/arch35/addr_tiling.cppCheckDtype要求 x1/x2/x3/beta/alpha/y 六者数据类型完全一致CheckShapes校验 x2、x3 为 1 维、x1 为 1 或 2 维且满足 broadcast、y 的 shape 严格等于[len(vec1), len(vec2)]、不允许空 Tensor计算前将 x2 reshape 为[n, 1]x2ReShape_与 x3 的[m]隐式按行广播相配即可完成外积DoOpTiling依据 x1 的数据类型float / fp16·bf16 / int8·bool / uint8分别读取 beta、alpha 常量并据其是否为 0 选择三种不同的 DAG 模板AddrWithoutAlphaCommon系列α 0y β·selfAddrWithoutBetaCommon系列β 0y α·(vec1⊗vec2)AddrWithBetaWithAlphaCommon系列α ≠ 0 且 β ≠ 0完整公式。通过tilingKey_GET_TPL_TILING_KEY(schMode, betaZeroOrNot, alphaZeroOrNot)把β/α 是否为 0这一信息编码进编译期分支。Kernel 侧 math/addr/op_kernel/addr.cpp 的addr入口函数正是依据betaZeroOrNot、alphaZeroOrNot两个模板参数在编译期选择AddrWithoutAlphaProc/AddrWithoutBetaProc/AddrWithBetaWithAlphaProc三条路径从而在零分支开销下执行计算。低精度与整型计算的精度处理math/addr/op_kernel/arch35/addr_dag.h 中定义的 DAG 揭示了各数据类型的计算策略fp16/bf16先Cast到 float32 计算Vec::Mulsfloat、Vec::Mulfloat、Vec::Addfloat最后以CAST_MODE_RINT四舍五入模式Cast回原类型避免低精度累加误差int8/uint8先Cast到 int32/uint32 计算且乘加后通过AndFF与 255 做按位与即取低 8 位处理再 cast 回 8 位类型——源码注释明确说明cast 成 uint8_t 避免转 int8_t 溢出饱和模式float直接在 float32 域内完成β·self α·(vec1⊗vec2)的乘加流水。测试用例验证仓库为 Addr 提供了覆盖全面的单元测试 math/addr/tests/ut/op_api/test_aclnn_addr.cpp这些用例即是对上文参数约束最直接的验证数据类型用例FLOAT16、BFLOAT16、FLOAT、INT8、INT32、INT64、UINT8、BOOL 全部返回ACL_SUCCESS非法类型用例COMPLEX64 类型返回ACLNN_ERR_PARAM_INVALID标量类型约束用例bool 类型beta/alpha搭配非 bool 输入张量返回ACLNN_ERR_PARAM_INVALID整型输入搭配浮点型beta/alpha同样返回ACLNN_ERR_PARAM_INVALIDShape 约束用例vec非 1 维、self维度超过 2、self无法 broadcast 到外积 shape 均返回ACLNN_ERR_PARAM_INVALID空 Tensor 用例vec1或vec2为空 Tensor 时返回ACL_SUCCESS对应 op_api 层的空 Tensor 短路处理空指针用例self/vec/out为nullptr返回ACLNN_ERR_PARAM_NULLPTR而beta/alpha为nullptr时返回ACL_SUCCESS按默认值 1 处理类型混合用例vec2使用 INT16 而其余为 FLOAT 时返回ACL_SUCCESS验证了类型提升Cast 到最高类型路径。此外 math/addr/tests/ut/op_host 下还包含 Tiling 与 InferShape 的 Host 侧单测可进一步验证 shape 推导与 tiling 逻辑。小结Addr 算子以out β·self α·(vec1⊗vec2)的融合语义将外积与加权求和合并为一次算子调用避免了中间矩阵的多次读写适合线性代数与矩阵分解类计算图的显存与带宽优化。在 CANN ops-math 仓库中其交付链路完整清晰接口层math/addr/op_api/aclnn_addr.cpp 提供两段式aclnnAddr/aclnnInplaceAddr接口与组合实现定义层math/addr/op_host/addr_def.cpp 完成算子注册与ascend950平台配置推导层math/addr/op_host/addr_infershape.cpp 实现输出 shape 推导含未知 rank 处理调度层math/addr/op_host/arch35/addr_tiling.cpp 依据 β/α 取值选择三套计算模板执行层math/addr/op_kernel/addr.cpp 与 math/addr/op_kernel/arch35/addr_dag.h 完成核上 DAG 计算并对低精度/整型类型做精度保护验证层math/addr/tests/ut/op_api/test_aclnn_addr.cpp 覆盖全数据类型与各类非法入参场景。读者可参考 math/addr/examples/test_aclnn_addr.cpp 直接运行验证并结合 math/addr/docs/aclnnAddraclnnInplaceAddr.md 查阅更完整的接口细节。赞分享算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math 算子指南aclnnAddr 与 aclnnInplaceAddr 向量外积融合算子外积-加-缩放使用详解CANN ops math 算子指南aclnnAddr 与 aclnnInplaceAddr 向量外积融合算子外积 加 缩放使用详解 导读 本文围绕 CA算子库人工智能CANNCANN ops-math 算子接口实战aclnnLinalgCross 向量叉乘Cross算子详解CANN ops math 算子接口实战aclnnLinalgCross 向量叉乘Cross算子详解 本篇文章以 CANN ops math 开源仓库中算子库人工智能CANNCANN ops-math 中 Crosslinalg.cross向量叉积算子的原理与 aclnn 调用实战CANN ops math 中 Crosslinalg.cross向量叉积算子的原理与 aclnn 调用实战 导读 本文围绕 CANN ops math 仓算子库人工智能CANN上一篇从零到实战PayloadsAllTheThings的Web安全payload库上手指南下一篇RAG-Challenge-2并行处理指南大幅提升PDF解析效率创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考