企业官网建设流程全解析

CANN ops-math 算子解析：PackV2 双张量维度堆叠（Stack）算子的实现与 aclnn 调用实战

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

PackV2 是 CANN ops-math 仓库（experimental/conversion/pack_v2）中实现的双输入张量维度堆叠算子：它将两个形状完全相同的张量按照指定维度d拼接为一个新张量，等价于 PyTorch 的torch.stack在双输入场景下的行为。本文基于该算子目录下的 README、Host 侧算子定义与 Tiling 逻辑、Kernel 侧 AscendC 实现以及 aclnn 调用样例，完整讲解 PackV2 的功能语义、参数约束、NPU 侧实现原理与可运行的调用示例，帮助开发者在 Atlas A2 系列产品上快速理解并使用该算子。

功能说明：两个张量的维度堆叠（Stack）

PackV2 算子实现的是张量的维度堆叠（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）：

在源码层面，该支持范围由 pack_v2_def.cpp 中的算子注册配置进一步印证：this->AICore().AddConfig("ascend910b", aicoreConfig)，即为ascend910b芯片（对应 Atlas A2 系列）注册 AICore 配置；同时 pack_v2_binary.json 的目录层级也直接以ascend910b为平台维度组织算子编译配置。

参数说明

PackV2 算子的接口参数如下（摘自 README）：

在 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 与动态 rank；
• Kernel 入口映射：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_DIM=8 核进行切分。

算子内部实现：从注册到 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.md

1. 算子注册（pack_v2_def.cpp）

pack_v2_def.cpp 定义了PackV2算子原语，声明输入x、y，输出z，属性d，并注册ascend910b平台的 AICore 配置。这是算子被图编译器（GE）识别、参与构图与编译的入口。

2. 形状推导（pack_v2_infershape.cpp）

pack_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.cpp）

pack_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完整记录了核间/核内切分参数、输入输出 shape（x1/x2/y1/y2/z2）、大核小核的行数与循环次数、每个核的x/y行区间以及堆叠维度d与维度数dimNum。

4. Kernel 实现（pack_v2.cpp / pack_v2.h / pack_v2l.h）

Kernel 入口 pack_v2.cpp 以schMode为模板参数，依据 Tiling Key 分派到两套实现：

• PACK_NORMAL分支（pack_v2.h 中的NsPackV2::PackV2）：处理非最后一维堆叠。核心思路是把张量视作"若干行"，CopyIn阶段根据当前全局行号对partnum取模判断该行来自x还是y，用DataCopyPad按行搬入inQueueX/inQueueY；CopyOut阶段再将各行写入输出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::vector<int64_t> selfXShape = {1, 2, 3, 4}; std::vector<DataType> selfXHostData(24, 7); // x 全 7 std::vector<int64_t> selfYShape = {1, 2, 3, 4}; std::vector<DataType> selfYHostData(24, 9); // y 全 9 std::vector<int64_t> outShape = {1, 2, 3, 8}; // 最后一维 4 -> 8

DataType定义为float（修改该处即可切换测试数据类型，源码注释明确提示"修改测试数据类型"）。

③ 调用第一段接口获取 workspace 与 executor：

uint64_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释放三个 aclTensor，aclrtFree释放 Device 内存（含 workspace），最后aclrtDestroyStream、aclrtResetDevice、aclFinalize。

由于示例输入x全为 7、y全为 9，在"最后一维堆叠"语义下，期望输出为 x 与 y 元素交替排列的序列（7 与 9 交替），读者可据此直接验证算子结果的正确性。若要验证非最后一维堆叠，可将axis改为 0~2，并按d处插入新维度（大小为 2）的方式调整outShape后重新构造输出 Tensor。

编译配置：Kernel 二进制与 simplified key

PackV2 的 Kernel 编译行为由op_host/config/ascend910b/下的两个文件控制：

• pack_v2_binary.json：声明算子类型为PackV2，列出 4 组op_list配置，每组对应一种 dtype（float32 / 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_mode=0处理等）。PackV2 的配置为default=0。

算子目录下的 CMakeLists.txt 通过add_subdirectory递归挂载子目录，ENABLE_TEST未开启时剔除tests目录，与仓库整体的算子构建体系（见 conversion/CMakeLists.txt）衔接。

总结

PackV2 是 CANN ops-math 在转换类算子中提供的双张量维度堆叠算子，面向 Atlas A2 系列产品（ascend910b），支持 ND 格式、动态 shape，dtype 覆盖 FLOAT、FLOAT16 及整数类型（源码注册集合）。其实现完整覆盖了算子注册（pack_v2_def.cpp）、形状推导（pack_v2_infershape.cpp）、多核 Tiling（pack_v2_tiling.cpp）与双缓冲 Kernel（pack_v2.h、pack_v2l.h）四层逻辑，并以aclnnPackV2GetWorkspaceSize+aclnnPackV2两段式接口对外提供服务。开发者在实际使用时应结合 README 语义与当前实现（最后一维翻倍、属性d默认 0、示例中axis=3）仔细核对 shape 推导结果，并参照 test_aclnn_pack_v2.cpp 完成端到端调用验证。

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