IBM 如何用 Inductor 扩展 Spyre 加速器的 torch.compile 后端

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 的编译器发展成能够深入理解该硬件的编译器。

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

1. 我们从 Inductor 中获得了实实在在的助力——但我们仍然需要在 Inductor 的较低层级和代码生成中做大量的树外扩展。这并非即插即用的替换。（这种助力具体到可以衡量——参见评估部分的代码量数据。）
2. Inductor 的大部分功能直接就能用。工作集中在少数几个地方，在这些地方 Spyre 的数据流设计与 Inductor 内置的 GPU 形态假设存在分歧——具体来说，就是 Spyre 如何移动数据、管理片上存储以及将工作分配到核心。
3. 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 注册了若干设备特定的操作分解，用于将 layernorm 等 PyTorch 操作映射为优化后的 Spyre 原语操作序列。随后它又注册了若干 lowering，以让 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 的每核心暂存器和共享设备内存

核心分配

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

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

关键编译器阶段

将工作分散到各核心

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

torch-spyre 的工作划分规划器通过三遍来完成这件事。第一遍是强制性的，对单个核心可寻址的设备内存量施加硬性限制，并提交所需的最小拆分，以将任何超大的张量降回该限制之下。第二遍是专门针对矩阵乘法的成本模型，枚举拆分该操作三个维度的可行方式，估算每种方式的运行时间，并选择最便宜的。第三遍处理其他所有情况。

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

所举矩阵乘法示例的前后对比：未拆分 vs. 完整 2×16 核心网格中的 2×4 切片，展示每个核心的 (K,N) 坐标如何选出其 A 行块和 B 列块

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

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

一个 4 MB 张量按行拆分到四个核心上，使每个核心的切片都能放入暂存区

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

暂存区规划

一旦工作划分就位，第二个主题就登场了。对于核心本地暂存区来说过大的中间张量会迫使做出选择：要么让它们通过设备内存往返，要么重新组织一系列操作以在张量块上工作，并用遍历块的循环包裹这样的序列。分块究竟属于 PyTorch 还是 Triton、Helion 等专用编程语言，仍是一个悬而未决的问题。不过在构建 torch-spyre 时，我们观察到 torch.scan 是一个足够强大的构造，可以表达分块循环，尽管有些笨拙。torch-spyre 自己的 for_each_tile 构造建立在 torch.scan 之上，使表达分块循环变得容易，让程序员能够控制时空权衡——迭代次数与内存占用——例如通过对矩阵乘法的归约维度进行分块。

for_each_tile 对 split-K 矩阵乘法的归约维度进行分块，因此每个操作数一次只有一个 K 切片驻留

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

ScratchpadAllocator 的 pass 流水线，具有可插拔的布局求解器选择

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

softmax 内核的暂存器感知调度四个阶段，每个阶段的总设备内存流量以张量大小为单位，从仅使用设备内存到完全优化

当两个决策相互冲突时

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

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

工作划分不匹配的代价：匹配的拆分使缓冲区免费留在暂存器上，不匹配的拆分将其两次溢出到设备内存

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

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

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

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

协同优化分配器将工作划分并入暂存器规划，包含一条剪枝的贪心路径和一条穷举的约束求解器路径

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

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

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

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

评估

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

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

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

总结

torch-spyre 转向 backend=inductor 带来了真正的杠杆：编译器比它所取代的小约 7 倍，并且每个函数的复杂度可测量地更低。尽管弥合三个特定的硬件差距——stick 粒度数据移动、仅有暂存区且无共享缓存的内存层次结构、静态核心分配——需要在 Inductor 自身代码树之外做大量工作。使这项工作可行的，是 Inductor 暴露了扩展点，让像 torch-spyre 这样的后端能够接入，而无需分叉底层的编译器。

这种关系是双向的。torch-spyre 需要注册特定于设备的分解——这是将 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 的硬件构建树外后端，我们希望听到你的声音——并且，我们正与上游一起努力，让我们走过的这条路在下次更容易被遵循。

文章 Extending Inductor for the IBM Spyre Accelerator 首次发表于 PyTorch。