From KTIR to SpyreCode: An MLIR Dataflow Backend for IBM Spyre

TL;DR Torch-spyre’s Inductor decides how each kernel’s work is divided across Spyre’s cores and manages tensor residency between on-chip scratchpad and off-chip memory, and the backend compiler turns those decisions into something the hardware can actually run. The contract between them is represented using KTIR (Kernel Tile IR), an MLIR representation built on a data-parallel view of the accelerator, in which every compute tile can see every memory and a single logical tensor can live spread across several cores’ scratchpads. A new open-source Dataflow Scheduler consumes KTIR and turns each tile’s work into a dataflow pipeline: it splits the kernel into programs, starts from a simple load → compute → store pipeline, routes every transfer through the memory hierarchy given in an architecture description, and expands that pipeline into a nested, double-buffered, multi-instance schedule. The output of the scheduler is represented using DataflowIR , an MLIR intermediate representation that Spyre’s device-specific stages use to produce SpyreCode , a symbolic SpyreBinary plus a job plan that the Spyre runtime prepares once and launches many times, filling in tensor addresses and shapes at launch. This post is the middle chapter of three. Extending Inductor for the IBM Spyre Accelerator described how torch-spyre’s Inductor extensions settle layout, tiling, work division and scratchpad placement, and Building Spyre as a Native PyTorch Device described how the runtime launches a compiled artifact whose symbolic values are only filled in at the last moment. This post covers how the decisions Inductor makes actually become a program the runtime can launch. Where the backend compiler sits in the Spyre software stack Three ideas run through the design. The frontend decides the parallel decomposition, and the backend makes it run. KTIR is designed to carry the mapping Inductor has already chosen (which core computes what, over which data, out of which memory) so that the backend never has to rediscover it. The backend’s job is dataflow scheduling inside that mapping. Scheduling means building a pipeline. Spyre is a dataflow machine with decoupled access and execution, so loads, stores and compute run on separate engines communicating via FIFOs. The scheduler starts from the simplest legal pipeline and refines it one transformation at a time, and every one of those steps is a readable MLIR rewrite. The boundary is an IR, not a library call. KTIR, the scheduler’s ScheduleIR and DataflowIR are all ordinary MLIR built on upstream dialects ( linalg , scf , arith , memref , tensor ), and the scheduler and its dialects are open source under Apache-2.0. Device-specific lowering stages for code generation start at the DataflowIR boundary. KTIR: the contract between Inductor and the backend Why a new interface One of the key requirements for efficient backend scheduling is a contract between torch-spyre’s Inductor and the backend compiler with the capability and flexibility to capture the complete data-parallel view of the accelerator, in which every compute tile can see every memory and a single logical tensor can live spread across several cores’ scratchpads. KTIR ( RFC 0682 , implemented in torch-spyre/ktir-mlir-frontend ) is a tile-based, block-structured IR built on open-source MLIR. Our goal was to make it expressive enough to capture data-parallel mappings from two quite different sources: Inductor on one side, and hand-written Spyre kernels in languages like Triton or Helion on the other. A data-parallel abstraction of the accelerator KTIR is written against a deliberately simple picture of the hardware: a set of cores, each with a compute engine and its own on-chip scratchpad, joined by an on-chip interconnect that also reaches one or more off-chip memories. The property that matters most is that every compute tile has a global view of every memory , which means its own scratchpad, every other core’s scratchpad, and device memory. Hardware abstraction of a multi-core accelerator embodied in KTIR (from RFC 0682) That picture matches Spyre’s shape closely: 32 cores on a ring, each with 2 MB of scratchpad, backed by off-chip device memory. Data moves between device memory and the scratchpads in sticks , fixed 128-byte chunks that hold 64 fp16 values. Intentionally, the picture leaves out details such as FIFOs, load/store units and decoupled execution engines. KTIR describes what each compute tile does and where its data lives , not how that data gets there , because working out how it gets there is exactly the scheduler’s job. KTIR is also compositional. Arithmetic, reductions and control flow come straight from upstream MLIR dialects ( arith , math , linalg , scf , tensor , memref ), and a dedicated dialect, ktdp , adds the two things those dialects can’t express: explicit memories and explicit accesses. A second, much smaller dialect, spyreop , covers Spyre’s own intrinsics. The core ktdp operations are: ktdp op What it expresses get_compute_tile_id Which compute tile this is; 1-D or multi-dimensional, over the grid attribute on the kernel construct_memory_view Interpret a region at some (possibly symbolic) address in a named memory space as a strided tensor; no allocation, no data movement construct_distributed_memory_view Compose several memory views, each covering part of a global coordinate space, into one logical tensor construct_access_tile A structured set of coordinates into a view: an integer set for the domain and an affine map for the traversal order, so tiles can be rectangular, strided, skewed or triangular construct_indirect_access_tile The same, with some dimensions indexed through another tensor, e.g. paged KV-cache gathers load / store The only data movement: consume an access tile and produce or write a data tile inter_tile_produce / inter_tile_reduce Communication between compute tiles, e.g. an all-reduce of partial sums Here is a complete kernel, an elementwise add over 96×64 fp16 tensors divided among 32 compute tiles of three rows each. #tile = affine_set<(d0, d1) : (d0 >= 0, -d0 + 2 >= 0, d1 >= 0, -d1 + 63 >= 0)> // 3x64 #full = affine_set<(d0, d1) : (d0 >= 0, -d0 + 95 >= 0, d1 >= 0, -d1 + 63 >= 0)> #id = affine_map<(d0, d1) -> (d0, d1)> // The tensors' start addresses arrive as arguments: PyTorch's allocator owns them. func.func @add(%A: index, %B: index, %C: index) attributes {grid = [32]} { %c0 = arith.constant 0 : index %c3 = arith.constant 3 : index %id = ktdp.get_compute_tile_id : index %row0 = arith.muli %id, %c3 : index // this tile's 3-row slab %A_view = ktdp.construct_memory_view %A, sizes: [96, 64], strides: [64, 1] { coordinate_set = #full, memory_space = #ktdp.memory_space } : memref<96x64xf16> // ... %B_view and %C_view alike %A_t = ktdp.construct_access_tile %A_view[%row0, %c0] { access_tile_set = #tile, access_tile_order = #id } : memref<96x64xf16> -> !ktdp.access_tile<3x64xindex> // ... %B_t and %C_t alike %a = ktdp.load %A_t : !ktdp.access_tile<3x64xindex> -> tensor<3x64xf16> %b = ktdp.load %B_t : !ktdp.access_tile<3x64xindex> -> tensor<3x64xf16> %e = tensor.empty() : tensor<3x64xf16> %c = linalg.add ins(%a, %b : tensor<3x64xf16>, tensor<3x64xf16>) outs(%e : tensor<3x64xf16>) -> tensor<3x64xf16> ktdp.store %c, %C_t : tensor<3x64xf16>, !ktdp.access_tile<3x64xindex> return } The kernel keeps four concerns apart. Views say how memory should be interpreted, access tiles say which coordinates a given tile touches, loads and stores are the only operations that move data, and linalg does the arithmetic. Notice also that the work division is already settled: grid = [32] , together with the arithmetic on %id , is the partitioning, and the backend is expected to honor it rather than second-guess it. Spyre intrinsics: the spyreop dialect Upstream arith and math cover most of the arithmetic a kernel needs, but Spyre’s compute units also offer operations that no generic dialect has a name for. A kernel that wants one of those s