Part II — Front end

The intermediate representation

Lesson 02 handed you a captured graph — Dynamo walked the bytecode and emitted an FX graph of the tensor ops. But "a graph" sells it short: what you actually have is a pile of framework objects, each holding references to the next, with shapes baked in by tracing and mutation hiding in views and in-place ops. You cannot safely rewrite that. To fold a constant, merge two ops, or cancel a pair of transposes — and to do it a thousand times in a row without breaking the math — you need a representation with rules: what a value is, what a type carries, and what makes a rewrite legal. This lesson builds that representation, and then makes the one move that defines an ML compiler: there is not one IR but a ladder of them.

The previous step left this broken

Capture gave you a graph, but a graph of what? In eager PyTorch a tensor is a mutable object: x.add_(1) rewrites it in place, a view aliases another tensor's storage, and two names can point at the same bytes. If your IR inherits that, every rewrite becomes a whole-program alias analysis — "is it safe to fold this constant?" depends on whether some far-away op mutates a tensor this one aliases. That does not compose: you cannot apply pass after pass to a fixed point when each pass must re-prove global safety. And even with clean semantics, a single flat list of framework ops (conv, softmax, layer_norm) is the wrong altitude for half the optimizations you want — it is too coarse to express loop tiling and too abstract to talk about vector registers. You need both a sound representation and the right level of it.

Linear position

Forced by: safe, composable, repeated transformation needs a principled IR, not a pile of mutable framework objects.
New idea: a good ML IR is a typed dataflow DAG in SSA / value-semantics form (each value defined once, no hidden mutation → every rewrite is a local, provable no-op) with a small op set; and crucially there are multiple levels of IR — high (framework ops) → mid (linalg / loop nests) → low (vector/scalar) — because no single abstraction can host every optimization. Lowering = descending one level, spending portability to buy concrete control. This is the abstraction ↔ specialization motif, stated here for the whole track.
Forces next: an IR with typed values and legality rules is an invitation — something has to actually match and rewrite subgraphs to a fixed point. That engine is lesson 04.

The plan

Five moves. (1) Why value semantics / SSA make rewriting local and sound — aliasing and mutation are the enemy. (2) The IR as a typed dataflow DAG, with shape + dtype + device living in the type. (3) The multi-level idea: MLIR dialects and XLA's HLO vs LLVM as progressive-lowering ladders, each rung able to express what the others cannot. (4) A worked example — the same softmax shown at high, mid, and low IR. (5) The abstraction ↔ specialization trade named as the central motif. Then drive the lowering ladder until you watch portability fall and optimizations open up.

1 · Value semantics and SSA: why rewriting must be local

The single property that makes an IR rewritable is value semantics: an operation takes input values and produces new output values, and nothing is ever mutated in place. A value, once defined, never changes. Pair that with SSA — static single assignment, the rule that every value is assigned exactly once and named uniquely — and you get the property compilers have leaned on since the 1980s: the definition of any value is unambiguous, so the uses of a value form an explicit, complete list, and a transform that touches one value cannot silently change another.

Contrast the two worlds with one line. In mutable eager land:

x = randn(1024)
x.add_(1.0)         # mutates x in place — the old x is gone
y = x[0:512]        # y aliases x's storage; writing y writes x
relu_(y)            # mutates y → mutates x → mutates anything aliasing it

To decide whether relu_(y) is safe to move, fold, or fuse, you must prove what y aliases and whether anything reads the old x afterward — a global question. Now the SSA / value-semantics form of the same computation:

%0 = randn() : tensor<1024xf32>
%1 = add(%0, 1.0) : tensor<1024xf32>     ; new value, %0 untouched
%2 = slice(%1, 0, 512) : tensor<512xf32> ; %2 reads %1; no aliasing
%3 = relu(%2) : tensor<512xf32>          ; new value

Here every edge is a pure data dependency: %3 depends on %2 depends on %1. There is no hidden channel. So "is it legal to replace relu(%2) with something cheaper?" is answered by looking only at %2's definition and %3's uses — a local question. That locality is what lets you stack hundreds of passes and iterate each to a fixed point: every rewrite is a provable no-op on the math, so composing them stays correct. This is the preserve-semantics motif made mechanical. The price is that PyTorch programs are full of mutation and views, so the compiler must first run functionalization — rewrite add_ / view / in-place ops into pure value-producing equivalents — before the IR is in this form. We return to functionalization when it pays off for autodiff (FILE 12).

2 · A typed dataflow DAG: shape, dtype, device in the type

With value semantics in hand, the IR is exactly a directed acyclic graph of typed values: nodes are operations, edges are values flowing producer → consumer, and there are no cycles (loops, when they exist, are explicit region ops, not back-edges). The thing that makes it an ML IR rather than a generic SSA IR is the type. For tensor code the type carries far more than "float": it carries shape, dtype, and often device / layout.

%w  : tensor<4096x4096xbf16, @cuda:0>     ; shape · dtype · device all in the type
%x  : tensor<8x4096x4096xbf16, @cuda:0>
%mm = matmul(%x, %w) : tensor<8x4096x4096xbf16>   ; shape inferred & checked

Putting shape and dtype in the type buys you three things. First, verification: a rewrite that produces a shape mismatch is rejected at the IR level, not at runtime. Second, cost: every later pass — fusion (FILE 05), memory planning (FILE 06), layout (FILE 07) — needs to know how many bytes a value is, and the type is where that lives; a 268 MB activation and a 4-byte scalar are both "tensors" but the compiler must tell them apart to price a round trip. Third, specialization: knowing a dimension is exactly 4,096 (not "some N") unlocks tiling and vectorization that a dynamic shape would block — which is also why dynamic dimensions are the headache of FILE 11, where the type carries a symbol s0 instead of a constant. A good ML IR also keeps the op set small and well-defined: a few dozen primitive ops with crisp semantics, rather than the thousands of overloads a framework's public API exposes. A small op set means a rewrite pass has a small surface to reason about and a small number of cases to handle correctly.

SSA / value semantics

Define once, never mutate. Makes "is this rewrite legal?" a local question, so passes compose to a fixed point.

typed dataflow DAG

Nodes = ops, edges = values. Shape · dtype · device live in the type → verifiable, priceable, specializable.

small op set

A few dozen primitives with crisp semantics beats thousands of framework overloads — less surface for a pass to break.

3 · One IR is not enough: the multi-level ladder

Here is the move that separates an ML compiler from a generic one. The high-level graph — conv, softmax, layer_norm as single ops — is the perfect altitude for algebraic optimizations: cancel adjacent transposes, fuse a bias into a matmul, recognize an attention pattern. But it is useless for loop optimizations: a single softmax op has no loops to tile, no memory accesses to reorder, no vector width to choose. Drop too low, though — to scalar instructions and registers — and the softmax has dissolved into hundreds of nameless adds and you can no longer recognize it as softmax to fuse it at all. No single abstraction hosts every optimization. The answer is a ladder of IRs, and a process called lowering that walks down it one rung at a time, each rung re-expressing the program at a lower level where a different class of optimization becomes possible.

Two real ladders dominate practice. MLIR — the compiler infrastructure under Triton, IREE, Torch-MLIR, and modern XLA — makes the ladder explicit through dialects. A dialect is a self-contained namespace of ops and types at one level of abstraction; lowering is rewriting from a higher dialect to a lower one, and the rewrite that crosses a level is called legalization (turning ops legal in dialect A into ops legal in dialect B). A typical descent:

high · frameworktosa / mhlo / stablehlo — whole-tensor ops (conv, softmax, dot). Portable across any backend; the level for algebraic rewrites and fusion.

mid · structuredlinalg — ops as named loop nests over tensors (a generic reduce, a generic matmul). The level for tiling, fusion-by-loops, and bufferization.

mid-low · loopsaffine / scf / vector — explicit loops with affine index maps and SIMD vector types. The level for loop reordering, unrolling, vectorization.

low · targetllvm / nvvm / gpu — scalar/vector instructions, registers, address spaces. Hands off to LLVM for register allocation and final codegen.

XLA draws the same ladder with fewer, fatter rungs: it represents the whole program in HLO (High-Level Optimizer IR) — a small, functional, value-semantics op set of a hundred-odd tensor primitives where almost all of its optimizations live (fusion, layout, GSPMD sharding) — and then lowers HLO to LLVM IR for the CPU/GPU backend to emit PTX. linalg is the linear-algebra dialect named above: a small set of ops, each defined as a loop nest with explicit indexing, that is deliberately the sweet spot for tiling and fusion. The names differ across stacks, but the shape is identical: a high portable rung, one or more structured loop-level rungs, a low target rung, and legalization passes carrying the program down between them.

4 · The same softmax at three levels

Make the ladder concrete with one op. At the high level, softmax over the last axis is a single op — maximally portable (every backend knows what softmax means), maximally compact, and the level at which a compiler can pattern-match "this is the softmax inside attention" and fuse it:

// HIGH IR (framework / HLO) — one op, fully portable
%p = softmax(%s, axis=-1) : tensor<8x4096xf32>

Lower one rung to the mid level and softmax decomposes into its definition — the numerically-stable five-step reduction — expressed as explicit loop nests over the tensor. Now the loops are visible, so a compiler can tile them, fuse the elementwise sub/exp into the reductions, and decide what stays in on-chip SRAM:

// MID IR (linalg / loop nests) — the math, as loops; tileable & fusible
%mx  = reduce_max(%s, axis=-1)              // for r: mx[r] = max_c s[r,c]
%sub = generic sub(%s, broadcast %mx)       // for r,c: t[r,c] = s[r,c] - mx[r]
%ex  = generic exp(%sub)                     // for r,c: e[r,c] = exp(t[r,c])
%sm  = reduce_add(%ex, axis=-1)             // for r: sm[r] = sum_c e[r,c]
%p   = generic div(%ex, broadcast %sm)      // for r,c: p[r,c] = e[r,c] / sm[r]

Lower again to the low level and the loops become a concrete schedule on this chip: a vector width, an unroll factor, register-resident accumulators. This rung can saturate the SIMD/SIMT lanes — but it is now written against this target's vector ISA and is no longer portable:

// LOW IR (vector / scalar, target-specific) — concrete, fast, locked to one chip
loop r in 0..4096:
  vmx = vmax.f32x8(load_row(%s, r))          // 8-wide vectorized row max
  vsm = vsplat(0.0)
  for c in 0..4096 step 8:
    ve  = vexp.f32x8(vsub(load8(%s,r,c), vbroadcast(vmx)))
    vsm = vadd.f32x8(vsm, ve)                 // register-resident accumulator
    store8(%tmp, r, c, ve)
  inv = vrcp(hsum(vsm))                       // reciprocal of the row sum
  for c in 0..4096 step 8:
    store8(%p, r, c, vmul(load8(%tmp,r,c), vbroadcast(inv)))

Read the three dumps top to bottom and the trade is visible in the line count alone: one op becomes five becomes a fully unrolled vector loop. Descending, you lose portability (the bottom rung names f32x8 vector ops that only exist on a particular target) and gain the ability to express concrete optimizations (vectorization, register accumulation) that the top rung literally has no vocabulary for. That is the entire point of the ladder.

5 · The motif: abstraction ↔ specialization

Step back and name what just happened, because it governs the rest of the track. A high IR is abstract: portable across backends, compact, easy to pattern-match and rewrite algebraically — but it cannot say anything about the memory hierarchy or the instruction set, so it cannot be made fast on its own. A low IR is specialized: it talks in tiles, vectors, and registers of one specific chip, so it can be made very fast — but it is locked to that target and too detailed to recognize high-level structure. Lowering is the act of spending portability to buy concrete optimization.

The abstraction ↔ specialization trade (the track's spine)

High / abstract: portable, compact, algebraically rewritable, hardware-blind — slow if run as-is. Low / specialized: tiled, vectorized, register-aware, hardware-optimal — locked to one target. Every pass in this compiler lives on this axis: descend only as far as the optimization you want requires, and no further, so the higher rungs stay available for the rewrites that need them. Each later lesson — fusion (05), layout (07), scheduling (08), codegen (09) — is one more step down this ladder, each forced by an optimization the rung above could not express.

This is why the ladder is progressive and not a single jump from Python to PTX. If you lowered straight to scalar instructions you would have a fast representation of one program on one chip and no ability to fuse, retile, or retarget. By keeping a high rung for the algebraic wins, a mid rung for the loop wins, and a low rung for the instruction wins, the compiler gets each class of optimization at the level where it is natural — and that is precisely the descent the lowering-ladder widget below lets you walk by hand.

6 · Drive it: the lowering ladder

The widget lowers softmax one rung at a time. At each level it shows the IR text and three readouts: the op count (how many operations the program has expanded into), whether the ops are hardware-specific (your portability), and which optimizations are now expressible at this rung. Step down and watch the trade: op count climbs, portability flips from "any backend" to "one chip," and the set of available optimizations grows from "algebraic only" to "loop + vector." There is no free rung — every step buys an optimization with a loss of portability.

The lowering ladder — softmax, high → low

Step the program down the IR levels. The knob that proves the trade: descending raises the op count and opens new optimizations, but flips the IR from portable to hardware-locked. There is no rung that is both portable and register-optimal — that is why a compiler keeps all four.

IR level

—

op count

—

portability

—

new optimizations

—

—

descend → +optimizations, −portability

The shape of the four readouts is the lesson. Op count and "optimizations now possible" rise together as portability falls — and crucially the optimizations available at a low rung (vectorization, register accumulation) are invisible at the high rung, while the high rung's win (recognize-and-fuse softmax) is gone once you lower. A compiler keeps every rung precisely so it can take each optimization at the level where it exists.

Failure modes & checklist

Failure modes

Rewriting an IR with hidden mutation. Folding or fusing on an IR that still has in-place ops / aliasing views. Signal: a "semantics-preserving" pass changes results, because some far-away op was mutating a tensor this one aliased — the missing functionalization step.
Putting too much in one IR level. Trying to do loop tiling on a high op-graph, or pattern-match softmax on scalar instructions. Signal: the pass has no handle to grab — no loops to tile up high, no recognizable op down low.
Lowering too early. Dropping to vector/scalar before fusion runs. Signal: softmax has dissolved into anonymous adds, so the fuser can no longer recognize it and the high-level win is permanently lost.
Shape/dtype not in the type. An untyped or "tensor of unknown shape" IR. Signal: later passes can't price bytes or verify rewrites; mismatches surface as runtime crashes instead of IR verification errors.
A sprawling op set. Mirroring every framework overload as a distinct IR op. Signal: every rewrite pass has hundreds of cases and quietly misses some, so optimizations fire inconsistently.

Checklist

Functionalize first. Get to value semantics / SSA before any rewrite, so legality is a local question.
Type carries shape · dtype · device. Enough in the type to verify, to price bytes, and to specialize on constant dims.
Keep the op set small. A few dozen crisp primitives, not thousands of overloads.
Pick the rung for the optimization. Algebraic/fusion up high, tiling in the middle, vectorization down low.
Descend as little as needed. Stay high enough that the next pass still has structure to work with.
Legalize, don't leap. Lower one dialect at a time so each level's passes get to run.

Checkpoint

Try it

Open the widget and step from the high rung to the mid rung. Note two things: the op count jumps from 1 to 5 (softmax's max/sub/exp/sum/div), and a new optimization appears — loop tiling / fusion — that simply did not exist at the high rung. Now step to low and watch portability flip to "one target (f32x8)." Question to answer for yourself: at which rung would you run a pass that cancels two adjacent transposes, and at which would you choose a SIMD vector width? (Transpose cancellation is an algebraic rewrite — high rung; vector width is an instruction choice — low rung. The fact that those two passes want different rungs is why the ladder exists.) Finally, predict what is lost if you lower straight from high to low, skipping mid: you lose the rung where loop tiling and op-level fusion are expressible, so you never get them.

Where this points next

You now have a representation worth optimizing: a typed dataflow DAG in SSA / value-semantics form, sitting on a ladder of levels you can lower through. But a representation with legality rules is inert — it describes which rewrites are allowed, not which ones happen. Something has to walk the graph, recognize "this subgraph matches a pattern I can improve," replace it with an equivalent, and iterate until nothing more fires. The raw captured graph is full of exactly such opportunities — constants waiting to be folded, repeated subexpressions, dead nodes, x·1 and back-to-back transposes — and they block the bigger wins until they are cleaned up. That machinery is the pass / pattern-rewrite engine, and the catalog of always-valid cleanups it runs, in 04 · Rewrites & canonicalization.

Takeaway

A captured graph is a pile of mutable framework objects you cannot safely rewrite. A good ML IR fixes that with value semantics + SSA (define once, never mutate, so every rewrite's legality is a local question and passes compose to a fixed point — preserve-semantics made mechanical), a typed dataflow DAG with shape · dtype · device in the type (so passes can verify, price bytes, and specialize), and a small op set. The defining ML move is that there is not one IR but a ladder of them — high framework ops → mid linalg/loop nests → low vector/scalar — because no single abstraction hosts every optimization; the same softmax is one op up high, a five-step reduction in the middle, and a vectorized register loop down low. Lowering (in MLIR, rewriting between dialects via legalization; in XLA, HLO → LLVM) walks down that ladder, and each step is the abstraction ↔ specialization trade — spending portability to buy concrete optimization. Descend only as far as the next optimization requires, no further. An IR with rules invites the engine that applies them: the pass / pattern-rewrite machinery of lesson 04.

Interview prompts

Why does SSA / value semantics make graph rewriting tractable? (§1 — define-once + no mutation means a value's uses are an explicit complete list and a transform can't change another value through a hidden alias, so "is this rewrite legal?" is a local question; that locality is what lets you stack passes to a fixed point without re-proving global safety.)
What does a tensor type carry in an ML IR, and why does each piece matter? (§2 — shape, dtype, device/layout: shape+dtype enable verification and let every later pass price bytes; constant dims unlock tiling/vectorization (specialization); device/layout drive placement and layout passes — a dynamic dim becomes the symbol that breaks specialization in FILE 11.)
Why have multiple levels of IR instead of one? (§3 — no single abstraction hosts every optimization: a high op-graph has no loops to tile, low scalar code has no recognizable softmax to fuse; you keep a high rung for algebraic/fusion wins, a mid rung for loop wins, a low rung for instruction wins, and lower between them.)
Define lowering, dialect, and legalization in the MLIR sense. (§3 — a dialect is a namespaced op/type set at one abstraction level; lowering is rewriting from a higher dialect to a lower one; legalization is the pass that turns ops legal in one dialect into ops legal in the next, e.g. tosa/mhlo → linalg → vector → llvm.)
What are HLO and linalg, and where do they sit? (§3 — HLO is XLA's high-level functional value-semantics tensor IR where most of its optimizations live before lowering to LLVM IR; linalg is MLIR's mid-level linear-algebra dialect expressing ops as loop nests, the sweet spot for tiling and fusion.)
State the abstraction ↔ specialization trade and one consequence. (§5 — high IR is portable/compact/algebraically-rewritable but hardware-blind; low IR is tiled/vectorized/register-aware but locked to one chip; lowering spends portability for concrete optimization, so you descend only as far as the optimization needs and keep higher rungs for rewrites that require structure.)
Why must PyTorch be functionalized before these passes run? (§1 — eager PyTorch has in-place ops and aliasing views, so the captured graph isn't in value-semantics form; functionalization rewrites mutation/views into pure value-producing ops so legality stays local — the SSA property the rewrites depend on.)