Topic — Collectives over messages: tree reduction and prefix scan

Needs: Lesson 9 — Broadcast and reduce along a row; topic Message passing: routing by address.

A collective is an operation that every PE of a group takes part in, such as a sum over all of them. Lesson 9 reduces along a row with static routes: each PE adds its value to a stream that passes through it. With message passing, the PEs can instead form a logical tree, in which each PE sends to a parent chosen by coordinates, wherever the parent sits. This topic builds the two standard tree collectives on one tree:

• a reduction: the root ends up with the sum of all the leaves' values;
• a prefix scan: every leaf ends up with the sum of the values of the leaves to its left.

The same kernels run on two different placements of the tree.

What you will be able to do

• Build a tree reduction from messages, in which each node waits for its children.
• Run a second, downward sweep that gives every leaf its prefix sum.
• Tag a message with the sender's role, because a message does not say who sent it.
• Place a logical tree on the fabric from a table in Paint, and keep the kernels independent of the placement.

The program

Four leaves l0 to l3 hold the values V0 to V3. Two inner nodes p0 and p1 and a root form a binary tree over them:

                 root
               /      \
             p0        p1
            /  \      /  \
          l0    l1  l2    l3
          10    20  30    40      (tutorial row)

Up-sweep (the reduction). Each leaf sends its value to its parent. Each inner node adds its two children's values and sends the sum to its parent. The root adds its children's sums: that is the total.

  p0 = 10 + 20 = 30      p1 = 30 + 40 = 70      root: total = 30 + 70 = 100

Down-sweep (the scan). Each node receives a prefix: the sum of every leaf left of its subtree. It sends that prefix unchanged to its left child, and the prefix plus its left child's up-sweep value to its right child. The root's own prefix is 0.

  root sends  0 -> p0,  0 + 30 = 30 -> p1
  p0   sends  0 -> l0,  0 + 10 = 10 -> l1
  p1   sends 30 -> l2, 30 + 30 = 60 -> l3

Each leaf stores the prefix it receives, the exclusive scan, and the prefix plus its own value, the inclusive scan:

Leaf	l0	l1	l2	l3
value	10	20	30	40
exclusive (excl)	0	10	30	60
inclusive (incl)	10	30	60	100

The exclusive scan is how PEs find where their part of a shared result starts: if each leaf's value is how many items it holds, a leaf's items start at its exclusive prefix.

test.rc runs the tree in two placements. The figure draws y growing upward, as lesson 4 does; row 0 is the north row:

  tutorial (SHAPE=0), one row:
     x=0  x=1  x=2  x=3  x=4  x=5  x=6
     l0   l1   l2   l3   p0   p1   root

  tworow (SHAPE=1), two rows:
            x=0   x=1   x=2   x=3
     y=1    l0    l1    l2    l3
     y=0    p0    root  .     p1        "." has no kernel: it only forwards messages

tworow uses the values 7, 300, 0 and 5, and expects the exclusive scan 0, 7, 307, 307, the inclusive scan 7, 307, 307, 312, and the total 312. The leaves' messages go north and the messages to the leaves go south, so both direction bits are used.

The code

The layout: bench.paint

The kernels never compute where their parent or children are. Paint looks the coordinates up in a table and passes them as parameters:

(* Column and row of l0 l1 l2 l3 p0 p1 root *)
let X = setif(SHAPE [0 1 2 3 0 3 1] [0 1 2 3 4 5 6])
let Y = setif(SHAPE [1 1 1 1 0 0 0] [0 0 0 0 0 0 0])
let W = setif(SHAPE 4 7)
let H = setif(SHAPE 2 1)

define cx(i) { arrayidx(X i) }
define cy(i) { arrayidx(Y i) }

setif(c a b) is a when c is not 0 and b otherwise. [0 1 2 ...] is a Paint array, and arrayidx(X i) its element i. Nodes are numbered 0 to 6 in the order of the comment.

define name(arguments) { ... } defines a Paint function. Its body runs when the function is called, and its last expression is the value it returns. The two functions below each make a one-PE region, attach a kernel to it, and return the region:

define leaf(i v side up) {
  let t = rect(cx(i) cy(i) 1 1)
  t : code('leaf [mp:mpc VAL:v SIDE:side UPX:cx(up) UPY:cy(up)])
  t
}

define node(i root side up l r) {
  let t = rect(cx(i) cy(i) 1 1)
  t : code('node [mp:mpc ROOT:root SIDE:side UPX:cx(up) UPY:cy(up)
                  LX:cx(l) LY:cy(l) RX:cx(r) RY:cy(r)])
  t
}

up, l and r are node numbers: the parent and the two children. side is 0 for a left child and 1 for a right child. The tree itself is seven lines:

let all = rect(0 0 W H)
all : mp(b0)

let l0 = name(leaf(0 V0 0 4) 'l0)
let l1 = name(leaf(1 V1 1 4) 'l1)
let l2 = name(leaf(2 V2 0 5) 'l2)
let l3 = name(leaf(3 V3 1 5) 'l3)
let p0 = name(node(4 0 0 6 0 1) 'p0)
let p1 = name(node(5 0 1 6 2 3) 'p1)
let root = name(node(6 1 0 6 4 5) 'root)

union(all l0 l1 l2 l3 p0 p1 root)

name(r 'l0) is the function form of r : name('l0): in Paint, r : f(...) stands for r = f(r ...). It returns the region with the name, so l0.excl is leaf 0's excl. all enables message passing on every PE of the bounding rectangle, including the PE at (2, 0) in tworow, which has no kernel but lies on the path from p1 to the root.

The module: message.w

It is the module of the message-passing topic: the union word and the function header(col, row, here, len). Every message here has one data wavelet: the value in w[0] and a tag in w[1].

The leaves: leaf.w

    hdr.u = header(UPX, UPY, row_id[0], 1);
    msg.w[0] = VAL;
    msg.w[1] = SIDE;
    mp[] <- hdr.u;
    mp[] <- msg.u;

    hdr.u <- mp[];
    msg.u <- mp[];
    t[1] <- tsc[0];
    excl[0] = msg.w[0];
    incl[0] = msg.w[0] + VAL;

The leaf sends its value, tagged with its side, and then waits for its prefix. t[0], read just before the send, and t[1] are lesson 3's timestamp counter.

The inner nodes and the root: node.w

    ∀c in [0, 2) {
        hdr.u <- mp[];
        msg.u <- mp[];
        if (msg.w[1] == 0) left = msg.w[0];
        else right = msg.w[0];
    }

Both children send to the same socket, and their messages arrive in either order. A message carries its destination, not its sender, so the node tells the children apart by the tag.

    xp prefix = 0;
    if (ROOT == 1) {
        total[0] = left + right;
    } else {
        hdr.u = header(UPX, UPY, here, 1);
        msg.w[0] = left + right;
        msg.w[1] = SIDE;
        mp[] <- hdr.u;
        mp[] <- msg.u;
        hdr.u <- mp[];
        msg.u <- mp[];
        prefix = msg.w[0];
    }

An inner node passes its sum up and waits for its prefix. The root keeps the total and starts the down-sweep with prefix 0. ROOT is a parameter, so the compiler keeps only one branch on each PE.

    hdr.u = header(LX, LY, here, 1);
    msg.w[0] = prefix;
    msg.w[1] = 0;
    mp[] <- hdr.u;
    mp[] <- msg.u;

    hdr.u = header(RX, RY, here, 1);
    msg.w[0] = prefix + left;
    mp[] <- hdr.u;
    mp[] <- msg.u;

Each inner node receives two messages from its children and then one from its parent. The root receives the two from its children, and each leaf receives one from its parent. So the up-sweep and down-sweep messages cannot be confused: a parent sends down only after both of its children have sent up.

The test

ref.py computes the scans with numpy.cumsum. The Makefile reads each leaf's excl and incl and the root's total into one file and compares it with the reference. Then cycles.py prints, for each leaf, the cycles from its send to the arrival of its prefix; it never fails.

What the hardware does

Twelve messages. The up-sweep sends six messages, one per tree edge, and the down-sweep six more. Each is two wavelets, the header and the data. A tree of P leaves needs about 2 log2 P message steps from the first send to the last receive. Lesson 9's row reduction needs a step per PE, because the stream passes through every PE in turn.

Placement sets the paths. A message goes east or west first, then north or south, so the placement decides which routers each message crosses and which messages share a link. In tutorial, l0's message to p0 crosses the routers of l1, l2 and l3, and shares them with l1's message. In tworow, no message travels more than two hops.

The cores' work dominates here. cycles.py prints, for each leaf, the cycles from its send to the arrival of its prefix:

Row	l0	l1	l2	l3
tutorial	222	231	231	236
tworow	210	224	228	236

The last prefix arrives after 236 cycles in both placements. With seven nodes, the time goes to building, sending and reading the messages, 31 cycles per message at the sender (see the message-passing topic), not to the distance they travel.

Order. Messages from different senders reach a node in no fixed order. In these runs the nearer child's message arrived first, and which child is nearer differs between the placements.

Budget. One message-passing bank: five colors on every PE of the bounding rectangle, and one input queue per node.

Run it

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

cd tutorial/ext-collectives/message-collectives
make TEST='*'

Success is a PASS line for each of tutorial and tworow, such as message-collectives : 0:02.50 PASS tutorial [ SHAPE=0 V0=10 V1=20 V2=30 V3=40 CYCLES=1500 ARCH=sdr ], and test-tutorial/make-log ends with the four lines of cycles.py. make deletes the simulator's memory image when a row finishes, so to see one leaf's results, rebuild it as lesson 4 does:

cd test-tutorial
make -f ../Makefile VPATH=.. TEST=tutorial bench.map bench-sim.core
miget image=bench-sim.core map=bench.map symbol=l2.incl decimal
cd ..

A failure prints FAIL; test-tutorial/make-log shows the cmp that failed. Run make clean after each edit.

Pitfalls

• Arrival order is not the tree's order. A node that takes its first message as its left child's gets the wrong prefixes. Tag every message that a receiver must attribute.
• A node must know how many messages it will receive. Here each inner node receives two messages from its children and one from its parent, the root two, and each leaf one. A node that waits for a message nobody sends hangs, and the run stops at its cycle limit with the leaves' results still 0.
• Every PE the messages cross needs mp(), including PEs with no kernel, such as (2, 0) in tworow. Enabling it on the bounding rectangle covers them all.
• setif evaluates both of its branches. It is safe for the constant tables here; the Paint reference explains why it cannot end a recursion.

Exercises

1. Make the nodes ignore the tag: in node.w, replace msg.w[1] == 0 with c == 0, so the first message counts as the left child's. Run both rows. Both fail: in tutorial, l1 gets the exclusive prefix 20 instead of 10; in tworow, l3 gets 312 instead of 307. Work out from the two placements which child's message reached p0, p1 and the root first.
2. Grow the tree to eight leaves: add l4 to l7 and four more inner nodes, one more level, so the tree has 15 PEs. Extend the tables X and Y, and the values V4 to V7 in test.rc. The kernels do not change; bench.paint, test.rc, the Makefile, ref.py and cycles.py do.
3. Lay the tree out so that every child sits as close to its parent as it can, for example in the order l0 p0 l1 root l2 p1 l3 on one row. Compare the cycles that cycles.py prints with the two placements here.

Next

This is the last topic of the group. The next group, Routing, starts with Topic — Message passing: routing by address, which routes each message by an address in its header, so destinations can come from data.

References

• Message passing in Tungsten
• Message passing in Paint
• Paint function definitions (define)
• Paint arrays and array primitives (arrayidx)