Week 5 · Lesson 4
Matmul and conv from movement
A matrix multiply is a few movement ops, one mul and one reduce. So is a convolution. Once your compiler handles movement ops, it can express the two operations that do almost all the work in a neural network, and the syllabus says that from here it can compile any model.
The syllabus GEMM
The public syllabus opens with a .uir program that multiplies two 256x256 matrices. In our format, where broadcasts are written out with expand, it reads:
%a = load %A # 65536 floats
%b = load %B # 65536 floats
%a3 = reshape %a shape=[256 256 1] # axes: k, m, 1
%b3 = reshape %b shape=[256 1 256] # axes: k, 1, n
%ax = expand %a3 shape=[256 256 256]
%bx = expand %b3 shape=[256 256 256]
%p = mul %ax %bx # p[k][m][n] = a[k][m] * b[k][n]
%c = reduce %p op=add pop=1 # c[m][n] = sum over k, in order k = 0, 1, ...Both inputs are lifted into a cube with axes (k, m, n). Each one is missing an axis, so expand repeats it along that axis: a is the same for every n, b the same for every m. The mul fills the cube with every product a[k][m] · b[k][n], and reduce pop=1 folds away the leading axis, k, leaving c[m][n]. Step through it on a 2x2 by 2x3 case.
reshape %a shape=[K M 1] reads the buffer as K rows of M. For the product to be A @ B, row k of that buffer must hold column k of A: the buffer holds A transposed. Lesson 0.5 met the same fact from the other side. If your A arrives row-major as M x K, the syllabus form silently computes Aᵀ @ B instead, with the right shape and wrong values.
The general form
When A arrives as M x K, row-major, add one permute to move k to the front. The reduce still folds the first axis, so whatever you want to sum over has to be first:
# %A, %B: loaded values (load the buffers first: memory can't be reshaped)
%a3 = reshape %A shape=[M K 1]
%aT = permute %a3 order=[1 0 2] # K M 1
%ax = expand %aT shape=[K M N]
%b3 = reshape %B shape=[K 1 N]
%bx = expand %b3 shape=[K M N]
%p = mul %ax %bx
%C = reduce %p op=add pop=1 # M NThe cube is never stored. It's a description: element (k, m, n) of %p is one multiply of two loaded values. Lesson 5.5 shows how rangeify turns this description into loops, with the cube gone: one loop over the outputs, one over k.
A matmul from movement ops
Slot 0 holds A, 2x3, row-major. Slot 1 holds B, 3x2, row-major. Store C = A @ B (2x2, row-major) into slot 2, using only load, movement ops, expand, one mul, one reduce and a store. No hand-written range. The starter is the syllabus form, which assumes A is stored transposed.
A convolution is shifted windows and a reduce
A 1D convolution (in the machine learning sense, which is a correlation) computes y[i] = x[i-1]·w[0] + x[i]·w[1] + x[i+1]·w[2], with zeros past the ends of x. Build it from movement: pad x so every window is full, cut one shifted window per tap with shrink, stack the windows, and reduce over the tap axis.
%p = pad %x pads=[1 1] # 7: one zero each side
%s0 = shrink %p 0 shape=[5] # x[i-1] for every i
%s1 = shrink %p 1 shape=[5] # x[i]
%s2 = shrink %p 2 shape=[5] # x[i+1]
%win = stack %s0 %s1 %s2 # 3 5: one row per tap
%w2 = reshape %w shape=[3 1]
%wx = expand %w2 shape=[3 5]
%pr = mul %win %wx
%y = reduce %pr op=add pop=1 # 5The stack of shifted windows is the im2col view: it turns the convolution into a multiply and a sum over taps, the same shape of computation as the matmul above. With x = 1..5 and w = [1 0 -1], the oracle gives y[0] = 0·1 + 1·0 + 2·(-1) = -2 and y[4] = 4·1 + 5·0 + 0·(-1) = 4.
A same-size conv1d
Slot 0 holds x (6 f32), slot 1 holds w (3 f32). Store y[i] = x[i-1]·w[0] + x[i]·w[1] + x[i+1]·w[2] into slot 2 (6 f32), with x read as zero outside its 6 elements. The starter pads on the wrong side.
The syllabus says that after week 5 the compiler "is now capable of compiling any model", splitting it into kernels and printing C. What's left is making it learn and making it fast: week 6 adds autodiff and starts timing kernels, weeks 7 and 8 bring threads, GPUs, memory hierarchies and fast GEMMs, and weeks 9 and 10 run real models. Every one of those builds on the ops in this lesson.
Check yourself
In the syllabus GEMM, %p has shape (256, 256, 256). What shape is %c = reduce %p op=add pop=1? Write the sizes separated by a space.
pop=1 drops one leading axis.
(256, 256): the leading axis, k, is folded away, leaving m and n.
The syllabus form reshapes slot 0 to [K M 1] and reduces over the first axis. What must slot 0 hold for the result to be A @ B?
Row k of the buffer becomes the k-th plane of the cube.
Plane k needs A[m][k] for every m, which is column k of A. Stored as rows, that's A transposed.
A = [[1 2] [3 4]] is stored transposed as [1 3 2 4], and B = [[5 6 7] [8 9 10]]. What is c[1][2] from the syllabus form?
It's row 1 of A dotted with column 2 of B.
3·7 + 4·10 = 61. The oracle agrees.
Why does our version of the syllabus program need expand before the mul?
Try mul on shapes (2, 2) and (1, 2).
The oracle rejects it: "ALU operands must share one shape (no implicit broadcast: use expand)". And expand copies nothing: it only changes which index each element reads.
x has 5 elements and the kernel has 3 taps. After pad %x pads=[1 1] and stacking the three length-5 shrinks, what shape is %win?
stack adds a new leading axis of size 3.
(3, 5): one row per tap, one column per output. With x = 1..5, summing over taps gives [3 6 9 12 9].
With x = 1 2 3 4 5, w = 1 0 -1 and pads=[1 1], what is y[0]?
The window for y[0] is 0, 1, 2: one padded zero, then the first two elements.
0·1 + 1·0 + 2·(-1) = -2.