scan.bend

bend-parallel@0.1.0.0 · 234 lines · 25 definitions

1# scan.bend: parallel prefix sums over an Array<U32>, in place.
2#
3#   import ./scan.bend as Scan
4#   Scan.exclusive(8n, a)   # (a, total): a[j] becomes a[0] + .. + a[j - 1]
5#   Scan.inclusive(8n, a)   # (a, total): a[j] becomes a[0] + .. + a[j]
6#
7# A call's result is taken apart by a def's parameter (a def of yours
8# taking r: Array<U32> & U32), not by a let or match on the call itself.
9#
10# Sums wrap at 2^32, as U32.add does. The first argument is the fork
11# depth b: the array is cut into 2^b blocks of equal size (b is lowered to
12# log2 of the size when larger). Every b gives the same bits; b only sets
13# how many lanes share the work. A call under `!` (Scan.exclusive!(b, a),
14# or a def of yours that calls it, marked with `!`) runs its forks on the
15# GPU; without `!` they run on the CPU threads.
16#
17# Three passes, with no atomics:
18#   1. per block, in parallel: the block's sum into s[i];
19#   2. s becomes its exclusive scan (the blocks' bases) and the total: on
20#      one lane below 2^12 sums; else by this same scan at half the depth
21#      (the 2^b sums in 2^(b/2) blocks), whose own pass 2 runs on one lane;
22#   3. per block, in parallel: the block's scan from its base, in place.
23# ref.bend's exclusive and inclusive are the sequential references; the
24# tests check both passes against them bit for bit at many sizes and b.
25#
26# Sharing. The blocks' lanes read and write one array through handles
27# from Array.fork (Base, @unsafe): two O(1) handles to one block, joined
28# back by Array.join. No def here is @unsafe itself: the fork trees below
29# recurse on their depth, which the checker verifies. But in Bend's theory
30# Array.fork(a) is (a, a) and Array.join(a, b) is a, so the theory drops
31# the writes made through the second handle: never evaluate these defs in
32# the checker (a main that returns a value); call them from an IO main,
33# which runs compiled. `bend` lists every def that relies on the fork.
34import Base
35
36# Loops (one lane)
37# ----------------
38
39# r holds the array and the entry x read at j: adds x and the n entries
40# after it to s. Returns the array and the sum.
41def sum.go(+n: Nat, +j: U32, +s: U32, r: Array<U32> & U32) -> Array<U32> & U32:
42  match n:
43    case 0n:
44      (a, +x) = r
45      (a, (s + x : U32))
46    case 1n+p:
47      (a, +x) = r
48      +q = (j + 1 : U32)
49      sum.go(p, q, (s + x : U32), a[q])
50
51# r holds the array and the entry x read at j: a[j] becomes s, the sum
52# before it, and so on for the n entries after it. Returns the array and
53# the running sum past the last entry.
54def ex.go(+n: Nat, +j: U32, +s: U32, r: Array<U32> & U32) -> Array<U32> & U32:
55  match n:
56    case 0n:
57      (a, +x) = r
58      (a[j] <- s, (s + x : U32))
59    case 1n+p:
60      (a, +x) = r
61      +q = (j + 1 : U32)
62      ex.go(p, q, (s + x : U32), (a[j] <- s)[q])
63
64# As ex.go, but a[j] becomes the sum through it.
65def in.go(+n: Nat, +j: U32, +s: U32, r: Array<U32> & U32) -> Array<U32> & U32:
66  match n:
67    case 0n:
68      (a, +x) = r
69      +t = (s + x : U32)
70      (a[j] <- t, t)
71    case 1n+p:
72      (a, +x) = r
73      +q = (j + 1 : U32)
74      +t = (s + x : U32)
75      in.go(p, q, t, (a[j] <- t)[q])
76
77# The scan of the n + 1 entries from j, from base s.
78def go(+inc: Bool, +n: Nat, +j: U32, +s: U32, r: Array<U32> & U32) -> Array<U32> & U32:
79  match inc:
80    case False{}:
81      ex.go(n, j, s, r)
82    case True{}:
83      in.go(n, j, s, r)
84
85# Handles
86# -------
87
88type Two is Type:
89  Two{a: Array<U32>, s: Array<U32>}
90
91def two.join(l: Two, r: Two) -> Two:
92  Two{la, ls} = l
93  Two{ra, rs} = r
94  Two{Array.join(U32, la, ra), Array.join(U32, ls, rs)}
95
96# A leaf's spare handles, joined back after its work.
97def two.rejoin(a2: Array<U32>, s2: Array<U32>, r: Two) -> Two:
98  Two{a, s} = r
99  Two{Array.join(U32, a, a2), Array.join(U32, s, s2)}
100
101# Pass 1
102# ------
103
104def sums.leaf.fin(+i: U32, s: Array<U32>, r: Array<U32> & U32) -> Two:
105  (a, +v) = r
106  Two{a, s[i] <- v}
107
108# Block i (2^e entries): its sum into s[i].
109def sums.leaf(+i: U32, +e: Nat, a: Array<U32>, s: Array<U32>) -> Two:
110  +at = U32.shln(i, e)
111  sums.leaf.fin(i, s, sum.go(U32.to_nat((U32.shln(1, e) - 1 : U32)), at, 0, a[at]))
112
113# Pass 1's forks: d levels down to the leaves under leaf index i. fa and
114# fs are two handles each to the data a and the sums s. Leaf i reads only
115# its block of a and writes only s[i], so no two lanes touch one slot.
116def sums(+d: Nat, +i: U32, +e: Nat, fa: Array<U32> & Array<U32>, fs: Array<U32> & Array<U32>) -> Two:
117  match d:
118    case 0n:
119      (a, a2) = fa
120      (s, s2) = fs
121      two.rejoin(a2, s2, sums.leaf(i, e, a, s))
122    case 1n+p:
123      (a0, a1) = fa
124      (s0, s1) = fs
125      +j = (i * 2 : U32)
126      l r = sums(p, j, e, Array.fork(U32, a0), Array.fork(U32, s0))
127        sums(p, (j + 1 : U32), e, Array.fork(U32, a1), Array.fork(U32, s1))
128      two.join(l, r)
129
130# Pass 3
131# ------
132
133def puts.leaf.fin(s: Array<U32>, r: Array<U32> & U32) -> Two:
134  (a, +t) = r
135  Two{a, s}
136
137def puts.leaf.at(+inc: Bool, +i: U32, +e: Nat, a: Array<U32>, r: Array<U32> & U32) -> Two:
138  (s, +base) = r
139  +at = U32.shln(i, e)
140  puts.leaf.fin(s, go(inc, U32.to_nat((U32.shln(1, e) - 1 : U32)), at, base, a[at]))
141
142# Block i: its scan in place, from its base s[i].
143def puts.leaf(+inc: Bool, +i: U32, +e: Nat, a: Array<U32>, s: Array<U32>) -> Two:
144  puts.leaf.at(inc, i, e, a, s[i])
145
146# Pass 3's forks, shaped as pass 1's. Leaf i reads s[i] and reads and
147# writes only its block of a: no two lanes touch one slot.
148def puts(+inc: Bool, +d: Nat, +i: U32, +e: Nat, fa: Array<U32> & Array<U32>, fs: Array<U32> & Array<U32>) -> Two:
149  match d:
150    case 0n:
151      (a, a2) = fa
152      (s, s2) = fs
153      two.rejoin(a2, s2, puts.leaf(inc, i, e, a, s))
154    case 1n+p:
155      (a0, a1) = fa
156      (s0, s1) = fs
157      +j = (i * 2 : U32)
158      l r = puts(inc, p, j, e, Array.fork(U32, a0), Array.fork(U32, s0))
159        puts(inc, p, (j + 1 : U32), e, Array.fork(U32, a1), Array.fork(U32, s1))
160      two.join(l, r)
161
162# Driver
163# ------
164
165# floor(d / 2).
166def half(+d: Nat) -> Nat:
167  match d:
168    case 0n:
169      0n
170    case 1n+p:
171      match p:
172        case 0n:
173          0n
174        case 1n+q:
175          1n+half(q)
176
177def run.fin(r: Two, +total: U32) -> Array<U32> & U32:
178  Two{a, s} = r
179  (a, total)
180
181# Pass 3, from pass 2's bases and total.
182def run.mid(+inc: Bool, +d: Nat, +e: Nat, a: Array<U32>, r: Array<U32> & U32) -> Array<U32> & U32:
183  (s, +total) = r
184  run.fin(puts(inc, d, 0, e, Array.fork(U32, a), Array.fork(U32, s)), total)
185
186# Pass 1 over a (2^lg entries) at depth d, into a new s.
187def run.pre(+d: Nat, +e: Nat, a: Array<U32>) -> Two:
188  sums(d, 0, e, Array.fork(U32, a), Array.fork(U32, Array.new(U32, d, 0)))
189
190# The scan with pass 2 on one lane.
191def run0.sums(+inc: Bool, +d: Nat, +e: Nat, r: Two) -> Array<U32> & U32:
192  Two{a, s} = r
193  run.mid(inc, d, e, a, ex.go(U32.to_nat((U32.shln(1, d) - 1 : U32)), 0, 0, s[0]))
194
195def run0(+inc: Bool, +b: Nat, r: Array<U32> & U32) -> Array<U32> & U32:
196  (a, +n) = r
197  +lg = U32.log2(n)
198  +d = Nat.min(b, lg)
199  +e = Nat.sub(lg, d)
200  run0.sums(inc, d, e, run.pre(d, e, a))
201
202# The scan with pass 2 by run0 at depth d / 2.
203def run1.sums(+inc: Bool, +d: Nat, +e: Nat, r: Two) -> Array<U32> & U32:
204  Two{a, s} = r
205  run.mid(inc, d, e, a, run0(False{}, half(d), Array.size(U32, s)))
206
207# Pass 2 by run1 from 2^12 sums up, else on one lane: on the GPU we
208# measured, the nested scan's forks cost more than one lane's walk below
209# that.
210def run.pick(lo: Bool, +inc: Bool, +d: Nat, +e: Nat, r: Two) -> Array<U32> & U32:
211  match lo:
212    case True{}:
213      run0.sums(inc, d, e, r)
214    case False{}:
215      run1.sums(inc, d, e, r)
216
217def run(+inc: Bool, +b: Nat, r: Array<U32> & U32) -> Array<U32> & U32:
218  (a, +n) = r
219  +lg = U32.log2(n)
220  +d = Nat.min(b, lg)
221  +e = Nat.sub(lg, d)
222  run.pick(U32.is_lt(U32.from_nat(d), 12), inc, d, e, run.pre(d, e, a))
223
224# API
225# ---
226
227# a[j] becomes a[0] + .. + a[j - 1] (a[0] becomes 0); returns the array
228# and the total, a[0] + .. + a[n - 1].
229def exclusive(+b: Nat, a: Array<U32>) -> Array<U32> & U32:
230  run(False{}, b, Array.size(U32, a))
231
232# a[j] becomes a[0] + .. + a[j]; returns the array and the total.
233def inclusive(+b: Nat, a: Array<U32>) -> Array<U32> & U32:
234  run(True{}, b, Array.size(U32, a))