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 in weeks 6 to 9 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 week 6 and week 8 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. 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).