Week 4 study guide

Week 4 lets you write whole-vector programs: elementwise ops on rank-1 values, reduce, expand and stage. Your compiler turns them back into week 3 loops by itself, a pass called rangeify. This page recaps the week, collects the rules worth memorizing, and walks the lab one step at a time.

The week in one page

4.1 Vectors and elementwise ops

ALU ops now accept rank-1 operands of one shared shape and work element by element. There's no implicit broadcasting: add %v 1.0 with a vector %v is a type error. Rank-1 values may not depend on a hand-written loop's open range.

4.2 Reduce and expand

reduce is a left fold that starts from the first element: x[0] op x[1] op ... op x[n-1], rounding after every step. The order is part of the spec because f32 results depend on it. expand broadcasts a scalar to a shape, and week 5 widens it to grow any size-1 axis.

4.3 Rangeify

Rangeify removes shapes by asking each vector for "your element at index i". An elementwise op at i is the op on its operands at i, a load is a scalar load at offset + i, an expand is just its scalar, a reduce becomes an accumulator loop, and a rank-1 store becomes a store loop.

4.4 Fuse or materialize

A computed value used twice can be recomputed at each use (fusion) or stored once in a temp buffer (materialization). Without stage your compiler may choose, and stage forces a temp. The results are bit-identical either way. Only speed differs, and speed is graded from week 6.

4.5 Rendering vectors

Rendered C must be at most 20,000 bytes, so a 4096-element vector has to become a loop. A whole x / sum(x) kernel is around 800 bytes: a reduce loop, a scalar computed once between loops, and a store loop.

4.6 Lab

Same commands as week 3. render may print for loops when the program has a range, a reduce, expand or stage, an elementwise op on a vector, or any rank-1 value over 32 elements. The suite has 77 tests.

Cheat sheet

op	operands	result	remember
ALU ops	one shared shape	that shape	elementwise, with no implicit broadcast
`expand x shape=[n]`	a scalar (week 5: also size-1 axes)	`(n,)`	every element is x, and growing a size-2 axis is always an error
`reduce x op=add\|mul\|max`	rank-1, i32 or f32	scalar	left fold from `x[0]`, and one element is just that element
`stage x`	any value	same	same value, but forces a temp buffer in render

reduce	i32	f32
`op=add`	wraps at every step	rounds at every step, order matters, and a single `-0.0` stays `-0.0`
`op=mul`	wraps at every step	rounds and can overflow to inf
`op=max`	plain max	any NaN gives NaN, and `max(-0.0, 0.0) = +0.0`

render may use for loops when the program has
a `range`, or a `reduce`, `expand` or `stage`
an elementwise ALU op on a rank-1 value
any rank-1 value with more than 32 elements
otherwise it must be loopless, and every render must be at most 20,000 bytes

Mistakes people make

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

Folding right to left. Caught by reduce_add_f32_order, where the other order gives 0.0 instead of the right answer.
Folding from zero instead of x[0]. 0.0 + -0.0 is +0.0, and 0 is wrong for max over negatives. Caught by reduce_add_negzero, reduce_max_i32 and reduce_max_signed_zero.
A max reduce that skips NaN. Any NaN makes the result NaN. Caught by reduce_max_nan.
Elementwise ops that read element 0 for every index. Caught by add_i32_vec, add_i32_wrap_vec and cast_vec.
An expand that produces zeros. Caught by expand_dtypes, chained_reduce and cmplt_where_vec.
Forgetting to rangeify something. If your kernel's size grows with n, a vector reached the printer unrolled. Caught by r_add_f32_2048 and every 4096-element test.
A plain large copy rendered loopless. A 4096-element load and store has no ALU op, but it still has to be a loop to fit in 20 KB. Caught by r_copy_4096.
A stage that reads the wrong element, or zeroes its temp. Caught by r_stage_twice_4096, r_used_twice_stage_2048 and normalize_small_stage.
Accepting implicit broadcasts or bad reduces. add %v 1.0, a reduce of a scalar, and a reduce of bools are all invalid. Caught by err_missing_expand, err_reduce_scalar and err_reduce_bool.

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 or render FILE).

Shapes on ALU ops. Done when this exits 2 with invalid: (no implicit broadcast), and every week 1 test still passes:
```
%v = stack 1.0 2.0 3.0
%w = add %v 1.0
%t = assert 1 1
```
Parse and type expand, reduce and stage. Done when the suite's err_reduce_scalar, err_reduce_bool and err_expand_no_shape exit 2.
Vectors in run. Done when these fold in the right order and from the right start:
```
%v = stack 1e8 1.0 -1e8 1.0 1.0
%s = reduce %v op=add
%t = assert %s 2.0
%z = stack -0.0
%q = reduce %z op=add
%r = recip %q
%u = assert %r -inf
```
If the first assert fails, you're not folding left to right (or not rounding each step). If the second fails, you started from 0.0 instead of x[0].
The element at an index. Done when a 64-element copy (load a buffer, store it into another) renders with one loop, while the same copy of 8 elements may stay loopless. If the 64-element version prints 64 stores, the size rule never triggered.
Stores and reduces become loops. Done when r_dot_4096 passes bit for bit. A reduce is week 3's accumulator: a one-element alloc seeded with element 0, then a loop over the remaining n - 1 elements.

Stage becomes a temp. Done when staged and unstaged versions agree. Check:

%x  = stack 1.0 2.0 3.0 4.0
%s  = reduce %x op=add
%e  = expand %s shape=[4]
%n  = mul %x %e
%xs = stage %x
%s2 = reduce %xs op=add
%e2 = expand %s2 shape=[4]
%n2 = mul %xs %e2
%t  = assert %n %n2

Check the size. Done when your largest kernel is around a kilobyte, whatever n is, and every render test passes.

When a render is bit-for-bit wrong

Print the failing element's index and both values as f32 bits. A difference in the last bit almost always means a different order of operations: a reduce folded the other way, or an a * b + c the compiler contracted into one rounding. The debugging guide shows how to print the bits.

Check yourself

What does reduce with op=add give for the f32 vector [1e8 1.0 1.0 -1e8]?

Fold left to right from the first element, rounding to f32 after every add.

1e8 + 1.0 rounds back to 1e8, so does the next + 1.0, and - 1e8 leaves 0.0. An order that adds the two 1.0s together first, (1.0 + 1.0) + (1e8 + -1e8), would give 2.0.

What does reduce with op=mul give for the i32 vector [65536 65536 3]?

65536 is 2¹⁶. What's 2³² in i32?

2¹⁶ · 2¹⁶ = 2³², which wraps to 0, and 0 · 3 = 0.

For the f32 vector [-0.0], %s = reduce %v op=add. What's recip %s?

A one-element reduce is just that element. What's recip of -0.0?

The fold starts from x[0] = -0.0 and has nothing else to add, so the sum is -0.0, and recip -0.0 = -inf. Starting from 0.0 would give +0.0 and inf.

What does reduce with op=max give for the f32 vector [-0.0 0.0 -0.0]?

What does the spec say max(-0.0, 0.0) is, in either order?

max(-0.0, 0.0) is +0.0 in either order, so once the +0.0 is folded in, it stays.

Which program must render with a for loop?

A plain copy has no ALU op. What else makes the loop rule apply?

Any rank-1 value with more than 32 elements must become a loop. The 8-element copy may stay loopless (the oracle prints 8 stores, about 600 bytes), but the 64-element one renders as one loop.

A program computes x / sum(x) twice: once as written, and once with stage on x. How do the outputs compare?

Does stage change any value, or only where values live?

Stage only decides whether a value is stored in a temp or recomputed. Elementwise ops are deterministic and every reduce folds in the same order, so the bits are identical. Only speed can differ.

Which line is a type error, with %v an f32 vector of shape (3,)?

Do both operands have the same shape?

%v has shape (3,) and 1.0 has shape (). Write %one = expand 1.0 shape=[3] first, then add.

More practice: the week 4 quiz review.

Go deeper

tinygrad and its documentation. Rangeify is tinygrad's name for this pass, and its scheduler decides fusion and materialization for real models.
Halide and its PLDI 2013 paper. The idea that what to compute and where to store it are separate decisions, which is exactly fuse versus materialize.
TVM: An Automated End-to-End Optimizing Compiler for Deep Learning. How a production ML compiler fuses operators and picks schedules.
Loop fission and fusion. The classic transformations behind lesson 4.4.
Pairwise summation and Kahan summation. Why the fold order matters, and how real libraries trade it for accuracy.
NumPy broadcasting. The implicit broadcast rules that this course makes explicit with expand.
itertools and functools. product walks indices in row-major order, and reduce folds from the first element, both exactly as the spec does.
Array. init, map2, and fold_left cover most vector work, and float arrays are stored unboxed.
Python and OCaml, deeper. This week's idioms in your language: shapes, folds, and counting uses.

Challenges

Find an order that matters. Build an f32 vector of 8 values where the left fold and a pairwise sum give different bits. Check the left fold with the oracle (reduce op=add) and the pairwise sum by hand or in Python with f32 rounding.
Count the memory traffic. For (x - max(x)) * recip(sum((x - max(x)) * (x - max(x)))), count how many times each element of x is read with no stage, and with stage on x - max(x). Render both and count the loops to check your reasoning.
Mean and variance, staged or not. Write a program that writes the mean and the variance of 2048 values into two slots. Render it with and without stage on the deviations x - mean, count the loops in each, and confirm with the oracle that both leave the same bits. (A one-pass formula like E[x²] - mean² uses fewer loops but different arithmetic, so its bits will differ: try it and see.)