Lab: Train MNIST

The week 6 lab

Teach your compiler the four math ops, the sugar, gradient, detach, bitcast and file buffers, and then let it train. You're done when every week 6 test passes, including the one that trains on MNIST and has to get 850 of 1000 digits right. Weeks 1 to 5 must still pass, and you write 10 new tests of your own.

Warm up

One step of gradient descent, by hand, on four points. The MNIST test is this, 50 times, with a bigger model.

Warm-up: one step of gradient descent

Slots 0 to 2 hold X (4x2), y (4 targets) and w (2 weights). With loss = 0.25 · Σ (X @ w - y)², update w in place to w - 0.1 · d(loss)/dw. The starter steps the wrong way.

The command-line contract

The same commands as before. What they accept grows.

command week 6 changes

your.exe run FILE New ops: exp2, log2, sqrt, sin (exactly as the spec defines them), bitcast, detach, gradient. New sugar: neg, sub, fdiv, relu, exp, log, floor, cmpeq, matmul. Buffers may take url= and sha256=, read relative to the program's folder. Op names may contain digits. New static errors (exit 2) include a math op on i32, cmpeq on f32, fdiv on i32, a gradient with a non-scalar loss or with respect to a literal, a call, load or reduce op=mul on a gradient path, a missing data file, a hash or size mismatch, and a .. in a url.

your.exe render FILE Math ops become their exact steps in C (inline, or one C function per op). bitcast uses memcpy and turns any NaN into 0x7fc00000. The harness fills url= buffers before calling your kernel. A store whose value reads the buffer it overwrites must still use the old contents.

command	week 6 changes
`your.exe run FILE`	New ops: `exp2`, `log2`, `sqrt`, `sin` (exactly as the spec defines them), `bitcast`, `detach`, `gradient`. New sugar: `neg`, `sub`, `fdiv`, `relu`, `exp`, `log`, `floor`, `cmpeq`, `matmul`. Buffers may take `url=` and `sha256=`, read relative to the program's folder. Op names may contain digits. New static errors (exit 2) include a math op on i32, `cmpeq` on f32, `fdiv` on i32, a gradient with a non-scalar loss or with respect to a literal, a call, load or `reduce op=mul` on a gradient path, a missing data file, a hash or size mismatch, and a `..` in a url.
`your.exe render FILE`	Math ops become their exact steps in C (inline, or one C function per op). `bitcast` uses `memcpy` and turns any NaN into `0x7fc00000`. The harness fills `url=` buffers before calling your kernel. A store whose value reads the buffer it overwrites must still use the old contents.

How the tests are graded

# expect: and # render tests work as before. Three headers are new:

# render check: slot S >= V: the grader reads element 0 of slot S after your kernel runs and needs it to be at least V. The MNIST test uses it for accuracy, and r_url_check uses it on a small example. Nothing else is compared, because no two compilers are required to train to the same bits.
# time: T: your compiler's render, the C compile and the kernel each get T seconds instead of 10. MNIST gets 120.
# max_bytes: B: the size limit on your C, raised for MNIST's 50 written-out steps.

The run tests check exact bits: the math ops against their definitions, each gradient against the rules and the share order, every sugar against its expansion. Several tests come in pairs: *_one_ulp_off asserts one value one ulp (one unit in the last place) away from the right one and is marked to fail. A compiler that is only close makes that assert pass, so the test catches it.

Running the suite

From the course repo root:

python tools/test.py --compiler "C:/path/to/your.exe" --week 6
python tools/test.py --compiler "C:/path/to/your.exe" --week 6 -k grad -v      # just the gradients, with details
python tools/test.py --compiler "C:/path/to/your.exe" --week 6 -k mnist -v     # the training run, about 30 s
python tools/test.py --compiler "C:/path/to/your.exe" --week 5                 # each earlier week is its own run
python tools/test.py --compiler "C:/path/to/your.exe" --week 4
python tools/test.py --compiler "C:/path/to/your.exe" --week 3
python tools/test.py --compiler "C:/path/to/your.exe" --week 2
python tools/test.py --compiler "C:/path/to/your.exe" --week 1

Pass the built executable, not dune exec.the same wrapper script as in week 1. Tests live in tests/week6/run/ and tests/week6/render/, and their data files in the data/ folder next to each.

Timing your kernels

python tools/bench.py --compiler "C:/path/to/your.exe" tests/week6/render/mnist_train.uir

This renders the program with your compiler, builds it at -O2 without the sanitizers, runs it, and prints the render time, the compile time, the kernel's time, and the useful work in FLOPs (floating-point operations): one per element of every elementwise op and size - result size per reduce. Divide one by the other and you have your speed in FLOP/s. The reference does the MNIST program's 0.17 GFLOP (a GFLOP is a billion FLOPs) in about 0.05 s. Week 6 grades only correctness and the time limits, but from here on these numbers are the scoreboard.

Write your 10 tests

Ten tests in path/to/your-repo/tests/week6/. Check them against the oracle:

python tools/test.py --check-tests path/to/your-repo/tests/week6

Good targets: an exp2 result in the subnormal range, log2 on both sides of the 1.414 fold, sin of a negative number of turns, and a bitcast of a NaN made by arithmetic. For gradients, try a node used three times with values whose sum depends on the order, a gradient of a gradient, and an x that the loss reaches only through an i32 cast. Then a file whose size is off by one byte, and an in-place update whose render would be wrong if fused.

A build order that works

Digits in op names, bitcast, detach

The parser accepts exp2. bitcast reinterprets 32 bits, mapping every NaN to 0x7fc00000. detach is the identity, for now.

Checkpoint: the bitcast_* and detach_* tests pass, and weeks 1 to 5 still pass.

floor, then exp2

Build floor as sugar from the spec's 12 lines, then exp2 exactly as written. Keep exp2 as one node, and evaluate it with its body.

Checkpoint: exp2_* and sugar_floor* pass, including the subnormal results.

log2, sqrt, sin

Same approach. Compare against python _reference/uirc.py run on your own one-line programs when a bit is off. A wrong constant or a swapped operand shows up immediately.

Checkpoint: every math run test passes, including the special values.

The rest of the sugar

neg, sub, fdiv, relu, exp, log, cmpeq, matmul, each expanded in the parser.

Checkpoint: every sugar_* test passes.

Gradient, elementwise

Compute which nodes reach x, check the path, then walk the reverse post-order, adding shares in arrival order. Start with add, mul and the math ops, then max, where, recip and the casts.

Checkpoint: the scalar grad_* tests and every err_grad_* test pass.

Gradient through shapes

Reshape, permute, pad, shrink, expand, reduce add and max, stack, index, stage.

Checkpoint: every week 6 run test passes.

File buffers

Resolve the path against the program's folder, reject .., check SHA-256 and size, decode little-endian. hashlib.sha256 is in the standard library.OCaml's standard library has no SHA-256. Write one (it's about 60 lines, and a good exercise), or run sha256sum.

Checkpoint: the url_* tests pass.

Render

Lower math ops to their bodies, render bitcast with memcpy, and stage a stored value that reads the buffer it overwrites.

Checkpoint: every render test except MNIST passes.

MNIST

Run it. If it's slow to render, profile your compiler: look for a whole-graph walk per node. If it's wrong, run a 2-step version (python tools/mnist.py program --steps 2) and compare buffers with the oracle.

Checkpoint: mnist_train passes, with 850 or more of 1000 right.

Checklist

The AI policy, still

Your repo's AGENTS.md hasn't changed: an AI assistant may explain concepts, draw diagrams, and help you debug by asking questions, but must never write or rewrite compiler code. "Why does my log2 come out one unit high for 1.4142137?" is a fine question. "Write my gradient pass" is not.

The oracle in _reference/ stays closed. Running it on your own programs to compare bits is fine, and the spec has every definition you need.

Check yourself

Why does the MNIST test check an accuracy threshold instead of every bit?

Is anything random in the program?

Nothing is random: the weights start at zero and the batches are fixed. The run tests check every rule bit for bit. The MNIST test checks that it all adds up to learning.

The update stores W - g into the buffer W was loaded from. Your renderer turns the store into a loop over 7840 elements. What can go wrong?

What does element 5000's gradient read?

The logits, which use every weight. The fix is to materialize the new values first and copy them in after, which is what the spec's semantics require: a store's value comes from the contents before it.

Your render of MNIST takes 60 seconds, but a 2-step version takes 0.1 s. What's the likely cause?

25 times more steps, 600 times more time.

That ratio is roughly quadratic (25² is 625). Typical culprits are a topological sort from each statement, or recomputing "which ranges are open" from scratch for every new node. Memoize, and share one "seen" set across walks.

The MNIST program does about 0.17 GFLOP. If the kernel takes 0.1 s, how many GFLOP/s is that?

Speed is work divided by time.

0.17 / 0.1 = 1.7 GFLOP/s. The reference does it in about 0.05 s, so about 3.4 GFLOP/s, which is still far below what one CPU core can do. Later weeks close that gap.