CANN PyPTOpypto.interleave接口详解Tensor 逐元素交织与双输出拆分原理、约束及实战用法【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读pypto.interleave是 CANN PyPTO 张量编程范式Parallel Tensor/Tile Operation中用于数据重排的核心接口它把两个 Shape 完全一致的输入 Tensor按最后一个维度逐元素交替交织成一条数据流再按中点拆分为两个输出 Tensor。该接口常用于 RoPE 旋转位置编码、奇偶通道重组、数据并行切分等场景是pypto.deinterleave的逆操作。读完本文你将掌握它的数学语义、函数原型、参数约束、TileShape 配置规则、完整调用示例以及从 Python 前端到 C 内核的底层实现链路。产品支持情况pypto.interleave并非所有昇腾产品均支持当前仓库明确标注的支持矩阵如下产品支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持在源码层面该结论同样可以得到印证framework/src/interface/operation/vector/interleave.cpp中通过INTERLEAVE_SUPPORTED_ARCHITECTURES {NPUArch::DAV_3510}限定了算子支持的 NPU 架构同时framework/src/interface/tileop/vector/interleave.h中的实现整体被#if defined PTO_NPU_ARCH_A5宏保护仅在对应 A5 架构即 950 系列下参与编译。功能说明交织Interleave语义将两个输入 Tensorinput和other按最后一个维度逐元素交织构造一条长度为原末维两倍的“交织流”interleaved[2 * k] input[k] interleaved[2 * k 1] other[k]即input的元素占据交织流的偶数位置other的元素占据交织流的奇数位置。中点拆分语义交织完成后将交织流按中点拆分为两个输出 Tensorout0保存交织流的前半部分out1保存交织流的后半部分。由于交织流的长度为2 * nn 为输入最后一维大小前半部分和后半部分各长n因此out0、out1与input、other的 Shape 完全一致。pypto.interleave是pypto.deinterleave的逆操作二者相互配合可实现“交织-反交织”的往返数据变换。从源码结构看底层实现通过TileInterleaveOperation递归按 TileShape 切分各维度在最内层对每一对输入 Tile 调用Opcode::OP_INTERLEAVE生成两个输出 Tile并将输入 Tile 的动态有效 ShapeDynValidShape同步传播给两个输出见 interleave.cpp。函数原型interleave(input: Tensor, other: Tensor) - Tuple[Tensor, Tensor]该接口在 Python 前端位于 joining.py通过op_wrapper装饰后调用底层pypto_impl.Interleave(input, other)C 绑定入口定义在 operation.cpp最终映射到npu::tile_fwk::Interleave并将返回的std::tuple转为 Python 元组。参数说明参数名输入/输出说明input输入第一个源操作数支持的类型为 Tensor。其元素在交织流中占据偶数位置。other输入第二个源操作数支持的类型为 Tensor。其元素在交织流中占据奇数位置。other的 Shape、数据类型必须与input一致。返回值说明返回一个二元组(out0, out1)out0保存交织流的前半部分out1保存交织流的后半部分out0和out1的 Shape、数据类型与input一致。约束说明使用pypto.interleave前必须满足以下约束否则会触发参数校验失败源码中的CheckInterleaveTileShape与CheckTensorFormat负责执行这些检查见 interleave.cppinput和other的数据类型、维度数、Shape 必须一致。支持的数据类型为DT_INT8、DT_UINT8、DT_INT16、DT_UINT16、DT_INT32、DT_UINT32、DT_FP16、DT_FP32、DT_BF16。当前支持 1 到 4 维 Tensor。最后一个维度的 Shape 必须为偶数源码通过viewShape[lastAxis] % NUM_VALUE_2 0校验保证交织流能被中点完整拆分。当 TileShape 有效配置时TileShape 维度应与输入 Tensor 维度一致且最后一维必须与输入 Tensor 的 Shape 最后一维相等即最后一维不允许切分其他维可按切分需求设置。源码对应的校验为vecTile.tile.size() viewShape.size()与vecTile.tile[lastAxis] viewShape[lastAxis]其中requireFullLastAxis true即代表交织操作要求末维整轴参与。Tensor 类型输入不支持TileOpFormat.TILEOP_NZ格式源码通过CheckTensorFormat(..., {TileOpFormat::TILEOP_NZ}, Interleave)对两个输入逐一检查。调用示例TileShape 设置示例输入input和other的 Shape 为[m, n]输出out0和out1的 Shape 均为[m, n]TileShape 设置为[m1, n1]则m1、n1分别用于切分m、n轴。由于交织沿最后一维进行n1必须等于n末维不切分m1可按需要设置以控制 Tile 粒度pypto.set_vec_tile_shapes(4, 16)set_vec_tile_shapes是 PyPTO 向量计算中设置各维度 Tile Shape 的全局配置接口定义于 _controller.py上述示例表示第一维以 4 为步长切分、第二维 Tile 大小为 16。接口调用示例input pypto.tensor([2, 4], pypto.DT_FP32) other pypto.tensor([2, 4], pypto.DT_FP32) out0, out1 pypto.interleave(input, other)结果示例如下输入数据 input: [[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0]] 输入数据 other: [[10.0, 11.0, 12.0, 13.0], [14.0, 15.0, 16.0, 17.0]] 输出数据 out0: [[0.0, 10.0, 1.0, 11.0], [4.0, 14.0, 5.0, 15.0]] 输出数据 out1: [[2.0, 12.0, 3.0, 13.0], [6.0, 16.0, 7.0, 17.0]]以第一行为例交织流为[0, 10, 1, 11, 2, 12, 3, 13]前半部分[0, 10, 1, 11]写入out0后半部分[2, 12, 3, 13]写入out1与文档语义完全一致。从源码看底层实现与验证Python 前端调用链pypto.interleave的完整调用链为pypto.interleave(input, other) - pypto_impl.Interleave(input, other) # joining.py - npu::tile_fwk::Interleave(self, other) # operation.cpp pybind 绑定 - InterleaveOperation # 生成 OP_INTERLEAVE 算子 - pto::TINTERLEAVE 内核 # interleave.h (PTO_NPU_ARCH_A5)Python 前端实现位于 joining.py同一模块中还提供了concat、deinterleave等张量拼接/重排接口方便统一查阅。C 算子注册与 Tile 切分npu::tile_fwk::Interleave是算子的宿主入口interleave.cpp其行为包括对两个输入逐一执行CheckTensorFormat拒绝 TILEOP_NZ 格式与CheckInterleaveTileShape末维必须为偶数、TileShape 秩必须与输入一致、末维 Tile 必须等于输入末维依据input的 Shape 与数据类型创建两个输出 Tensor并继承输入的有效 Shape调用InterleaveOperation在当前的 Function 中登记算子。Tile 切分逻辑由TileInterleaveOperation实现interleave.cpp它按vecTile递归遍历除末维外的每个维度为每个 Tile 计算 offset 与 valid shape最终调用function.AddOperation(Opcode::OP_INTERLEAVE, {inputTile1, inputTile2}, {resultTile1, resultTile2})生成内核指令。算子通过REGISTER_OPERATION_TILED_FUNC(OP_INTERLEAVE, ...)注册到算子框架中。在设备侧interleave.h 中InterleaveComputeImpl最终落到pto::TINTERLEAVE(dst1, dst0, src1, src0)单条硬件指令上并由InterleaveCompute按 Tile 布局循环展开执行。测试用例验证仓库为该算子提供了完整的 ST 测试支撑Python 侧端到端测试 test_interleave_onboard.py构造(8, 64)的随机 FP32 输入经pypto.view切 view、双层循环、set_vec_tile_shapes(2, 64)、pypto.interleave与pypto.assemble完整流水后用 Torch 参考实现校验interleaved torch.stack((input0, input1), dim-1).flatten(-2)再按中点切分为 golden0 / golden1并以rtol3e-3, atol3e-3比对输出测试用例标注pytest.mark.soc(950)与文档支持矩阵一致。C/CSV 测试用例表 Interleave_st_test_cases.csv覆盖 1 到 4 维、FP32/FP16/BF16/INT8/UINT8/INT16/UINT16/INT32/UINT32 等多种数据类型输入格式均为 ND其中末维统一取 64保证偶数且 Tile 可整分。C 测试主体 test_interleave_operation.cpp驱动上述用例执行算子并对拍黄金数据。如果你在 950 系列环境上本地验证可按仓库docs/zh/install/build_and_install.md的流程构建 PyPTO 并运行python/tests/st/operation/vector/test_interleave_onboard.py。常见使用场景提示RoPE 旋转位置编码PyPTO 内部在实现 Interleave 风格 RoPE 时使用了交错/反交错数据流仓库测试 test_interleave_rope.py 及其黄金参考 interleave_rope_golden.py 展示了pypto.set_vec_tile_shapes(1, 1, 64, 64)配合交错数据的典型 kernel 写法可作为动手实践的参考。数据并行切分/重组当需要将两组按奇偶位置组织的特征流合并并均衡分配到两个输出时pypto.interleave提供了一条指令级的解决方案且通过 TileShape 控制切分粒度以适配不同算子Vector流水布局。相关接口pypto.deinterleave将交织流反交织回偶数位置和奇数位置的元素流是pypto.interleave的逆操作。它支持双输入形式接收前半/后半部分与单输入形式接收完整交织流输出末维减半两种调用方式可在需要还原原始数据布局时配合使用。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考