Optimizing matrix multiplication on modern CPUs
going from naive matrix multiplication to avx2 microkernels, openmp, cache blocking, prefetching, compiler flags, and eventually comparing pico with openblas and intel mkl.
i have spent a very unhealthy amount of time staring at matrix multiplication code.
not because i think i'm going to beat intel mkl in a weekend, i'm not insane. but because gemm is one of those problems where the simple version and the fast version look like completely different species of code.
the naive version is one of the first things you write when learning arrays:
for(int i = 0; i < M; i++)
for(int j = 0; j < N; j++)
for(int k = 0; k < K; k++)
C[i][j] += A[i][k] * B[k][j];and then somehow production blas libraries are doing the same math but 50x, 100x, sometimes 200x faster, and that made me very curious, also paired with the fact that gemm might be one of the most important algorithms today due to the fact that it makes up for most of the flops powering ai models today, i decided to take my time with it.
so inside pico, my tiny c tensor library, i decided to build the matmul kernel up from the stupid version into something that at least looks like it belongs in the same conversation as openblas and intel mkl.
the goal was not really to beat blas. the goal was to understand the complexity. if i can explain why the fast code is fast, then the project worked.
the current matmul code lives in src/kernels/matmul, while the benchmark harness lives in bench/matmul. the benchmark logs are also saved in bench/results, including the cache blocking retrospective.
the baseline was very normal and very bad
the first implementation was just normal triple loop matmul. no simd, no cache tricks, no threading, no nothing.
for i
for j
for k
c[i][j] += a[i][k] * b[k][j]correct, readable, and slow.
the first issue is memory access. assuming row-major tensors, A[i][k] walks across a row, which is nice, but B[k][j] walks down a column if you use the classic ijk order. that means every load from B jumps by the row stride.
the cpu hates that.
modern cpus don't fetch one float at a time from memory. they fetch cache lines. if your code walks through memory contiguously, the cpu can bring in nearby values and your next loads are already waiting for you. if your code jumps around, the cache becomes less helpful and you spend more time waiting.
so the first meaningful change was loop ordering.
ikj is the first "oh wait" moment
instead of:
i -> j -> kyou can do:
i -> k -> jwhich turns the inner loop into:
for(int i = 0; i < M; i++) {
for(int k = 0; k < K; k++) {
float a_val = A[i][k];
for(int j = 0; j < N; j++) {
C[i][j] += a_val * B[k][j];
}
}
}this looks like a small change, but it changes the personality of the loop. now, for a fixed i and k, we load one scalar from A, then walk across a row of B and a row of C.
that means:
B[k][j]is contiguousC[i][j]is contiguousA[i][k]gets reused across many columns
same math, much nicer memory behavior. the cpu gets to stream through the data instead of jumping around like it forgot where everything lives.
this is the first lesson i had to internalize: algorithmic complexity is not the whole story.
O(n^3)can still have wildly different performance depending on how carefully you walk through memory.
then came register accumulation
the next annoying thing was writing partial sums back to memory too often.
for each output value, the naive version conceptually does:
load c
add a*b
store c
load c again
add a*b
store c again
...that is rude.
if the cpu has registers, use them. so the better pattern is:
load c once
accumulate in a register for the whole k loop
store c oncethis is one of those optimizations that sounds obvious after you understand it, but it matters a lot because memory is slow compared to registers.
in matrix multiplication, the output tile is precious. you want to keep that tile close, update it many times, and only write it back when you're done.
this idea later becomes the heart of the microkernel.
avx2
after the scalar loop started making sense, i moved to avx2.
with avx2, one __m256 register holds 8 floats.
so instead of doing:
c[j] += a * b[j]one float at a time, the kernel does:
c[j..j+7] += a * b[j..j+7]the pattern is:
broadcast one scalar from A
load 8 contiguous floats from B
fma into 8 C valuesroughly:
__m256 a_vec = _mm256_broadcast_ss(&A[i][k]);
__m256 b_vec = _mm256_loadu_ps(&B[k][j]);
acc = _mm256_fmadd_ps(a_vec, b_vec, acc);the important instruction here is fma: fused multiply-add. conceptually it is just acc = (a * b) + acc, but the cpu performs the multiply and the add as one fused operation, which is exactly the shape of matrix multiplication's inner loop.
and yes, this still has to be assigned back to acc. fma doesn't magically mutate your variable in c. the intrinsic returns the new vector value, and you keep it.
this was one of those tiny things that made me laugh because the hardware instruction is doing multiply-add, but the c intrinsic still behaves like a function. you don't escape c semantics just because the instruction is sexy.
the first avx family in pico was the x8 family: 1x8, 2x8, 4x8, and 8x8.
the files are here:
each kernel computes some number of rows and 8 columns at a time. for example, the 8x8 idea is:
8 rows of C
8 columns of C
64 output values
kept in avx accumulators while reducing over Kthis is where matmul starts feeling less like a loop and more like a tiny machine.
microkernels are just register budgeting
at first i thought "bigger tile should be faster", because more work per kernel call sounds good.
yh, no.
avx2 gives you 16 ymm registers. that is really not a lot unfortunately.
if your microkernel wants too many accumulators, plus broadcast registers, plus b vectors, plus address temporaries, the compiler starts spilling registers to the stack. now your "optimized" kernel is secretly doing memory traffic again.
that was the big register pressure lesson.
the early 8x8 kernel was useful, but it was not obviously the best shape. so i started reading other small blas-like kernels and found a different layout:
6 rows x 16 columnsthat sounds larger, but the register usage is cleaner.
each row gets two 8-wide accumulators:
row 0: acc0_0, acc0_1
row 1: acc1_0, acc1_1
...
row 5: acc5_0, acc5_1so 6 rows x 2 vectors = 12 accumulator registers, then you still have a few registers left for broadcast/load work. that spare room matters because the inner loop still needs to load two B vectors and broadcast values from A; if the accumulators already eat the entire register file, the compiler has nowhere comfortable to put the rest of the computation.
the actual shape of the 6x16 hot loop became something like this:
__m256 b_vec_0 = _mm256_loadu_ps(&B[k][j + 0]);
__m256 b_vec_1 = _mm256_loadu_ps(&B[k][j + 8]);
__m256 a0 = _mm256_set1_ps(A[i + 0][k]);
__m256 a1 = _mm256_set1_ps(A[i + 1][k]);
acc0_0 = _mm256_fmadd_ps(b_vec_0, a0, acc0_0);
acc0_1 = _mm256_fmadd_ps(b_vec_1, a0, acc0_1);
acc1_0 = _mm256_fmadd_ps(b_vec_0, a1, acc1_0);
acc1_1 = _mm256_fmadd_ps(b_vec_1, a1, acc1_1);the nice thing here is that the same two B vectors feed every row accumulator. for one k, load B[k][j..j+15] once, broadcast each needed A scalar, and update 12 accumulators. this is the point where the kernel started looking closer to the small examples i had seen in blas tutorials, not because the code was identical, but because the data movement pattern finally made sense.
this became the 16x family:
this was one of the biggest architecture changes in the kernel. instead of trying to make the old 8x8 family win forever, the code now had separate families so they could be benchmarked properly.
one of the later avx_kernels benchmark runs made the direction pretty obvious:
| shape | 8x8-family | 16x-family | winner |
|---|---|---|---|
256^3 | 194.12 GFLOP/s | 260.79 GFLOP/s | 16x-family |
512^3 | 236.69 GFLOP/s | 247.92 GFLOP/s | 16x-family |
1024x128x64 | 153.60 GFLOP/s | 239.48 GFLOP/s | 16x-family |
64x128x1024 | 112.07 GFLOP/s | 162.41 GFLOP/s | 16x-family |
70x70x70 | 29.18 GFLOP/s | 46.11 GFLOP/s | 16x-family |
the benchmark is noisy, and tiny shapes can move around depending on the run, but the trend was enough for me. the 16x family gave the compiler a better register story and usually gave the cpu more useful work per trip through the inner loop.
the funny thing is that "6x16" looks weirder than "8x8", but it maps better to the register file. the hardware really does not care about your sense of symmetry and shape.
tails are where clean code goes to suffer
real matrices are not always divisible by 6 or 16.
so after the main 6x16 kernel, the remaining rows/columns still have to be handled.
the rough strategy became:
try the largest clean microkernel first
then smaller vector tails
then scalar fallbackso for columns, if a matrix has 31 columns:
first 16 columns -> 6x16
remaining 15 -> 6x8 + scalar tailand for rows, if there are not enough rows for 6, the dispatcher tries smaller row kernels before falling back to scalar. this is not glamorous, but it's the difference between "fast for nice benchmark sizes" and "actually works for arbitrary tensor shapes".
then threading happened
once a single thread was doing useful work, the next obvious question was: why use one core? matrix multiplication is convenient here because different output rows are independent. if thread 0 writes rows 0..255 and thread 1 writes rows 256..511, they don't need to synchronize while computing. they both read A and B, and they write separate parts of C.
so the first threading experiment used pthreads and later a small thread pool.
that worked, but it also created a lot of overhead and complexity:
- job submission
- waiting
- mallocs around task args
- threadpool lifecycle
- more code than the matmul idea deserved
eventually i moved the matmul parallelism to openmp. not because openmp is magic, but because this exact pattern:
#pragma omp parallel for
for(row blocks)
compute independent output rowsis basically what openmp is good at. also openmp already manages worker threads internally, which means pico does not need to pretend its tiny thread pool is smarter than a mature runtime.
i still like the threadpool because it taught me what the overhead looked like. but for this kernel, openmp is the better tool.
the cache blocking confusion
cache blocking was the most humbling part. i had read enough to know that serious gemm implementations use blocking, so i tried to add:
for ii in row blocks
for jj in column blocks
for kk in k blocks
compute blockconceptually, this is correct. you break the matrices into smaller chunks so the cpu can reuse data while it is still in cache. but my first implementation had a very simple bug: the inner loops compared against the block starts instead of the block ends.
so for the first block:
ii = 0
jj = 0
kk = 0and the loop conditions were basically:
i < ii
j < jj
k < kkwhich means the loops did not run. C stayed empty. very funny. very painful.
the corrected shape is:
i_end = min(ii + block_size, rows)
j_end = min(jj + block_size, columns)
k_end = min(kk + block_size, k_dim)then every loop walks from the block start to the clamped block end.
there was also another subtle bug later: after adding k_start and k_end, one path still restarted k from zero. this passed small tests because small matrices only had one k block. then bigger shapes exposed the lie.
this is why performance work without shape-specific tests is unserious. the bug will politely wait until the benchmark size changes.
i wrote the retrospective here: cache blocking matmul retrospective
what cache blocking is actually doing
the easiest way i understand it now is this: the microkernel is a small register block, while cache blocking is a larger memory block. they solve different levels of the same problem.
the microkernel says:
while computing this tiny C tile,
keep the outputs in registers
reuse A and B as much as possiblecache blocking says:
while computing this region of C,
reuse a region of A and B while they are still in cache
before moving to far away memoryso no, 6x16 is not "cache blocking" by itself. it is register blocking. it blocks the work at the register level, but the macro loop can still stream across a massive region of B and then come back later. proper cache blocking controls the outer traversal so a panel of B and a block of A get reused before they fall out of cache.
the thing that made it tricky in pico is that the original kernel had a nice property:
load C tile once
accumulate across full K
store C tile onceif you split K into blocks carelessly, you can accidentally turn that into:
load C tile
accumulate partial K
store C tile
load C tile again
accumulate next partial K
store C tile againwhich means cache blocking can make things worse if the rest of the macro-kernel isn't designed around it.
eventually, after fixing the loop structure properly and combining it with the newer 16x family, prefetching, and openmp, cache blocking became useful. but the first version losing was not surprising anymore. it was the wrong interaction with the microkernel.
prefetching was annoyingly useful
i also tried __builtin_prefetch. the intuitive version was to prefetch inside the hot k loop:
__builtin_prefetch(&B[k + distance][j], 0, 3);that is easy to understand:
while computing with the current row of B, ask the cpu to start fetching a future row of B.
but putting prefetch inside the hottest loop can be expensive too. if you add extra instructions to the loop that already runs billions of times, you better be sure those instructions are mad useful.
the version that worked better was panel prefetching before entering the microkernel work. instead of prefetching every single k iteration, the kernel gives the cpu an early hint that a whole future region of B will be needed soon.
roughly:
before computing this j panel,
touch the B rows for the upcoming k block and column rangethis is less obvious than hot-loop prefetching, but it avoids adding prefetch instructions directly into the tightest fma loop.
this one is going in the "things i did not expect to help this much" folder. prefetching felt like fake optimization until the numbers moved.
the compiler flags were embarrassing
the most annoying discovery was that i spent all this time thinking about microkernels, prefetch distance, cache blocks, thread counts, and openblas comparisons, then realized the benchmark was not being built aggressively enough.
the serious command became:
taskset -c 0-11 env \
OMP_NUM_THREADS=4 \
OPENBLAS_NUM_THREADS=4 \
OMP_DYNAMIC=false \
OPENBLAS_DYNAMIC=0 \
make -C bench matmul_focused_openblas \
CFLAGS="-std=c11 -O3 -march=native -ffast-math -I ../src -Wall -pthread -fopenmp"-O3, -march=native, and -ffast-math matter here.
especially -march=native, because the compiler is then allowed to target the actual cpu instead of emitting a generic binary. when you're writing avx2/fma code and then asking the compiler to optimize around it, the surrounding scalar/index/address code matters too.
imagine doing cpu optimization and forgetting to let the compiler optimize for your cpu. wallahi i was fighting ghosts.
the go-all-out log is here: matmul_go_all_out_2026-07-25.md
best saved openblas comparison from that run:
| shape | pico-avx | openblas | pico/openblas |
|---|---|---|---|
256^3 | 141.08 | 147.00 | 96.0% |
512^3 | 156.56 | 171.38 | 91.4% |
768^3 | 126.24 | 139.20 | 90.7% |
1024^3 | 104.78 | 143.01 | 73.3% |
512x1024x2048 | 84.67 | 134.70 | 62.9% |
2048x1024x512 | 112.93 | 119.70 | 94.3% |
not beating openblas everywhere, but close enough in some shapes that the implementation stopped feeling like a toy.
then i compared with intel mkl
after openblas, i added a separate benchmark target for intel mkl: bench_matmul_mkl.c. the mkl bench reuses the same focused matmul harness, so the shapes and timing structure stay identical.
the command looked like:
taskset -c 0-11 env \
MKLROOT=/opt/intel/oneapi/mkl/latest \
OMP_NUM_THREADS=4 \
MKL_NUM_THREADS=4 \
OMP_DYNAMIC=false \
MKL_DYNAMIC=false \
make -C bench matmul_mkl \
CFLAGS="-std=c11 -O3 -march=native -mavx2 -mfma -I ../src -Wall -pthread -fopenmp"one of the mkl runs looked like this:
| shape | pico-avx | intel mkl | pico/mkl |
|---|---|---|---|
256^3 | 138.25 | 155.64 | 88.8% |
384^3 | 147.47 | 144.97 | 101.7% |
512^3 | 154.84 | 169.55 | 91.3% |
640^3 | 155.55 | 164.98 | 94.3% |
768^3 | 147.03 | 169.44 | 86.8% |
960^3 | 149.54 | 171.41 | 87.2% |
1024^3 | 125.07 | 162.13 | 77.1% |
1280^3 | 136.00 | 171.84 | 79.1% |
512x1024x2048 | 109.24 | 155.16 | 70.4% |
2048x1024x512 | 153.01 | 162.51 | 94.2% |
these numbers are noisy, so i would not overclaim from one run. but still, getting a tiny c tensor library's matmul kernel within shouting distance of intel mkl on some shapes was insane to see, especially because the path there was not magic:
better loop order
register accumulation
avx2/fma
microkernels
tail kernels
openmp
cache blocking
panel prefetch
compiler flags
benchmark disciplinethat's it. painful, but not mystical.
why mkl is still better
this is the part where i have to be honest: intel mkl is not just "a faster microkernel". it has years of engineering around:
- packing panels of
AandB - architecture-specific kernels
- better thread scheduling
- cache hierarchy tuning
- many matrix shapes
- transposed inputs
- different datatypes
- edge cases i have not even thought about yet
pico does not have that. pico currently has a decent avx2 path, some cache blocking, openmp parallelism, and a focused benchmark harness. that is already a lot, but it is not a full blas implementation.
the worst pico shapes are still informative:
- large square matrices like
1024^3 - very wide outputs like
512x1024x2048
those shapes likely need a more serious macro-kernel and packing strategy, not just more tiny tweaks.
packing is probably the next real gemm frontier: copy panels of B or A into a layout that the microkernel can consume with less addressing pain and better cache reuse.
i tried (asked codex to speedrun) a quick packed-b experiment, and it was correct but slower. that does not mean packing is bad. it means cheap packing bolted onto the side is bad.
real packing has to be part of the macro-kernel design.
the final shape of the lesson
the biggest thing i learnt is that high performance matmul is not one trick. it's layers. you don't "add simd" and suddenly become blas.
you fix memory traversal. then you keep partial sums in registers. then you choose a register tile that fits the architecture. then you handle tails without destroying the fast path. then you parallelize over independent output regions. then you manage cache reuse. then you compile it properly. then you benchmark it honestly and realize half your assumptions were fake, tbh and that's what made this fun.
the final implementation is still simple enough that i can understand it. it is not openblas. it is not mkl. but it is no longer naive, and it is no longer pretending, and my friends cannot call it a toy anymore.
for pico, that is probably good enough for now, there's alot of work on the platform itself that i can't keep postponing if i want to achieve my small goal of training an llm on pico itself.
i'll probably come back later for packing, better macro-kernels, and maybe avx-512 if i hate myself enough.
thank you for taking the time to read this.
links
- pico repo: https://github.com/sijirama/pico
- matmul kernels: https://github.com/sijirama/pico/tree/main/src/kernels/matmul
- matmul benches: https://github.com/sijirama/pico/tree/main/bench/matmul
- focused benchmark: https://github.com/sijirama/pico/blob/main/bench/matmul/bench_matmul_focused.c
- intel mkl benchmark target: https://github.com/sijirama/pico/blob/main/bench/matmul/bench_matmul_mkl.c