Operator Lowering and IRs

A PyTorch matmul call carries no information about cache lines, warp occupancy, or register pressure. The CUDA kernel that eventually runs on the GPU does. Something has to bridge that gap, and that something is a lowering pipeline: a chain of intermediate representations (IRs) where each step either optimises at one abstraction level or discards abstraction to expose the level below. Getting this chain right is why torch.compile can give 2-3x throughput gains on the same hardware without changing a single line of user code.

What an IR actually is

An IR is a data structure that represents a computation in a form that is simultaneously easy to analyse and easy to transform. Unlike source code, an IR is designed for machines to read. Unlike binary instructions, it retains enough structure to permit rewrites that would be impossible or illegal at a lower level.

A useful mental model: each IR level answers a different question.

Lowering moves a program from the top row toward the bottom. At each step the compiler can apply passes that are only sound at that abstraction level. Loop tiling is meaningless in a graph IR; algebraic simplification of exp(log(x)) is hard to spot in PTX.

The MLIR dialect stack

MLIR (Multi-Level IR) formalises this idea by making dialect switching explicit. Every operation belongs to a named dialect; a lowering pass converts operations from one dialect into operations in another. The Toy tutorial in the MLIR documentation (chapters 3 and 5) demonstrates this concretely: a high-level toy.transpose is first canonicalised within the Toy dialect, then partially lowered to Affine and MemRef operations for loop-level optimisation, and finally fully lowered to LLVM IR for code generation.

A simplified trace of a matrix multiply through the stack looks like this:

linalg.matmul ins(%A, %B) outs(%C)   # structured op: knows it's a matmul
  -> affine.for loops with affine.load/store  # explicit loops, affine bounds
    -> memref.load/store + arith.mulf         # concrete memory, scalar arith
      -> llvm.load + llvm.fmul                # LLVM, ready for NVPTX backend
        -> PTX fma.rn.f32                     # hardware instruction

Each arrow is one or more MLIR conversion passes. Because the conversion framework enforces legality constraints, you cannot accidentally leave a linalg.matmul dangling in an LLVM IR module - the compiler will error at the conversion stage, not silently miscompile.

The full list of MLIR dialects - including linalg, affine, llvm, and gpu - is maintained in the official MLIR dialect reference.

torch.compile and the FX graph

PyTorch's torch.compile (introduced in 2.0) builds its own two-stage lowering pipeline on top of MLIR and Triton.

1. TorchDynamo traces user Python into an FX graph. This is the graph IR: nodes are ATen operations, edges are tensor dependencies. Dynamo handles control flow by "guarding" on shapes and dtypes and re-tracing if guards break.

2. TorchInductor takes the FX graph and lowers it through its own IR (a loop-level representation called "loops IR") to either Triton kernel source (for GPU) or C++ with OpenMP (for CPU). Inductor decides which operations to fuse - it will fold an add into the epilogue of a preceding mm when memory bandwidth is the bottleneck.

The key insight in Inductor's design is that it generates Triton source code, not PTX directly. This means the Triton compiler then handles the GPU-specific lowering: tiling, vectorisation, shared memory allocation, and register pressure management. The two compilers compose cleanly because Triton's own IR (also built on MLIR dialects) speaks the same language as the hardware abstractions Inductor targets.

Operator fusion as a lowering decision

Fusion - merging two operators into one kernel - is almost always a lowering-time decision, not a graph-time one. At graph IR level you can identify fusion candidates (producer-consumer pairs with no intervening consumers of the intermediate tensor), but you cannot perform the fusion until you know the loop structure. Fusion only makes sense once you are at a level where you can see that two operations iterate over the same index space.

This is why XLA's lowering strategy matters. XLA uses HLO (High Level Operations) as its graph IR and lowers through a series of passes before reaching either LLVM or a hardware-specific library call. Its fusion pass runs on HLO, not on PTX - it inspects shapes and element-wise access patterns at the HLO level, marks fusible clusters, then emits a fused loop nest during the code-generation lowering step. The arxiv.org/abs/2301.13062 analysis of XLA fusion confirms that fusion decisions are made on the HLO graph and that the reinforcement-learning environment they build operates at that same abstraction level.

A fused vs. unfused kernel for relu(layernorm(x)) illustrates the difference:

• Unfused: layernorm kernel writes a full intermediate tensor to HBM; relu kernel reads it back. Two memory round trips.
• Fused: Single kernel computes layernorm values in registers, applies relu in the same thread, and writes only the final result. One memory round trip.

On an A100 where HBM bandwidth is around 2 TB/s but compute throughput is tens of TFLOP/s, the fused version is memory-bandwidth-limited at a much higher throughput.

When it falls down

Shape specialisation breaks generalisation. Dynamo traces per (shape, dtype) combination. A model that receives variable-length sequences - common in production serving - will retrace frequently, and each retrace pays the compilation latency again. The fix is either to bucket input lengths or to use torch.compile(dynamic=True), which generates guards over symbolic sizes. Symbolic shapes make the lowering harder: Inductor can no longer statically tile loops when bounds are unknown at compile time.

Dialect gaps cause silent fallbacks. When an MLIR conversion pass encounters an operation it does not know how to lower, it may fall back to a library call (cuBLAS, cuDNN) or refuse to compile. The silent-fallback case is dangerous: your kernel runs correctly but skips all the fusion and tiling work you expected. Profiling with torch._inductor.config.debug = True or with Nsight Systems is the only reliable way to see whether Inductor actually generated a Triton kernel or silently called into ATen.

Affine analyses assume affine access patterns. The Affine dialect's polyhedral analyses are exact only for loops with affine bounds and subscripts. Gather/scatter operations, indirect indexing, and dynamic slices fall outside this fragment. At that point the compiler either inserts a conservative memory model (blocking fusion) or hands off to a runtime library.

Multi-backend divergence. A lowering pipeline tuned for NVIDIA A100 will produce different tiling parameters, different warp counts, and different shared-memory layouts than one targeting AMD MI300X or Google TPU v5. The abstraction gap between "linalg dialect" and "hardware-optimal kernel" is not architecture-neutral. XLA manages this with backend-specific lowering passes; Triton manages it with a hardware-specific register-file and SIMD width abstraction. The key implication: a kernel generated by torch.compile for CUDA may or may not be near-optimal when you switch hardware vendors.

Quantised operators break standard lowering assumptions. INT8 or FP8 matmuls require type-aware lowering: the accumulation type differs from the storage type, and scaling factors must be injected at specific points in the loop nest. Standard linalg lowering assumes a single arithmetic type. Quantised kernels either require custom dialect extensions or hand-written Triton kernels that bypass the standard lowering pipeline entirely.

Further reading

• MLIR: A Compiler Infrastructure for the End of Moore's Law - Lattner et al., the original MLIR paper covering the dialect system and progressive lowering design.
• MLIR Dialect Reference - authoritative list of all core dialects including linalg, affine, memref, llvm, and gpu.
• MLIR Toy Tutorial Ch. 5: Partial Lowering to Lower-Level Dialects - hands-on walkthrough of multi-stage dialect conversion with concrete before/after IR.
• XLA Architecture Overview - describes the HLO IR, target-independent passes, and backend-specific code generation stages.