04 · Matmul & GEMM

GEMM — GEneral Matrix Multiply — is the single kernel that consumes >90% of the compute in any transformer forward pass. If you can write a competent GEMM, you can write anything numerical. If you cannot, everything downstream (attention, convolutions, quantized inference) is a black box you cannot open. This file walks the exact ladder every OpenBLAS/BLIS contributor climbed — naive → cache-blocked → SIMD micro-kernel → threaded — and locks the target: a 4× speedup over -O3 scalar on your machine, benchmarked against your system’s OpenBLAS.

The problem, precisely

C = A × B + C, where A is M×K, B is K×N, C is M×N. All FP32, row-major, contiguous. Work = 2·M·N·K FLOPs (each output cell does K multiplies + K adds). For M=N=K=1024, that’s ~2.1 GFLOPs per call.

Rung 1: naive triple loop

for (int i = 0; i < M; i++)
    for (int j = 0; j < N; j++) {
        float acc = C[i*N + j];
        for (int k = 0; k < K; k++)
            acc += A[i*K + k] * B[k*N + j];
        C[i*N + j] = acc;
    }

On a 3 GHz core: expect ~0.5–2 GFLOPS for M=N=K=1024. Peak of the ALU is ~50–100 GFLOPS single-thread. You are at 1–4% of peak. Where does the time go? Every B[k*N + j] walks down a column of B in memory. B is row-major, so each access is N*4 = 4096 bytes apart — a new cache line every iteration. K misses per output cell × M·N outputs = catastrophe.

Rung 2: loop reorder (i-k-j)

Swap the two inner loops:

for (int i = 0; i < M; i++)
    for (int k = 0; k < K; k++) {
        float a_ik = A[i*K + k];
        for (int j = 0; j < N; j++)
            C[i*N + j] += a_ik * B[k*N + j];
    }

Now the inner loop walks along a row of B (contiguous) and a row of C (contiguous). One line of B feeds N/16 vector accumulates. Expect 5–15 GFLOPS — an easy 5–10× win from a single change no one taught you in school. The compiler will auto-vectorize this cleanly with -O3 -march=native. Verify with -fopt-info-vec.

Rung 3: cache blocking (tiled matmul)

Even with the reorder, once B doesn’t fit in L2, you re-read it from DRAM for every row of A. Fix: process A, B, C in tiles that fit in cache.

Three loop nests deep:

#define MC 256   // rows of A per L2 block
#define KC 256   // shared dim per L2 block
#define NC 1024  // cols of B per L3 block
#define MR 8     // rows in register block
#define NR 8     // cols in register block

for (int jc = 0; jc < N; jc += NC)
  for (int pc = 0; pc < K; pc += KC)
    for (int ic = 0; ic < M; ic += MC)
      // pack A block into contiguous buffer
      // pack B block into contiguous buffer
      for (int jr = 0; jr < NC; jr += NR)
        for (int ir = 0; ir < MC; ir += MR)
          // micro-kernel: MR × NR × KC using registers only

This is the Goto algorithm (Kazushige Goto, 2008; the intellectual DNA of GotoBLAS → OpenBLAS → BLIS). Tune MC, KC, NC to your cache sizes. On an M-class Apple Silicon core: 128 KB L1D, 4 MB L2 per core cluster. On a Ryzen: 48 KB L1D, 1 MB L2. Read them with sysctl hw.l1dcachesize (Mac) or lscpu (Linux).

Expected: 20–40 GFLOPS on one core — 30–60% of peak. This is where you fight OpenBLAS.

Rung 4: SIMD micro-kernel

The innermost MR × NR block should live entirely in vector registers. On AVX2 with 16 ymm registers of 8 floats each, MR=8, NR=8 fits perfectly: 8 accumulator vectors (__m256), plus room for one row of A broadcast and one row of B loaded. Pattern:

__m256 c0 = _mm256_load_ps(&C[0*N]);
// ... c1..c7
for (int p = 0; p < KC; p++) {
    __m256 b = _mm256_load_ps(&Bp[p*NR]);
    __m256 a0 = _mm256_broadcast_ss(&Ap[0*KC + p]);
    c0 = _mm256_fmadd_ps(a0, b, c0);
    // ... a1..a7 broadcast, fmadd into c1..c7
}
_mm256_store_ps(&C[0*N], c0);
// ... c1..c7

On NEON (Apple Silicon) the analog is vfmaq_f32 with float32x4_t; MR/NR are usually 4×8 or 8×8 to match 32 128-bit v registers.

Expected: 50–80 GFLOPS single-thread on modern desktop. This is roughly 70–85% of OpenBLAS on the same core. The remaining gap is packing efficiency, instruction scheduling, and prefetch hints — things OpenBLAS spent 15 years tuning.

The four things OpenBLAS/BLIS still do that you probably won’t

  1. Assembly micro-kernels per micro-architecture. BLIS has hand-written kernels for Haswell, Skylake, Zen, Zen2/3/4, Cortex-A76, Neoverse-N1, Apple Firestorm. Look at blis/config/ in the BLIS repo.

  2. Runtime CPU dispatch. One binary, N kernels; a startup probe picks the right one via cpuid or getauxval(AT_HWCAP).

  3. NUMA-aware packing. On multi-socket boxes, pack buffers on the same NUMA node as the thread.

  4. Perfect prefetch schedules. Tuned to hide the L2→L1 latency of the next block while computing the current block.

Do not chase these in M10–M11. Match steps 1–3 and note the gap. That is the honest engineering.

Benchmarking rules (non-negotiable)

  • Warm up: run the kernel 3× before timing.

  • Time only the compute, not allocation.

  • Report GFLOPS = 2·M·N·K / seconds / 1e9. Round to one decimal.

  • Vary matrix size: sweep 64, 128, 256, 512, 1024, 2048. Plot on log-x, linear GFLOPS.

  • Report against OpenBLAS on your machine (brew install openblas on Mac, apt install libopenblas-dev on Ubuntu). Call cblas_sgemm.

  • Report against numpy.dot, which usually calls the same OpenBLAS — useful sanity check.

  • Pin to one core with taskset -c 0 ./bench (Linux) so you’re not measuring scheduler noise.

Threaded GEMM (only after single-thread is good)

Parallelize the outer jc loop with OpenMP: #pragma omp parallel for schedule(static). Expect near-linear scaling up to ~4 cores, then diminishing returns from L3 contention and memory bandwidth saturation. On a 16-core Ryzen you’ll typically max at ~8–10× not 16×. This is a bandwidth wall, not a compute one — back to the roofline in 03_cache_and_memory_hierarchy.md.

What most people get wrong about this

They think the win comes from SIMD. It comes from cache-blocking. The SIMD micro-kernel is the last 2×; the packing and blocking are the first 20×. If your naive scalar is 1 GFLOPS and your final SIMD-blocked is 60 GFLOPS, the SIMD contributed a factor of ~3, the packing contributed a factor of ~4, and the blocking contributed a factor of ~5. Add SIMD to a non-blocked matmul and you get maybe a 2× improvement over an already-slow baseline — still slow. Blocking is the fulcrum. Read Kazushige Goto’s original 2008 paper “Anatomy of High-Performance Matrix Multiplication” (TOMS) before you start writing.


Return to README.md · Next: 05_quantization_kernels_in_c.md