PyPTO Tensor.assemble 详解：将分块计算结果写回大张量指定区域

PyPTO Tensor.assemble 详解将分块计算结果写回大张量指定区域【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptoPyPTOParallel Tensor/Tile Operation的assemble接口用于把源 Tensor通常是大输出张量的一个 tile 分片或中间计算结果按offsets指定的偏移位置写回目标输出 Tensorout的对应区域是 PyPTO 编程范式中分块计算、聚合写回的关键一环。本文基于仓库文档 pypto-Tensor-assemble.md 与其关联文档 pypto-assemble.md 展开并结合 operation.py、operation_impl.cpp 等源码从函数原型、参数约束、并行写回语义到底层实现与实战示例进行完整剖析。读完本文你将掌握单源与批量多源两种 assemble 调用方式、动态偏移的使用场景以及parallel参数背后的调度语义。一、产品支持情况根据文档assemble接口在当前仓库支持的产品形态如下Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持注以上支持范围以仓库文档中标注的产品矩阵为准实际可用性请以所部署环境的驱动与固件版本为准。二、两种调用形式与函数原型1. Tensor 成员方法形式在 tensor.py 中Tensor类提供了assemble成员方法assemble(self, input: Tensor, offsets: List[Union[int, SymbolicScalar]]) - None该方法将自身视为目标输出把input写入自身指定偏移位置等价于pypto.assemble(input, offsets, self)s pypto.tensor((16, 16), pypto.DT_FP32) a pypto.tensor((2, 2), pypto.DT_FP32) s.assemble(a, [0, 0]) # 等价于 pypto.assemble(a, [0, 0], s)2. 顶层函数形式在 operation.py 中定义了带overload的顶层函数支持单源与批量多源两种调用# 单源形式将 input 写入 out 的 offsets 位置 assemble(input: Tensor, offsets: List[Union[int, SymbolicScalar]], out: Tensor, parallel: bool False) - None # 批量多源形式一次性将多个 (源Tensor, 偏移) 写入同一个 out assemble(inputs: List[Tuple[Tensor, List[Union[int, SymbolicScalar]]]], out: Tensor, parallel: bool False) - None运行时入口 operation.py#L121-L149 根据位置参数个数分发传入 3 个位置参数(src, offsets, dst)先做类型校验src必须是Tensoroffsets必须是int或SymbolicScalar序列dst必须是Tensor再将 offsets 通过to_syms()统一转为符号标量序列调用底层pypto_impl.Assemble(...)传入 2 个位置参数(srcs, dst)校验srcs为二元组Tensor, offsets的列表/元组若为空列表则直接返回空批量输入是合法空操作再调用底层批量接口pypto_impl.Assemble(srcs, dst, parallel)参数个数不对时抛出TypeError并给出如果是单源形式则缺少 offsets 参数的提示。从源码结构可以推断该分发逻辑的存在是为了让单源与批量多源共用同一命名空间的同时保持调用语义清晰、便于静态类型检查。三、参数说明参数名输入/输出说明input输入源操作数。支持 PyPTO 支持的数据类型不支持空 TensorShape Size 不大于 2147483647即 INT32_MAX。inputs输入由源操作数 输出偏移二元组组成的列表。单个元素支持 PyPTO 支持的数据类型不支持空 TensorShape Size 不大于 2147483647即 INT32_MAX。offsets输入相对于目标输出的偏移需要保证offsets小于out的 Shape即每个维度的偏移值 对应写入范围不得越界。out输出目的操作数需要与input的维度数量一致。支持 PyPTO 支持的数据类型不支持空 TensorShape Size 不大于 2147483647即 INT32_MAX。parallel输入是否允许并行写回。默认值为False当写回区域互不重叠、可安全并行时应显式传入parallelTrue。其中offsets的每个元素可以是int也可以是SymbolicScalar符号标量后者用于循环索引等编译期未知、运行期才确定的动态偏移场景典型用法见下文跨 loop 动态偏移示例。返回值说明无返回值接口会直接修改out对应的底层存储。四、约束说明valid shape 由用户保证输出 Tensorout的 valid shape 需由用户在调用assemble前确保正确该接口不会自动推导。这是因为assemble只是把源数据按偏移嵌入目标不负责计算目标张量的逻辑形状。维度数量一致输入张量input与输出张量out的维度数量必须一致。底层实现 operation_impl.cpp#L1584-L1587 有显式校验CHECK_OP(dstShape.size() srcShape.size()) Assemble: src and dest requires same shape; CHECK_OP(dstShape.size() dynOffset.size()) Assemble: dynOffset and dest requires same shape;同时还会校验 src 各维度 shape 大于 0、src 与 dest 的 Format 与 DataType 一致否则直接报错见 operation_impl.cpp#L1589-L1595。写后写依赖与串行语义当多个assemble对同一out的重叠区域存在写后写WAR/WAW依赖且这些写回分布在不同 loop 迭代或不同 function 中时默认parallelFalse已保证写回按依赖顺序串行执行框架会在对应 outcast 上标记NORMAL供后续调度按串行写处理。若写回区域互不重叠、可安全并行应显式传入parallelTrue。五、调用示例与运行结果文档给出了完整的可运行示例x pypto.tensor([2, 2], pypto.DT_FP32) out pypto.tensor([4, 4], pypto.DT_FP32) offsets [0, 0] pypto.assemble(x, offsets, out) y pypto.tensor([2, 2], pypto.DT_FP32) pypto.assemble([(x, offsets), (y, [2, 2])], out)结果示例如下输出数据x: [[1, 1] [1, 1]] 输入数据out: [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]] 输出数据out: [[1, 1, 0, 0], [1, 1, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]] 输出数据out1: [[1, 1, 0, 0], [1, 1, 0, 0], [0, 0, 1, 1], [0, 0, 1, 1]]可以看到单源assemble(x, [0, 0], out)把 2x2 的x写到了out左上角随后批量多源调用又把x与y分别写到(0,0)与(2,2)两个互不重叠的区域最终out呈对角两块 1 的形态。在 operation.py 的 docstring 中还给出了等价的成员方法调用方式与相同的结果可直接对照运行验证。六、跨 loop / 跨 function 的串行写回当同一输出在不同 loop 迭代或不同 function 之间存在写后写依赖、必须串行 assemble 时默认行为即为串行写回也可显式传入parallelFalse# 跨loop后一次写依赖前一次写的结果默认串行 for i in pypto.loop(0, n, nameSEQ_WRITE): tile ... pypto.assemble(tile, [i * tile_m, 0], out) # 跨function下游function继续写同一out的重叠区域时同样默认串行 pypto.assemble(partial, [offset_m, offset_n], out)说明单 tensor 与批量多源 assemble 在未传parallel时均默认为False需要并行写回时显式传入parallelTrue。若各次 assemble 写回区域不重叠且需并行调度应设置parallelTrue。值得注意的是上例中[i * tile_m, 0]的偏移是循环变量i的表达式这正是offsets支持SymbolicScalar的典型场景——偏移在编译期未知由运行期循环索引决定。七、底层实现原理从 Python 到算子级 IR1. pybind11 绑定层顶层函数最终会进入 pybind11 绑定 operation.cpp#L737-L753。批量形式把 Python 侧的(Tensor, offsets)列表转换为std::vectornpu::tile_fwk::AssembleItem再调用npu::tile_fwk::Assemble(items, dest, parallel)单源形式直接调用npu::tile_fwk::Assemble(tensor, dynOffset, dest, parallel)。2. 框架核心实现C 侧核心实现在 operation_impl.cpp#L1580-L1610先做维度数量、offsets 长度、shape 非空、Format/DataType 一致性等前置校验通过TensorDInnerAssemble在当前 Function 中插入 assemble 语义操作将offsets与输入 Tensor 封装为AssembleOpAttribute含静态 offset 与动态 dynOffset 两部分生成OP_ASSEMBLE或 SSA 变体算子最后通过Program::GetInstance().GetTensorSlotManager()-TensorWrite(dest, SlotProperty::ASSEMBLE_DST)向 Tensor Slot 管理器登记对dest的写以便后续调度/依赖分析识别写后写关系——这从源码层面印证了默认串行、框架跟踪依赖的文档描述。3. 解释器/计算内核在解释器侧 calc_common.cpp#L27-L56 注册了ExecuteOpAssemble它把输出数据视图oop按输入 shape 与求值后的 offset 建立View然后执行calc::Copy(ret, iop)完成数据搬移若该 op 带atomic_add属性则执行calc::Add。其中还有一处值得注意的优化当 InplaceProcess 之后 Assemble 的输入输出共享同一底层 RawTensorRawMagic 相同时该 op 不会翻译成任何硬件指令直接跳过 View/Copy即输入即输出区域时 assemble 退化为空操作。OP_ASSEMBLE、OP_CONTRACT、OP_ASSEMBLE_SSA三种 opcode 都复用该实现见 calc_common.cpp#L54-L56。八、真实场景分块计算的聚合写回assemble最常见的用途是在 kernel 内按 tile 循环分块计算后把每个分块结果写回大输出张量。仓库测试 test_add_onboard.py 给出了标准范式——以 2D 和 4D 两种 kernel 分别演示def add_2d_kernel(a: pypto.Tensor(), b: pypto.Tensor(), output: pypto.Tensor(), config: AddConfig): for row_index in pypto.loop((config.output_shape[0] row_view - 1) // row_view): for column_index in pypto.loop((config.output_shape[1] column_view - 1) // column_view): row_offset row_index * row_view column_offset column_index * column_view result ... # 对当前分块执行 add得到 row_view x column_view 的结果 pypto.assemble(result, [row_offset, column_offset], output)同样模式还出现在 matmul 场景如 test_scaled_mm_mxfp8.py 中pypto.assemble(out_view, [m_offset, n_offset], out_tensor)、分块归约场景如 test_argmax_onboard.py 中pypto.assemble(block_result, [b_idx * view_shape[0]], dst_tensor)等大量测试中。框架内部实现如 operation_impl.cpp 中的TiledAssemble、TiledInnerAssemble与各 cube 算子实现中的批量写回调用也大量使用该接口完成 tiling 结果的聚合说明assemble是 PyPTO tile 编程范式中的基础构件。九、注意事项小结调用前务必自行确保out的 valid shape 正确接口不做形状推导offsets不能越界每个维度的偏移需要满足偏移 源 shape out shapeinput与out维度数必须一致底层还会校验 Format 与 DataType 一致存在写后写依赖时保持默认parallelFalse串行写回仅当写回区域互不重叠且需并行调度时才显式传parallelTrue空批量输入空inputs列表是合法调用直接返回不产生任何操作动态偏移如循环索引请使用SymbolicScalar类型元素实现编译期未知偏移的写回。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考