企业官网建设流程全解析

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], dim=0)得到形状[2, S, D]的输出；dim=1则得到[S, 2, D]。它与cat的本质区别在于：

• cat沿已有维度拼接，输出各维之和增大，各输入在拼接轴上分片；
• stack沿新插入的维度堆叠，输入张量本身被原样保留，新增轴的长度等于输入个数。

在 PyPTO 算子对标手册 references/torch-pypto-op-mapping.md 中，stack被明确归类为「形状变换」类组合方案：

即从数学语义上，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, name="batch", 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, name="batch", 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）区分两个输入各自占据的「一片」。

name="batch"为 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_shapes：Vector 单元 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)：新轴位置取0；
• b写到输出偏移[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 对占位符的约定，本骨架的轴处理可归纳为一张表：

这种设计把一个「多输入、多轴」的 stack 问题，降维成「单输入片搬运」的循环：每次迭代只处理两个输入的第i个 batch 片，搬运逻辑完全对称，代码可读性与后端优化空间都更好。对于输入个数更多（如 4 个输入）的场景，只需按同样模式增加对应的view/assemble调用并分别赋予新轴坐标0,1,2,3。

四、与相邻搬运类骨架的对比

在 examples/ 目录下，stack与cat、unbind、repeat同属「metadata/搬运语义」（view 取片 + assemble 拼回、无逐元素计算），但轴切分方式不同，对比有助于加深理解：

• stack vs cat：cat 的多次assemble共享同一组偏移、仅在拼接轴上连续递增；stack 的多次assemble在新增轴上取不同坐标0/1，这正是「unsqueeze + concat」组合方案的代码形态。
• stack vs repeat：repeat 是单输入在复制轴上重复写同一片；stack 是多输入在新增轴上各写各的片，输入源不同。
• stack vs unbind：unbind 是 stack 的逆操作，把一个大张量沿拆分轴拆成多个输出；stack 则是把多个输入合并到带新轴的大张量。

五、使用注意事项与调优要点

按 examples/README.md 的声明，examples/下每个<op>.md均为kernel 参考骨架，仅展示接口组合与轴切分模式（哪些轴 loop、哪些轴整块），不作为标准模板，且未逐一经 NPU 编译验证。实际使用时需注意：

1. 占位符替换：sl（输入 shape，如[B, S, D]）、ol（输出 shape）、pypto_dtype（如pypto.DT_FP32）、batch（loop 外层轴长，通常sl[0]）、inner（内层 shape，如sl[1:]）均需按实际算子替换，最小可运行 setup 可参考 examples/README.md 中的示例。
2. loop 轴与 tile shape 调优：unroll_list、tile shape（set_vec_tile_shapes(1, *inner)）需按实际 shape / dtype 与平台约束确定并调优；TileShape 每维 > 0、最多 4 维。
3. 动态 shape 风险：若 stack 输入含动态轴，需注意 SKILL.md 中记录的硬约束——归约类等计算 API 在编译期需要 concrete shape；搬运类骨架若依赖具体偏移计算，同样应在 API 可行性报告中评估动态轴影响。
4. 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