The IR ladder, and the MLIR fork in the road

Because $N$ operators times $M$ backends is a maintenance disaster, and because large operators hide the structure an optimiser needs.

Gained. A new backend implements ~70 Prims (or ~250 core ATen) rather than 2000 ATen ops with their overloads and in-place variants — the difference between a feasible project and an impossible one. Optimisation also gets easier: a fused gelu is opaque, while erf, multiplies and adds are things a fusion scheduler can reason about, reorder and combine with neighbours. And precise Prims semantics remove the ambiguity that accumulated across two thousand hand-written operators.

Lost, and it is real. Decomposing destroys the knowledge that this cluster of primitives was a gelu or a softmax, and hand-written kernels for those patterns are often far better than anything a scheduler reconstructs. Worse, a decomposition can be numerically different: computing softmax naively from its primitive parts loses the max-subtraction trick unless the decomposition encodes it.

So the ladder is not uniformly downward. Inductor keeps pattern matching to recognise decomposed structures and swap in a library kernel — recovering scaled_dot_product_attention and routing it to FlashAttention rather than emitting the quadratic materialised form. And some operators are deliberately not decomposed, staying as opaque calls into cuBLAS or cuDNN, because nothing generated will beat them.

The general principle: decompose to expose structure to the optimiser, then pattern-match to recover the cases where a human already solved it better. Both directions are necessary.

It removes mutation and aliasing, which are precisely the two things that make reordering unsound.

With mutation, the meaning of a graph depends on execution order in a way the dataflow edges do not express:

y = x.view(-1)     # y aliases x
x.add_(1)          # mutates both
z = y.sum()        # depends on the add — but no dataflow edge says so

A scheduler looking at dataflow sees z depending on y depending on x, and is free to move the add_ after the sum. That is wrong, and to avoid it the compiler would need full alias analysis — which is expensive, conservative, and in the presence of arbitrary views not reliably decidable.

Functionalization rewrites this into pure form: x_1 = x_0 + 1, and y becomes a view of the appropriate version. Now every dependency is a dataflow edge, so anything that respects dataflow is legal. Fusion, reordering, dead-code elimination, buffer reuse and recomputation all become straightforward local decisions.

Two practical notes. Functionalization can appear to increase memory, since in-place became out-of-place — but the buffer reuse pass afterwards recovers it, now safely, because it has exact lifetimes rather than a guess. And mutations visible to the caller (an updated parameter, a modified input) must be reapplied at the graph boundary, which is why the compiled artifact has an epilogue copying results back.

This is the same argument SSA form makes in a classical compiler, and for the same reason.

Because the interesting optimisation lives between them, and separate compilation cannot see it.

The decision is which forward intermediates to save for backward and which to recompute. Saving costs memory and a write plus a read; recomputing costs FLOPs. Given the joint graph, that is a min-cut problem: cut the graph so the total size of tensors crossing the boundary is minimised subject to the recompute cost. AOTAutograd’s partitioner solves exactly this.

What that yields:

• Activation checkpointing, automatically. The thing people hand-annotate with torch.utils.checkpoint falls out of the partition. Recomputing a cheap elementwise chain instead of storing it is almost always right, and a human rarely annotates at that granularity.
• Fusion across the boundary. A backward op can fuse with a recomputed forward op, which is impossible if the two graphs were compiled independently.
• Correct handling of decompositions in backward. Backward for a decomposed op is derived from the decomposition, so you do not need a hand-written backward for every operator.

The cost is that you must trace backward ahead of time, which requires knowing the graph is fixed — so anything that breaks capture (Module 4) or varies (Module 10) affects backward too. And there is a subtlety worth knowing: the partitioner optimises a proxy for memory and runtime, so its choices are occasionally worse than a knowledgeable human’s on unusual models, which is why manual checkpointing still exists.

Take it when your target is not a GPU with a mature CUDA backend.

Concretely: TPU-like accelerators, NPUs in phones and laptops, and custom silicon where the vendor has an MLIR-based compiler and no realistic path to supporting Triton. That is why IREE uses torch-mlir as its PyTorch frontend and why vendors maintain forks. It is also the right path when you want an ahead-of-time, serialisable artifact — StableHLO is versioned with compatibility guarantees, so you can compile now and deploy later, which Inductor’s Triton output is not designed for.

What you give up:

• Dynamism. The MLIR/StableHLO route is fundamentally an export-then-compile flow, closer to TF1’s model. torch.export requires a whole graph, so Module 4’s graph breaks become hard errors and Part M’s dynamism becomes a much bigger problem.
• The ecosystem’s default path. Inductor gets the tuning, the bug reports and the kernel coverage, because that is where the users are.
• Triton’s escape hatch. On the default path you can hand-write a Triton kernel and drop it in. On the MLIR path you are writing dialect passes.

The honest summary: these are not competitors so much as different destinations. Inductor optimises for the GPU everyone has, with Python-level dynamism intact. The MLIR route optimises for heterogeneous hardware and stable deployable artifacts, and it inherits both the strengths and the rigidity of the Google compiler lineage — XLA’s fusion quality and StableHLO’s stability, along with XLA’s longstanding difficulty with shapes that are not known in advance.