Part III — Middle end

Rewrites & canonicalization

Lesson 03 handed you a typed SSA dataflow graph: every value defined once, every tensor carrying its shape, dtype, and device, every op drawn from a small, well-defined set. That representation has rules — which is precisely what makes it safe to rewrite. But a freshly captured graph is a mess: constants you could have folded at compile time, the same subexpression computed twice, nodes whose results nobody reads, and dozens of x + 0 / x · 1 spelled forty different ways. This lesson builds the engine that cleans all of that up — match a subgraph, replace it with an equivalent one, repeat until nothing changes — and the catalog of cheap, always-valid passes it drives. None of these are the big win. They are the setup: every one of them shrinks the graph and normalizes its shape so the high-leverage pass next door — fusion — has clean, matchable patterns to grab.

1 · The rewrite engine: match, replace, iterate to a fixed point

Every middle-end optimization in this lesson is the same shape. A rewrite rule is a pair: a pattern — a small subgraph template, possibly with wildcards (“any value a times a constant 1”) — and a replacement — the equivalent subgraph to splice in (“just a”). The engine walks the IR, finds where a pattern matches, swaps in the replacement, and rewires the SSA edges so every consumer of the old result now reads the new one. Because the IR is single-assignment with value semantics (lesson 03), that rewire is purely local: no aliasing, no hidden mutation, nothing downstream silently changes meaning.

The catch is that one rewrite enables another. Fold a constant and a subexpression two nodes over suddenly becomes a duplicate that CSE can collapse; collapse it and a third node loses its last consumer and becomes dead. So the engine doesn't run each pass once — it iterates the whole rule set until a full sweep changes nothing. That terminal state is a fixed point: applying any rule produces a graph identical to the one you have. f(g) = g. As long as every rule strictly simplifies (fewer nodes, or a canonical form that's never undone), the iteration is guaranteed to terminate — there's a finite amount of graph to shrink.

This is not exotic. It is exactly what MLIR's PatternRewriter and applyPatternsAndFoldGreedily do over a set of RewritePatterns; what XLA's HLO passes do as an algebraic-simplifier fixpoint; and what torch.compile's Inductor does with its declarative pattern matcher (register_replacement) before it lowers to Triton. Different IRs, identical engine: match a subgraph, replace it with an equivalent, iterate.

2 · The staples: constant folding, CSE, DCE

Constant folding evaluates, at compile time, any op whose inputs are all known constants, and replaces the op with its result. The trace baked 1 / √64 as a live computation; folding turns it into the literal 0.125 so it costs nothing at run time. It cascades: fold one node and its constant output may make the next node foldable too.

before                              after constant folding
  %a = const 64                       %a = const 64        (now dead — see DCE)
  %b = sqrt %a        ; = 8.0         %scale = const 0.125
  %s = div 1.0, %b    ; = 0.125
  %y = mul %x, %s                     %y = mul %x, 0.125

Common-subexpression elimination (CSE) finds two nodes that compute the same op on the same inputs and keeps one, redirecting all consumers of the second to the first. Trivial in SSA: hash each node by (op, input-ids, attrs); equal hashes with equal inputs are the same value. The residual block that recomputed x + bias twice now computes it once.

before                              after CSE
  %h1 = add %x, %bias                 %h1 = add %x, %bias
  %h2 = add %x, %bias    ; dup        ; %h2 removed; its uses point to %h1
  %p  = relu %h1                       %p = relu %h1
  %q  = gelu %h2                       %q = gelu %h1

Dead-code elimination (DCE) deletes any node whose result has no consumers and no side effect — the leftover print-branch subgraph, the constant %a = 64 that folding orphaned, the second add that CSE just disconnected. In an SSA DAG this is one reverse sweep: anything not reachable backward from the graph's outputs is dead. DCE is what makes folding and CSE pay off — they disconnect nodes; DCE actually removes them so the kernel count and memory plan shrink.

before                              after DCE
  %a = const 64        ; no uses      ; %a removed
  %dbg = sum %x        ; no uses      ; %dbg removed
  %y = mul %x, 0.125                  %y = mul %x, 0.125
  return %y                           return %y

The three run as a fixpoint loop, not once each, because they feed each other: fold → orphans constants → DCE; CSE → disconnects a node → DCE; DCE → exposes a now-single-use node → in-place fusion later. Run them to a fixed point and the graph contains exactly the live, distinct computation and nothing else.

3 · Algebraic simplification — and the float-associativity trap

Algebraic simplification rewrites a subgraph into a cheaper algebraically-equal one. These are the identity rules every compiler ships:

before                              after algebraic simplify
  %t1 = transpose %w  (0,1)→(1,0)     ; both transposes removed
  %t2 = transpose %t1 (1,0)→(0,1)     %y = mul %x, %w        ; (uses %w directly)
  %r  = div %a, sqrt(%b)              %r = mul %a, rsqrt(%b)
  %z  = mul %r, 1.0                   ; mul-by-1 removed

Every rule above is an exact identity on real numbers — that's the license to apply it. But tensors are floating-point, and floating-point is not the real numbers. The one rule that looks just as innocent and is not safe is reassociation: (a + b) + c → a + (b + c). Over reals these are equal; over IEEE-754 they are not, because each + rounds.

Same three numbers, different grouping, answers that differ by 1.0 — because in the second grouping the 1.0 is swallowed by the rounding of a number 20 orders of magnitude larger. Reassociation can also create or destroy overflow/NaN/Inf. So a correct compiler treats + and · on floats as non-associative and non-distributive by default: it will not reorder a reduction, balance a sum tree, or distribute a·(b+c) → a·b + a·c unless you opt in. That opt-in is fast-math — --ffast-math, XLA's xla_allow_excess_precision / reduced-precision-reassociation, the LLVM fast / reassoc / contract flags, PyTorch's TF32 and reduced-precision-reduction toggles. Fast-math says “I accept that the compiler may reassociate and contract, and my result may differ in the last bits.” It is the explicit boundary of the preserve-semantics invariant: every other rule in this lesson is bit-exact; reassociation trades a few ulps of accuracy for the freedom to reorder. Know which side of that line a rewrite sits on, because a silent reassociation is how a model's loss curve quietly diverges from the reference run.

4 · Canonicalization: one normal form so patterns match

Here is the move that makes everything above worth more than the sum of its parts. Canonicalization rewrites every equivalent spelling of a thing into one chosen canonical form — a single normal form. Commutative ops get their operands sorted into a fixed order (constant on the right: add(2, x) → add(x, 2)); x - c becomes add(x, -c); div-by-constant becomes mul by the reciprocal; subtraction folds into addition; a chain of reshapes collapses to one. The point is not that the canonical form is faster — it's that it is unique.

Why does uniqueness matter so much? Because every later pass is a pattern matcher, and a pattern matcher only fires when the IR looks like its template. Without canonicalization, a fusion rule that wants “matmul followed by add(bias)” would need to also spell out add(bias, matmul), matmul then add-via-sub-of-negative, and every other variant — a combinatorial explosion of near-duplicate patterns, most of which someone forgets to write. Canonicalize first and you write the pattern once, against the one form the IR is guaranteed to be in. Canonicalization doesn't optimize; it makes optimization matchable. In MLIR it's literally a built-in pass (-canonicalize) that ops register getCanonicalizationPatterns into; XLA folds it into the algebraic simplifier; Inductor normalizes before its pattern matcher runs. The discipline: canonicalize → fold → CSE → DCE, iterate, and only then reach for the passes that actually move bytes.

5 · Phase ordering, and the e-graph escape

There's a worm in the greedy fixpoint. Passes interact, and the order you run them in changes the result — this is the phase-ordering problem. A rewrite that looks like a win locally can block a much bigger win two passes later. Classic case: greedily fold a·b + a·c down by some local rule and you may destroy the a·(b+c) factoring that fusion wanted; canonicalize a transpose away and you might forfeit a layout that a later pass could have exploited. Greedy rewriting commits to the first improvement it sees and can never take a step that's temporarily worse to reach a better global form. No fixed pass order is optimal for all graphs, and finding the best sequence is, in general, intractable.

The principled escape is equality saturation over an e-graph (equality graph). Instead of destructively replacing the matched subgraph — throwing away the original — you add the equivalent form and record that the two are equal. An e-graph compactly represents all the equivalent programs you've discovered at once: nodes are grouped into equivalence classes (“e-classes”), and a rewrite just merges classes rather than mutating the graph. You apply every rule everywhere, repeatedly, until the e-graph saturates (no rule adds anything new — a fixed point over the whole space of rewrites), and only then run a cost-model-driven extraction that picks the single cheapest program out of the exponentially many the e-graph encodes. Because nothing was ever destroyed, you sidestep phase ordering entirely: the engine never had to choose an order, it kept everything.

This is the idea behind TASO and PET (superoptimizing tensor graphs by exploring equivalent subgraphs) and the egg library (fast, general e-graphs and equality saturation) now used in several ML and DSL compilers. The cost is real — the e-graph can blow up in memory, and extraction with a realistic cost model (one that accounts for fusion, layout, and traffic, not just node count) is itself hard. So today's production stacks are mostly greedy-fixpoint with carefully tuned phase orders, and equality saturation is the frontier you reach for when the greedy order leaves measurable speed on the table. Either way, the engine's job is done: a clean, canonical, redundancy-free graph — handed to the pass that actually collapses the bytes.

6 · Drive it: the rewrite playground

The graph below is a deliberately messy little program with one of each problem: a constant expression (1/√64), a repeated subexpression (x + bias twice), a dead node (an unused debug sum), and a no-op mul by 1.0. Apply the passes in any order and watch #nodes, #ops (the kernels you'd launch), and est. cost shrink toward a fixed point — when every button is disabled, the graph is saturated. Then flip “reassociate floats?”: it offers an extra reduction-reordering rewrite, and warns you that it may change the result in the last bits. That toggle is the preserve-semantics line drawn as a switch.

Notice what the scoreboard cannot do: it bottoms out with one live op per surviving node. Constant fold deleted compute, CSE deleted a duplicate, DCE deleted the dead, algebraic deleted the no-ops — but matmul, add(bias), and gelu all remain as separate ops, and in eager execution that's three kernels and three HBM round trips. Cheap rewrites shrink the graph; they don't merge the survivors. That's the wall this lesson hands off.

Failure modes & checklist

Checkpoint

Where this points next

The engine did its job: a typed SSA graph is now canonical, constant-folded, free of duplicates and dead nodes, and algebraically tightened — and crucially, it's in the one normal form that later pattern matchers expect. But look at what survives. Every cleanup in this lesson deleted work; none of them merged work. The graph still runs as one kernel per op, and from lesson 01 we know that's the expensive part: each surviving memory-bound op pays its own HBM round trip, loading and storing the whole activation just to do a few FLOPs per element. Constant folding and CSE removed nodes that never needed a kernel; they did nothing about the kernels that remain. The single highest-leverage rewrite left is categorically different — instead of removing a node, fuse a chain of them into one kernel that reads from HBM once, does all the math while the data is hot in registers, and writes once. Everything in this lesson was setup: a clean, canonical graph is exactly what makes fusion patterns findable. That is lesson 05, Operator fusion — the central win.

Interview prompts

• What is a rewrite rule and when does the engine stop? (§1 — a pattern (subgraph template + wildcards) paired with a provably-equivalent replacement; the engine matches, splices in the replacement, rewires SSA uses, and iterates the whole rule set in a worklist loop until a fixed point — a full sweep that changes nothing — guaranteed to terminate if every rule strictly simplifies.)
• Why is CSE almost free in an SSA/value-semantics IR? (§2 — no aliasing or hidden mutation, so two nodes with the same op, inputs, and attrs are provably the same value; hash by (op, input-ids, attrs), redirect all consumers of the duplicate to the original, then let DCE remove it.)
• Give a concrete case where reassociating floats changes the answer, and how compilers gate it. (§3 — (1e20 + −1e20) + 1 = 1 vs 1e20 + (−1e20 + 1) = 0: the +1 is swallowed by rounding; floats are non-associative, so reassociation/distribution is off by default and only enabled under fast-math flags, which trade ulps of accuracy for reorder freedom.)
• What does canonicalization buy you if it doesn't make the code faster? (§4 — uniqueness: it maps every equivalent spelling to one normal form, so every downstream pattern (fusion, layout) is written once against the guaranteed form instead of as a combinatorial set of near-duplicate patterns; it makes optimization matchable.)
• What is the phase-ordering problem and how do e-graphs sidestep it? (§5 — passes interact and no fixed order is globally optimal because greedy rewriting destroys the matched form and can't take a temporarily-worse step; equality saturation over an e-graph adds equivalent forms instead of replacing, saturates the whole rewrite space, then extracts the cheapest by cost model — order never has to be chosen.)
• When is it unsafe to constant-fold, and when is DCE wrong? (§2, Failure modes — don't fold an op whose input only looks constant (dynamic/shaped value, externally mutated buffer); don't DCE a node that has no SSA consumer but does have a side effect — RNG advance, a collective, an in-place write — or you silently break correctness.)
• After all these cleanups, why is the graph still slow? (Where this points next — every pass here only deleted nodes; it never merged the survivors, so the graph still runs one kernel per op and each memory-bound op pays its own HBM round trip. The high-leverage rewrite is fusion, lesson 05.)