MLIR Empowers IBM Spyre: Dataflow Compilation from KTIR to SpyreCode

A recent PyTorch official blog post revealed the key design of the torch-spyre compilation stack: its Inductor component is responsible for deciding how each kernel's tasks are partitioned across Spyre chip cores and for managing tensor residency strategies between on-chip scratchpad and off-chip memory, while the backend compiler translates these decisions into executable SpyreCode. This pipeline is built on MLIR as the intermediate representation foundation, forming a dataflow compilation path from KTIR to SpyreCode.

Why Spyre Needs a Dedicated Dataflow Backend

Spyre is an accelerator designed by IBM for AI inference workloads, and its architecture emphasizes distributed on-chip storage and explicit dataflow scheduling. General-purpose GPU compilers typically assume unified addressing and implicit cache hierarchies, whereas Spyre relies on software to explicitly manage data location—on-chip scratchpad capacity is limited and independent per core, and off-chip memory access latency is significant. Therefore, the compiler's core task is no longer simple instruction selection, but rather deciding "what data is in which core's local storage at what time."

torch-spyre's Inductor takes on this planning responsibility: it partitions the computation graph into kernels, assigns cores to each kernel, and inserts the necessary data movement operations. The backend compiler then converts this planning result into SpyreCode, which contains core-level parallel instructions and explicit DMA descriptions. This division of labor decouples high-level scheduling decisions from low-level code generation, making it easier to optimize each separately.

scratchpad is a type of on-chip storage explicitly managed by software, unlike the cache hierarchy in GPUs that is automatically controlled by hardware. On architectures that use scratchpad, the compiler must explicitly decide when each block of data is brought in, how long it resides, and when it is evicted; the hardware does not make cache hit decisions on the programmer's behalf. Each Spyre core has its own independent scratchpad with limited capacity, which makes "data location planning" a first-class task for the compiler: if the plan is wrong, a core may spin idle waiting for DMA transfers. This is also why such architectures are often called "dataflow architectures"—computation proceeds driven by data arrival, rather than by a unified program counter.

DMA (Direct Memory Access) descriptions refer to data movement instructions explicitly declared in the compilation output, used to transfer tensor tiles between off-chip memory and each core's scratchpad. Unlike CPU/GPU systems where caches automatically prefetch, transfers on Spyre must be inserted by the compiler at appropriate points and synchronized with compute instructions. Inductor's core partitioning and residency decisions ultimately manifest as a set of DMA descriptions and their dependencies, and the backend compiler must ensure overlap between transfers and computation to hide off-chip access latency.

The Role of MLIR in the Compilation Stack

MLIR's multi-level dialect system provides a flexible abstraction ladder for this pipeline. KTIR (Kernel Tensor IR), as the upstream representation, carries the core partitioning and residency decisions made by Inductor; through layer-by-layer lowering, it ultimately generates SpyreCode. The value of MLIR in this process lies in the fact that each dialect layer can independently define verification rules and transformation passes—for example, inserting synchronization primitives at the dataflow layer, and performing buffer allocation and reuse at the memory layer.

This design also means that debugging and performance analysis can be done in stages—developers can inspect whether the dataflow graph matches expectations at the MLIR intermediate representation level, without having to directly face the final machine code. For non-mainstream architectures like Spyre, observability is often the key to whether a compilation stack can be accepted by the ecosystem.

MLIR dialect is a customizable layer of intermediate representation in MLIR, where each dialect layer defines its own operation set, type system, and verification rules. The compiler uses layer-by-layer lowering to convert higher-level dialects into lower-level dialects, down to the target code. This "abstraction ladder" allows different concerns to be handled at different layers: the dataflow layer cares about synchronization and parallelism, the memory layer cares about buffer allocation and reuse, and the instruction layer cares about core-level encoding. For non-mainstream architectures like Spyre, custom dialects can also preserve upstream scheduling decisions in the IR in a structured form, facilitating verification and analysis.

Inductor is the default compilation backend for torch.compile in PyTorch 2.x, responsible for partitioning FX graphs into kernels, generating Triton or C++ code, and providing a pluggable backend interface for different accelerators. torch-spyre reuses Inductor's scheduling capabilities for Spyre by implementing this interface, allowing PyTorch models to run on the chip without hand-written kernels.

Practical Significance for Developers

As PyTorch ecosystem support for IBM Spyre, torch-spyre's compilation stack design draws on Inductor's experience in integrating with other accelerators, but it has been specially tailored for Spyre's explicit memory hierarchy. For application developers, this means PyTorch models can be executed on Spyre through the torch.compile path without hand-written kernels; for compiler researchers, this KTIR→MLIR→SpyreCode pipeline demonstrates a feasible pattern for separating dataflow scheduling decisions from code generation.

Currently, publicly available information focuses on the architectural description level, and specific performance data or comparisons with hand-written kernels have not yet been disclosed. Worth watching going forward is the backend's end-to-end performance on real models, as well as the actual benefits of MLIR passes in optimizing data movement.