Week 1 · Lesson 1

The UOp: one node type for everything

Your whole compiler is built from one tiny data structure: an op name, a tuple of source nodes, and an argument. Programs are graphs of these. This lesson is about why that's enough, and the trick (hash consing) that makes it fast.

Three fields

UOp(op, src: tuple[UOp, ...], arg)

Read src: tuple[UOp, ...] as "src is a tuple of any number of UOps". Op names are written in capitals (ADD) when we talk about the compiler and in lowercase (add) in .uir text. They're the same op.

op: what this node does, like ADD, CONST or WHERE.
src: the nodes it reads. Their order matters: CMPLT a b is a < b, not b < a.
arg: a non-node payload. The number for a CONST, the slot for a PARAM, the target dtype for a CAST. Most ops have none.

Every node also carries a dtype (void, bool, i32, f32) and a shape, computed from its sources when it's built. Inside the compiler there are no variables or blocks. A program is one expression, stored as a graph.

Why a graph and not a tree

In (2 + 3) * (2 + 3), the tree has the 2 + 3 part twice. In a graph it's one node with two arrows coming out. That's a DAG (directed acyclic graph): arrows only point from sources to users, with no loops. Sharing is the reason ML compilers use DAGs. In a model, one intermediate value often feeds many later ops.

Try it: type .uir, see the graph

A quick primer, since lesson 2 has the full format: each line names one node, %name = op operands. Plain numbers are constants. assert a b checks that two values are equal. Edit the program. Sources sit on top and results flow down. Identical statements collapse into one node. Try adding %e = add 2 3 and watch the node count stay the same.

live graph

Hash consing

Hash consing means never build the same node twice. Every node you create goes through one cache, keyed on (op, src, arg), where dtype and shape are part of the key. If the key is already there, you get the existing node back. It's the same object, not just an equal copy.

Because sources are already unique, the key can compare sources by identity (their address, or a unique id). Building a node is then one hash lookup, no matter how big the expression underneath is.

What you get for free

Equality takes one step, however big the expression. "Are these two expressions the same?" becomes "are these the same object?". Rewrite rules like x xor x → 0 need exactly that check, and it doesn't walk any subtrees.
Sharing is automatic. Write 2 + 3 a hundred times and you still have one node.
Memoization has an easy key. Caching "the value of node n" by node identity means each node is computed at most once. (Once call arrives in lesson 5, this cache has to be per call, because the same param node means different values in different calls.)

Structure, not meaning

Hash consing is structural. add 2 3 and add 3 2 are different nodes, even though they're equal as numbers, because the src tuple is in a different order. Making them meet is the job of rewrite rules (lesson 3), which can put the operands of order-independent ops like add (commutative ops) into one fixed (canonical) order.

Code it: a hash-consing table

Build the cache yourself, on a toy node type in your language. The key is the op plus the ids of the sources, which is exactly why sources being unique already makes the lookup cheap.

Hash cons a toy node

Write create (). It returns a make function with its own private table. make op src must return the same node (the same object: ==is) when it has already built a node with that op and those exact source nodes. Otherwise it builds a new node whose id is the next number: 0, 1, 2, and so on.

Two different create () calls must not share a table. Source order matters: an "add" of x then y and an "add" of y then x are different nodes.

The shape: let create () = let tbl = Hashtbl.create 16 in let next = ref 0 in fun op src -> .... The table lives inside each create () call, and the function it returns can see it. Key the Hashtbl on (op, list of source ids).

The shape: inside create(), make a dict and a counter, then define make and return it. Each create() call gets its own dict, and the make it returns can see it. Key the dict on (op, tuple of source ids), and compare nodes with is, not ==.

The doubling DAG

Here's a program the tests really use. Each line adds the previous value to itself:

%x0 = const 1
%x1 = add %x0 %x0
%x2 = add %x1 %x1
%x3 = add %x2 %x2      # ... keep going to %x300

x0 to x3 as a DAG: 4 nodes

Both arrows into each add come from the same node.

levels	nodes as a DAG	nodes if you expand it into a tree
3	4	15
30	31	2,147,483,647
300	301	about 4 × 10⁹⁰, more than the atoms in the universe

An evaluator that recurses on sources without remembering results visits %x0 2³⁰⁰ times. One that caches each node's value visits each node once, so 301 steps. The spec requires this case to run in milliseconds, so "I'll add memoization later" isn't an option.

Count distinct nodes

Write dag_size t: the number of physically distinct nodes reachable from t. A node reached twice counts once. dag_size (dbl 3) is 4, and dag_size (dbl 60) has to finish, so walking the tree is out.

Remember what you've visited. List.memq tests membership with ==.A set of id(node) values tests membership by identity.

x to the 8th in 3 ops

The starter computes x⁸ correctly with 7 muls. Rewrite %out to use at most 3 ops. Sharing is the whole trick: a node can feed both sides of a mul. %x = param 0 dtype=i32 is the input the tests fill in (lesson 2 explains params), and const and param lines don't count as ops.

i32 wraps, so 16⁸ is 0. Your program gets that for free.

Properties come from sources

A UOp's dtype and shape aren't stored by hand. They're computed when the node is built, from its op and its sources. This is the type checker. A node that can't get a valid dtype (say, add of an i32 and an f32) is rejected right there.

node	dtype	shape	rule
`add a b`	a's	`()`	a and b must share a dtype, rank 0 (no axes: a single value)
`cmplt a b`	bool	`()`	compares, so the result is always bool
`stack a0 .. a(n-1)`	a0's	`(n,)`	adds a leading axis
`index x i`	x's	`()`	pops x's leading axis
`assert a b`	void	`()`	a root, not a value

Preview: bounds and cost (lesson 4)

The syllabus lists two more recursive properties. bounds is the smallest and largest value a node could ever have: max %p 0 is in [0, 2³¹-1], and mod of that by 8 is in [0, 7]. Bounds let you prove cmplt (mod (max p 0) 8) 8 is always true without knowing p. cost estimates how expensive a subgraph is, so rewrites can tell if they made things better.

Check yourself

How many nodes does this program have after hash consing? Count constants as nodes.

%a = add 2 3
%b = add 2 3
%c = mul %a %b

Literals are nodes too. Two statements with the same op and the same sources collapse into one.

const 2, const 3, one add (shared by %a and %b), and one mul. %c evaluates to 25.

Are %a = add 2 3 and %b = add 3 2 the same UOp?

Hash consing compares the op, the source tuple in order, and the arg. It never computes values.

Hash consing looks at structure, not values. A rewrite rule that sorts operands can make them meet.

%p = param 0 dtype=i32 and later %q = param 0 dtype=i32. How many nodes?

A node is keyed on (op, src, arg), with dtype and shape as part of the key. Compare each of those for %p and %q.

Same op, same (empty) src, same arg, same dtype: one node. What a param means depends on which call is evaluating it (lesson 5), so sharing it is safe.

Start the doubling DAG at %x0 = const 1. What is %x31? Remember i32 wraps.

%x31 is 2 to the 31st. The largest i32 is 2 to the 31st minus 1.

One past INT_MAX wraps around to the most negative i32.

%x31 is 2³¹, one more than INT_MAX, so it wraps to INT_MIN = -2147483648. One more doubling gives 2³², which wraps to 0, and it stays 0 all the way to %x300.

An evaluator caches each node's value by node identity. How many node evaluations does %x300 take?

How many distinct nodes does the chain %x0 to %x300 have, if each is computed only once?

Each distinct node is computed once. The second arrow into an add hits the cache.

What dtype is cmplt 1 2?

A comparison answers a yes or no question, whatever it compares.

Comparisons always produce bool, whatever their operands are. Its value is true.

What's the shape of stack 1 2 3?

stack adds one new leading axis, with one entry per source.

stack adds a new leading axis with one entry per source. index pops it back off: index (stack 1 2 3) 2 is 3.