IBM Spyre 加速器的 torch.compile 后端从 sendnn 迁移到 Inductor

TL;DR

The PyTorch-based software stack for IBM’s Spyre accelerator recently moved its torch.compile backend from sendnn, a separate compilation stack built entirely from scratch, to backend=inductor, an open-source backend built by extending PyTorch’s own Inductor. Most of Inductor carried over untouched. The real work concentrated in three places where Spyre’s hardware asks a fundamentally different question of the compiler than a GPU does. Inductor turned out to be exactly the right substrate for this work, and the result is a compiler roughly 7x smaller than what it replaced.

IBM designed Spyre as a clean-sheet AI accelerator for enterprise inference workloads, prioritizing power efficiency. It’s a 75-watt, PCIe-attached accelerator that requires no custom power delivery or cooling.

Spyre is a dataflow accelerator architected for the reduced-precision arithmetic of inference. It’s built from 32 small cores wired together over a ring, each with its own local 2 MB scratchpad rather than a shared cache. It moves data around in fixed-size chunks called sticks – 128 bytes, or 64 values in fp16.

Torch-spyre is the PyTorch integration that targets it, and this post is about growing torch-spyre’s compiler, on top of Inductor, into something that understands that hardware well.

Three things stood out as we made that move:

1. We gained real leverage from Inductor – but we still needed substantial out-of-tree extensions at Inductor’s lower levels and in code generation. This wasn’t a drop-in swap. (The leverage is concrete enough to measure – see the code-size numbers in Evaluation.)
2. Most of Inductor just worked. The effort concentrated on a small number of places where Spyre’s dataflow design diverges from the GPU-shaped assumptions built into Inductor – specifically, how Spyre moves data around, manages on-chip storage, and distributes work to cores.
3. Inductor’s extension points made this tractable. We’re now working with the upstream PyTorch project to formalize the points we leaned on most heavily, so other out-of-tree hardware backends can benefit from the same leverage.

The rest of this post walks through the backend’s overall structure, the hardware differences that drove its design, how the compiler bridges them, and what we’ve measured so far.

Overall structure of the Spyre Inductor backend

Torch-spyre’s compiler implementation consists of two stages. The first stage uses Inductor extensions (living in the torch-spyre repository) to lower a torch.compile’d graph down to an intermediate representation called SuperDSC. The second stage uses a proprietary compiler called DeepTools to translate SuperDSC into the instructions that run on the Spyre hardware.

Everything in this post focuses on the first stage: the Inductor extensions.

The Inductor side plugs into the standard torch.compile pipeline – Dynamo traces the model, AOTAutograd partitions it, and Inductor lowers the resulting FX graph into its LoopLevelIR before generating code. Torch-spyre’s extensions augment this pipeline.

Torch-spyre registers a number of device-specific operation decompositions to map PyTorch operations such as layernorm to optimized sequences of primitive Spyre operations. It then registers a number of lowerings to teach Inductor about these primitive operations.

Most of the interesting work happens in CustomPreSchedulingPasses – it’s where tiling, layout assignment, work division, and scratchpad allocation actually get decided. It’s also the point that is, today, wired in via a monkey-patch on Inductor’s GraphLowering.codegen() rather than a first-class extension hook – the concrete formalization we’re pursuing with upstream, which we come back to at the end of this post.

LoopLevelIR is the right level for that work, but Inductor’s own hook there, _pre_fusion_custom_pass, comes too late: it only runs after LoopLevelIR has already been turned into SchedulerNodes, and Inductor doesn’t support rewriting the graph at scale once those nodes exist. Torch-spyre’s passes need exactly that kind of large-scale rewriting, which is why CustomPreSchedulingPasses hooks GraphLowering one step earlier – before Inductor has committed to SchedulerNodes and closed off large-scale transformation.

While these Spyre-specific extensions are the focus of this post, we leverage other torch.compile extension hooks too, for instance, to pool many on-device tensors into few memory allocations or to parse and propagate work division hints from source code to related compiler passes.

Three hardware differences that drive the design

Inductor was built against a mental model of the hardware underneath it – one shaped, understandably, by GPUs. Most of that model travels well to Spyre; three pieces of it don’t, and those three pieces account for nearly everything that had to be built out-of-tree.

Tensor layouts and sticks

Spyre doesn’t let a kernel address memory at the individual-element granularity that GPUs do. On a GPU, a kernel can request an individual element, even though the hardware may satisfy that request by fetching a larger cache line. The GPU exposes individual elements as the programming abstraction, while Spyre exposes the stick – a 128-byte aligned chunk of memory – as the fundamental unit of data movement and computation. A Spyre kernel must move data in sticks, with no equivalent path for moving an individual element on its own.

For a compiler that inherited GPU-shaped assumptions, that’s not a minor footnote. The index arithmetic Inductor uses to determine which tensor element a given loop iteration touches now has to reason in terms of sticks instead of elements. Beyond that indexing problem, sticks create additional challenges that GPU-focused compilers never had to face.

Each operation has specific constraints on how input and output tensor elements can be assembled into sticks. A pointwise op, for instance, requires all of its inputs and its output to agree on one shared stick dimension – the innermost dimension. When an op’s inputs arrive with incompatible sticks, the compiler has to insert a restickify: a data-movement op that physically reorders the tensor elements not just by permuting sticks but also reforming the sticks themselves. The restickify op is the one exception to the stick-immutability law. Finding an assignment of stick dimensions across a whole graph that minimizes the total cost of data movement is a global optimization problem.

Even once every stick constraint is satisfied, the on-device tensor layouts torch-spyre picks also determine performance. Suboptimal layouts have a direct impact on wall-clock time. On a batched matmul with all tensors already resident in device memory, changing how the input and output tensors are tiled cut kernel time by 29-69% across a range of batch sizes and shapes.

Together, these constraints turn layout into a graph-wide optimization problem – one that Inductor’s GPU backend was never designed to solve.

Scratchpad memories

The second difference has to do with memory hierarchies. Each of Spyre’s cores carries its own small local scratchpad – fast, but modest. Behind it there’s nothing but the global device memory. Nothing moves on its own: every trip from device memory into scratchpad, and every trip from scratchpad into a core’s compute engines, happens because the compiler emitted an explicit instruction to move it there. There’s no hardware-managed layer deciding on its own what’s worth keeping close – if the compiler doesn’t plan for a value to stay in scratchpad, it has to issue the instructions to write it back out to device memory and fetch it again the next time it’s needed. Workloads on Spyre are frequently memory-bound rather than compute-bound, so that round trip is very often where the time actually goes.

That’s a difference of degree from a GPU, not of kind – GPUs also have memory hierarchies, and plenty of the burden of using them well already falls on the programmer or compiler. But Spyre pushes further in the same direction with no automatic fallback at all. The compiler has to decide ahead of time when and where to move data to support the computation. That’s the job of the scratchpad-planning machinery described below.

Spyre’s per-core scratchpad and shared device memory

Core assignment

There’s a related, third difference: Spyre’s cores execute as a static dataflow machine. No GPU-style runtime scheduler balances load across streaming multiprocessors on the fly. That means the division of work across cores has to be decided at compile time too – which is the first theme below.

In GPU terminology, this means that Spyre kernels are naturally persistent kernels: the division of work is fixed a priori, but the scratchpad of each core can persist valuable data across many rounds of computation.

Key compiler phases

Spreading work across the cores

The first theme is dividing a single operation’s work across Spyre’s cores. Because Spyre runs as a static dataflow machine rather than relying on a runtime scheduler to balance load, this division has to be decided at compile time, per operation, before anything runs.

Torch-spyre’s work-division planner does this in three passes. First, a mandatory pass enforces a hard per-core limit on how much device memory a single core can address, committing the minimum splits needed to bring any oversized tensor back under that limit. Second, a cost-model pass specifically for matrix multiplications enumerates plausible ways to split the operation’s three dimensions, estimates the runtime of each, and picks the cheapest. A third pass handles everything else.

Take a matmul with a [8192, 32768] input, a [32768, 4096] weight, and a [8192, 4096] output, all fp16. Left unsplit, the input tensor’s memory span would exceed what a single core is allowed to address. So the first pass immediately commits a 2-way split along the reduction dimension. From there, the second pass estimates that a 16-way split along the output dimension is the most beneficial. The planner isn’t just picking a valid division. It’s picking one that minimizes the projected execution time.

Before/after for the worked matmul example: unsplit vs. a 2×4 slice of the full 2×16 core grid, showing how each core’s (K,N) coordinate picks out its A row-tile and B column-tile

The planner prioritizes matrix multiplications because they typically dominate the end-to-end computation time.

Splitting work across cores also splits the data those cores touch. A [1024, 2048] fp16 tensor is 4 MB – too big for any single core’s scratchpad – but split by rows across four cores, each core’s slice is a comfortably-sized 1 MB.

A 4 MB tensor split by rows across four cores so each core’s slice fits in scratchpad

Multi-core scratchpad residency isn’t free, though. It depends on every operation touching that tensor or tile (a buffer) agreeing on how it’s split – and nothing so far has guaranteed that agreement.

Scratchpad planning

The second theme picks up once that division of work is in place. Intermediate tensors too large for a core’s local scratchpad force a choice. Either round-trip them through device memory, or reorganize a sequence of ops to work on tensor tiles, and wrap such sequences with loops iterating over tiles. Whether tiling belongs to PyTorch or specialized programming languages such as Triton and Helion is an open question. While building torch-spyre, though, we observed that torch.scan is a powerful enough construct to express tiled loops, albeit somewhat awkwardly. Built on top of torch.scan, torch-spyre’s own for_each_tile construct makes it easy to express tiled loops, giving the programmer control about the space-time trade-off – iteration count vs. memory footprint – for instance by tiling the reduction dimension of a matrix multiplication.

for_each_tile tiling the reduction dimension of a split-K matmul, so only one K-slice of each operand is ever resident at once

Once a kernel has been tiled and tiles have been assigned to cores, the scratchpad allocator decides what actually stays resident in the core’s scratchpad and at what address. It offers a choice of solvers, from a simple greedy algorithm up to CP-SAT and simulated-annealing options for harder cases.

ScratchpadAllocator’s pass pipeline, with a pluggable choice of layout solver

The payoff is measurable. Torch-spyre’s documentation walks through a single-core softmax over a modestly sized fp16 tensor – the lowered sequence is max → sub → exp → sum → div – and shows what happens as scratchpad-aware scheduling is added one capability at a time: first pinning small reduction outputs on-core, then reusing scratchpad addresses in place for pointwise steps, then cloning the input into scratchpad once instead of re-fetching it from device memory for every reader. End to end, median kernel runtime on this example drops from 32.51 µs to 23.67 µs — a 27% reduction, purely from deciding what stays on-chip and when.

Four stages of scratchpad-aware scheduling for a softmax kernel, with total device-memory traffic per stage in units of the tensor size, from device-memory-only to fully optimized

When the two decisions fight each other

Scratchpad allocation and work division are described above as two separate passes, and torch-spyre’s first implementation ran them that way: work division picked each operation’s core split independently, and scratchpad allocation placed buffers afterward, treating those splits as fixed. That turned out to be limiting, which is why torch-spyre is moving toward the co-optimization approach described in detail below.

The problem shows up whenever two adjacent operations that share a buffer pick different core splits. Work division optimizes each operation on its own terms: a matmul wants the split that keeps its compute engines busy; a neighboring reduction or pointwise op wants whatever split matches its own shape. Nothing forces those choices to agree. When they don’t, the buffer between them is sliced differently by its producer than by its consumer – one core’s scratchpad has a piece of the data, the next op expects a different piece. Torch-spyre calls this a work-division mismatch. Historically the only way to handle it was to disqualify the buffer from scratchpad residency outright: the producer writes it out to device memory, and the consumer reads it back from there, paying for the round trip in both directions – exactly the traffic scratchpad allocation intends to avoid.

Cost of a work-division mismatch: matching splits keep a buffer on scratchpad for free, mismatched splits spill it to device memory twice over

Next, torch-spyre explored a promising alternative: decide optimal work division for each matrix multiplication in isolation but derive work division for intermediate ops from the preceding matrix multiplication, hence the decoupling of the second and third work-division passes. Device memory round trips remained the fallback when mismatches could not be avoided.

That’s no longer the only option. Spyre’s scratchpad hardware supports a core-to-core relayout – a shuffle that moves data directly from one core’s scratchpad slice to another’s, without ever touching device memory. Torch-spyre has begun turning that capability on under the name lx_relayout. A relayout doesn’t make a mismatch free: reshuffling data between cores still costs real time. And because the old layout and the new one have to coexist in scratchpad while the shuffle happens, it costs scratchpad space too. What it changes is the shape of the choice. A mismatch used to mean an all-or-nothing fallback to device memory; now it can instead mean a reasonably priced, on-chip shuffle, and the compiler gets to decide.

A better shuffle is still a shuffle, though. There is another option: moving from independent work division decisions to coordinated decisions accurately accounting for both the benefits and drawbacks of consistent scratchpad allocation across adjacent ops. Torch-spyre’s new co-optimizing allocator folds the standalone work-division pass into scratchpad planning instead of treating its output as fixed. This allocator compares the data shuffle cost resulting from optimal but mismatched work division decisions across adjacent ops to the cost of aligned but suboptimal work division decisions.

In the end, the co-optimizer has many decisions to make and many options to consider; two solver backends do that search differently. The pruned, greedy path narrows the space with hardware-specific heuristics – for a matmul, anchoring on the divisions its own cost-model pass already favors – and then places buffers with best-fit packing, keeping the search fast enough for every graph. An exhaustive constraint-programming path (torch-spyre uses OR-Tools CP-SAT) instead considers the full space of candidate divisions and relayout choices in one pass: which candidate division each operation takes, and whether any given producer-consumer pair pays for a relayout between them, tied together by hard constraints – does a buffer’s chosen division match every consumer’s, directly or via a paid-for relayout; does everything actually fit in scratchpad, including both copies while a relayout is live – with a cost to minimize once those constraints are satisfied. That relayout cost is a plain term built from the same hardware-fitted cost model that prices ordinary device-memory traffic, so the solver is never comparing a real cost against an assumed-free one. The exhaustive path can afford to fully subsume work division precisely because it isn’t pruning; the greedy path still needs the cost-model pass’s heuristics to keep its own search tractable.

The co-optimizing allocator folding work division into scratchpad planning, with a pruned greedy path and an exhaustive constraint-solver path

The payoff shows up concretely on real kernels. On a fused MLP-style kernel, letting a pointwise operation adopt a neighboring matmul’s core split – rather than insisting on its own locally optimal one – lifted core utilization from roughly two-thirds to about 80% and cut kernel time by around 17%. That translates into roughly a 2x improvement over the equivalent sendnn-compiled kernel.

On a multi-head-attention kernel, two matmuls around a softmax contract along different axes, so they can never simply adopt each other’s split outright. Offering both of them – and the softmax reduction sitting between them – a shared batch-and-output-dimension tiling let the whole chain converge on one consistent division. That pinned what would otherwise have been a device-memory round trip for the attention scores matrix.

And on a small synthetic graph built specifically to put work division and scratchpad placement in tension, letting the solver choose a relayout instead of forcing a spill replaced an 80 microsecond device-memory round trip with a 1.1 microsecond on-chip shuffle. The old strict-match gate would have paid 70x more than the problem actually required.

This is still a young part of the compiler. The search space grows quickly with the kernel size. The exhaustive solver considers that whole space. Therefore, we bound its execution time, and continue to explore smart pruning heuristics that accelerate the search for solutions without jeopardizing quality.

Evaluation

The software-engineering case for this move is easiest to see in raw size. sendnn, the compilation stack that this Inductor backend replaced, was built independently of torch-spyre and of Inductor alike. Measured with lizard (non-comment lines of code, tests excluded), it totals roughly 273,000 lines of C++ and Python across 6,170 functions, with an average per-function cyclomatic complexity around 10. The Inductor extension that replaces it, living under torch_spyre/_inductor, comes in at roughly 38,600 lines of Python across 1,878 functions, averaging a cyclomatic complexity of 5.6 per function: about 7x fewer lines, a third as many functions, and simpler ones at that.

That’s not quite an apples-to-apples “we wrote less code to do the same thing from scratch” comparison, since the Inductor extension leans on all of Inductor’s own machinery underneath it in a way sendnn never could. But that asymmetry is the point: a 7x smaller, less complex codebase is what “real leverage from an existing compiler” looks like in practice.

On performance, the fuller story is still being assembled, but two real, measured data points show up at opposite ends of the memory hierarchy. On-core, the softmax example above gives us a 27% drop in median kernel time purely from better use of scratchpad. Off-core, the matmul layout example gives us a 29-69% cut in kernel time purely from how tiling determines whether sticks land contiguously in device memory. Neither number depends on the other. Together they’re evidence that Spyre’s unique architecture poses real demands on the compiler, and that Inductor and torch-spyre can successfully tackle this challenge.

Wrapping up

Torch-spyre’s move to backend=inductor bought real leverage: a compiler roughly 7x smaller than what it replaced, and measurably less complex per function. That’s despite the fact that bridging three specific hardware gaps – stick-granularity data movement, a scratchpad-only memory hierarchy with no shared cache, a static core allocation – took substantial work outside Inductor’s own tree. What made that work tractable at all is that Inductor exposes extension points that let a backend like torch-spyre plug in without forking the compiler underneath it.

That relationship runs in both directions. Torch-spyre’s need to register device-specific decompositions – a concrete requirement that fell out of adapting Inductor to Spyre’s op set – helped drive improvements to out-of-tree Inductor backend support that landed upstream in PyTorch 2.12. We proposed the change as an RFC, and it shipped as two merged PRs: #176811 and #185909. Two design documents describing further extension points that torch-spyre relies on are also public – one proposing a tiled-layout abstraction for stick-granularity memory, one on a new MLIR-based intermediate representation (KTIR) – and we’ve already shared and discussed them with upstream Inductor maintainers: RFC 0047 – tiled tensors, RFC 0682 – KTIR. And the extension point most central to this post, CustomPreSchedulingPasses, is next in line. Turning it from today’s monkey-patch into a supported hook is an active, ongoing conversation with upstream.

We think torch-spyre is a useful proof point that Inductor’s extension-point model generalizes well beyond GPU-shaped hardware. If you’re building an out-of-tree backend for hardware that doesn’t look like a GPU, we’d like to hear from you – and, together with upstream, we’re working to make the path we took easier to follow next time.

The post Extending Inductor for the IBM Spyre Accelerator appeared first on PyTorch.