Caches and tiling: reuse what you've loaded

The memory hierarchy

Lesson 0.2 introduced caches. Here are the levels with rough sizes and costs on a current desktop processor. The exact numbers vary, the ratios don't:

Memory moves in cache lines of 64 bytes, 16 floats. Reading one float brings its 15 neighbors along for free. Two things make a kernel fast: using all 16 floats of every line you load (walking memory in order), and using each value many times while it's still close (reuse).

Where a plain GEMM wastes time

The plain week 3 GEMM computes one output at a time: for each (i, j), a loop over k reads row i of A and column j of B. Every multiply needs two loads. Reading down a column of B jumps N floats per step, so each load touches a new cache line and uses one float of it. And row i of A, just loaded, is used once and then reloaded for the next j.

Tiling

Compute a tile of outputs at once, say 4 rows by 8 columns. For each k, load 4 values of A (one per row of the tile) and 8 of B (one per column), and do all 32 multiply-adds they make possible. That's 12 loads for 32 multiply-adds instead of 64 loads. The 8 values of B sit next to each other, so they come from one or two cache lines. And the 32 running sums can live in registers for the whole k loop.

Why the sums don't change

Each output C[i][j] is still its own sum over k in the same order: A[i][0]·B[0][j] first, then + A[i][1]·B[1][j], and so on. Tiling only changes which outputs are worked on side by side. The order of the adds inside any one sum is untouched, so every result is bit for bit the same as the plain loops'. Reordering k itself, say summing in blocks, would change the bits, and the spec doesn't allow it.

A tiled GEMM in .uir

The week 7 tests write the tile out explicitly. Ranges ib and jb pick the tile, ki and kj (4 and 8) index inside it, and a 32-element alloc holds the tile's sums:

%acc = alloc 32 dtype=f32
%ib  = range 192            # tile row: rows 4ib .. 4ib+3
%jb  = range 96             # tile column: columns 8jb .. 8jb+7
# start each sum from its first product, A[i][0]·B[0][j]
...
# then for k = 1 .. 767: for each (ki, kj) in the tile, acc[ki*8 + kj] += A[i][k]·B[k][j]
%k0  = range 767
%k   = add %k0 1
%ki  = range 4
%kj  = range 8
...
# finally write the tile to C

Each tile's sums start from the tile's own first product, not from zero, for two reasons. First, it's the same order reduce folds in, so the result equals the matmul sugar's bit for bit. Second, it's a placement rule from week 3: a store that doesn't depend on jb runs outside the jb loop. A reset to 0.0 would run once for the whole program instead of once per tile. The first product depends on the tile's indices, so it lands inside the loop where it belongs.

On its own, writing a tile as loops already helps a little. The big win comes in lesson 7.3, when the loops inside the tile are unrolled and the 32 sums become registers.