Upcast: unrolling loops into registers

axis=upcast

A range can now say how its loop runs:

%kj = range 8 axis=upcast

An upcast loop means the same thing as a plain loop, and run treats it exactly like one. The difference is in render: instead of a for loop, emit the body 8 times, with kj replaced by the literal 0, then 1, up to 7. The bound must be a literal from 1 to 32, because the body is copied that many times.

Here is the innermost loop of the tiled GEMM before and after, in the C your compiler emits (simplified):

/* plain loops: acc[] is indexed by variables, so it stays in memory */
for (int32_t ki = 0; ki < 4; ki++)
  for (int32_t kj = 0; kj < 8; kj++)
    acc[ki * 8 + kj] = acc[ki * 8 + kj] + A[(4*ib + ki) * K + k] * B[k * N + 8*jb + kj];

/* upcast: every index is a literal */
float a0 = A[(4*ib + 0) * K + k];   /* one A value per tile row: a0 .. a3 */
acc[0] = acc[0] + a0 * B[k * N + 8*jb + 0];
acc[1] = acc[1] + a0 * B[k * N + 8*jb + 1];
...
acc[31] = acc[31] + a3 * B[k * N + 8*jb + 7];

With constant indices and an array that never escapes, the C compiler replaces acc[0] to acc[31] with 32 variables, and those go in registers. The eight B values for one row of the tile sit next to each other in memory, so the compiler can also use SIMD instructions: one instruction that does 4 or 8 multiplies (or adds) on adjacent floats. SIMD works on separate outputs side by side, so every sum still adds in its own order, and the bits don't change.

What it's worth

Measured on the 768×768 GEMM, single-threaded, against the same loops without axes:

The two compilers differ a lot. Both get faster, and neither changes a bit of the result. The speed test for upcast alone, s_gemm_768_upcast, asks for at least 0.65 of the reference's speed. A renderer that turns upcast loops into plain for loops lands around 0.25 with clang and 0.55 with gcc, and fails.

Upcast and the accumulator rule

Week 3's rule for loads still decides where things run. In the tiled GEMM, the load of acc[ki*8 + kj] sits inside the k loop, because the k loop stores into the version it reads. Unrolling doesn't change that: it happens after linearizing, when you print C. Unroll in the renderer, not in the graph, and the rules from earlier weeks keep working untouched.