CANN ops-math 算子解析PackV2 双张量维度堆叠Stack算子的实现与 aclnn 调用实战【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathPackV2 是 CANN ops-math 仓库experimental/conversion/pack_v2中实现的双输入张量维度堆叠算子它将两个形状完全相同的张量按照指定维度d拼接为一个新张量等价于 PyTorch 的torch.stack在双输入场景下的行为。本文基于该算子目录下的 README、Host 侧算子定义与 Tiling 逻辑、Kernel 侧 AscendC 实现以及 aclnn 调用样例完整讲解 PackV2 的功能语义、参数约束、NPU 侧实现原理与可运行的调用示例帮助开发者在 Atlas A2 系列产品上快速理解并使用该算子。功能说明两个张量的维度堆叠StackPackV2 算子实现的是张量的维度堆叠Stack操作沿用户指定的维度将两个形状完全相同的输入张量堆叠成一个新的张量。以 README 中的示例为例输入selfX、selfY均为 shape 为[2, 3, 4, 5]的 Tensor若dim 1可取 0、1、2、3、4则输出 shape 为[2, 2, 3, 4, 5]若dim 3则输出 shape 为[2, 3, 4, 2, 5]。即输出张量的维度数等于输入维度数加 1在dim处插入一个新的维度大小为 2其余维度与输入完全一致。由于只有两个输入插入的新维度大小恒为 2因此输出总元素数为单个输入元素数的 2 倍。值得注意的是README 描述的语义是新增一个维度Stack 语义而仓库中 pack_v2_infershape.cpp 的形状推导实现的是在最后一维翻倍yShape-SetDim(i, dim * 2)仅对最后一维生效这与 README 及调用样例示例中outShape {1, 2, 3, 8}即最后一维由 4 翻倍为 8中的实际行为一致。读者在使用时应以实际算子行为最后一维拼接、d取值为最后一维索引为准并注意 README 与当前实现的差异。产品支持情况PackV2 算子的支持产品与芯片平台如下表所示见 README产品是否支持Atlas A2 训练系列产品 / Atlas 800I A2 推理产品 / A200I A2 Box 异构组件√在源码层面该支持范围由 pack_v2_def.cpp 中的算子注册配置进一步印证this-AICore().AddConfig(ascend910b, aicoreConfig)即为ascend910b芯片对应 Atlas A2 系列注册 AICore 配置同时 pack_v2_binary.json 的目录层级也直接以ascend910b为平台维度组织算子编译配置。参数说明PackV2 算子的接口参数如下摘自 README参数名输入/输出/属性描述数据类型数据格式x输入张量需要进行维度堆叠的输入张量。见下方NDy输入张量需要进行维度堆叠的输入张量。见下方NDout输出维度为 4 维shape 由 dim 和原 selfx 的 shape 共同决定dtype 需要与 selfx 一致。同 xNDd属性堆叠的目标维度即 dim默认值为 0。int—在 Atlas A2 训练系列产品 / Atlas 800I A2 推理产品 / A200I A2 Box 异构组件上数据类型支持FLOAT、FLOAT16。源码中的参数定义佐证从 Host 侧算子定义 pack_v2_def.cpp 可以看到算子接口的实际定义比 README 描述更完整输入x与y均为必选输入ParamType(REQUIRED)支持DT_FLOAT、DT_INT32、DT_INT16、DT_FLOAT16四种数据类型格式限定为FORMAT_ND并配置了AutoContiguous()内存自动连续化未知形状时同样限定为FORMAT_ND输出z同样为必选输出数据类型与输入保持一致属性d通过this-Attr(d).AttrType(OPTIONAL).Int(0)声明为可选属性默认值为0动态特性DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)、PrecisionReduceFlag(true)说明算子支持动态 shape 与动态 rankKernel 入口映射ExtendCfgInfo(opFile.value, pack_v2)将算子与 Kernel 侧源文件名pack_v2.cpp关联。在 pack_v2_binary.json 中四个二进制配置分别对应float32、int32、int16、float16四种 dtype 组合输入输出均标记为 ND 格式、shape: [-2]动态维度即支持任意 shape、format_match_mode: FormatAgnostic。需要说明的是README 的支持 FLOAT、FLOAT16是产品支持情况的简化表述而源码实际注册的 dtype 集合为 FLOAT、INT32、INT16、FLOAT16读者以源码注册为准。约束说明README 中 PackV2 的约束说明为无。结合源码进一步补充可推断的约束信息来自 pack_v2_tiling.cpp 的 dtype 校验逻辑输入数据类型必须是{DT_FLOAT, DT_INT32, DT_INT16, DT_FLOAT16}之一否则 Tiling 阶段直接返回失败invalid dtype输入x与y形状需要保持一致示例代码中selfXShape与selfYShape相同当前 Kernel 实现pack_v2l.h中PackV2L分支以float模板实例化 GlobalTensor且 Tiling 阶段根据 UB 容量约束了单核分块大小预留 5% 余量见FitsUB判断大 tensor 会按 BLOCK_DIM8 核进行切分。算子内部实现从注册到 Kernel 的四层结构PackV2 遵循 CANN 自定义算子开发的经典四层结构仓库目录组织如下experimental/conversion/pack_v2/ ├── examples/test_aclnn_pack_v2.cpp # aclnn 调用示例 ├── op_host/ │ ├── pack_v2_def.cpp # 算子原语注册IR 定义 │ ├── pack_v2_infershape.cpp # 形状推导InferShape │ ├── pack_v2_tiling.cpp # Tiling 策略多核切分 │ └── config/ascend910b/ │ ├── pack_v2_binary.json # 二进制 Kernel 编译配置 │ └── pack_v2_simplified_key.ini # simplified key 模式配置 ├── op_kernel/ │ ├── pack_v2.cpp / pack_v2.h # 常规非最后一维场景 Kernel │ ├── pack_v2l.h # 最后一维PACK_LAST场景 Kernel │ ├── pack_v2_tiling_data.h # Tiling 数据结构 │ └── pack_v2_tiling_key.h # Tiling Key模板调度 ├── CMakeLists.txt └── README.md1. 算子注册pack_v2_def.cpppack_v2_def.cpp 定义了PackV2算子原语声明输入x、y输出z属性d并注册ascend910b平台的 AICore 配置。这是算子被图编译器GE识别、参与构图与编译的入口。2. 形状推导pack_v2_infershape.cpppack_v2_infershape.cpp 实现了InferShapePackV2读取输入x的 shape将输出yShape的维度数设置为与输入一致并将最后一维翻倍if (i dim_num - 1) dim * 2;。从该实现看当前形状推导仅对最后一维做翻倍处理即输出 shape 与调用示例中outShape {1, 2, 3, 8}一致。3. Tiling 策略pack_v2_tiling.cpppack_v2_tiling.cpp 是 PackV2 性能调度的核心主要完成三件事获取平台信息通过platform_ascendc::PlatformAscendC读取 UB 大小与 AI Core 数量并在异常时返回失败dtype 与 shape 校验校验输入数据类型是否在支持集合内读取输入总元素数多核切分与 UB 分块以BLOCK_DIM 8为核数基准将输入按行x1 除最后一维外的总行数均匀划分到各核计算每个核上x、y的起始行、结束行与行数startX/endX/rowsX、startY/endY/rowsY再根据 UB 容量做核内二次分块core_tile_x1通过倍增 回溯确定FitsUB校验双缓冲下总占用不超过 UB 的 95%生成 Tiling 数据与 Tiling Key将切分结果写入PackV2TilingData结构体定义见 pack_v2_tiling_data.h并根据d是否为最后一维设置PACK_LAST或PACK_NORMAL两种 Tiling Key见 pack_v2_tiling_key.h供 Kernel 模板实例化时分流。Tiling 数据结构PackV2TilingData完整记录了核间/核内切分参数、输入输出 shapex1/x2/y1/y2/z2、大核小核的行数与循环次数、每个核的x/y行区间以及堆叠维度d与维度数dimNum。4. Kernel 实现pack_v2.cpp / pack_v2.h / pack_v2l.hKernel 入口 pack_v2.cpp 以schMode为模板参数依据 Tiling Key 分派到两套实现PACK_NORMAL分支pack_v2.h 中的NsPackV2::PackV2处理非最后一维堆叠。核心思路是把张量视作若干行CopyIn阶段根据当前全局行号对partnum取模判断该行来自x还是y用DataCopyPad按行搬入inQueueX/inQueueYCopyOut阶段再将各行写入输出zGm的对应位置。Process按tileNum大循环 tailNum尾块处理全程使用双缓冲BUFFER_NUM 2。PACK_LAST分支pack_v2l.h 中的NsPackV2L::PackV2L处理最后一维堆叠的扁平化场景直接把两输入视为一维数据流CopyIn分别搬入x、y的一段连续数据Compute阶段在 UB 中做交错写入zLocal.SetValue(2 * i, real); zLocal.SetValue(2 * i 1, imag);CopyOut一次性写出长度为processDataNum * 2的结果同样以双缓冲流水化执行。调用说明通过 aclnnPackV2 接口调用PackV2 算子通过aclnn 接口aclnnPackV2对外提供调用这是 CANN 上层框架PyTorch、TensorFlow 等经 ACL 调用底层算子的标准途径。完整可编译的示例见 test_aclnn_pack_v2.cpp其调用流程可分为以下 9 步① 初始化 ACL 环境固定写法aclInit→aclrtSetDevice→aclrtCreateStream② 构造输入输出 aclTensor通过aclrtMalloc申请 Device 内存、aclrtMemcpy将 Host 数据拷入再调用aclCreateTensor创建 ND 格式、连续 strides 的aclTensor。示例中std::vectorint64_t selfXShape {1, 2, 3, 4}; std::vectorDataType selfXHostData(24, 7); // x 全 7 std::vectorint64_t selfYShape {1, 2, 3, 4}; std::vectorDataType selfYHostData(24, 9); // y 全 9 std::vectorint64_t outShape {1, 2, 3, 8}; // 最后一维 4 - 8DataType定义为float修改该处即可切换测试数据类型源码注释明确提示修改测试数据类型。③ 调用第一段接口获取 workspace 与 executoruint64_t workspaceSize 0; int32_t axis 3; // 堆叠维度示例取最后一维 aclOpExecutor* executor; ret aclnnPackV2GetWorkspaceSize(selfX, selfY, axis, out, workspaceSize, executor);注意示例中的形参名为axis对应算子定义中的属性d此处取值为 3最后一维索引。若返回的workspaceSize 0则需aclrtMalloc申请对应大小的 workspace 内存。④ 调用第二段接口执行算子ret aclnnPackV2(workspaceAddr, workspaceSize, executor, stream);⑤ 同步等待执行结束aclrtSynchronizeStream(stream)⑥ 取回结果aclrtMemcpy将 Device 侧输出拷贝到 Host按 48 个元素打印输出总元素 24 × 2。⑦ 释放资源aclDestroyTensor释放三个 aclTensoraclrtFree释放 Device 内存含 workspace最后aclrtDestroyStream、aclrtResetDevice、aclFinalize。由于示例输入x全为 7、y全为 9在最后一维堆叠语义下期望输出为 x 与 y 元素交替排列的序列7 与 9 交替读者可据此直接验证算子结果的正确性。若要验证非最后一维堆叠可将axis改为 0~2并按d处插入新维度大小为 2的方式调整outShape后重新构造输出 Tensor。编译配置Kernel 二进制与 simplified keyPackV2 的 Kernel 编译行为由op_host/config/ascend910b/下的两个文件控制pack_v2_binary.json声明算子类型为PackV2列出 4 组op_list配置每组对应一种 dtypefloat32 / int32 / int16 / float16的输入输出组合并指定编译产出的bin_filename如PackV2_a1532827238e1555db7b997c7bce2928、动态 shape-2、ND 格式与FormatAgnostic匹配模式。该文件指导 opc 工具为不同 dtype 组合生成对应的 Kernel 二进制。pack_v2_simplified_key.ini配置--simplified_key_mode的取值文件内注释详细说明了配置规则默认 mode、按平台差异化配置、缺省时 AscendC 算子按simplified_key_mode0处理等。PackV2 的配置为default0。算子目录下的 CMakeLists.txt 通过add_subdirectory递归挂载子目录ENABLE_TEST未开启时剔除tests目录与仓库整体的算子构建体系见 conversion/CMakeLists.txt衔接。总结PackV2 是 CANN ops-math 在转换类算子中提供的双张量维度堆叠算子面向 Atlas A2 系列产品ascend910b支持 ND 格式、动态 shapedtype 覆盖 FLOAT、FLOAT16 及整数类型源码注册集合。其实现完整覆盖了算子注册pack_v2_def.cpp、形状推导pack_v2_infershape.cpp、多核 Tilingpack_v2_tiling.cpp与双缓冲 Kernelpack_v2.h、pack_v2l.h四层逻辑并以aclnnPackV2GetWorkspaceSizeaclnnPackV2两段式接口对外提供服务。开发者在实际使用时应结合 README 语义与当前实现最后一维翻倍、属性d默认 0、示例中axis3仔细核对 shape 推导结果并参照 test_aclnn_pack_v2.cpp 完成端到端调用验证。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考