Week 0 · Lesson 5

Matrix multiply is a pile of dot products

You only need one piece of linear algebra for this course: the dot product, and matrix multiply as a grid of them. Most of the compute in a neural network is matrix multiply, so most of the speed work, once kernels are timed from week 6, is about this one loop.

The dot product

Two vectors of the same length. Multiply them element by element, then add up the products.

[1, 2, 3] · [4, 5, 6] = 1*4 + 2*5 + 3*6 = 32

That's a multiply followed by a reduce (sum). Hold on to that pattern. Later in this lesson it turns into two operations in the course's IR: a broadcast multiply and a reduce.

A strided dot product

Write dot(a, sa, b, sb, n): the sum of a[i*sa] * b[i*sb] for i from 0 to n-1. With a stride, the same function can dot a row (stride 1) with a column (stride = row width) of a flat matrix, which is exactly what one cell of a matmul does.

Matrix times matrix

If A has shape (M, K) and B has shape (K, N), then C = A @ B has shape (M, N). (@ means matrix multiply. It's the Python and NumPy symbol for it.) Each output cell is one dot product:

C[i][j] = sum over k of  A[i][k] * B[k][j]

That's row i of A dotted with column j of B. The two inner sizes (both K) have to match, and they disappear from the result.

The triple loop

In C, with all three matrices stored row-major as flat arrays (lesson 4):

void matmul(int M, int N, int K, const float *A, const float *B, float *C) {
  for (int i = 0; i < M; i++)
    for (int j = 0; j < N; j++) {
      float acc = 0.0f;                   // 0.0f is a float literal; acc is the running sum
      for (int k = 0; k < K; k++)
        acc += A[i*K + k] * B[k*N + j];   // A walks a row, B walks a column
      C[i*N + j] = acc;
    }
}

It does M * N * K multiplies and the same number of adds. Look at the index into B: as k goes up it jumps N floats at a time, which is a stride, so it's cache-unfriendly for large N. Reordering the loops, tiling them, and keeping values in registers is how later weeks make this fast. The math doesn't change, only the order.

Count it

A 256 by 256 matmul is 256 * 256 * 256 = 16,777,216 multiply-adds. People quote it as about 33.5 million FLOPs (floating point operations), because each multiply-add is 2. That FLOP count is how you'll grade your own speed later: FLOPs divided by seconds.

Count the FLOPs

Write gemm_flops(M, N, K): the FLOP count of an (M, K) @ (K, N) matmul, counting a multiply-add as 2.

The starter is the obvious one-liner. Try it on 2048.

The syllabus GEMM, made small

The course syllabus writes a 256 by 256 matrix multiply (a "GEMM") with no loops at all:

%0 = buffer 65536 dtype=f32 slot=0
%1 = buffer 65536 dtype=f32 slot=1
%2 = reshape %0 [256 256 1]
%3 = reshape %1 [256 1 256]
%4 = mul %2 %3
%5 = reduce %4 op=add pop=1

These ops arrive in weeks 2 to 5, so don't worry about the syntax yet. This is the syllabus's shorthand. The course's own .uir spells it out more: it loads each buffer before reshaping, writes the shape as shape=[...], and broadcasts with an explicit expand. The idea is lesson 4's broadcasting. Reshape both buffers to 3 axes, with a size-1 axis in a different place in each. Then mul broadcasts each to the full (K, M, N) cube of products, and reduce sums away the leading axis k. That's the multiply-then-sum of the dot product, done for every (i, j) at once.

Below is the same thing with K = 2, M = 2, N = 3, so you can see every number.

Read the axes carefully

Because the summed axis comes first, slot 0 holds A transposed (K by M). Broadcast-multiply, then reduce, is a matmul.

Code it

The syllabus GEMM, in C

Write the matmul that the syllabus program computes. A arrives stored transposed: At is K x M, so A[i][k] is At[k*M + i]. B is K x N and C is M x N, all flat and row-major.

Fill C = A @ B. Overwrite C: the tests hand you a C full of junk.

The triple loop above is your starting point. Only the index into A changes. In the inner loop over k, B was already strided, and now A is strided too.

Check yourself

[1, 2, 3] · [4, 5, 6] = ?

Multiply matching positions, then add the three products.

4 + 10 + 18 = 32.

What's the shape of (2, 3) @ (3, 5)?

The inner sizes have to match, and they are the ones that disappear.

The outer sizes survive and the matching inner size 3 is summed away.

With A = [[1, 2, 3], [4, 5, 6]] and B = [[7, 8], [9, 10], [11, 12]], what is C[1][0]?

Row 1 is the second row of A. Column 0 is the first column of B.

Dot them: three products, then the sum.

Row 1 of A dotted with column 0 of B: 4*7 + 5*9 + 6*11 = 28 + 45 + 66 = 139.

Which of these is a valid matrix multiply?

Write each as (M, K) @ (K, N): the two inner numbers must be equal.

The result has the two outer sizes.

The inner sizes must match. (4, 1) @ (1, 4) is an "outer product": every cell is one multiply, with K = 1.

How many multiplies does the triple loop do for M = N = K = 256?

The innermost line runs once for every (i, j, k) triple. How many triples are there?

256³ = 16,777,216, which is 2²⁴. Counting the adds too gives about 33.5 million FLOPs.

In the syllabus GEMM, what's the shape of %4 = mul %2 %3, before the reduce?

Broadcasting stretches every size-1 axis to the other side's size.

Line up [256 256 1] and [256 1 256] axis by axis.

Each size-1 axis broadcasts to match the other side. That's a cube of 16.7 million products, which the reduce collapses to [256 256]. A good compiler never actually stores that cube. It fuses the multiply into the sum (week 4).