OPTIMIZATION LAB #01 — matmul: 1 GFLOP/s → 40 GFLOP/s, and proving which step did it

TRACK: optimization | DIFFICULTY: 3 | NEEDS: mlir-opt, mlir-translate, clang, taskset FILES: matmul.mlir (payload), bench_matmul.c (harness), sched/gen.sh (schedule generator), run_matmul.sh (driver)

What this teaches

The usual shape of a perf task on a production MLIR backend: "kernel K is fast with X, slow without X — find X, explain it, make the compiler always do X." This lab runs that loop end to end on the most important kernel there is:

1. Ablate. Build the SAME kernel with the optimisation stages switched on one at a time. Each row of the table differs from the previous one in exactly ONE thing, so each speedup has exactly one cause.
2. Explain each step with evidence from the artifacts, not with a story: instruction counts, spill counts, the loop structure in the IR.
3. Sweep the knob (register tile size), and learn to tell a real effect from noise.

The stages (all driven by a transform script, lab 03_linalg/03):

STEP 1 — Look at the schedule the driver generates

$ sched/gen.sh T V H TM=4 TN=16 | grep -E 'tile_using_for|vectorize|hoist|lower_contraction'
    %t, %loops:3 = transform.structured.tile_using_for %mm tile_sizes [4, 16, 1]
    transform.structured.vectorize %t : !transform.any_op
    %f2 = transform.structured.hoist_redundant_vector_transfers %func : (!transform.any_op) -> !transform.any_op
      transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"

sched/gen.sh prints the transform script for any subset of stages. Read the full output once: it's 40 lines and it is the whole optimiser.

STEP 2 — The ablation table

./run_matmul.sh

Real run on this machine (i7-1260P, one P-core pinned, -march=native = AVX2+FMA):

A_naive                      min   230.79 ms   p50   235.93 ms     1.16 GFLOP/s   maxerr 0.0e+00 OK
B_tile_only                  min    84.43 ms   p50    86.88 ms     3.18 GFLOP/s   maxerr 0.0e+00 OK
C_tile_vec                   min    17.97 ms   p50    18.42 ms    14.94 GFLOP/s   maxerr 0.0e+00 OK
D_tile_vec_hoist             min     8.98 ms   p50    10.21 ms    29.89 GFLOP/s   maxerr 0.0e+00 OK
E_full_noFMA                 min    16.07 ms   p50    16.48 ms    16.70 GFLOP/s   maxerr 0.0e+00 OK

(Representative: your times will differ; the ratios between rows are what matters.) Every row prints maxerr ... OK. A speedup without a correctness check is not a result.

STEP 3 — Explain every row with evidence

A → B (3×): tiling alone, and NO SIMD. Count packed (...ps) vs scalar (...ss) float instructions:

$ for v in A_naive B_tile_only D_tile_vec_hoist; do printf "%-18s packed=%s scalar=%s\n" $v $(clang -O2 -march=native -S out/$v.ll -o - -Wno-override-module | grep -cE '^\s+v\w+ps\s') $(clang -O2 -march=native -S out/$v.ll -o - -Wno-override-module | grep -cE '^\s+v\w+ss\s'); done
A_naive            packed=0 scalar=66
B_tile_only        packed=0 scalar=64
D_tile_vec_hoist   packed=180 scalar=16

A and B are both 100 % scalar. So B's 3× comes from memory order, not from vectors. Look at B's loops:

$ mlir-opt out/B_tile_only.t.mlir --convert-linalg-to-loops | grep -E 'scf.for|memref.load' | sed 's/ :.*//' | awk '{$1=$1};1' | tail -9
scf.for %arg3 = %c0 to %c512 step %c8 {
scf.for %arg4 = %c0 to %c512 step %c16 {
scf.for %arg5 = %c0 to %c512 step %c1 {
scf.for %arg6 = %c0 to %c8 step %c1 {
scf.for %arg7 = %c0 to %c16 step %c1 {
scf.for %arg8 = %c0 to %c1 step %c1 {
%0 = memref.load %subview[%arg6, %arg8]
%1 = memref.load %subview_0[%arg8, %arg7]
%2 = memref.load %subview_1[%arg6, %arg7]

In A (naive: m, n, k innermost) every k step walks down a column of B: a new cache line per multiply, 512 lines per dot product. In B, k sits outside an 8×16 tile, so each k step reads 16 contiguous floats of one B row, and the 8×16 C tile stays in L1. Same FLOPs, far fewer cache misses.

B → C (4.7×): vectorization in MLIR. 180 packed instructions appear. The tile body became vector.fma on vector<16xf32> (2 × ymm). LLVM didn't have to guess anything: the vector shapes came from the tile sizes.

C → D (2×): hoisting. Without H, every k iteration loads the C tile from memory, does one FMA per row, and stores it back: a load+store per FMA. With H the k loop carries C in registers:

$ grep -m1 -E 'scf.for .*iter_args' out/D_tile_vec_hoist.t.mlir | sed 's/^ *//'
%2 = scf.for %arg5 = %c0 to %c512 step %c1 iter_args(%arg6 = %1) -> (vector<8x16xf32>) {
$ grep -c 'iter_args' out/C_tile_vec.t.mlir
0

iter_args(... -> (vector<8x16xf32>): the accumulator IS the loop-carried value. That's the "microkernel" shape every BLAS uses.

D → E (1.8× slower): same IR, weaker target. -march=x86-64 = SSE2: 4-wide vectors and no FMA instruction (fma is emulated as mul + add). Identical MLIR, different machine flags. On AArch64 the same trap is forgetting +sve, +i8mm, +bf16, or -mcpu (lab 05_aarch64).

STEP 4 — The knob: register tile size

./run_matmul.sh sweep

Real run (GFLOP/s, min of 5):

TM\TN     8      16     32     64
 1       8.95  16.22  25.66  26.00
 2      16.75  26.62  34.15  23.51
 4      24.37  33.09  22.77  33.66
 8      35.98  30.45  33.94  17.45
12      (vectorize failed — see STEP 5)
16      26.26  41.42  32.85  27.18

Now the honest part: this table is mostly noise. Re-running four entries three times each:

tile_16x16   26.92 28.43 41.19      <- 1.5x spread on the SAME binary
tile_4x16    42.22 43.87 46.39
tile_4x32    25.97 26.25 29.46
tile_8x8     40.31 40.41 40.82

A laptop with turbo and a powersave governor moves 30–50 % between runs. Rules (same as ~/debug_lab/perf):

• Trust ratios ≥ 2× on one run. Anything under ~1.3× needs repeated, interleaved runs, or a noise-free metric.
• Get a noise-free metric from the artifact. Count spills:

$ for v in tile_4x16 tile_8x16 tile_16x16 tile_8x64 tile_16x64; do printf "%-12s spills=%s\n" $v $(clang -O2 -march=native -S out/$v.ll -o - -Wno-override-module | grep -c 'Spill$'); done
tile_4x16    spills=0
tile_8x16    spills=18
tile_16x16   spills=84
tile_8x64    spills=232
tile_16x64   spills=486

AVX2 has 16 ymm registers. A TM×TN f32 tile needs TM×TN/8 accumulators, plus a register for the B row and one for the broadcast A element:

• 4×16 → 8 accumulators: fits, 0 spills, and consistently the fastest (42–46).
• 8×16 → 16 accumulators: already over budget.
• 16×64 → 128 accumulators: 486 spills, i.e. the "registers" live on the stack.

On AArch64 NEON there are 32 × 128-bit registers (4 f32 each), so the sweet spot moves: 8×8 or 8×12 tiles are common. With SVE the tile is [8, [4]], "4 × vscale" wide (lab 05_aarch64). This is why IREE's KernelDispatch.cpp picks tile sizes per target, and why a production MLIR backend will need the same logic.

STEP 5 — A tile that doesn't divide the problem

$ ./run_matmul.sh one tile_12x16 "T V H TM=12 TN=16" 2>&1 | head -2
matmul.mlir:7:3: error: Attempted to vectorize, but failed
  linalg.matmul ins(%A, %B : memref<512x512xf32>, memref<512x512xf32>)

512/12 isn't whole, so the last tile is 8 rows: a dynamic shape. Plain vectorize needs static shapes. Three standard fixes, all used in real compilers:

1. Masking: transform.structured.vectorize %t vector_sizes [12, 16, 1] (vector ops with a mask for the ragged edge; SVE does this in hardware).
2. Peeling: split off the remainder loop (transform.loop.peel, IREE's LLVMCPUPeel.cpp).
3. Padding: pad the tile to 12 (transform.structured.pad).

Triton kernels hit this constantly: BLOCK_SIZE rarely divides the runtime n.

The method, to reuse on any kernel

1. Get a correct baseline and a harness that checks results (bench_*.c).
2. Write the optimisation as independent, switchable stages.
3. Ablate: turn stages on one at a time. The step with the jump is your X.
4. Explain each jump from artifacts (asm counts, spills, IR loop structure, perf stat counters), never from timing alone.
5. Sweep the parameters, then distrust differences smaller than your noise.

Try this yourself

1. Add a cache-level K tile: ./run_matmul.sh one kc128 "T V H TM=4 TN=16 KC=128". Does it help at 512? Change N in both files to 2048 and try again.
2. Use perf stat -e cycles,instructions,L1-dcache-load-misses on A_naive vs B_tile_only. Does the miss count explain the 3×?
3. Cross-compile D for AArch64 (lab 02_lowering/04) and count NEON fmla in the inner loop. Then change TN to 8 and TM to 8 and count again.