Week 5 · Lesson 5
Rangeify movement ops
Movement ops never compute anything, so rangeify can make them disappear. Each one becomes a little index arithmetic in front of a load: a multiply, a divide, a clamp. After rangeify, a matmul written as a broadcast cube is two loops (one over the flattened outputs, one over k) and a few integer ops, and the cube was never there.
Every op is a rule for indices
Week 4's rangeify asked each value one question: what is your element at index i? Week 5 asks the same question with an index tuple (i0, i1, ...), one entry per axis. A movement op answers by rewriting the tuple and passing the question to its operand:
| op | element at (i0, ..., ik) becomes |
|---|---|
reshape | flatten the index in the new shape (f = i0·stride0 + i1·stride1 + ..., row-major strides), then unflatten f in the old shape: axis k is (f div stride_k) mod size_k |
permute order=p | the operand at the index whose axis p_k is i_k: just a reordering, no arithmetic |
pad pads=[l r ...] | the operand at the clamped index clamp(i - l, 0, n - 1), then where(valid, that, 0), with valid = l ≤ i < l + n |
shrink | the operand at (i0 + o0, ...) |
expand | the operand with every size-1 axis read at 0 |
stack | item i0 at the rest of the index: a chain of wheres on i0, or a plain C array index when the items are scalars |
reduce pop=m | one accumulator loop over all m leading axes together: a flat counter t, unflattened into the m indices with div and mod |
Eventually every question reaches a load of a rank-1 buffer, which answers with a scalar load at offset + i. Everything above it was index arithmetic.
Pad: a clamp and a mask
Here is the oracle's C for pad %x pads=[1 2] on a 3-element x, stored into a 6-element buffer:
for (int32_t r9 = 0; r9 < 6; r9++) {
bool v11 = (0 < r9);
bool v12 = (r9 < 4);
bool v13 = (v11 ? v12 : false); /* valid: 1 <= r9 < 4 */
int32_t v14 = (int32_t)((uint32_t)r9 + (uint32_t)-1);
int32_t v15 = (v14 > 0 ? v14 : 0); /* clamp below */
int32_t v16 = (int32_t)((uint32_t)v15 * (uint32_t)-1);
int32_t v17 = (v16 > -2 ? v16 : -2);
int32_t v18 = (int32_t)((uint32_t)v17 * (uint32_t)-1); /* clamp above: min(v15, 2) */
float v20 = b0[v18];
float v21 = (v13 ? v20 : 0x0.0p+0f);
b1[r9] = v21;
}There's no min op, so the upper clamp is written as -max(-x, -(n-1)). Step through the loop to see what each iteration reads.
The where throws away the loaded value outside the real elements, so it's tempting to drop the clamp and read b0[r9 - 1]. But the load still runs, every iteration: at r9 = 0 it reads b0[-1], and at r9 = 5 it reads b0[4] of a 3-element buffer. Reading out of bounds is undefined behavior in C, even if the value is discarded. The clamp is what keeps every read inside the buffer.
Pad: clamp and mask
Write the two halves of a padded read. For a pad with l zeros on the left of an n-element x: pad_valid(i, l, n) is 1 when position i lands on a real element and 0 otherwise, and pad_src(i, l, n) is the index to read, i - l clamped into [0, n-1]. Positions may be anywhere, including negative.
The movement GEMM, rendered
And here is the syllabus GEMM from lesson 5.4 (K = 2, M = 2, N = 3), after rangeify. The store of the [2 3] result became one loop over its 6 elements, and the reduce pop=1 became a seeded accumulator and a loop over the remaining K - 1 = 1 values of k:
float a4[1] = {0};
for (int32_t r11 = 0; r11 < 6; r11++) { /* output element: (m, n) = (r11 / 3, r11 % 3) */
int32_t v13 = (r11 / 3);
int32_t v14 = (v13 % 2);
int32_t v15 = (v14 % 4);
float v17 = b0[v15]; /* a at k = 0 */
int32_t v18 = (r11 % 3);
int32_t v19 = (v18 % 6);
float v21 = b1[v19]; /* b at k = 0 */
float v22 = (v17 * v21);
a4[0] = v22; /* seed with the k = 0 product */
for (int32_t r26 = 0; r26 < 1; r26++) { /* k = 1 */
float v27 = a4[0];
int32_t v28 = (int32_t)((uint32_t)r26 + (uint32_t)1);
int32_t v29 = (v28 % 2);
...
float v40 = (v34 * v39);
float v41 = (v27 + v40);
a4[0] = v41;
}
float v45 = a4[0];
b2[r11] = v45;
}The [K M N] cube from lesson 5.4 is gone. Every element of it was a multiply of two loads, and rangeify only ever asked for the elements the reduce needed, one k at a time.
Cleaning up the index math
Look at v15 = (v14 % 4). v14 is (r11 / 3) % 2, which is always 0 or 1, so the % 4 does nothing. The same goes for v19 = (v18 % 6). These come from composing an unflatten with the next flatten, and week 1 already gave you the tool to remove them: bounds analysis proves v14 is in [0, 1], and the rule mod(x, c) → x when 0 ≤ x < c deletes the op. Constant folding (lesson 1.3) handles any arithmetic on constants that's left. The output is correct either way. Cleaner index math is smaller, and from week 6 it's faster.
The shape after a chain of movement ops
Given a starting shape and a list of movement ops, compute the final shape. Reshape and Expand give the new shape directly, Permute reorders axes, Pad adds a left and right amount per axis, and Reduce m drops the first m axes. The inputs are always valid.
Check yourself
A [6] value is reshaped to [2 3]. Element [1][2] of the result reads which flat position of the original?
Row-major strides for [2 3] are 3 and 1.
1·3 + 2·1 = 5.
A flat index f is unflattened in shape [2 3 4]. What is the coordinate on axis 1?
Axis 1's stride is the product of the sizes after it.
The stride of axis 1 is 4, so divide by 4, then take the coordinate mod its size, 3. For f = 23: (23 div 4) mod 3 = 2, and element 23 of a [2 3 4] is [1][2][3].
In the rendered pad loop above, which index of b0 does iteration r9 = 5 read?
5 - 1 = 4, then clamp into [0, 2].
2, the last real element. The mask then replaces it with 0, because position 5 is past the end of x.
Why does a rendered pad clamp the index when the mask already zeroes the value outside x?
Does where stop the load from happening?
No: both branches are evaluated, as in week 1. Without the clamp the loop reads b0[-1] and past the end, which is undefined behavior even if the value is thrown away.
In the rendered syllabus GEMM with K = 2, how many times does the inner accumulator loop run per output element?
The accumulator is seeded with the k = 0 product.
Once: r26 < 1. The reduce folds K values. The first seeds the accumulator, and the loop folds the remaining K - 1.
How does rangeify answer "your element at (k, m, n)" for expand %b3 shape=[K M N], where %b3 has shape [K 1 N]?
Which axis of %b3 has size 1?
The size-1 axis is read at 0, so m is ignored. No memory is involved: expand is pure index rewriting.