PyPTO Gym 算子参考：用 view + assemble 在 NPU 上实现 torch.stack 搬运内核

PyPTO Gym 算子参考用 view assemble 在 NPU 上实现 torch.stack 搬运内核【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym本文围绕 pypto-api-explore/examples/stack.md 展开讲解如何基于 PyPTO 的「搬运语义」pypto.view取片 pypto.assemble拼回在 NPU 上实现torch.stack并给出 batch 轴 loop 切分、新 stack 轴逐输入占片的完整内核骨架。读者阅读后可掌握 stack 类多输入搬运算子的 PyPTO 编码范式、Tiling 配置要点以及它与cat/unbind/repeat等相邻算子的切分差异可直接套用于 PyPTO 算子开发与 API 可行性评估。一、算子背景torch.stack 与 PyPTO 组合方案torch.stack的作用是在新增的维度上把多个张量堆叠起来。例如两个形状均为[S, D]的输入torch.stack([a, b], dim0)得到形状[2, S, D]的输出dim1则得到[S, 2, D]。它与cat的本质区别在于cat沿已有维度拼接输出各维之和增大各输入在拼接轴上分片stack沿新插入的维度堆叠输入张量本身被原样保留新增轴的长度等于输入个数。在 PyPTO 算子对标手册 references/torch-pypto-op-mapping.md 中stack被明确归类为「形状变换」类组合方案Torch 算子PyPTO 组合方案参考实现stackunsqueezeconcatstack.md即从数学语义上stack等价于「先对每个输入做unsqueeze插入新轴再沿该新轴做concat」。而 examples/stack.md 给出的内核骨架正是把这一语义落地为「每个输入在输出新轴的不同偏移上各写一片」的搬运实现——两个输入分别写[0, ...]与[1, ...]两条偏移带。二、内核骨架stack_kernel 逐行解析examples/stack.md 给出的完整参考骨架如下pypto.frontend.jit(runtime_options{run_mode: pypto.RunMode.NPU}) def stack_kernel(a: pypto.Tensor(sl, pypto_dtype), b: pypto.Tensor(sl, pypto_dtype), out: pypto.Tensor(ol, pypto_dtype)): for i in pypto.loop(batch, namebatch, unroll_list[1]): a_s pypto.view(a, [1] inner, [i] [0] * len(inner)) b_s pypto.view(b, [1] inner, [i] [0] * len(inner)) pypto.set_vec_tile_shapes(1, *inner) pypto.assemble(a_s, [0, i] [0] * (len(ol) - 2), out) pypto.assemble(b_s, [1, i] [0] * (len(ol) - 2), out)2.1 装饰器与函数签名pypto.frontend.jit(runtime_options{run_mode: pypto.RunMode.NPU})通过pypto.frontend.jit将内核函数编译为 NPU 可执行程序runtime_options中run_mode指定为pypto.RunMode.NPU表明该内核目标是 NPU 运行时而非 CPU 仿真/模拟模式。三个形参均为pypto.Tensor描述的张量输入a、b形状为sl输入 shape 列表输出out形状为ol输出 shape 列表元素 dtype 统一由占位符pypto_dtype表示如pypto.DT_FP32。2.2 batch 轴 loop外轴切片搬运for i in pypto.loop(batch, namebatch, unroll_list[1]):这是本骨架的核心切分策略——参考 examples/README.md 的 Note 约定batch是被 loop 的外层轴长度通常为sl[0]。其设计意图原文 Note 概括为输入 batch 轴 loop 切分新增 stack 轴上各输入占一片此处 2 输入输入内轴整块。也就是说loop 的轴输入张量的 batch 轴首个轴长度为batch整块的轴输入张量的内轴inner如sl[1:]每次迭代整块搬运不切分新增 stack 轴不参与 loop而是通过两次assemble的偏移0与1区分两个输入各自占据的「一片」。namebatch为 loop 命名便于调试与后端分析unroll_list[1]表示该 loop 按迭代粒度 1 展开保证每次迭代处理一片 batch 数据。2.3 view 取片从输入中切出单片a_s pypto.view(a, [1] inner, [i] [0] * len(inner)) b_s pypto.view(b, [1] inner, [i] [0] * len(inner))pypto.view是元数据视图操作不搬数据、不产生逐元素计算仅描述「从大张量中取一个子片」的 shape 与起始偏移第一个参数为源张量第二个参数[1] inner是取片后的形状batch 轴上取长度为 1 的一片内轴保持inner整块第三个参数[i] [0] * len(inner)是取片的起始偏移batch 轴偏移i第i个 batch内轴偏移全为 0。第i次迭代取出a、b各自的第i个 batch 片形状均为[1] inner。2.4 set_vec_tile_shapesVector 单元 Tiling 配置pypto.set_vec_tile_shapes(1, *inner)本算子属于纯搬运 / metadata 语义无逐元素计算、无归约、无 matmul因此走 Vector 通道而非 Cube 通道只需配置set_vec_tile_shapes无需set_cube_tile_shapes。该调用的第一个参数1是 batch 片维度上的 tile 长度*inner将内轴各维依次展开为 tile shape 的其余维度。它告诉后端每次 Vector 搬运以[1] inner的整块为粒度处理与view取出的片形状保持一致。按 SKILL.md 中的硬约束速查TileShape 每维必须 0 且最多 4 维实际取值需结合具体 shape 与平台约束确定。2.5 assemble 拼回按新轴偏移写输出pypto.assemble(a_s, [0, i] [0] * (len(ol) - 2), out) pypto.assemble(b_s, [1, i] [0] * (len(ol) - 2), out)pypto.assemble与view互为逆操作把一片数据写入大张量在指定偏移处的区域。第二个参数是写入输出的起始偏移这里两个输入各自携带「stack 新轴坐标」a写到输出偏移[0, i] [0] * (len(ol) - 2)新轴位置取0b写到输出偏移[1, i] [0] * (len(ol) - 2)新轴位置取1。其中len(ol)是输出张量的维度数。以输出形状ol [2, B, S, D]即对[B, S, D]的两个输入在 dim 0 做 stack为例[0, i] [0] * 2展开为[0, i, 0, 0][1, i] [0] * 2展开为[1, i, 0, 0]——新轴坐标0/1区分两个输入batch 坐标取当前迭代i内轴坐标全为 0。两次assemble合起来正好完成一次迭代内两个输入片在输出新轴上的「拼接」等价于把两个unsqueeze后的子张量沿新轴concat。三、切分策略为何「batch 轴 loop 新轴整块」结合 examples/README.md 对占位符的约定本骨架的轴处理可归纳为一张表轴处理方式原因batch 轴输入首轴长度batchpypto.loop切分每次迭代取一片控制单次搬运的数据量适配 Vector tile 能力输入内轴inner如sl[1:]整块view时偏移为 0内轴整块搬运避免切碎导致 assemble 偏移计算复杂新增 stack 轴整块不 loop靠assemble偏移0/1区分该轴长度即输入个数此处为 2规模小直接按片落位这种设计把一个「多输入、多轴」的 stack 问题降维成「单输入片搬运」的循环每次迭代只处理两个输入的第i个 batch 片搬运逻辑完全对称代码可读性与后端优化空间都更好。对于输入个数更多如 4 个输入的场景只需按同样模式增加对应的view/assemble调用并分别赋予新轴坐标0,1,2,3。四、与相邻搬运类骨架的对比在 examples/ 目录下stack与cat、unbind、repeat同属「metadata/搬运语义」view 取片 assemble 拼回、无逐元素计算但轴切分方式不同对比有助于加深理解算子参考骨架loop 的轴写入输出的偏移规律stackstack.mdbatch 轴各输入在新增轴的不同坐标0/1各写一片catcat.mdbatch 轴各输入沿concat 轴顺序搬运到输出对应偏移[i][0]*...unbindunbind.md拆分轴每片对应一个输出view取片 assemble拼到单个输出repeatrepeat.md输出新增复制轴每次迭代把整块输入搬到输出第i份偏移[i][0]*...stack vs catcat 的多次assemble共享同一组偏移、仅在拼接轴上连续递增stack 的多次assemble在新增轴上取不同坐标0/1这正是「unsqueeze concat」组合方案的代码形态。stack vs repeatrepeat 是单输入在复制轴上重复写同一片stack 是多输入在新增轴上各写各的片输入源不同。stack vs unbindunbind 是 stack 的逆操作把一个大张量沿拆分轴拆成多个输出stack 则是把多个输入合并到带新轴的大张量。五、使用注意事项与调优要点按 examples/README.md 的声明examples/下每个op.md均为kernel 参考骨架仅展示接口组合与轴切分模式哪些轴 loop、哪些轴整块不作为标准模板且未逐一经 NPU 编译验证。实际使用时需注意占位符替换sl输入 shape如[B, S, D]、ol输出 shape、pypto_dtype如pypto.DT_FP32、batchloop 外层轴长通常sl[0]、inner内层 shape如sl[1:]均需按实际算子替换最小可运行 setup 可参考 examples/README.md 中的示例。loop 轴与 tile shape 调优unroll_list、tile shapeset_vec_tile_shapes(1, *inner)需按实际 shape / dtype 与平台约束确定并调优TileShape 每维 0、最多 4 维。动态 shape 风险若 stack 输入含动态轴需注意 SKILL.md 中记录的硬约束——归约类等计算 API 在编译期需要 concrete shape搬运类骨架若依赖具体偏移计算同样应在 API 可行性报告中评估动态轴影响。lint 门禁优先骨架只是 API 用法参考当写法与 lint / 门禁冲突时以 lint 为准参见 SKILL.md 对用法参考的定位说明。综上stack.md 为 stack 类算子提供了「batch 轴 loop 新轴按输入落位」的通用搬运骨架view负责切、assemble负责拼、set_vec_tile_shapes负责 Vector 通道 Tiling三者组合即可在 NPU 上复现torch.stack的 unsqueeze concat 语义是编写多输入搬运类 PyPTO 内核时可复用的最小范式。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考