Week 1 · Lesson 3

Rewrites and constant folding

A compiler pass is mostly a list of small rules like add(x, 0) → x, applied over and over until none of them fire. This lesson shows how that loop works, one rule at a time.

A rule is a pattern and a replacement

A rewrite rule has two halves. The pattern describes a shape of graph to look for. The replacement says what to build instead. We'll write rules like this:

add(x, 0)          → x                  # adding zero does nothing (i32)
mul(x, 1)          → x                  # multiplying by one does nothing
add(c1, c2)        → const(c1 + c2)     # both inputs known: compute it now
mul(c, x)          → mul(x, c)          # move the constant to the right

x and y match any node. Names starting with c (c, c1, c2) match only const nodes. add(x, 0) is shorthand for the .uir node add %x 0, and in the OCaml exercise below the same pattern is written Add (x, Const 0). A rule fires on a node when the pattern matches it, and then the node is replaced by whatever the right side builds.

The engine itself is tiny, and you'll write one this week:

Sources first. Visit nodes in topological order, sources before their users. By the time you look at a node, its sources have already been simplified.
Greedy. At each node, try the rules. If one fires, try them all again on the result, until nothing matches.
To a fixpoint. Keep running whole passes until a pass changes nothing. That graph is your answer.

Watch it, one rule per step

This is the program being simplified. Use the buttons or the arrow keys. A bold outline means a pattern matched. Green means the rule just built that node. Faded nodes are no longer used by anything.

%x   = param 0 dtype=i32
%a   = add 2 3
%b   = add 4 %a
%c   = mul %b %x
%out = add %c 0

Notice the sources-first order doing its job. add 4 %a can't fold until %a has become the constant 5. Because sources are visited first, that has already happened by the time the engine reaches %b. Visiting results before their sources, you'd need an extra pass for every level of the graph.

Be the simplifier

The starter has five ops that compute ((x + 3) * 1 + 0) + 2 * 5. Write a %out with the same value for every x in at most 1 op.

Is (x + 3) + 10 really x + 13 when x + 3 wraps? For i32, yes. Lesson 4 says why.

Constant folding

Constant folding is the rule family "every input is a const, so compute the answer now." For each arithmetic or logic op (the spec calls these ALU ops) it's one rule, and it evaluates the op with exactly the same semantics as your interpreter: i32 wraps around, f32 rounds to nearest-even, division truncates toward zero.

That last part matters. If your folder computes mul 1073741824 4 with 64-bit integers, it gets 4294967296, which isn't an i32. The interpreter would give 0. A simplifier that changes the answer is a bug, so share one implementation of each op between folding and interpreting.

Don't fold errors into values

div 1 0 has two constant inputs, so the folding rule matches it. But there's no i32 that is "one divided by zero". It's a runtime error, and a program whose asserts depend on it must exit 2. If your folder catches the error and returns something like const 0, the program silently "works". The right move is to not fire: leave div 1 0 in the graph, and let evaluation report it.

Canonicalization: one way to write it

add 3 %x and add %x 3 mean the same thing. Without help, every rule would need two versions, one for each side. Instead we add one canonicalization rule for each commutative op (add, mul, max, xor):

op(c, x) → op(x, c)      when c is a const and x is not

max is only commutative because the spec pins down its one ambiguous case: max(-0.0, 0.0) is +0.0 in either order. A naive if a > b then a else b returns whichever comes second, and recip can see the difference.

Now every other rule can assume constants sit on the right. Canonical forms also help hash consing, because two equivalent expressions are more likely to become the same node.

Why hash consing makes this cheap

In a hash-consed graph, building a node that already exists hands you back the existing one. Two statements with the same (op, src, arg), where dtype and shape are part of the key, are a single UOp. Three things follow:

Shared work is rewritten once. If add %x 1 appears ten times, it's one node, simplified once.
"Did anything change?" is an identity check. If the rule returns the very same object, nothing happened. No deep comparison of trees.
Rules like xor(x, x) → 0 are a pointer comparison. "Both sources are the same node" is just "the two sources are the same node" (OCaml's physical equality, ==).

Above: %a = add %x 1 and %b = add %x 1 collapse into one add, and mul %a %b points at it twice. Simplifying %out here leaves exactly 2 ops.

Termination: every rule must make progress

"Repeat until nothing changes" only stops if each rule moves the graph toward something smaller or more canonical. A rule like

add(x, y) → add(y, x)        # always swap

matches its own output, so the engine swaps forever. The canonicalization rule above is safe because it only fires when the const is on the left, and its output has the const on the right, where the rule no longer matches. A good habit: for each rule, name what it shrinks (number of ops, number of non-const sources, constants on the left).

Code it: a toy rewriter

Before rewriting UOps in your own compiler, write the loop on a toy expression type. Same engine: sources first, greedy, to a fixpoint.

A toy rewriter

Write simplify e. Simplify both children first, then apply these rules to the node, and keep re-applying them to the result until none fires. Integers are i32: fold with 32-bit wraparound.

Add (Const a, Const b)        → Const (a + b)
Mul (Const a, Const b)        → Const (a * b)
Add (Const c, x)              → Add (x, Const c)     (x not a Const)
Mul (Const c, x)              → Mul (x, Const c)     (x not a Const)
Add (x, Const 0)              → x
Mul (x, Const 1)              → x
Mul (x, Const 0)              → Const 0
Add (Add (x, Const a), Const b) → Add (x, Const (a + b))
Mul (Mul (x, Const a), Const b) → Mul (x, Const (a * b))

Leave everything else alone. Add (Const 1, Add (Const 2, Var "y")) becomes (y + 3), which takes three rules in a row.

Key idea

A simplifier is a small set of local, obviously-correct rules plus a dumb loop. All the cleverness lives in choosing rules, and every rule has to give exactly the same answer as the interpreter on every input.

Check yourself

After simplifying, what does %out = add 3 %x become? Write it as op %src %src.

No rule folds this, because %x isn't a constant. One rule still fires: where should constants sit?

Canonicalization moves the constant to the right: add %x 3. Nothing else applies, so that's the fixpoint.

How many ops (non-const, non-param nodes) are left after simplifying this?

%x   = param 0 dtype=i32
%a   = add 2 3
%b   = mul %a %x
%out = add %b 0

Apply the rules in order: fold the all-constant node, canonicalize, then drop the identity.

add 2 3 folds to 5, mul 5 %x canonicalizes to mul %x 5, and add(.., 0) disappears. One op: mul %x 5.

Your folder sees mul 1073741824 4 (that's 2³⁰ × 4). What i32 constant should it produce?

1073741824 is 2 to the 30th, so the product is 2 to the 32nd. Keep only the low 32 bits.

2³² wraps around to 0 in 32-bit two's complement. Folding has to wrap exactly like the interpreter does.

The folding rule matches div 1 0. What should the simplifier do?

Is div 1 0 always an error, or only when an assert needs its value?

The simplifier only sees one node at a time. It can't know whether anything depends on it.

Only programs that depend on the error are invalid, and the simplifier can't know that by looking at one node (an unused div 1 0 is fine). Leaving the node in place keeps the error exactly where the interpreter will find it.

Predict the exit code of run on this program.

%d = div 1 0
%t = assert 1 1

Which nodes do the asserts actually depend on?

Exit 0. The roots are the asserts, and assert 1 1 doesn't depend on %d, so %d is never evaluated. (It is still type-checked. div 1 0 is well typed.)

Why visit nodes sources first?

Think about what add 4 %a needs before it can fold, and when %a becomes a constant.

In the stepper, add 4 %a could only fold because %a had already become 5. Top down, you'd need one extra pass per level of depth.

You add the rule add(x, y) → add(y, x) with no condition. What happens?

After the swap, does the result still match the rule's pattern?

The output matches the pattern again, so the rule always fires. Rules must make progress. Guard it as "only when the const is on the left".

How many ops are left after simplifying this?

%x   = param 0 dtype=i32
%a   = add %x 1
%b   = add %x 1
%out = mul %a %b

Are %a and %b different nodes after hash consing?

Two. %a and %b are the same hash-consed node, so the graph is one add and one mul that uses it twice.