理解 pypto.assemblePyPTO 中按偏移量写回大 Tensor 的核心接口【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto本篇围绕 CANN pypto 仓库中的 pypto.assemble API 文档 展开讲解pypto.assemble的单源写回与批量多源写回两种调用形态、parallel并行写回参数的语义以及框架在底层如何把 assemble 展开为具体的OP_ASSEMBLE操作节点。读完后你将能够在大 Tensor 分块计算、循环分块写回、跨 function 累积写等场景正确使用 assemble并理解串行/并行写回在 IR 层的实现差异。产品支持情况pypto.assemble在以下产品上受支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持功能说明assemble的核心语义是按偏移量写回以offsets指定的out索引位置为基准将输入 Tensorinput的值赋值到输出 Tensorout的对应区域。它不是拼接出一个新 Tensor 的算子而是对已有目标 Tensorout的就地in-place区域写入因此天然适合分块计算后写回全局结果、在循环迭代中逐块填充输出等编程模式。从 Python 前端入口 可以看到pypto.assemble是一个重载函数最终都会调用底层 C 绑定的pypto_impl.Assemble单源形态assemble(input, offsets, out, parallelFalse)批量多源形态assemble([(x, offsets), (y, [2, 2])], out, parallelFalse)函数原型assemble(input: Tensor, offsets: List[Union[int, SymbolicScalar]], out: Tensor, parallel: bool False) - None assemble(inputs: List[Tuple[Tensor, List[Union[int, SymbolicScalar]]]], out: Tensor, parallel: bool False) - None参数说明参数名输入/输出说明input输入源操作数。支持的数据类型为PyPto支持的数据类型。不支持空TensorShape Size不大于2147483647即INT32_MAX。inputs输入源操作数和输出偏移组成的Tuple列表。单个支持的数据类型为PyPto支持的数据类型。不支持空TensorShape Size不大于2147483647即INT32_MAX。offsets输入相对于目标输出的偏移。需要保证offsets小于out的Shape。out输出目的操作数需要和input的维度数量一致。支持的数据类型为PyPto支持的数据类型。不支持空TensorShape Size不大于2147483647即INT32_MAX。parallel输入是否允许并行写回。默认值为False当写回区域互不重叠、可安全并行时应显式传入parallelTrue。几点补充说明依据 前端入口实现offsets中的元素可以是int也可以是 SymbolicScalar即在编译期未知的动态偏移由运行时解析单源形态要求恰好 3 个位置参数批量形态要求 2 个位置参数参数个数不符或类型不符会抛出TypeError批量形态中若srcs为空列表会直接返回no-op若误把单源的第一个参数Tensor当作srcs传入错误信息会额外提示缺少 offsets 参数降低误用成本。返回值说明无返回值会直接对out进行修改。这与底层实现一致在 C 侧Assemble实现 中函数会把OP_ASSEMBLE操作挂到当前 Function 上并将dest标记为SlotProperty::ASSEMBLE_DST写点即out本身作为被写对象参与后续版本推导。约束说明输出Tensor out的valid shape需由用户在调用assemble前确保正确该接口不会自动推导。输入张量input和输出张量out的维度数量需要一致。当多个assemble对同一out的重叠区域存在写后写依赖且这些写回分布在不同loop迭代或不同function中时默认parallelFalse已保证写回按依赖顺序串行执行框架会在对应outcast上标记NORMAL供后续调度按串行写处理。若写回区域互不重叠、可安全并行应显式传入parallelTrue。从源码结构看这些约束在 C 实现 中都有对应的运行期检查src and dest requires same shape维度数量一致dynOffset and dest requires same shapeoffsets 与 out 的维度数量一致shape of src tensor requires integersrc 各维不能为空src and dest requires same format/requires same dtype与参数表支持的数据类型一致性的底层校验。也就是说文档中维度数量需一致不支持空 Tensor等约束不是文档约定而是框架在图构建阶段的硬性检查。调用示例单源写回x pypto.tensor([2, 2], pypto.DT_FP32) out pypto.tensor([4, 4], pypto.DT_FP32) offsets [0, 0] pypto.assemble(x, offsets, out)批量多源写回在单源写回之后可以继续用批量形态把另一个小 Tensor 写到不同区域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, [0,0]), (y, [2,2])]则同时完成左上角与右下角两个区域的一次性写回得到out1。串行与并行写回的底层实现这是parallel参数最有工程价值的部分。在 C 侧批量Assemble实现 中两种模式的 IR 展开方式完全不同串行模式parallelFalse默认按items顺序逐个处理每一步先对前一步结果做View在offsets处取区域视图再插入一个OP_ASSEMBLE_SSA操作把新结果作为下一次迭代的输入preResult curResult.GetStorage()。从源码结构看这形成了一条链式的 SSA 版本序列assemble_seq_out0 - assemble_seq_out1 - ...天然保证了写后写依赖按序执行。并行模式parallelTrue所有 item 直接对同一个src基版本做ViewOP_ASSEMBLE_SSA共同写向同一个输出版本最后通过SetSameMemId将src与result绑定到同一物理内存。由于各写回区域互不重叠这种展开允许调度器将多次写回并行化。单源形态同样会把parallel作为属性记录到操作节点上见TensorDInnerAssemble中的op.SetAttribute(OpAttributeKey::parallel, parallel)供后续 pass 判定写回语义。跨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)这正是分块 GEMM、分块 softmax 等 kernel 中每轮算一个 tile、再写回全局 out的标准写法由于相邻迭代写回区域可能重叠且语义上要求顺序默认parallelFalse会按依赖顺序串行执行框架在对应 outcast 上标记NORMAL属性供后续调度按串行写处理避免跨迭代写回被错误并行化。使用要点小结单 tensor 与批量多源 assemble 在未传parallel时均默认为 False需要并行写回时显式传入parallelTrue。若各次 assemble 写回区域不重叠且需并行调度应设置parallelTrue一旦存在区域重叠的写后写依赖务必保持默认的串行语义否则会破坏正确性。offsets支持SymbolicScalar动态偏移配合pypto.loop可实现编译期未知偏移的循环分块写回调用 assemble 前需自行保证out的 valid shape 正确该接口不会自动推导。与 assemble 相关的其他接口可参考同目录文档pypto.scatter、pypto.atomic_add、pypto.gatherPython 端类型标注见 pypto_impl 接口声明C 绑定入口见 operation.cpp。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
