CANN ops-math 算子 aclnnExpand 使用指南：广播扩展接口的原理解析与编程实践

算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载本指南以 CANN ops-math 数学算子库中的 aclnnExpand 算子文档 为核心系统讲解 Expand广播扩展算子的功能语义、两段式 aclnn 接口原型、参数约束、返回码以及完整调用示例并深入其所在算子目录 math/expand 的源码实现帮助你理解该接口从入参校验、算子下发到 NPU/AICPU 内核执行的完整链路能够独立编写并运行调用 aclnnExpand 的应用代码。算子功能与数学定义Expand 算子将输入张量self广播broadcast成指定 shape 的张量。其语义与 PyTorch 的Tensor.expand一致当输入张量的某一维度为 1 时可以在对应维度上重复该维数据以扩展至目标 shape维度不足时则在头部补 1 再参与广播。扩展后张量 Y 与原始张量 X 的元素映射关系为$$Y_{i_1,\dots,i_k,\dots,i_n} X_{i_1,\dots,\lfloor i_k / m \rfloor,\dots,i_n}$$其中 Y 为 expand 后的张量m 为第 k 维的扩展倍数。从算子仓库的 README 可得到直观示例输入 tensor 的 shape 为(1, 4)指定 size 为(2, 4)则输出是 shape 为(2, 4)的 tensor——第 0 维由 1 扩展为 2。产品支持情况根据 aclnnExpand.md 与 README当前支持的产品如下产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品支持Atlas 训练系列产品支持两段式接口与函数原型aclnnExpand 采用 CANN aclnn 算子库通用的两段式接口设计必须先调用aclnnExpandGetWorkspaceSize获取计算所需 workspace 大小以及包含了算子计算流程的执行器再调用aclnnExpand接口执行计算。aclnnStatus aclnnExpandGetWorkspaceSize( const aclTensor* self, const aclIntArray* size, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnExpand( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)从源码 op_api/aclnn_expand.cpp 可以看到两段接口的实现分工第一段aclnnExpandGetWorkspaceSize依次完成创建 OpExecutor、参数校验CheckParams、空 tensor 处理然后通过l0op::Contiguous将输入self转成连续 tensor调用l0op::Expand构图计算再通过l0op::ViewCopy将结果拷贝到输出outout可以是非连续 tensor最后通过uniqueExecutor-GetWorkspaceSize()汇总计算所需的 workspace 大小第二段aclnnExpand直接调用CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成实际计算这也是所有 aclnn 接口统一的执行入口。这里还包含两个值得注意的特殊处理空 tensor 短路当self或out为空 tensor 时workspaceSize直接置 0 并返回成功无需真正下发计算0 维标量特例当selfDimNum 0 size-Size() 0即 0 维标量时直接使用l0op::ViewCopy完成拷贝而非常规的 Contiguous Expand 流程。aclnnExpandGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度shape非连续张量 Tensorself输入表示待广播的目标张量公式中的 selfshape 与 size 满足 broadcast 关系FLOAT16、FLOAT、UINT8、INT8、INT32、INT64、BOOL、BF16ND0-8√size输入广播时指定的 size-INT---out输出广播后的张量shape 需要满足 self 的 shape 根据 size 的推导结果与 self 一致ND-√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小-----executor输出返回 op 执行器包含了算子计算流程-----BF16 平台差异说明Atlas 训练系列产品Ascend 910与Atlas 推理系列产品Ascend 310P不支持 BFLOAT16 数据类型。这一限制在源码中得到印证——aclnn_expand.cpp 中定义了按 NPU 架构区分的支持列表默认的ASCEND910_DTYPE_DTYPE_SUPPORT_LIST仅包含 FLOAT16、FLOAT、UINT8、INT8、INT32、INT64、BOOL 七种类型而DAV_2201910B与DAV_3510350/950 系列架构额外支持DT_BF16。此外CheckFormat会对self的存储格式进行检查若为 NZFORMAT_FRACTAL_NZ格式会打印警告日志提示该格式可能导致精度问题见 aclnn_expand.cpp。返回值与错误码aclnnStatus 为函数返回状态码具体含义参见 aclnn 返回码说明。第一段接口完成入参校验以下场景会报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self、size、out 是空指针ACLNN_ERR_PARAM_INVALID161002self、out 的数据类型或数据格式不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self 与 out 的数据类型不一致ACLNN_ERR_PARAM_INVALID161002self 或 out 的 shape 和 size 不匹配size 个数小于 self 的 dimNumself 任意维度不等于 size 偏移后对应维度且不等于 1偏移量为 size 与 self 的 dimNum 差值output 的 shape 不等于预期 shapesizeACLNN_ERR_PARAM_INVALID161002self 最大维度超过 8这些校验在 aclnn_expand.cpp 的CheckParams中依次执行CheckNotNull空指针检查、CheckDtypeValiddtype 一致性 支持列表检查、CheckShape广播关系检查、CheckMaxDimension最大 8 维限制对应源码中MAX_SUPPORT_DIM 8。其中CheckShape还支持 size 中维度值为-1的语义-1表示该维继承 self 对应维的大小见 aclnn_expand.cpp。aclnnExpand 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnExpandGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream第二段接口同样返回 aclnnStatus 状态码参见 aclnn 返回码说明。约束说明确定性计算aclnnExpand 默认确定性实现同一输入多次执行结果可复现。除上述参数校验约束外无其他额外约束。调用示例以下示例代码摘自 aclnnExpand.md仓库中还提供了可直接运行的完整样例 test_aclnn_expand.cpp使用 RAII 智能指针管理资源以及图模式调用样例 test_geir_expand.cpp。具体编译和执行过程请参考编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_expand.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 {4, 2}; std::vectorint64_t outShape {4, 2}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclIntArray* size nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectorfloat outHostData {0, 0, 0, 0, 0, 0, 0, 0}; int64_t sizeValue[2] {4, 2}; size aclCreateIntArray((sizeValue[0]), 2); // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); 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; // 调用aclnnExpand第一段接口 ret aclnnExpandGetWorkspaceSize(self, size, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnExpandGetWorkspaceSize 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); } // 调用aclnnExpand第二段接口 ret aclnnExpand(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnExpand 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 resultSize GetShapeSize(outShape); std::vectorfloat resultData(resultSize, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, resultSize * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy resultData from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i resultSize; i) { LOG_PRINT(resultData[%ld] is: %f\n, i, resultData[i]); } // 6. 释放aclTensor和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(out); aclDestroyIntArray(size); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例程序遵循 aclnn 调用的标准七步流程① device/stream 初始化 → ② 构造输入输出 aclTensor通过aclCreateIntArray创建 size、aclCreateTensor创建 self/out→ ③ 两段式接口调用先aclnnExpandGetWorkspaceSize拿 workspace 大小并申请内存再aclnnExpand执行→ ④aclrtSynchronizeStream同步 → ⑤ 将结果从 device 侧拷回 host 侧并打印 → ⑥ 销毁 tensor/intArray → ⑦ 释放 device 资源并aclFinalize。源码级实现纵深从算子定义到内核执行算子定义与 Infershape算子注册在 op_host/expand_def.cpp 中输入x与输出y支持 FLOAT、FLOAT16、INT32、UINT8、INT8、BOOL、BF16、INT64 共 8 种数据类型格式均为 ND输入shape为 INT32/INT64 的常量张量且声明了ValueDepend(OPTIONAL)该输入参与 shape 推导属于值依赖输入。AICore 配置声明了动态 shape、动态 rank 支持并指定内核文件expand_apt。形状推导逻辑在 op_host/expand_infershape.cpp 的InferShape4Expand中实现其广播规则与 PyTorch 语义对齐shape维数不能小于x的维数否则报错多出的头部维度视为对x补 1目标维度为-1时输出该维取x对应维的大小目标维度为 1 而x对应维不为 1 时输出该维取x的维度不做双向广播广播合法性检查x维度必须为 1 或与目标维度相等否则报cannot be broadcast错误。同时该文件通过InputsDataDependency({1})声明 shape 输入为值依赖框架会先取到 shape 的常量值再执行推导。Tiling 与 KernelAscend 350/950 架构arch35的 tiling 实现在 op_host/arch35/expand_tiling_arch35.cpp 中从源码结构看它复用了广播类算子的 tiling 基类brcto::BroadcastToTilingAscendC来自 conversion/broadcast_to 模块先通过AdjustShapesToSameDimNum将输入输出对齐到相同维数再依次执行广播规则校验、DeleteOneSizeAxis删除维度为 1 的轴与MergeAxis轴合并等优化最后调用DoTiling生成计算分片信息。内核实现 op_kernel/expand_apt.cpp 更为直接——expand内核函数直接调用broadcast_to_impl(x, shape, y, workspace, tiling)即 Expand 在 NPU 侧完全复用 broadcast_to 的向量内核实现。AICPU 兜底实现op_kernel_aicpu/expand_aicpu.cpp 提供了 AICPU 侧的兜底实现ExpandCpuKernel其计算流程为空 tensor 处理HandleEmptyTensor在 shape 元素数为 0标量场景时直接拷贝输入到输出shape 归一化NormalizeExpandShape对 target_shape 做合法性与广播规则校验-1继承输入维度、维度 1 保持、输入非 1 且不等于目标则报错逐层扩展ExpandByLayer从最低维开始找到第一个输入与目标不同的维度作为 break_axis计算出copy_size需要整块复制的元素数与expand_factor扩展倍数通过CalculateOutIndex按块复制生成中间结果迭代直到所有维度对齐。内核注册通过OPS_MATH_REGISTER_CPU_KERNELV2(kExpand, ExpandCpuKernel)完成并在 expand_aicpu_def.cpp 中补充对应定义。测试与验证仓库为 Expand 提供了多层次验证API 层单测tests/ut/op_api/test_aclnn_expand.cppInfershape 单测tests/ut/op_host/test_expand_infershape.cpp覆盖广播合法性、-1维度继承等规则Tiling 单测tests/ut/op_host/arch35/test_expand_tiling.cppAICPU 单测tests/ut/op_kernel_aicpu/test_expand.cppST 用例tests/st/aclnnExpand/atk_aclnnExpand.json 与配套的 executor_aclnnExpand.py。常见问题与建议shape 与 size 不匹配self任意维度必须等于 size 偏移后对应维度或为 1且out的 shape 必须等于 size 推导结果size的维数不能小于self的维数。数据类型不一致self与out的数据类型必须完全一致且需落在当前产品架构支持的 dtype 列表中注意 Ascend 910 / 310P 不支持 BF16。维度上限self最大支持 8 维超出会返回ACLNN_ERR_PARAM_INVALID。空指针与内存self、size、out传空指针会返回ACLNN_ERR_PARAM_NULLPTRworkspace 只有在workspaceSize 0时才需要申请申请与释放均应使用aclrtMalloc/aclrtFree成对完成。非连续张量self与out均支持非连续 tensor接口内部通过Contiguous与ViewCopy自动完成连续性处理调用方无需手动转换。赞分享算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math 算子接口指南aclnnClampMaxTensor 与 aclnnInplaceClampMaxTensor 的使用与原理CANN ops math 算子接口指南aclnnClampMaxTensor 与 aclnnInplaceClampMaxTensor 的使用与原理 本指南算子库人工智能CANNCANN ops-math 算子开发实战aclnnExpandv 接口深度解析与 NPU 广播实现原理CANN ops math 算子开发实战aclnnExpandv 接口深度解析与 NPU 广播实现原理 导读 本文以 CANN ops math 开源仓库中的算子库人工智能CANNCANN ops-math 中 Roll 算子的 aclnnRoll 接口使用指南与实现原理CANN ops math 中 Roll 算子的 aclnnRoll 接口使用指南与实现原理 本文以 CANN ops math 开源仓库中 experimen算子库人工智能CANN上一篇Pokémon卡片CSS全息特效部署指南从开发到生产的完整流程下一篇res-downloader3步捕获视频号、抖音无水印资源创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考