Week 6 study guide

Week 6 makes your compiler able to train. Four math functions are built from adds and multiplies, nine sugar ops expand into core ops, gradient walks the chain rule backward, and buffers can come from files checked by hash. The lab trains a digit classifier on MNIST. This page recaps the week, collects the rules worth memorizing, and walks the lab one step at a time.

The week in one page

6.1 Math from adds

exp2, log2, sqrt and sin are each defined by an exact program in the spec: reduce the input to a small range, then a short polynomial. bitcast reads an f32's bits as an i32 and back, which is how exp2 builds 2^n and log2 reads an exponent. Matching the bits means doing exactly the spec's steps.

6.2 Sugar

neg, sub, fdiv, relu, exp, log, floor, cmpeq and matmul each mean their expansion. Expand them in the parser and nothing downstream changes. fdiv is a · recip(b), not correctly rounded division, and div stays i32-only.

6.3 Automatic differentiation

gradient y x is built as more graph. Seed y with 1.0, visit y's graph in reverse post-order, add each node's shares in arrival order, and send shares to its f32 operands on a path to x. detach stops a gradient. A call, load or loop on the path from x to y is invalid, and so is a literal as x.

6.4 Gradients through shapes

Every movement op has a partner: reshape and reshape back, permute and the inverse order, pad and shrink, expand and reduce add. reduce max splits the gradient between tied maxima. matmul needs no rule, and a staged value's gradient is staged.

6.5 Data and MNIST

url= and sha256= fill a buffer from a raw little-endian file, rejected if the hash or size is wrong. The MNIST program trains a softmax regression for 50 steps of 100 digits and counts how many of 1000 test digits it gets right.

6.6 Lab

The suite has 147 tests. One renders and runs the whole training program and needs 850 of 1000 right within 120 seconds. Weeks 1 to 5 must stay green. tools/bench.py times your kernels.

Cheat sheet

Mistakes people make

Each of these is a bug that was planted in the reference compiler to prove the suite catches it. The week 6 test guide lists every test with the bugs known to fail it.

1. A polynomial one term short, or a different one. The answer is close and the bits are wrong. Caught by r_exp2, r_exp_log, r_grad_softmax.
2. Rounding n down instead of to nearest in exp2. f lands in [0, 1), where the same polynomial is less accurate. Caught by r_exp2, r_exp_log, r_grad_softmax.
3. Building 2^n as one bit pattern. Below 2^-126 there is no such pattern, so the subnormal results come out wrong. Caught by r_exp2, r_exp_log, exp2_integers.
4. Skipping log2's fold at 1.414, or its subnormal scaling. Caught by r_exp_log, r_log2, r_sqrt.
5. Leaving out sqrt's Newton step. The guess alone is a few units off. Caught by r_sqrt, sqrt_values.
6. A bitcast that leaks which NaN you had. x86 makes a negative NaN for 0 · inf. Caught by bitcast_nan_from_bits, r_bitcast.
7. A floor that truncates. floor -2.5 is -3. Caught by r_exp2, r_exp_log, r_floor.
8. Mixing up a mul's shares. a gets g · b, not g · a. Caught by mnist_train, r_grad_regression_step, r_grad_softmax.
9. Adding shares in a different order. f32 addition isn't associative. Caught by grad_accumulation_order_inner.
10. Ignoring detach, or letting gradients flow into integers. Caught by mnist_train, r_grad_softmax, grad_detach.
11. Getting a movement op's partner backward. The inverse permutation, the side of a pad, the offsets of a shrink. Caught by grad_permute, r_grad_movement, grad_pad.
12. Not splitting reduce max between ties. Caught by grad_reduce_max_ties.
13. Trusting a file without its hash or size. Caught by err_url_sha_mismatch, err_url_file_too_long.
14. Fusing a store that reads its own buffer. Later elements see earlier ones' new values. The week 5 test r_inplace_transpose checks this, and r_grad_regression_step and mnist_train update their weights in place.

Stuck on the lab?

Build in this order, and don't move on until the check passes. Save each check as a file and run it through your compiler and the oracle (python _reference/uirc.py run FILE).

1. bitcast and floor. Done when this passes:

   %a = bitcast 1.0 dtype=i32
   %t1 = assert %a 1065353216
   %f = floor -2.5
   %t2 = assert %f -3.0
   %n = bitcast nan dtype=i32
   %t3 = assert %n 2143289344

2. exp2, exactly. Integers are exact, so start there: exp2 3.0 is 8.0 and exp2 -149.0 is 1.4e-45. Then compare exp2 0.1 bit for bit with the oracle using bitcast: if it's one unit off, check the constants and the order of the Horner steps against the spec.
3. log2, sqrt, sin. log2 8.0 is 3.0, sqrt 16.0 is 4.0, sin 0.0 is 0.0. Then the suite's *_bits tests.
4. Sugar. Done when sugar_* passes. fdiv 10.0 3.0 must give 3.3333335.
5. Gradient on scalars. Done when this passes:

   %b = buffer 1 dtype=f32 slot=0 data=[2.0]
   %v = shrink %b 0 shape=[]
   %x = load %v
   %c = mul %x %x
   %y = mul %c %x
   %g = gradient %y %x
   %t = assert %g 12.0

   If you get 4, each share is overwriting the last instead of adding to it: x receives three shares of 4.
6. Gradient through shapes. The suite's grad_expand_*, grad_permute, grad_pad and grad_matmul_* tests, one rule at a time.
7. Files, then render, then MNIST. url_* first. For a slow MNIST render, look for any walk over the whole graph that runs once per node.

Check yourself