Topic — Matrix multiplication across PEs

Needs: Lesson 9 — Broadcast and reduce along a row, Lesson 10 — Switches: PEs that take turns on one color.

This topic computes C = A × B, with A of size M×K and B of size K×N, on an M×N grid of PEs. Each PE computes one element of C while rows of A stream in from the left and columns of B from the top. This layout is called an output-stationary systolic array: the results stay put and the inputs flow past them. The topic then compares it with three other ways to lay out the same product.

What you will be able to do

• Stream two inputs across a 2-D grid, one along the rows and one along the columns, with routes that deliver and forward each wavelet.
• Give each sender of a row or column its own data with miset, and read a result from every PE of a grid in one miget.
• Explain why the PEs need no clock-aligned input schedule.
• Choose between four matrix-multiply layouts by what each costs in PEs, colors and time.

The program

The grid PE in column x = 1 + j and row y = i computes C[i][j] = Σₖ A[i][k] · B[k][j]. The A sender in row i holds row i of A, and the B sender above column 1 + j holds column j of B. Each sender streams its K values.

  y=M              bsend          bsend               bsend         column j of B,
                     |              |                   |           streamed down on bcolor
                     v              v                   v
  y=M-1   asend --> C[M-1][0] --> C[M-1][1] --> ... --> C[M-1][N-1]
                     |              |                   |
  ...      ...       v              v                   v
  y=0     asend --> C[0][0] ----> C[0][1] ----> ... --> C[0][N-1]   row i of A,
                                                                    streamed right on acolor
          x=0       x=1            x=2                 x=N

The tutorial row multiplies the 2×2 example

A = [[1, 2],    B = [[5, 6],    C = A × B = [[19, 22],
     [3, 4]]         [7, 8]]                 [43, 50]]

and leaves C[0][0] = 19 on the PE at (1, 0) and C[1][1] = 50 on the PE at (2, 1). The other rows use random integers from −4 to 4:

Row	M	K	N	PEs
tutorial	2	2	2	8
rect	3	5	4	19
square	8	8	8	80

All of them are exact in HP: every product and partial sum is an integer far below 2048, the largest range in which HP holds every integer. So the check does not depend on the order in which the PE rounds.

The code

send.w

xp param K;

hp v[n=K];

socket tx;

main {
    ∀k in [0, K) tx[] <- v[k];
}

Every sender runs the same kernel. v has no initializer: miset writes each sender's values after compilation (lesson 9: never copy a kernel to change its data).

pe.w

xp param K;

socket arx;
socket brx;

hp c[n=1];

main {
    hp acc = 0.0;
    ∀k in [0, K) {
        hp a = arx[];
        hp b = brx[];
        acc = acc + a * b;
    }
    c[0] = acc;
}

In step k the PE takes the k-th value from each socket and adds their product to acc. Carbide compiles acc + a * b as one fused multiply-add.

bench.paint

Three regions, each with one code() for all its PEs:

let asend = rect(0 0 1 M 'x 'y)
asend : code('send [tx:acolor K:K])
asend : paint(acolor [C>R])
asend : name('asend)

let bsend = rect(1 M N 1 'x 'y)
bsend : code('send [tx:bcolor K:K])
bsend : paint(bcolor [C>D])
bsend : name('bsend)

let grid = rect(1 0 N M 'x 'y)
grid : code('pe [arx:acolor brx:bcolor K:K])
grid : paint(acolor [L>CR|R])
grid : paint(bcolor [U>CD|D])
grid : name('grid)

union(asend bsend grid)

• acolor [L>CR|R] delivers each A value to the PE's core and forwards it right. |R drops the forward on the grid's right edge, so the last column only receives.
• bcolor [U>CD|D] does the same downwards: from the PE above (U, +y) to the core and to the PE below. |D ends it on the bottom row.
• The top-left corner (0, M) belongs to no region and is left empty.

The rules sit on one rect() each, which is where | works as intended (boundary operators).

The Makefile: data in, results out

bench.img: bench.obj a.bin bt.bin bench.map
	cp $< $@
	miset image=$@ out=$@ map=bench.map symbol=asend.v order="y n" data=a.bin
	miset image=$@ out=$@ map=bench.map symbol=bsend.v order="x n" data=bt.bin

c-%.bin: bench-%.core bench.map
	miget image=$< map=bench.map symbol=grid.c order="y x n" data=$@

The dimension names 'x 'y on each rect() become dimensions of its symbols in bench.map. For the rect row:

asend.v                 [ y<3 x<1 n<5      ] ( 0x0+n        0   0+y )
bsend.v                 [ y<1 x<4 n<5      ] ( 0x0+n        1+x 3   )
grid.c                  [ y<3 x<4 n<1      ] ( 0x1          1+x 0+y )

a.bin is A in row-major order, so with order="y n" the A sender in row y gets row y of A. asend.v also has x<1; a dimension of size 1 may be left out of order. bt.bin is B transposed, so with order="x n" the B sender above grid column x gets column x of B. grid.c comes back with order="y x n": C in row-major order, which ref.py wrote as c-ref.bin.

What the hardware does

No input schedule is needed. A textbook systolic array is clocked in lockstep. It delays each input stream by a different amount, so that A[i][k] and B[k][j] reach each PE in the same cycle. Here each socket has its own input queue, and a receive waits until its queue has a value. The PE pairs the k-th value of one queue with the k-th value of the other, whenever each arrives. Flow control does the rest: a route that delivers and forwards waits until both outputs have taken the wavelet (lesson 4), so a slow PE holds back its row and column, and nothing is lost.

The streams run at the speed of the PE loop. The listing of pe.w (test-square/pe_artefacts/pe*.yeti) shows the loop:

	setjt JTR0 = main_bb3;  //@../pe.w:13

main_bb1: 
	imov16 R1 = [T02];  //@../pe.w:13
main_bb2: 
	fmov16 R2 = [T11];  //@../pe.w:14
	fmov16 R3 = [T03];  //@../pe.w:15
	fmach R0 = R0, R2, R3;  //@../pe.w:16
	jmp main_bb1;  //@../pe.w:13
main_bb3: 
	st16 [c] = R0;  //@../pe.w:18

As in lesson 3, T02 is a descriptor that only counts: after K uses the CE goes on at main_bb3, the address setjt stored before the loop, which stores acc into c. Each step reads one HP value from each input queue ([T11] and [T03] are the fabric operands for arx and brx) and does one fused multiply-add in the register R0. A link moves a wavelet per cycle, so the fabric is not the limit; the loop is.

Time grows with K plus the grid's size. The last memory write of each run:

M	K	N	Last write (cycle)
1	8	1	61
2	2	2	39
3	5	4	55
8	8	8	75
2	8	16	90
16	8	2	89
16	16	16	123
8	32	8	171
8	64	8	299

These fit roughly 25 + 4·K + 2·max(M, N) cycles: about 4 cycles per step, and about 2 cycles per hop for the first values to reach the far corner. Making the grid larger adds hops, not steps.

Budget. Two colors of 24. Each grid PE receives on two input queues and sends on none. The senders and the grid PEs each hold a few words of memory.

Four layouts for one product

Three other layouts compute the same product. They are not part of this topic; the figure and the table compare how they are built:

Broadcast B, one result color per PE         Broadcast B, one shared result color
  send --> mul0 --> mul1      (B, bcast)       send --> mul0 --> mul1       (B, bcast)
           mul0 --> mul1 --> recv  (res0)               mul0 --> mul1 --> recv  (res)
                    mul1 --> recv  (res1)      mul1's switch: [L>R ; C>R]; mul0's ADV moves it

K-split with the fabric adder                Output-stationary systolic (this topic)
  mul0 --> mul1 --> recv     (res)             A from the left, B from the top,
  mul1's router adds its partial product       one PE per element of C

	Broadcast, a color per PE	Broadcast, one shared color	K-split, fabric adder	Systolic
Each compute PE holds	a row of A	a row of A	column k of A, row k of B	nothing at the start
Each compute PE computes	a row of C, one dot product per column of B	the same	the partial products for one k	one element of C
How it computes	a task per arriving B value, by dispatch (lesson 7)	the same	one product per element of C	a loop of K steps
Where C ends up	a receiver PE	a receiver PE	a receiver PE	in the grid
Compute PEs for M×N	M	M	K	M·N
Colors	1 + M	2	1	2

How they differ:

• Broadcast with a color per PE. Each compute PE sends its row of C on its own color, so the rows drain in parallel. The colors run out as M grows.
• Broadcast with one shared color. The compute PEs take turns on one result color with switch positions, as in lesson 10: mul0 sends its row and then a control(ADV) that moves mul1's switch from passing mul0's data through to sending its own. Two colors serve any M, but the rows drain one after another. With more than two compute PEs, one advance would move every switch on its path at once. Each switch must then remove the advance once it has used it, so that the switches further on stay where they are; the route option pop=adv does that, and Topic — Switch protocols teaches it.
• K-split with the fabric adder. PE k computes A[i][k]·B[k][j] for every (i, j), and the routers add the partial products as they pass, so C arrives at the receiver complete. It uses one color, and only K compute PEs, but each PE does M·N products in sequence. Topic — Reductions in the router: the fabric adder teaches the adder and its limits.
• Systolic. M·N PEs work at once, each for K steps, and no PE ever holds a whole row or column. C stays distributed, which suits a computation whose next step uses C where it is.

Run it

In your copy of the tutorial (see Lesson 1):

cd tutorial/ext-algorithms/matrix-multiply
make TEST=tutorial       # the 2x2 example
make TEST='*'            # every row of test.rc

A passing row prints one line, such as matrix-multiply : 0:02.25 PASS tutorial [ M=2 K=2 N=2 ARCH=sdr ], and test-tutorial/make-log shows the matrices ref.py used. A failing row prints FAIL, and its make-log ends with the first byte where c-ref.bin and the simulator's C differ. Run make clean after you edit a source file: make reuses an existing test-*/ directory.

Pitfalls

• A missing |R or |D passes the test. Without them the last column forwards every A value out of its right side and the bottom row every B value out of its bottom. Here nothing lies beyond the grid, and the rect row still passes (checked). Next to another region, those copies enter that region's routers.
• Both streams need K values. A PE that receives fewer values on one socket than its loop reads waits forever, and its row and column stop behind it.
• B is sent by column. The sender above grid column x streams B[0][x], B[1][x], …, so the Makefile loads B transposed. Loading B in row-major order without the transpose gives each B sender K consecutive values of B's rows instead of a column. When B is square, the result is A × Bᵀ.
• Real data does not stay exact. With fractional values, the PE rounds acc to HP after every step, in k order. A reference must round the same way, or use a tolerance.

Exercises

1. Add a row M=4 K=16 N=12 to test.rc. From the table, estimate when the last PE finishes, and check that the row passes.
2. Make the grid accumulate in SP: declare sp acc, add (sp)a * (sp)b, and store into sp c[n=1]. miget then needs order="y x n _word", and ref.py must write c-ref.bin as float32. Which rows' results change?

Next

Topic — Sorting networks along a row pairs neighbouring PEs in phases to sort data spread along a row.

References

• Rectangles and dimension names
• Boundary operators
• Receiving data
• Floating-point evaluation
• The CE's queues
• Colors