Lesson 9 — Broadcast and reduce along a row

Each PE of a row holds a vector v of M sp values. After this program, every PE holds the element-wise sum of all K vectors. This collective is an allreduce: a reduce combines one value from every PE into a total, and a broadcast delivers the total back to every PE. Here the sum is computed on the stream as it passes each PE, and the routers copy the total back to every other PE.

What you will be able to do

• Compute on a stream: add a PE's own data to values passing through it, in one instruction, without storing them first.
• Send one result back along a row with a multicast route.
• Give each role in a row its own kernel and name, and load 32-bit data per PE with _word in order.
• Predict how a row collective's time grows with the number of PEs and the length of the vector.

The program

The row has K PEs in three roles:

• the head, at x=0, starts the running sum with its v and then receives the total;
• each mid, at x=1 to K-2, adds its v to the sum as it streams past, then keeps a copy of the total;
• the tail, at x=K-1, adds its v, keeps the total, and sends it back.

reduce, on r (colorpair r/rp), left to right:

   head          mid            mid                  tail
   v0  ---r--->  +v1  ---r--->  +v2  ---> ... --->  +v(K-1)  -> s = total

broadcast, on b, right to left:

   s  <---b---   s  <---b---    s   <--- ... <---   (sends s)
 b [R>C]       b [R>CL]       b [R>CL]             b [C>L]

Every PE's s must equal the sum of all v. ref.py writes each PE's v and the expected s of every PE, and cmp compares the result byte for byte. The values are multiples of 1/8, so every partial sum is exact in sp and the order of the additions cannot change it.

The rows of test.rc are tutorial (K=4, M=8), k3 (K=3, M=8, a single mid), k8 (K=8, M=8), k16 (K=16, M=8) and k16m64 (K=16, M=64). The same files run every row.

The code

Arithmetic on a stream: mid.w

xp param M;

socket r;
socket b;

sp v[n=M];
sp s[n=M];

{
    let j in [0, M);
    ∀j r[] <- r[] + v[j];
    ∀j s[j] <- b[];
}

∀j r[] <- r[] + v[j]; reads one value from the fabric, adds v[j] from memory, and sends the sum on, M times. It is a single instruction. The listing, test-tutorial/mid_artefacts/mid2.yeti, shows it with its three operands:

	fadds [T22] = [T01], [T12];  //@../mid.w:14

The declarations above the task name the three DSDs (lesson 3): T22 is a send DSD on r of length 8, T01 a receive DSD on r, and T12 a memory vector over v. The running sum never lands in the mid's memory. The colorpair of lesson 8 makes the one socket r both the input from the left and the output to the right.

After the reduce, the mid receives the M values of the total into s.

The head and the tail

head.w sends its v to start the sum and receives the total. It also reads the timestamp counter of lesson 3 before its first send and after its last receive:

    let j in [0, M);
    t[0] <- tsc[0];
    ∀j r[] <- v[j];
    ∀j s[j] <- b[];
    t[1] <- tsc[0];

head.w enables the counter just before this block, so t[0] comes right after the enable, the case lesson 3 warns about, and reads 0 in every row. In the simulator that does no harm: a test kernel timed a 32-element sp addition at 18 cycles whether its first read came right after the enable or 63 cycles after it. Nothing guarantees that in general, so in your own kernels enable the counter earlier, as lesson 3 does.

tail.w adds its own v into memory, which gives it the total, and sends the total back:

    let j in [0, M);
    ∀j s[j] <- r[] + v[j];
    ∀j b[] <- s[j];

The tail keeps its own result and never receives on b: its route for b is [C>L], with no output to its own core.

The layout: one region per role

let r  = color()
let rp = color()
let b  = color()

let row = rect(0 0 K 1)
row : paint(rp [CXR|R])
row : paint(r  [L|LXC])
row : colorpair(r rp)

let head = rect(0 0 1 1)
head : code('head [r:r b:b M:M])
head : paint(b [R>C])
head : name('head)

let mids = rect(1 0 K-2 1 'x)
mids : code('mid [r:r b:b M:M])
mids : paint(b [R>CL])
mids : name('mids)

let tail = rect(K-1 0 1 1)
tail : code('tail [r:r b:b M:M])
tail : paint(b [C>L])
tail : name('tail)

union(row head mids tail)

row carries the reduce, with the same two colorpair rules as lesson 8's beach. The broadcast color b has one rule per role:

• the tail sends it left, [C>L];
• each mid takes it from the right and sends it both to its core and on to the left, [R>CL], a multicast;
• the head takes it from the right for its core, [R>C].

mids is one rectangle with one code() for all K-2 mids, as in lesson 8. Its x dimension is named 'x. Without a name Paint also calls it x, but leaves it out of the map when the rectangle is one PE wide (K = 3), and the Makefile's order="x n _word" then fails with order: 'x' is not a name or number. A name keeps the dimension at any width; the k3 row checks that.

Each role region has its own name(). The map then has one entry per symbol and role:

head.v                  [ n<8 _word<2      ] ( 0x0+2*n+_word  0   0 )
mids.v                  [ x<2 n<8 _word<2  ] ( 0x0+2*n+_word  1+x 0 )
tail.v                  [ n<8 _word<2      ] ( 0x0+2*n+_word  3   0 )

(These are the tutorial row's entries, K=4.) Pitfalls shows what one name for all three regions does.

One kernel, different data: _word

Every mid runs the same mid.w; only v differs. Do not copy a kernel to change its data: K-2 copies of mid.w, each with its own constants, would each be compiled and kept in step by hand. Paint builds mid.w for the whole mids rectangle, every mid runs the same compiled code, and the Makefile loads each PE's v from one file:

	miset image=$@ out=$@ map=bench.map symbol=head.v order="n _word" data=v-head.bin
	miset image=$@ out=$@ map=bench.map symbol=mids.v order="x n _word" data=v-mids.bin
	miset image=$@ out=$@ map=bench.map symbol=tail.v order="n _word" data=v-tail.bin

Memory words are 16 bits, so an sp value takes two words, and the map gives every sp array a last dimension, _word<2. order must name it last: order="x n _word" reads v-mids.bin as one block per mid, n values per block, two words per value, low word first. NumPy writes float32 arrays in that layout. The results come back the same way, with miget and the same orders.

What the hardware does

The reduce runs in the fabric operands

The mid's fadds takes one operand from its input queue and writes its result to its output queue, so the add runs at the speed of the stream. Each partial sum is one sp value, one wavelet, and a link carries one wavelet per cycle in each direction.

The broadcast runs in the routers

[R>CL] makes each mid's router copy every wavelet of the total to two outputs: its own core and the left neighbour. The cores do no forwarding. As lesson 4 showed, the router holds each wavelet until both outputs have taken it, so a mid that stopped reading b would stop the broadcast for every PE to its left.

The row uses 3 of the 24 colors: r, rp and b. Carbide gave each mid two input queues, for r and b, and one output queue, for r; the iqueue and oqueue lines at the top of the listing show them.

Cost, measured

cycles.py prints how long the head waited, from before its first send to after its last receive. It never fails.

Row	K	M	Cycles
k3	3	8	54
tutorial	4	8	68
k8	8	8	124
k16	16	8	236
k16m64	16	64	362

Each extra PE adds 14 cycles: the sum passes one more PE on the way out, and the total one more on the way back. Each extra value adds about 2.25 cycles. The tail receives the whole sum before it sends the total back, so the M values cross the row twice, one after the other, at a little over one cycle per value each way. The time grows with K + M, not with K × M: the PEs work on different values at the same time.

The sum here is a chain: the head's value passes through every PE. Its time grows with the length of the row. A tree of messages, or the adders in the routers themselves, can do better; see the topics Collectives over messages: tree reduction and prefix scan and Reductions in the router: the fabric adder.

Run it

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

cd tutorial/main-3-rows/09-broadcast-and-reduce
make TEST=tutorial
cat test-tutorial/make-log

Success is a line ending in PASS tutorial [ K=4 M=8 ARCH=sdr ]. make-log shows the time, such as K=4 M=8: the head had the whole total 68 cycles after its first send. A wrong sum makes cmp print out-ref.bin out-sim.bin differ: byte ..., and make prints FAIL. make TEST='*' runs all five rows.

Pitfalls

• One name for several kernels. A single name('row) on union(head mids tail) looks tidier, but the three kernels then put three entries named row.v and three named row.s in the map. miget symbol=row.s warns symbol [row.s] has 3 exact map matches; using first entry, and reads only the head's s. Name each region that runs a different kernel.
• _word left out. With order="x n", miset expects one word per value and stops with v-mids.bin: too long.
• A receive the route never feeds. A receive on b in tail.w would wait forever, because the tail's b route does not reach its core.
• At least one mid. The Makefile loads mids.v, so the layout needs K ≥ 3.

Exercises

1. From the table, predict the time for K=32 and M=8. Add a row k32 K=32 M=8 ARCH=sdr to test.rc, run make TEST=k32, and compare.
2. Turn the sum into a maximum: write max(r[], v[j]) in place of r[] + v[j] in mid.w and tail.w, and np.maximum(total, v[p]) in ref.py. Run make clean and make TEST='*'. Then look at the listing: which instruction replaced fadds, and are its operands still three DSDs?

Next

Lesson 10 — Switches: PEs that take turns on one color sends the values of many PEs to one receiver on a single color.

References

• Color pairs and routing
• Naming
• What dimension names are for, including _word
• miset and miget
• Colors in the architecture notes