Week 3 · Lesson 3
A GEMM by hand
The syllabus says that with loops "we can write gemms and convs in this language". This lesson writes the GEMM: C = A @ B on flat row-major buffers, three nested ranges, and one accumulator per output element. Then it shows the C your renderer should produce, and why the order of the loops matters even though every order gives the same answer.
The plan
A is M-by-K, B is K-by-N and C is M-by-N, all flat and row-major, as in lesson 0.5. Element (r, c) of an R-by-C matrix sits at r*C + c. For every output element:
C[i*N + j] = sum over k of A[i*K + k] * B[k*N + j]We'll use M = 2, K = 3, N = 2, the same matrices as lesson 0.5: A = [[1 2 3] [4 5 6]], B = [[7 8] [9 10] [11 12]], so C = [[58 64] [139 154]]. Each output element needs K multiply-adds. Step through which elements one iteration touches:
Building it, one piece at a time
The accumulator is C itself. For each (i, j), reset C[i*N + j] to 0, then let the k loop add into it. The reset stores into a view at offset i*N + j, so it depends on both ranges and runs once per output element. That's the trick from lesson 3.2.
The whole program checks out with the oracle, and its C, run on real arrays, prints 58 64 139 154. Here's the rendered kernel, with the wraparound casts trimmed for reading:
for (int32_t r6 = 0; r6 < 2; r6++) { /* i */
int32_t v7 = r6 * 2; /* i*N */
int32_t v8 = r6 * 3; /* i*K */
for (int32_t r9 = 0; r9 < 2; r9++) { /* j */
int32_t v10 = v7 + r9; /* i*N + j */
b2[v10] = 0x0.0p+0f; /* reset C[i][j] */
for (int32_t r15 = 0; r15 < 3; r15++) { /* k */
float v16 = b2[v10]; /* loop-carried load */
int32_t v17 = v8 + r15;
float v19 = b0[v17]; /* A[i][k] */
int32_t v21 = r15 * 2 + r9;
float v23 = b1[v21]; /* B[k][j] */
float v24 = v19 * v23; /* two statements, as the oracle prints them: */
float v25 = v16 + v24; /* a*b + c on one line invites an FMA (lesson 2.4) */
b2[v10] = v25;
}
}
}Every index computation sits in the loop of its deepest range: i*K outside the j loop, i*N + j outside the k loop. You get that hoisting for free from open ranges.
The view shrink %A %ik needs i*3 + k to be provably in [0, 5]. With i in [0, 1] and k in [0, 2], week 1's interval rules give i*3 in [0, 3] and i*3 + k in [0, 5]. Make A five elements long and the same program is rejected: "can't prove the view is in bounds: offset in [0, 5], 1 element(s), parent has 5".
C = A @ B transposed
A and B are both 2-by-3 i32 matrices, row-major, in slots 0 and 1. Compute C = A @ B^T (2-by-2) into slot 2: C[i*2 + j] is the dot product of row i of A and row j of B. The starter is the plain GEMM from this page, which indexes B as if it were 3-by-2.
Same answer, different speed
The order of the ends picks the loop order. Here's a 2-by-3 fill written two ways. Both store i*3 + j into element i*3 + j, and both leave the buffer holding 0 to 5:
%ej = end %s %j # j inner
%ei = end %ej %i # i outer
for (i ...)
for (j ...)
b[i*3 + j] = ... /* 0 1 2 3 4 5 */%ei = end %s %i # i inner
%ej = end %ei %j # j outer
for (j ...)
for (i ...)
b[i*3 + j] = ... /* 0 3 1 4 2 5 */The left version walks memory in order: addresses 0, 1, 2, 3, 4, 5. The right one jumps by 3 each step. Lesson 0.2 showed what that costs: in-order access uses every byte of each cache line it loads, and strided access can miss on nearly every step once the arrays are bigger than the cache. In this GEMM, the inner k loop walks A along a row (stride 1) but B down a column (stride N). Reordering the loops so the innermost index is the one each array is contiguous in is one of the first speedups in week 6. For now, the grader only checks correctness.
Dot product
The smallest GEMM: a 1-by-4 times a 4-by-1. Put the dot product of slots 0 and 1 (four i32s each) into slot 2's only element. The starter resets the output and stops.
Check yourself
With %i = range 2 and %k = range 3, what interval does the bounds check give i*3 + k? (low and high)
A range's interval is [0, n-1]. Multiply, then add, end by end.
i is in [0, 1], so i*3 is in [0, 3]. Adding k in [0, 2] gives [0, 5], which fits a six-element buffer exactly.
For the matrices on this page, what does C[1*2 + 1] end up holding?
Row 1 of A is [4 5 6]. Column 1 of B is [8 10 12].
4*8 + 5*10 + 6*12 = 32 + 50 + 72 = 154, the same answer lesson 0.5 got with the triple loop in C.
Suppose the accumulator is a one-element alloc that's never reset, and each (i, j) copies it into C after its k loop. What does C hold?
The alloc starts at 0 once, not once per output element.
The sum carries over from one output element to the next: 58, then 58 + 64 = 122, then 122 + 139 = 261, then 261 + 154 = 415. That's why the reset has to depend on i and j.
Change A to five elements, leaving the rest of the GEMM alone. What exit code does run give?
The interval of i*3 + k doesn't change.
The offset is still in [0, 5], and a five-element buffer can't hold a view at 5, so the bounds check fails and the program is invalid: exit 2.
How many multiply-adds does an M-by-K times K-by-N GEMM do?
How many output elements are there, and how many multiply-adds does each one take?
There are M*N outputs, and each is a sum of K products, so the innermost body runs M*N*K times. That's the trip count of the three nested ranges.
Swapping the order of two ends in the fill example changes...
Each element is written exactly once in both versions.
Each element gets the same value either way, so the result is identical and the bounds don't change. What changes is the loop nest the renderer emits, and so the access pattern: stride 1 versus stride 3. That decides cache behavior, which week 6 starts measuring.