1 # scan.bend: parallel prefix sums over an Array, 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), 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. 34 import 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. 41 def sum.go(+n: Nat, +j: U32, +s: U32, r: Array & U32) -> Array & 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. 54 def ex.go(+n: Nat, +j: U32, +s: U32, r: Array & U32) -> Array & 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. 65 def in.go(+n: Nat, +j: U32, +s: U32, r: Array & U32) -> Array & 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. 78 def go(+inc: Bool, +n: Nat, +j: U32, +s: U32, r: Array & U32) -> Array & 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 88 type Two is Type: 89 Two{a: Array, s: Array} 90 91 def 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. 97 def two.rejoin(a2: Array, s2: Array, 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 104 def sums.leaf.fin(+i: U32, s: Array, r: Array & U32) -> Two: 105 (a, +v) = r 106 Two{a, s[i] <- v} 107 108 # Block i (2^e entries): its sum into s[i]. 109 def sums.leaf(+i: U32, +e: Nat, a: Array, s: Array) -> 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. 116 def sums(+d: Nat, +i: U32, +e: Nat, fa: Array & Array, fs: Array & Array) -> 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 133 def puts.leaf.fin(s: Array, r: Array & U32) -> Two: 134 (a, +t) = r 135 Two{a, s} 136 137 def puts.leaf.at(+inc: Bool, +i: U32, +e: Nat, a: Array, r: Array & 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]. 143 def puts.leaf(+inc: Bool, +i: U32, +e: Nat, a: Array, s: Array) -> 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. 148 def puts(+inc: Bool, +d: Nat, +i: U32, +e: Nat, fa: Array & Array, fs: Array & Array) -> 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). 166 def 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 177 def run.fin(r: Two, +total: U32) -> Array & U32: 178 Two{a, s} = r 179 (a, total) 180 181 # Pass 3, from pass 2's bases and total. 182 def run.mid(+inc: Bool, +d: Nat, +e: Nat, a: Array, r: Array & U32) -> Array & 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. 187 def run.pre(+d: Nat, +e: Nat, a: Array) -> 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. 191 def run0.sums(+inc: Bool, +d: Nat, +e: Nat, r: Two) -> Array & 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 195 def run0(+inc: Bool, +b: Nat, r: Array & U32) -> Array & 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. 203 def run1.sums(+inc: Bool, +d: Nat, +e: Nat, r: Two) -> Array & 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. 210 def run.pick(lo: Bool, +inc: Bool, +d: Nat, +e: Nat, r: Two) -> Array & 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 217 def run(+inc: Bool, +b: Nat, r: Array & U32) -> Array & 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]. 229 def exclusive(+b: Nat, a: Array) -> Array & U32: 230 run(False{}, b, Array.size(U32, a)) 231 232 # a[j] becomes a[0] + .. + a[j]; returns the array and the total. 233 def inclusive(+b: Nat, a: Array) -> Array & U32: 234 run(True{}, b, Array.size(U32, a))