scan.bend source
scan.bend on the hub · documented module
# scan.bend: parallel prefix sums over an Array<U32>, in place.## import ./scan.bend as Scan# Scan.exclusive(8n, a) # (a, total): a[j] becomes a[0] + .. + a[j - 1]# Scan.inclusive(8n, a) # (a, total): a[j] becomes a[0] + .. + a[j]## A call's result is taken apart by a def's parameter (a def of yours# taking r: Array<U32> & U32), not by a let or match on the call itself.## Sums wrap at 2^32, as U32.add does. The first argument is the fork# depth b: the array is cut into 2^b blocks of equal size (b is lowered to# log2 of the size when larger). Every b gives the same bits; b only sets# how many lanes share the work. A call under `!` (Scan.exclusive!(b, a),# or a def of yours that calls it, marked with `!`) runs its forks on the# GPU; without `!` they run on the CPU threads.## Three passes, with no atomics:# 1. per block, in parallel: the block's sum into s[i];# 2. s becomes its exclusive scan (the blocks' bases) and the total: on# one lane below 2^12 sums; else by this same scan at half the depth# (the 2^b sums in 2^(b/2) blocks), whose own pass 2 runs on one lane;# 3. per block, in parallel: the block's scan from its base, in place.# ref.bend's exclusive and inclusive are the sequential references; the# tests check both passes against them bit for bit at many sizes and b.## Sharing. The blocks' lanes read and write one array through handles# from Array.fork (Base, @unsafe): two O(1) handles to one block, joined# back by Array.join. No def here is @unsafe itself: the fork trees below# recurse on their depth, which the checker verifies. But in Bend's theory# Array.fork(a) is (a, a) and Array.join(a, b) is a, so the theory drops# the writes made through the second handle: never evaluate these defs in# the checker (a main that returns a value); call them from an IO main,# which runs compiled. `bend` lists every def that relies on the fork.import Base# Loops (one lane)# ----------------# r holds the array and the entry x read at j: adds x and the n entries# after it to s. Returns the array and the sum.def sum.go(+n: Nat, +j: U32, +s: U32, r: Array<U32> & U32) -> Array<U32> & U32: match n: case 0n: (a, +x) = r (a, (s + x : U32)) case 1n+p: (a, +x) = r +q = (j + 1 : U32) sum.go(p, q, (s + x : U32), a[q])# r holds the array and the entry x read at j: a[j] becomes s, the sum# before it, and so on for the n entries after it. Returns the array and# the running sum past the last entry.def ex.go(+n: Nat, +j: U32, +s: U32, r: Array<U32> & U32) -> Array<U32> & U32: match n: case 0n: (a, +x) = r (a[j] <- s, (s + x : U32)) case 1n+p: (a, +x) = r +q = (j + 1 : U32) ex.go(p, q, (s + x : U32), (a[j] <- s)[q])# As ex.go, but a[j] becomes the sum through it.def in.go(+n: Nat, +j: U32, +s: U32, r: Array<U32> & U32) -> Array<U32> & U32: match n: case 0n: (a, +x) = r +t = (s + x : U32) (a[j] <- t, t) case 1n+p: (a, +x) = r +q = (j + 1 : U32) +t = (s + x : U32) in.go(p, q, t, (a[j] <- t)[q])# The scan of the n + 1 entries from j, from base s.def go(+inc: Bool, +n: Nat, +j: U32, +s: U32, r: Array<U32> & U32) -> Array<U32> & U32: match inc: case False{}: ex.go(n, j, s, r) case True{}: in.go(n, j, s, r)# Handles# -------type Two is Type: Two{a: Array<U32>, s: Array<U32>}def two.join(l: Two, r: Two) -> Two: Two{la, ls} = l Two{ra, rs} = r Two{Array.join(U32, la, ra), Array.join(U32, ls, rs)}# A leaf's spare handles, joined back after its work.def two.rejoin(a2: Array<U32>, s2: Array<U32>, r: Two) -> Two: Two{a, s} = r Two{Array.join(U32, a, a2), Array.join(U32, s, s2)}# Pass 1# ------def sums.leaf.fin(+i: U32, s: Array<U32>, r: Array<U32> & U32) -> Two: (a, +v) = r Two{a, s[i] <- v}# Block i (2^e entries): its sum into s[i].def sums.leaf(+i: U32, +e: Nat, a: Array<U32>, s: Array<U32>) -> Two: +at = U32.shln(i, e) sums.leaf.fin(i, s, sum.go(U32.to_nat((U32.shln(1, e) - 1 : U32)), at, 0, a[at]))# Pass 1's forks: d levels down to the leaves under leaf index i. fa and# fs are two handles each to the data a and the sums s. Leaf i reads only# its block of a and writes only s[i], so no two lanes touch one slot.def sums(+d: Nat, +i: U32, +e: Nat, fa: Array<U32> & Array<U32>, fs: Array<U32> & Array<U32>) -> Two: match d: case 0n: (a, a2) = fa (s, s2) = fs two.rejoin(a2, s2, sums.leaf(i, e, a, s)) case 1n+p: (a0, a1) = fa (s0, s1) = fs +j = (i * 2 : U32) l r = sums(p, j, e, Array.fork(U32, a0), Array.fork(U32, s0)) sums(p, (j + 1 : U32), e, Array.fork(U32, a1), Array.fork(U32, s1)) two.join(l, r)# Pass 3# ------def puts.leaf.fin(s: Array<U32>, r: Array<U32> & U32) -> Two: (a, +t) = r Two{a, s}def puts.leaf.at(+inc: Bool, +i: U32, +e: Nat, a: Array<U32>, r: Array<U32> & U32) -> Two: (s, +base) = r +at = U32.shln(i, e) puts.leaf.fin(s, go(inc, U32.to_nat((U32.shln(1, e) - 1 : U32)), at, base, a[at]))# Block i: its scan in place, from its base s[i].def puts.leaf(+inc: Bool, +i: U32, +e: Nat, a: Array<U32>, s: Array<U32>) -> Two: puts.leaf.at(inc, i, e, a, s[i])# Pass 3's forks, shaped as pass 1's. Leaf i reads s[i] and reads and# writes only its block of a: no two lanes touch one slot.def puts(+inc: Bool, +d: Nat, +i: U32, +e: Nat, fa: Array<U32> & Array<U32>, fs: Array<U32> & Array<U32>) -> Two: match d: case 0n: (a, a2) = fa (s, s2) = fs two.rejoin(a2, s2, puts.leaf(inc, i, e, a, s)) case 1n+p: (a0, a1) = fa (s0, s1) = fs +j = (i * 2 : U32) l r = puts(inc, p, j, e, Array.fork(U32, a0), Array.fork(U32, s0)) puts(inc, p, (j + 1 : U32), e, Array.fork(U32, a1), Array.fork(U32, s1)) two.join(l, r)# Driver# ------# floor(d / 2).def half(+d: Nat) -> Nat: match d: case 0n: 0n case 1n+p: match p: case 0n: 0n case 1n+q: 1n+half(q)def run.fin(r: Two, +total: U32) -> Array<U32> & U32: Two{a, s} = r (a, total)# Pass 3, from pass 2's bases and total.def run.mid(+inc: Bool, +d: Nat, +e: Nat, a: Array<U32>, r: Array<U32> & U32) -> Array<U32> & U32: (s, +total) = r run.fin(puts(inc, d, 0, e, Array.fork(U32, a), Array.fork(U32, s)), total)# Pass 1 over a (2^lg entries) at depth d, into a new s.def run.pre(+d: Nat, +e: Nat, a: Array<U32>) -> Two: sums(d, 0, e, Array.fork(U32, a), Array.fork(U32, Array.new(U32, d, 0)))# The scan with pass 2 on one lane.def run0.sums(+inc: Bool, +d: Nat, +e: Nat, r: Two) -> Array<U32> & U32: Two{a, s} = r run.mid(inc, d, e, a, ex.go(U32.to_nat((U32.shln(1, d) - 1 : U32)), 0, 0, s[0]))def run0(+inc: Bool, +b: Nat, r: Array<U32> & U32) -> Array<U32> & U32: (a, +n) = r +lg = U32.log2(n) +d = Nat.min(b, lg) +e = Nat.sub(lg, d) run0.sums(inc, d, e, run.pre(d, e, a))# The scan with pass 2 by run0 at depth d / 2.def run1.sums(+inc: Bool, +d: Nat, +e: Nat, r: Two) -> Array<U32> & U32: Two{a, s} = r run.mid(inc, d, e, a, run0(False{}, half(d), Array.size(U32, s)))# Pass 2 by run1 from 2^12 sums up, else on one lane: on the GPU we# measured, the nested scan's forks cost more than one lane's walk below# that.def run.pick(lo: Bool, +inc: Bool, +d: Nat, +e: Nat, r: Two) -> Array<U32> & U32: match lo: case True{}: run0.sums(inc, d, e, r) case False{}: run1.sums(inc, d, e, r)def run(+inc: Bool, +b: Nat, r: Array<U32> & U32) -> Array<U32> & U32: (a, +n) = r +lg = U32.log2(n) +d = Nat.min(b, lg) +e = Nat.sub(lg, d) run.pick(U32.is_lt(U32.from_nat(d), 12), inc, d, e, run.pre(d, e, a))# API# ---# a[j] becomes a[0] + .. + a[j - 1] (a[0] becomes 0); returns the array# and the total, a[0] + .. + a[n - 1].def exclusive(+b: Nat, a: Array<U32>) -> Array<U32> & U32: run(False{}, b, Array.size(U32, a))# a[j] becomes a[0] + .. + a[j]; returns the array and the total.def inclusive(+b: Nat, a: Array<U32>) -> Array<U32> & U32: run(True{}, b, Array.size(U32, a))