Week 3 study guide

Week 3 adds loops. A range is a loop counter, an end repeats an effect for every value of it, and your compiler learns to place each value at the right loop depth and print for loops. This page recaps the week, collects the rules worth memorizing, and walks the lab one step at a time.

The week in one page

3.1 Ranges and end

%i = range 4 is a loop counter that takes the values 0, 1, 2, 3, named by its statement. end %s %i runs the effect %s once for every value of %i. Every node has a set of open ranges: the ranges it depends on that no end has closed yet. Asserts, calls and range bounds may not depend on an open range.

3.2 Memory in loops

A value is computed inside every loop whose range it uses, and no deeper. A load is the one exception: if the innermost loop holding every use of it stores into the version it reads, the load moves into that loop and sees the current contents on each iteration. That is the accumulator. A use counts as inside a loop only if it is placed there itself.

3.3 A matmul by hand

Three ranges i, j, k. Reset C[i*N + j] to zero with a store that depends on i and j, then let the k loop add into it. Bounds come free: a range's interval is [0, n-1], so i*K + k is provably inside a buffer of M·K elements.

3.4 A convolution by hand

A valid convolution sums x[o + k] * w[k]. Padding is a clamped index (so the read stays in bounds) plus a where that zeroes the taps outside the input. A 2D convolution is four loops on flat buffers.

3.5 Rendering loops

Give every node a path of loops, outermost first, from its open ranges. Linearize into nested blocks where sources come first and loads of a version come before the store that overwrites it, then print each block as a for loop. Values that depend on no range end up outside the loops for free.

3.6 Lab

run and render keep their week 2 contracts, and render may now print for loops (never while, goto or recursion). The suite has 90 tests, and weeks 1 and 2 must stay green.

Cheat sheet

op	operands	result	remember
`range n`	an i32 scalar with no open ranges	i32	values 0 to n-1, `n ≤ 0` runs zero times, and the statement name is the loop's name
`end x1 ... xk r`	effects (stores or ends), then a range	effect	runs the effects for r = 0, 1, ..., closes r, and must use r
`after m e`	memory, then an effect	memory	may sit inside loops, but only an after with no open ranges is an output

rule	what it says
open ranges	a range is open in itself, `end ... r` removes r, and everything else takes the union of its sources
no open ranges allowed in	`assert`, `call` (and call bodies), range bounds, buffer and alloc sizes, outputs
one end per range	each range is closed by exactly one end, and the ends' nesting is the loop nesting
placement	a value sits in every loop whose range it uses, no deeper
load exception	if the innermost loop holding every use is deeper and stores into the load's version, the load moves in and sees current contents
bounds	a range's interval is `[0, hi(n) - 1]`

python tools/test.py --compiler YOUR_COMPILER --week 3           # grade
python tools/test.py --compiler YOUR_COMPILER --week 2           # keep week 2 green
python _reference/uirc.py render FILE                            # the oracle's C

Mistakes people make

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

Never moving loads into loops. An accumulator load whose only use is inside the k loop is then read once, and each iteration overwrites the sum instead of adding to it. Caught by r_conv1d_padded and r_conv1d_valid.
Moving a load to its deepest use instead of the loop holding all of its uses. A load used both inside and after a loop must run once, before it. Caught by split_use_across_asserts.
Deciding load placement from one assert at a time. Where a load runs depends on every use in the program, so run and render must agree. Caught by split_use_across_asserts and r_split_use_across_asserts.
Off-by-one loops. Printing <=, starting at 1, or running one extra trip. Caught by r_conv2d_valid, r_bounds_from_ranges and bounds_from_ranges.
Running a zero-trip loop once. range 0 and negative bounds run zero times. Caught by zero_trip and negative_trip.
Running a loop backwards. f32 sums depend on the order: 1e8 + 1 + -1e8 + 1 is 1.0 left to right. Caught by reduce_sum_f32_order and prefix_sum.
Too-tight or missing range bounds. A range's interval must be [0, n-1], no tighter, or an overrunning view slips through. Caught by err_range_overrun and err_wrong_stride.
Treating an after inside a loop as an output. Only top-level afters are outputs. Caught by inloop_after_not_output in week 4's suite.
Skipping the static loop rules. A range closed twice, a loop never closed, an end that doesn't use its range, an assert that reads a range: all exit 2. Caught by err_range_two_ends, err_loop_never_closed, err_end_unused_range_closed and err_assert_closed_range.
Hoisting a value out of a loop it depends on. Caught by r_bounds_from_ranges and r_conv1d_padded.

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), comparing exit codes.

Parse and type range and end. Done when this exits 0: a loop that writes b[i] = i, then reads b[3].

%b  = buffer 4 dtype=i32 slot=0
%i  = range 4
%v  = shrink %b %i shape=[]
%s  = store %v %i
%e  = end %s %i
%b1 = after %b %e
%v3 = shrink %b1 3 shape=[]
%x  = load %v3
%t  = assert %x 3

If it exits 1, check the trip count and the start value.

Open ranges and the static rules. Done when closing a range the effect doesn't use exits 2 with invalid:. Check: the same program with an extra %j = range 2, an %ej = end %e %j wrapped around the i loop, and %b1 = after %b %ej. Nothing inside uses %j, so the program must exit 2.
Bounds through ranges. Done when range 5 indexing a 4-element buffer exits 2 (the offset is in [0, 4]), while shrink %b (mul %k 2) shape=[2] with %k = range 5 on a 10-element buffer exits 0 (offset in [0, 8], plus 2, fits in 10). If both pass or both fail, look at the range case of your interval analysis.

Loop placement and the load rule. Done when this sums 1 to 4 into o:

%x   = buffer 4 dtype=i32 slot=0 data=[1 2 3 4]
%o   = buffer 1 dtype=i32 slot=1
%ov  = shrink %o 0 shape=[]
%z   = store %ov 0
%o0  = after %o %z
%k   = range 4
%xv  = shrink %x %k shape=[]
%xk  = load %xv
%ov0 = shrink %o0 0 shape=[]
%cur = load %ov0
%sum = add %cur %xk
%s   = store %ov0 %sum
%e   = end %s %k
%o1  = after %o0 %e
%ov1 = shrink %o1 0 shape=[]
%r   = load %ov1
%t   = assert %r 10

If you get 4, %cur was loaded once before the loop: its only use, %sum, depends on %k, and the loop stores into its version, so it must move in.

Linearize. Done when the suite's old_version_after_loop passes: a load of the version before a loop, used after the loop, must run before the loop overwrites it.
Execute in run. Done when every week 3 run test passes and weeks 1 and 2 still do.
Render. Done when python _reference/uirc.py render and your render agree on the program from step 4 (same results, not necessarily the same text), and every render test passes.

The pinned load

If no use of a load depends on the inner range, the load stays outside the inner loop. In quiz 3 below, %cur is read through a view that depends only on i, and its use %nx = add %cur 1 doesn't depend on j either, so the load runs once per i iteration, not once per j iteration. The fix is to make the view (and so the load) depend on the inner range.

Check yourself

How many times does the store run, with %n = add 0 -3 and %i = range %n?

What does a range do when its bound is zero or less?

A negative bound runs zero times, just like range 0. The buffer keeps its data.

With %k = range 5, is shrink %b (mul %k 2) shape=[2] provably inside a 10-element buffer? (Written as two statements.)

What's the interval of %k, and what does mul by 2 do to it?

%k is in [0, 4], so the offset is in [0, 8]. A 2-element view at offset 8 ends at 10, which fits.

This program counts with two nested loops (3 and 4 trips). What does %r end up as?

%o   = buffer 1 dtype=i32 slot=0
%i   = range 3
%j   = range 4
%z   = mul %i 0
%ov  = shrink %o %z shape=[]
%cur = load %ov
%nx  = add %cur 1
%jj  = mul %j 0
%nx2 = add %nx %jj
%s   = store %ov %nx2
%ej  = end %s %j
%ei  = end %ej %i
%o1  = after %o %ei
%ov1 = shrink %o1 0 shape=[]
%r   = load %ov1

Which loop is %cur placed in? Its use %nx depends on %i but not on %j.

So %cur runs once per i, and the j loop stores the same value four times.

3. The view depends only on %i, and %cur's only use (%nx) depends only on %i, so the load runs once per i iteration. The j loop writes cur + 1 four times. Making the view's offset depend on %j too (add %z (mul %j 0)) moves the load into the j loop and gives 12.

An f32 accumulator starts at x[0] and adds x[1], x[2], x[3] in a loop over k. With x = [1e8 1.0 -1e8 1.0], what's the result?

1e8 + 1.0 in f32: is 100000001 representable?

1e8 + 1.0 rounds back to 1e8 (f32 has about 7 decimal digits). Then - 1e8 gives 0.0, and + 1.0 gives 1.0.

Which program is invalid?

One of these has an end that closes a range its effects don't depend on.

An end must close a range its effects use. A computed bound is fine (it has no open ranges), one end may repeat several effects, and an inner loop body may ignore an outer range: it just reruns.

A load of buffer version %b is used both inside a loop that stores into %b and in an assert after the loop. Where does it run?

What is the innermost loop that contains every use?

The uses are split, so the innermost loop holding all of them is the outer block. The load runs once, before the loop, and the use inside the loop sees the same pre-loop value every iteration.

In a looped matmul over i < 3, j < 4 and k < 5, how many times does the reset store C[i*4 + j] = 0 run?

The reset depends on i and j but not on k. Which loops is it inside?

It depends on i and j, so it runs once per (i, j) pair: 3 · 4 = 12 times, before each k loop.

More practice: the week 3 quiz review.

Go deeper

tinygrad and its documentation. RANGE and END come from tinygrad's IR, and its codegen places values by their ranges just as lesson 3.5 does.
Loop-invariant code motion. The classic name for what placement gives you for free: values that don't depend on a loop leave it.
Loop nest optimization. Tiling and interchange, the transformations that make a looped matmul fast (week 6 onward).
Fast multidimensional matrix multiplication on CPU from scratch. Takes exactly this week's triple loop and speeds it up step by step.
What Every Programmer Should Know About Memory (Drepper). Why loop order changes speed but not results. Read the cache chapters.
CS231n: Convolutional Neural Networks. The convolution arithmetic of lesson 3.4, with pictures, from the machine learning side.
collections. defaultdict and Counter, the two containers that make grouping nodes by loop short.
Set.Make. Immutable sets of range ids, with the union and subset tests that placement needs.
Python and OCaml, deeper. This week's idioms in your language: sets of ranges, bucketing nodes by loop, and emitting nested loops.

Challenges

Loop order without changing results. Write the 2×3 by 3×2 matmul twice, once with the k loop innermost and once with j innermost. Check with the oracle that both leave identical bits. Then explain why an f32 matmul would change bits if you moved k outermost and accumulated partial sums differently.
A strided view. Write a loop that copies every other element of an 8-element buffer into a 4-element one. Prove the offsets in bounds in your head first, then confirm the oracle accepts it, and that changing the stride to 3 is rejected.
A 2D running max. For a 3×4 buffer, write each row's running maximum into an output of the same shape, with a reset per row. Compare your expected values (computed by hand or in Python) with python _reference/uirc.py run.