PyTorch Team Details MLIR Dataflow Backend That Turns KTIR Into IBM Spyre Code

KTIR: the contract between Inductor and the backend

According to the post, torch-spyre's Inductor decides how each kernel's work is split across Spyre's cores and manages where tensors reside, either in on-chip scratchpad or off-chip memory. The backend compiler then turns those decisions into something the hardware can run. The contract between the two is KTIR, short for Kernel Tile IR, an MLIR representation built on a data-parallel view of the accelerator.

In that view, every compute tile can see every memory, and a single logical tensor can be spread across the scratchpads of several cores. KTIR is described in RFC 0682 and implemented in torch-spyre/ktir-mlir-frontend. The authors say they wanted it expressive enough to capture mappings from two different sources: Inductor, and hand-written Spyre kernels in languages such as Triton or Helion.

The hardware picture behind KTIR is intentionally simple: cores with a compute engine and their own scratchpad, joined by an on-chip interconnect that also reaches off-chip memory. The post says this matches Spyre closely, with 32 cores on a ring, each with 2 MB of scratchpad, backed by off-chip device memory. Data moves in sticks, fixed 128-byte chunks that hold 64 fp16 values. Details such as FIFOs, load/store units and decoupled execution engines are left out on purpose, because KTIR describes what each tile does and where its data lives, not how the data gets there.

How KTIR is built

KTIR reuses upstream MLIR dialects for arithmetic, reductions and control flow: arith, math, linalg, scf, tensor and memref. A dedicated dialect, ktdp, adds explicit memories and explicit accesses. Its core operations include get_compute_tile_id, construct_memory_view, construct_distributed_memory_view, construct_access_tile, construct_indirect_access_tile, load and store, and inter_tile_produce and inter_tile_reduce for communication between tiles such as an all-reduce of partial sums.

The post walks through an elementwise add over 96×64 fp16 tensors divided among 32 compute tiles of three rows each. Views say how memory is interpreted, access tiles say which coordinates a tile touches, loads and stores are the only data movement, and linalg does the arithmetic. The work division is already fixed by grid = [32] and the arithmetic on the tile id, and the backend is expected to honor it rather than second-guess it.

A second, much smaller dialect, spyreop, exposes Spyre intrinsics as ordinary MLIR operations. It covers elementwise and activation functions such as exp, gelu, sigmoid, silu, softplus, sqrt, rsqrt, reciprocal and realdiv, building blocks for normalization, integer and address arithmetic including idx32toaddr, and compare and select. It also defines df16, IBM's 16-bit DLFloat format, which splits its bits 1-6-9 between sign, exponent and mantissa, compared with 1-5-10 for IEEE half precision.

The Dataflow Scheduler and SpyreCode

A new open-source Dataflow Scheduler consumes KTIR and turns each tile's work into a dataflow pipeline. It splits the kernel into programs and starts from a simple load, compute, store pipeline. It then routes every transfer through the memory hierarchy given in an architecture description and expands the pipeline into a nested, double-buffered, multi-instance schedule.

The post explains that Spyre is a dataflow machine with decoupled access and execution, so loads, stores and compute run on separate engines that communicate through FIFOs. The scheduler begins with the simplest legal pipeline and refines it one transformation at a time, each a readable MLIR rewrite.

The scheduler's output is DataflowIR, an MLIR intermediate representation. Spyre's device-specific stages use it to produce SpyreCode, which is a symbolic SpyreBinary plus a job plan. The Spyre runtime prepares this once and launches it many times, filling in tensor addresses and shapes at launch. KTIR, the scheduler's ScheduleIR and DataflowIR are all ordinary MLIR built on upstream dialects, and the scheduler and its dialects are open source under Apache-2.0. Device-specific lowering for code generation starts at the DataflowIR boundary.

What it means for developers

The post is the middle of a three-part series. The earlier posts cover Inductor's extensions for layout, tiling and scratchpad placement, and how the runtime launches compiled artifacts. For developers, the design means the parallel decomposition is decided once in the frontend and not rediscovered by the backend. Because the boundaries are MLIR IRs rather than library calls, and the scheduler is Apache-2.0 licensed, the stages can be inspected and reused. KTIR is also meant to accept hand-written kernels from Triton or Helion alongside Inductor output.