MLIR赋能IBM Spyre：从KTIR到SpyreCode的数据流编译实践

PyTorch官方博客近期披露了torch-spyre编译栈的关键设计：其Inductor组件负责决定每个内核在Spyre芯片各核心间的任务划分，并管理张量在片上暂存器与片外内存之间的驻留策略，而后端编译器则将这些决策翻译为可执行的SpyreCode。这一流程以MLIR为中间表示基础，构建了一条从KTIR到SpyreCode的数据流编译通路。

为什么Spyre需要专属数据流后端

Spyre是IBM面向AI推理负载设计的加速器，其架构强调分布式片上存储与显式数据流调度。通用GPU编译器通常假设统一寻址与隐式缓存层次，而Spyre依赖软件显式管理数据位置——片上scratchpad容量有限且各核心独立，片外内存访问延迟显著。因此，编译器的核心任务不再是简单的指令选择，而是决定“什么数据在什么时刻位于哪个核心的本地存储中”。

torch-spyre的Inductor承担了这一规划职责：它将计算图切分为内核，为每个内核分配核心，并插入必要的数据搬运操作。后端编译器随后把这一规划结果转换为SpyreCode，其中包含核心级并行指令与显式DMA描述。这种分工使得高层调度决策与底层代码生成解耦，便于分别优化。

scratchpad（暂存器） 是一种由软件显式管理的片上存储，与GPU中由硬件自动控制的缓存层次不同。在采用scratchpad的架构上，编译器必须明确决定每块数据何时被搬入、驻留多久、何时被换出，硬件不会替程序员做缓存命中判断。Spyre各核心拥有独立的scratchpad且容量有限，这使“数据位置规划”成为编译器的第一等任务：一旦规划失误，核心可能因等待DMA搬运而空转。这也是此类架构常被称为“数据流架构”的原因——计算的推进由数据到达驱动，而非由统一的程序计数器驱动。

DMA（直接内存访问）描述 指编译产物中显式声明的数据搬运指令，用于在片外内存与各核心的scratchpad之间传输张量分块。与CPU/GPU上由缓存自动预取不同，Spyre上的搬运必须由编译器在合适位置插入，并与计算指令做同步。Inductor的分核与驻留决策最终就体现为一组DMA描述及其依赖关系，后端编译器需保证搬运与计算的重叠以隐藏片外访问延迟。

MLIR在编译栈中的角色

MLIR的多级dialect体系为这一流程提供了灵活的抽象阶梯。KTIR（Kernel Tensor IR）作为上游表示，携带Inductor做出的分核与驻留决策；经过逐层lowering，最终生成SpyreCode。MLIR在此过程中的价值在于：每一层dialect可以独立定义验证规则与变换pass，例如在数据流层插入同步原语、在内存层做缓冲区分配与复用。

这种设计也意味着调试与性能分析可以分阶段进行——开发者可以在MLIR中间表示上检查数据流图是否与预期一致，而无需直接面对最终机器码。对于Spyre这类非主流架构，可观测性往往是编译栈能否被生态接受的关键。

MLIR dialect（方言） 是MLIR中可自定义的一层中间表示，每层dialect定义自己的操作集、类型系统与验证规则。编译器通过逐层lowering把高层dialect转换为更低层的dialect，直至目标代码。这种“抽象阶梯”使不同关注点可以分层处理：数据流层关心同步与并行，内存层关心缓冲区分配与复用，指令层关心核心级编码。对Spyre这类非主流架构，自定义dialect还能让上游调度决策以结构化形式保留在IR中，便于验证与分析。

Inductor 是PyTorch 2.x中torch.compile的默认编译后端，负责把FX图切分为内核、生成Triton或C++代码，并为不同加速器提供可插拔的后端接口。torch-spyre正是通过实现该接口，把Inductor的调度能力复用到Spyre上，从而让PyTorch模型无需手写内核即可落到该芯片执行。

对开发者的实际意义

torch-spyre作为PyTorch生态对IBM Spyre的后端支持，其编译栈设计借鉴了Inductor对接其他加速器的经验，但针对Spyre的显式内存层次做了专门处理。对应用开发者而言，这意味着PyTorch模型可以通过torch.compile路径落到Spyre上执行，而无需手写内核；对编译器研究者而言，这条KTIR→MLIR→SpyreCode的流水线展示了一种将数据流调度决策与代码生成分离的可行模式。

目前公开信息集中在架构描述层面，尚未披露具体性能数据或与手写内核的对比。后续值得关注的是该后端在真实模型上的端到端表现，以及MLIR pass在数据搬运优化上的实际收益。