TL;DR

IBM Spyre 加速器基于 PyTorch 的软件栈最近将其 torch.compile 后端从 sendnn（一个完全从零构建的独立编译栈）迁移到了 backend=inductor，后者是通过扩展 PyTorch 自有的 Inductor 构建的开源后端。Inductor 的大部分内容未经修改即可直接使用。真正的工作集中在三个地方，在这些地方 Spyre 的硬件对编译器提出了与 GPU（图形处理器） 根本不同的问题。事实证明，Inductor 正是这项工作的理想基础，最终得到的编译器比被替代的方案小了约 7 倍。

IBM 将 Spyre 设计为一款面向企业推理工作负载的全新 AI 加速器，优先考虑能效。它是一款 75 瓦的 PCIe 附加加速器，无需定制供电或散热。

Spyre 是一款数据流加速器，专为推理所需的低精度算术而设计。它由 32 个小核心通过环形网络连接而成，每个核心拥有自己本地的 2 MB 暂存器而非共享缓存。它以称为 stick 的固定大小块来移动数据——128 字节，即 fp16 格式下的 64 个值。

Torch-spyre 是针对它的 PyTorch 集成，本文讨论的是如何在 Inductor 之上将 torch-spyre 的编译器发展成能够深入理解该硬件的方案。

在迁移过程中，有三点尤为突出：

• 我们从 Inductor 中获得了实实在在的收益——但我们仍然需要在 Inductor 的较低层级和代码生成部分进行大量的树外扩展。这不是简单的即插即用替换。（收益具体到可以量化——参见评估部分的代码规模数据。）

• Inductor 的大部分功能直接可用。工作集中在少数几个地方，在这些地方 Spyre 的数据流设计与 Inductor 内置的 GPU 形态假设存在分歧——具体来说，就是 Spyre 如何移动数据、管理片上存储以及将工作分配到核心。

• Inductor 的扩展点使这项工作变得可行。我们目前正在与上游 PyTorch 项目合作，将我们最依赖的扩展点正式化，以便其他树外硬件后端也能从同样的收益中获益。

本文的其余部分将介绍该后端的整体结构、驱动其设计的硬件差异、编译器如何弥合这些差异，以及我们目前所测量的结果。

Spyre Inductor 后端的整体结构

Torch-spyre 的编译器实现包含两个阶段。第一阶段使用 Inductor 扩展（位于 torch-spyre 仓库中）将 torch.compile 编译后的图降低为一种称为 SuperDSC 的中间表示。第二阶段使用名为 DeepTools 的专有编译器将 SuperDSC 转换为在 Spyre 硬件上运行的指令。

本文的所有内容都聚焦于第一阶段：Inductor 扩展。

Inductor 侧接入标准的 torch.compile 流水线——Dynamo 追踪模型，AOTAutograd 对其进行分区，Inductor 将生成的 FX 图降级为其 LoopLevelIR，然后再生成代码。torch-spyre 的扩展增强了这一流水线。

| 扩展点 | 阶段 | torch-spyre 在此处的工作 |
| --- | --- | --- |
| Decompositions | FX 图 | Spyre 特定的算子分解 |
| Lowerings | FX 图到 LoopLevelIR | Spyre 特定的算子降级 |
| CustomPreSchedulingPasses | LoopLevelIR | Spyre 特定编译器的主体：布局分配、分块、工作划分、暂存器分配 |

torch-spyre 注册了大量设备特定的算子分解，将 layernorm 等 PyTorch 算子映射为 Spyre 原语算子的优化序列。随后它注册了大量降级规则，让 Inductor 了解这些原语算子。

大部分有趣的工作发生在 CustomPreSchedulingPasses 中——分块、布局分配、工作划分和暂存器分配实际上都是在这里决定的。这也是目前通过 monkey-patch Inductor 的 GraphLowering.codegen() 接入的点，而非一等扩展钩子——我们正在与上游推进的具体形式化方案，将在本文末尾再次讨论。

LoopLevelIR 是这项工作的正确层级，但 Inductor 自身的钩子 _pre_fusion_custom_pass 来得太晚：它只在 LoopLevelIR 已经被转换为 SchedulerNodes 之后才运行，而一旦这些节点存在，Inductor 就不支持大规模重写图。torch-spyre 的 pass 恰恰需要这种大规模重写，这就是为什么 CustomPreSchedulingPasses 提前一步钩住 GraphLowering——在 Inductor 确定 SchedulerNodes 并关闭大规模变换之前。

虽然这些 Spyre 特定的扩展是本文的重点，我们也利用了其他 torch.compile 扩展钩子，例如将许多设备上的张量池化为少量内存分配，或从源代码中解析并传播工作划分提示到相关的编译器 pass。

驱动设计的三个硬件差异

Inductor 是围绕其底层硬件的思维模型构建的——这个模型可以理解地由 GPU 塑造。该模型的大部分都能很好地迁移到 Spyre；但有三部分不行，而这三部分几乎解释了所有必须在树外构建的内容。

张量布局与 stick

Spyre 不允许内核以 GPU 那样的单元素粒度访问内存。在 GPU 上，内核可以请求单个元素，即使硬件可能通过获取更大的缓存行来满足该请求。GPU 将单个元素作为编程抽象暴露出来，而 Spyre 将 stick——128 字节对齐的内存块——作为数据移动和计算的基本单位。Spyre 内核必须以 stick 为单位移动数据，没有单独移动单个元素的等效路径。

对于一个继承了 GPU 式假设的编译器来说，这不是一个小注脚。Inductor 用于确定给定循环迭代触及哪个张量元素的索引算术，现在必须以 stick 而非元素来推理。除了索引问题之外，stick 还带来了 GPU 导向编译器从未面临过的额外挑战。

每个算子对输入和输出张量元素如何组装成 stick 都有特定约束。例如，逐点算子要求其所有输入和输出在一个共享的 stick 维度上保持一致——即最内层维度。当算子的输入以不兼容的 stick 到达时，编译器必须插入 restickify：一种数据移动算子，不仅通过置换 stick 来物理重排张量元素，还重新形成 stick 本身。restickify 算子是 stick 不可变性法则的唯一例外。在整个图中找到一种 stick 维度分配，使数据移动的总成本最小化，是一个全局优化问题。

即使每个 stick 约束都得到满足，torch-spyre 所选择的设备端张量布局也决定了性能。次优布局会直接影响实际运行时间。在一个所有张量都已驻留在设备内存中的批量矩阵乘法中，改变输入和输出张量的分块方式，在一系列批量大小和形状上将内核时间缩短了 29-69%。

这些约束共同将布局问题转化为一个图级别的优化问题——而 Inductor 的 GPU 后端从未被设计来解决这个问题。

暂存存储器

第二个差异与内存层次结构有关。Spyre 的每个核心都带有自己的小型本地暂存存储器——速度快，但容量有限。其背后只有全局设备内存。没有任何数据会自行移动：从设备内存到暂存存储器的每一次传输，以及从暂存存储器到核心计算引擎的每一次传输，都是因为编译器发出了显式指令将其移动到那里。没有硬件管理的层级自行决定哪些数据值得保留在近处——如果编译器没有计划让某个值留在暂存存储器中，它就必须发出指令将其写回设备内存，并在下次需要时重新获取。Spyre 上的工作负载通常是内存受限而非计算受限，因此这种往返传输往往才是时间真正消耗的地方。

这是与 GPU 在程度上的差异，而非本质上的差异——GPU 也有内存层次结构，而且用好它们的大量负担已经落在程序员或编译器身上。但 Spyre 在同一方向上走得更远，完全没有自动回退机制。编译器必须提前决定何时以及在哪里移动数据以支持计算。这就是下面描述的暂存存储器规划机制的工作。

核心分配

还有一个相关的第三个差异：Spyre 的核心作为静态数据流机器执行。没有 GPU 风格的运行时调度器在流式多处理器之间动态平衡负载。这意味着跨核心的工作划分也必须在编译时决定——这是下面的第一个主题。

用 GPU 术语来说，这意味着 Spyre 内核本质上是持久内核：工作划分是先验固定的，但每个核心的暂存存储器可以在多轮计算中持久保存有价值的数据。

关键编译阶段

将工作分散到各核心

第一个主题是将单个操作的工作划分到 Spyre 的各个核心上。由于 Spyre 作为静态数据流机器运行，而不是依赖运行时调度器来平衡负载，这种划分必须在编译时、在任何东西运行之前、针对每个操作来决定。

Torch-spyre 的工作划分规划器通过三个步骤来完成这项工作。首先，一个强制步骤对单个核心可以寻址的设备内存量施加硬性限制，执行必要的最小拆分以使任何超大张量回到该限制以下。其次，一个专门针对矩阵乘法的成本模型步骤枚举拆分该操作三个维度的可行方式，估算每种方式的运行时间，并选择最便宜的。第三个步骤处理其他所有情况。

考虑一个 matmul，输入为 [8192, 32768]，权重为 [32768, 4096]，输出为 [8192, 4096]，全部为 fp16。如果不进行拆分，输入张量的内存跨度将超过单个核心允许寻址的范围。因此，第一遍立即在归约维度上提交 2 路拆分。在此基础上，第二遍估计沿输出维度进行 16 路拆分最为有利。规划器不仅仅是在选择一个有效的划分，它选择的是使预计执行时间最小化的划分。

规划器优先处理矩阵乘法，因为它们通常主导端到端计算时间。

将工作拆分到多个核心上，也会拆分这些核心所接触的数据。一个 [1024, 2048] 的 fp16 张量为 4 MB——对于任何单个核心的暂存器来说都太大——但按行拆分到四个核心上，每个核心的切片大小舒适地降至 1 MB。

然而，多核暂存器驻留并非没有代价。它取决于每个接触该张量或瓦片（缓冲区）的操作都对其拆分方式达成一致——而到目前为止，还没有任何机制保证这种一致性。

暂存器规划

第二个主题在工作划分确定后展开。对于核心本地暂存器而言过大的中间张量迫使做出选择：要么通过设备内存进行往返，要么重新组织一系列操作以在张量瓦片上工作，并用遍历瓦片的循环包裹这样的序列。瓦片化究竟属于 PyTorch 还是 Triton、Helion 等专用编程语言，仍是一个开放问题。然而，在构建 torch-spyre 的过程中，我们观察到 torch.scan 是一个足够强大的构造，可以表达瓦片化循环，尽管方式略显笨拙。基于 torch.scan 构建，torch-spyre 自有的 for_each_tile 构造使表达瓦片化循环变得容易，让程序员能够控制时空权衡——迭代次数与内存占用——例如通过对矩阵乘法的归约维度进行瓦片化。

一旦内核被瓦片化，且瓦片被分配给核心，暂存器分配器就会决定哪些内容实际驻留在核心的暂存器中以及驻留在什么地址。它提供多种求解器选择，从简单的贪心算法到用于更难情况的 CP-SAT 和模拟退火选项。

收益是可衡量的。torch-spyre 的文档展示了一个在中等大小 fp16 张量上的单核 softmax——降级后的序列为 max → sub → exp → sum → div——并展示了随着暂存器感知调度逐项添加能力时发生的情况：首先将小的归约输出固定在核心上，然后为逐点步骤就地复用暂存器地址，然后将输入克隆到暂存器一次，而不是为每个读取者从设备内存重新获取。端到端来看，此示例上的内核运行时间中位数从 32.51 µs 降至 23.67 µs——降低了 27%，纯粹来自决定什么留在片上以及何时留在片上。

当两个决策相互冲突时

上文将暂存器分配和工作划分描述为两个独立的遍，torch-spyre 的第一个实现就是这样运行的：工作划分为每个操作独立选择核心拆分，暂存器分配随后放置缓冲区，并将这些拆分视为固定不变。事实证明这具有局限性，这就是为什么 torch-spyre 正在转向下文详细描述的协同优化方法。

当两个共享缓冲区的相邻操作选择了不同的核心拆分时，问题就会出现。工作划分以各自的视角优化每个操作：matmul 希望采用能让其计算引擎保持忙碌的拆分；相邻的归约或逐点操作希望采用与其自身形状匹配的拆分。没有任何机制强制这些选择保持一致。当它们不一致时，它们之间的缓冲区被生产者切片的方式与消费者不同——一个核心的暂存器拥有数据的一部分，而下一个操作期望的是另一部分。torch-spyre 将此称为工作划分不匹配。历史上，处理此问题的唯一方法是直接取消该缓冲区驻留暂存器的资格：生产者将其写出到设备内存，消费者从那里读回，双向支付往返代价——这正是暂存器分配意图避免的流量。

接下来，torch-spyre 探索了一个有前景的替代方案：为每个矩阵乘法单独决定最优工作划分，但中间操作的工作划分从前一个矩阵乘法推导而来，从而将第二和第三工作划分遍解耦。当不匹配无法避免时，设备内存往返仍然是后备方案。

这不再是唯一的选择。Spyre 的 scratchpad 硬件支持核间重布局——一种将数据直接从一个核的 scratchpad 切片移动到另一个核的切片、而无需触碰设备内存的 shuffle。torch-spyre 已开始以 lx_relayout 之名启用这一能力。重布局并不能让不匹配变得免费：在核之间重新洗牌数据仍然要耗费真实时间。而且由于在 shuffle 进行期间旧布局和新布局必须在 scratchpad 中共存，它也会占用 scratchpad 空间。它改变的是选择的形态。过去不匹配意味着全有或全无地回退到设备内存；现在它可以改为一次价格合理的片上 shuffle，由编译器来决定。

不过，更好的 shuffle 仍然是 shuffle。还有另一种选择：从独立的工作划分决策转向协调决策，准确权衡相邻算子之间一致的 scratchpad 分配所带来的收益与代价。torch-spyre 新的协同优化分配器将独立的工作划分 pass 并入 scratchpad 规划，而不是将其输出视为固定不变。该分配器会比较相邻算子之间因最优但不匹配的工作划分决策所产生的数据 shuffle 成本，与对齐但次优的工作划分决策的成本。

最终，协同优化器要做出许多决策并考虑许多选项；两个求解器后端以不同方式完成这一搜索。剪枝的贪心路径通过硬件特定的启发式方法缩小空间——对于 matmul，锚定在其自身成本模型 pass 已经偏好的划分上——然后用最佳适配打包放置缓冲区，使搜索对每个图都足够快。穷举的约束编程路径（torch-spyre 使用 OR-Tools CP-SAT）则一次性考虑候选划分和重布局选择的整个空间：每个算子采用哪个候选划分，以及任意给定的生产者-消费者对之间是否要为一次重布局付费，这些由硬约束绑定在一起——缓冲区的选定划分是否与每个消费者的划分匹配，无论是直接匹配还是通过付费重布局匹配；所有内容是否真的能放进 scratchpad，包括重布局生效期间的两份副本——一旦这些约束满足，就有一个要最小化的成本。该重布局成本是一个普通项，由与普通设备内存流量定价相同的硬件拟合成本模型构建，因此求解器绝不会拿真实成本与假定免费的成本作比较。穷举路径之所以能完全吞并工作划分，正是因为它在做剪枝；贪心路径仍然需要成本模型 pass 的启发式方法来保持自身搜索的可处理性。

收益在真实内核上具体显现。在一个融合的 MLP 风格内核上，让一个逐点算子采用相邻 matmul 的核划分——而不是坚持其自身局部最优的划分——将核利用率从大约三分之二提升到约 80%，并将内核时间缩短了约 17%。这相当于比等效的 sendnn 编译内核提升了约 2 倍。

在一个多头注意力内核上，softmax 两侧的两个 matmul 沿不同轴收缩，因此它们永远无法直接采用彼此的划分。为它们两者——以及位于它们之间的 softmax 归约——提供共享的批次和输出维度分块，使整条链收敛到一种一致的划分。这固定住了原本会成为注意力分数矩阵的一次设备内存往返。

而在一个专门为让工作划分与 scratchpad 放置产生张力而构建的小型合成图上，让求解器选择一次重布局而不是强制溢出，将一次 80 微秒的设备内存往返替换为一次 1.1 微秒的片上 shuffle。旧的严格匹配门控会付出比问题实际所需多 70 倍的代价。

这仍是编译器中年轻的一部分。搜索空间随内核规模迅速增长。穷举求解器会考虑整个空间。因此，我们限制其执行时间，并继续探索智能剪枝启发式方法，以在不损害质量的前提下加速解的搜索。

评估

这一举措的软件工程理由在原始规模上最容易看出。sendnn，即这个 Inductor 后端所取代的编译栈，是独立于 torch-spyre 和 Inductor 构建的。用 lizard 测量（非注释代码行数，不含测试），它在 6,170 个函数中总计约 273,000 行 C++ 和 Python，平均每函数圈复杂度约为 10。取代它的 Inductor 扩展位于 torch_spyre/_inductor 下，约 38,600 行 Python，分布在 1,878 个函数中，平均每函数圈复杂度为 5.6：代码行数约少 7 倍，函数数量为其三分之一，而且函数更简单。

这并不完全是“我们从零编写更少代码做同样的事”的同类比较，因为 Inductor 扩展在底层依赖 Inductor 自身的全部机制，而 sendnn 永远无法做到这一点。但这种不对称正是重点：一个规模小 7 倍、复杂度更低的代码库，就是“从现有编译器获得真正杠杆”在实践中的样子。

在性能方面，更完整的情况仍在整理中，但两个真实、经过测量的数据点出现在内存层次结构的两端。在核上，上面的 softmax 示例仅凭更好地利用 scratchpad 就使内核时间中位数下降了 27%。在核外，matmul 布局示例仅凭分块决定 stick 是否连续落在设备内存中，就使内核时间减少了 29-69%。这两个数字互不依赖。它们共同证明，Spyre 独特的架构对编译器提出了真实的要求，而 Inductor 和 torch-spyre 能够成功应对这一挑战。

总结

Torch-spyre 转向 backend=inductor 带来了实实在在的收益：编译器体积大约缩小到被替换版本的 1/7，且每个函数的复杂度也显著降低。尽管弥合三个特定的硬件差距——stick 粒度的数据移动、仅有暂存器（scratchpad）而无共享缓存的内存层级、静态核心分配——需要在 Inductor 自身代码树之外做大量工作。使这些工作变得可行的关键在于，Inductor 暴露了扩展点，让 torch-spyre 这样的后端能够接入，而无需分叉底层的编译器。

这种关系是双向的。Torch-spyre 需要注册设备特定的分解（decomposition）——这是将 Inductor 适配到 Spyre 算子集时产生的一个具体需求——推动了树外 Inductor 后端支持的改进，并已合入上游 PyTorch 2.12。我们以 RFC 形式提出了该变更，并作为两个合并的 PR 发布：#176811 和 #185909。描述 torch-spyre 所依赖的进一步扩展点的两份设计文档也已公开——一份提出针对 stick 粒度内存的平铺布局（tiled-layout）抽象，另一份关于新的基于 MLIR 的中间表示（KTIR）——我们已经与上游 Inductor 维护者分享并讨论过：RFC 0047 – tiled tensors，RFC 0682 – KTIR。而本文最核心的扩展点 CustomPreSchedulingPasses 是下一个。将其从目前的 monkey-patch 转变为受支持的钩子，是与上游正在进行的积极讨论。

我们认为 torch-spyre 是一个有用的例证，表明 Inductor 的扩展点模型可以很好地推广到 GPU 形态之外的硬件。如果你正在为不像 GPU 的硬件构建树外后端，我们期待听到你的声音——并且，我们正与上游一起努力，让我们走过的这条路在下次更容易被遵循。

转载声明：本文为转载整理，原文版权归原作者及原站点所有。

原文来源：PyTorch Blog

原创地址：https://pytorch.org/blog/extending-inductor-for-the-ibm-spyre-accelerator/