Python and OCaml, deeper

Ops as an enum

An Enum gives each op one object, so a typo is an AttributeError instead of a string that silently never matches. The parser can look ops up by name.

from enum import Enum, auto

class Op(Enum):
  CONST = auto()
  PARAM = auto()
  ADD = auto()
  MUL = auto()
  # ...

Op["cmplt".upper()]   # Op.CMPLT, or a KeyError for an unknown op

Hash consing: a frozen dataclass and a dict

Make the node immutable with frozen=True, and keep one table from a key to the node already built for it. Pass eq=False too. Then == and hash use identity, which is exactly right once nodes are shared, and constant time.

from dataclasses import dataclass

@dataclass(frozen=True, eq=False)
class UOp:
  op: Op
  dtype: str
  src: tuple["UOp", ...]
  arg: object

_table: dict[tuple, UOp] = {}

def uop(op, dtype, src=(), arg=None) -> UOp:
  key = (op, dtype, src, arg_key(arg))
  ...   # return the node stored under key, or build one and store it

Rewrite rules with match

Python 3.10's match reads like the rule it implements. Keyword patterns work on any dataclass, and nested patterns match a whole subtree at once. Return None when no rule fires, so the engine knows to stop.

def rule(u: UOp) -> UOp | None:
  match u:
    case UOp(op=Op.ADD, dtype="i32", src=(x, UOp(op=Op.CONST, arg=0))):
      return x                       # add(x, 0) -> x, for i32 only
    case UOp(op=Op.MUL, dtype="i32", src=(x, UOp(op=Op.CONST, arg=1))):
      return x
  return None

32-bit integers and f32 floats, by hand

Python ints never overflow and Python floats are doubles. Neither matches the course's types, so give every arithmetic op one small helper and never use bare + on a value.

• i32. After each op, keep the low 32 bits (mask with 0xFFFFFFFF) and read the top bit as a sign. ctypes.c_int32(n).value does the same and makes a good reference for a test.
• Division. // floors and % takes the divisor's sign, and int(a / b) goes through a double. Build truncating division from the quotient of the absolute values.
• f32. Round a double to f32 by packing and unpacking it: struct.unpack("<f", struct.pack("<f", x))[0]. Do it after every op and on every literal, not once at the end.

Memoize per call, and know what functools.cache does

functools.cache on an evaluator is one global table that lives forever. That's fine for a pure property of a node, like its dtype or shape. It's wrong for values that depend on a call's arguments: the same body called with 3 and then with 4 returns the first answer twice. Use a plain dict created at the start of each call for those.

An explicit worklist instead of deep recursion

Python stops at about 1000 nested calls (sys.getrecursionlimit()), and a long chain of ops is deeper than that. Raising the limit works until the C stack overflows and the process dies with no message. Walk the graph with your own stack instead. The trick for "sources before the node" is to push each node twice:

stack, done = [(root, False)], set()
while stack:
  u, expanded = stack.pop()
  if u in done: continue
  if expanded:
    ...                      # every source of u is finished: handle u here
    done.add(u)
  else:
    stack.append((u, True))
    stack.extend((s, False) for s in u.src)

Exit codes

An uncaught Python exception prints a traceback and exits 1, which the grader reads as "an assert failed". Catch errors at the top of main, print a line starting invalid:, and sys.exit(2). Catch RecursionError there too, so a deep graph fails loudly instead of looking like a failed assert.

type op = Const | Param | Add | Mul | Cmplt | Where (* ... *)
type dtype = I32 | F32 | Bool
type arg = No_arg | Int of int32 | Float of float | Slot of int

type uop = { id : int; op : op; dtype : dtype; src : uop list; arg : arg }

module Key = struct
  type t = op * dtype * int list * arg_key (* src ids, and arg with floats as bits *)
  let equal = ( = )                        (* cheap now: no subtrees inside *)
  let hash = Hashtbl.hash
end
module Table = Hashtbl.Make (Key)

let table : uop Table.t = Table.create 1024
let make op dtype src arg = ...            (* find, or build with a fresh id and add *)

let rule u =
  match u.op, u.dtype, u.src with
  | Add, I32, [ x; { op = Const; arg = Int 0l; _ } ] -> Some x
  | Mul, I32, [ x; { op = Const; arg = Int 1l; _ } ] -> Some x
  | _ -> None   (* fine here: most nodes have no rule *)

An f-string emitter

Collect C lines in a list and join them at the end. Name each value once, in a dict from node to C name, so a shared node is computed once in the C too.

lines: list[str] = []
names: dict[UOp, str] = {}           # identity keys, thanks to eq=False

def emit(u: UOp) -> None:
  a, b = (names[s] for s in u.src)
  names[u] = f"v{len(names)}"
  lines.append(f"  int32_t {names[u]} = {a} ... {b};")   # the wrap-safe form is yours to write

Two f-string details bite in C. Literal braces must be doubled (f"void kernel(...) {{"), and a quote inside the braces needs the other quote style before Python 3.12.

Exact float literals with float.hex

str(x) does round-trip, but formats like f"{x:f}" and %g drop digits, and a decimal literal is something you have to argue is exact. x.hex() prints the bits as a C99 hex float, exact by construction: 0.1 rounded to f32 prints 0x1.99999a0000000p-4, and C reads it back with an f suffix to the same float. It doesn't cover everything, though: float("inf").hex() is inf, which C doesn't accept, so infinities and NaN need their own spelling.

Bounds with unbounded ints

Here Python's big ints help. Represent an interval as a pair of plain ints and compute lo and hi exactly, with no overflow along the way. Only at the end ask whether the result fits the buffer, or i32. In a 32-bit language you'd have to reason about the bounds' own overflow.

The runtime: subprocess, tempfile, struct

• tempfile.TemporaryDirectory() in a with block gives a fresh folder and deletes it afterward, even on an error.
• subprocess.run([cc, "-O2", src, "-o", exe], capture_output=True, text=True, timeout=10) takes a list, so paths with spaces need no quoting, and timeout saves you from a kernel that hangs. Check returncode and show stderr.
• Read a float back by its bits, never by parsing a decimal: int(h, 16), then struct.pack("<I", n) and struct.unpack("<f", ...). str.split() with no argument also strips the \r of Windows line endings.

let b = Buffer.create 4096
let name = Hashtbl.create 64                 (* node id -> "v12" *)

let emit u =
  let v = Printf.sprintf "v%d" (Hashtbl.length name) in
  Hashtbl.replace name u.id v;
  Printf.bprintf b "  int32_t %s = ...;\n" v   (* the wrap-safe form is yours to write *)

module type PASS = sig
  val name : string
  val run : uop -> uop
end

module Debug (P : PASS) : PASS = struct     (* a functor: wraps any pass *)
  let name = P.name
  let run u = let u' = P.run u in (* print u' as .uir here *) u'
end

Sets of ranges, computed once

Placement asks, for each node, which loops it depends on. That's a set, and frozenset is hashable and cheap to union. Compute it bottom up with the worklist from week 1 and store it in a dict keyed by node, or with functools.cache, since it's a pure property of the node.

Bucketing nodes by loop

collections.defaultdict(list) keyed by a loop gives the body of each loop in visit order, which is already the order sources must come first. Emit with an indentation depth you pass down, as " " * depth + line, rather than fixing whitespace afterward.

Testing a matmul by hand

Before you trust your C, compute the expected matrix in plain Python with f32 rounding after each multiply and add, in the same k order as your loop. Nested list comprehensions and zip make the reference short enough to read at a glance.

Shapes as tuples

Use tuples for shapes, never lists: they're hashable, so a shape can sit in a hash-consing key, and a list there fails loudly. math.prod(shape) is the element count, and math.prod(()) is 1, which is right for a scalar.

Row-major order for free

itertools.product(*(range(d) for d in shape)) yields every index tuple in row-major order, last axis fastest. That's the order the spec uses for reduce, so an interpreter for vector values can loop over it directly.

A left fold that starts from the first element

functools.reduce(f, xs) with no initial value starts from the first element and folds left, which is the spec's order. Passing 0 as an initial value changes the answer for -0.0 and for max. Remember to round to f32 inside f.

Counting uses

Materialization starts from how many times each value is used. collections.Counter(s for u in nodes for s in u.src) counts them in one line.

Strides and unraveling

Row-major strides are suffix products of the shape: each stride is the product of the sizes after it. list(itertools.accumulate(reversed(shape[1:]), operator.mul, initial=1))[::-1] builds them ((12, 4, 1) for (2, 3, 4)), and divmod unravels a flat index one axis at a time. Write both as small functions and test them against itertools.product on random shapes.

Index expressions are UOps

When you rangeify a movement op, build the new index with the same smart constructor you use everywhere else, not as a string or a Python int. Then hash consing shares it and your week 1 simplifier cleans it up, which keeps the rendered C small.

Property tests for permutations

Permute and its inverse should cancel, and reshape there and back should be the identity. These are perfect for a property test: generate a random shape and order, and check the round trip. random.sample(range(r), r) gives a random order.

Profiling

python -m cProfile -s cumulative main.py run big.uir prints where time goes by function. py-spy record -o profile.svg -- python main.py run big.uir samples a running program with almost no slowdown and draws a flame graph. The usual finding in this course is a missing memo, a recursive __hash__, or an O(n) list lookup inside a loop.

Testing

With pytest, any function named test_* in a file named test_*.py is a test, and a plain assert reports both sides when it fails. Use it for small facts about your helpers, and keep using tools/test.py for the compiler as a whole. Hypothesis generates inputs for you and shrinks a failure to the smallest one. It's a natural fit for i32 helpers:

from hypothesis import given, strategies as st
import ctypes

i32 = st.integers(-2**31, 2**31 - 1)

@given(i32, i32)
def test_add_wraps(a, b):
  assert add_i32(a, b) == ctypes.c_int32(a + b).value

The same idea checks a rewrite rule: evaluate the program before and after on random arguments and assert the values agree. That's what the grader's simplify tests do.

; test/dune
(test
 (name test_compiler)
 (libraries mycompiler alcotest))

let%expect_test "fold add" =
  print_string (to_uir (simplify (parse "...")));
  [%expect {| %out = const 3 |}]

• tinygrad's uop/ops.py. The real UOp: a dataclass with eq=False, a hash-consing cache keyed on its op, dtype, sources and arg, and UPat and PatternMatcher for rewrite rules. Read it after week 1 and compare with yours.
• tinygrad documentation. The overview of how the pieces in that file fit into a full ML framework.
• PEP 636: Structural Pattern Matching tutorial. The friendliest walk through match, including class patterns and guards.
• The match statement reference. The exact rules, such as why a literal pattern uses ==.
• dataclasses. What frozen, eq, and slots generate, and how they interact with __hash__.
• struct. Format codes and byte order for reading and writing f32 bits.
• float.hex and float.fromhex. The exact hex float format your renderer prints.
• Floating-Point Arithmetic: Issues and Limitations. Short, and it explains why Python's floats are doubles and what repr promises.
• functools. cache, reduce, and singledispatch, each useful somewhere in a compiler.
• Enum HOWTO. Beyond auto(): members with values, and iterating over all ops.
• subprocess. Running clang and your kernel, with timeouts and captured output.
• Crafting Interpreters (Robert Nystrom). A free book on scanners, parsers, and tree-walking interpreters, written in a style that maps well onto Python.
• Fluent Python (Luciano Ramalho, 2nd edition). The book on how Python's data model works: hashing, equality, dataclasses, and pattern matching.

Tools mentioned above: pytest, Hypothesis, py-spy, and cProfile.