1 # histogram.bend: a parallel histogram of U32 keys into K buckets, with no 2 # atomics. 3 # 4 # import ./histogram.bend as Hist 5 # Hist.count(8n, 256, keys, Hist.Bufs.new()) # (keys, bufs) 6 # Hist.Bufs{counts, table} = bufs # counts[c]: the keys in bucket c 7 # 8 # k is the bucket count K (0 counts as 1), any U32. A key at or past K 9 # counts as K - 1. counts holds 2^ceil(log2 K) entries: the K counts, then 10 # zeros. b is the fork depth, as in scan.bend: 2^b blocks of keys (b is 11 # lowered to log2 of the key count when larger), and every b gives the 12 # same counts. A call under `!` runs its forks on the GPU. 13 # 14 # Bufs carries the two buffers from call to call: the count table (2^b 15 # rows of 2^ceil(log2 K) entries, so 2^(b + ceil(log2 K)) U32s) and the 16 # counts. A call keeps a buffer that is big enough and allocates a new one 17 # otherwise. Array.new fills on one lane, which under `!` costs about 76 ms 18 # per million entries on the GPU we measured (see README), so size the 19 # buffers once with Bufs.sized outside the `!` and hand them back each 20 # call. 21 # 22 # Two passes: 23 # 1. per block of keys, in parallel: zero the block's row of the table, 24 # then count its keys into it (the block owns its row); 25 # 2. the column sums. Below 2^12 rows, per group of buckets, in 26 # parallel: a bucket's count is the sum of its column over the rows 27 # (the group owns its counts). From 2^12 rows up, in two steps: the 28 # 2^b rows are cut into 2^(b/2) groups of rows; a) per (group of 29 # rows, group of buckets), in parallel: each bucket's sum over the 30 # group's rows into the group's last row; b) per group of buckets, in 31 # parallel: a bucket's count is the sum of those. Every lane walks 32 # 2^(b/2) rows at most, not 2^b. 33 # After a call the table holds scratch (from 2^12 rows up, partial sums 34 # in the groups' last rows), not per-block counts. ref.bend's histogram 35 # is the sequential reference. 36 # 37 # Sharing: as in scan.bend, lanes share arrays through Base's @unsafe 38 # Array.fork and Array.join; no def here is @unsafe itself, and none may 39 # be evaluated in the checker (call them from an IO main). 40 import Base 41 import ./scan.bend as Scan 42 43 # Buffers 44 # ------- 45 46 type Bufs is Type: 47 Bufs{counts: Array, table: Array} 48 49 # Empty buffers: the first call allocates what it needs. 50 def Bufs.new() -> Bufs: 51 Bufs{[0 : U32^0n], [0 : U32^0n]} 52 53 def kbits.of(big: Bool, +k: U32) -> Nat: 54 match big: 55 case False{}: 56 0n 57 case True{}: 58 1n+U32.log2((k - 1 : U32)) 59 60 # ceil(log2 K), for K >= 1. 61 def kbits(+k: U32) -> Nat: 62 kbits.of(U32.is_lt(1, k), k) 63 64 # Buffers for 2^lg keys at fork depth b and k buckets, allocated now. 65 def Bufs.sized(+b: Nat, +k: U32, +lg: Nat) -> Bufs: 66 +kb = kbits(U32.max(k, 1)) 67 Bufs{Array.new(U32, kb, 0), Array.new(U32, Nat.add(Nat.min(b, lg), kb), 0)} 68 69 def fit.pick(ok: Bool, +d: Nat, t: Array) -> Array: 70 match ok: 71 case True{}: 72 t 73 case False{}: 74 Array.new(U32, d, 0) 75 76 # t when it holds exactly 2^d entries, else a new one. 77 def fit.eq(+d: Nat, r: Array & U32) -> Array: 78 (t, +n) = r 79 fit.pick(U32.is_eq(n, U32.shln(1, d)), d, t) 80 81 # t when it holds at least 2^d entries, else a new one. 82 def fit.ge(+d: Nat, r: Array & U32) -> Array: 83 (t, +n) = r 84 fit.pick(U32.is_ge(n, U32.shln(1, d)), d, t) 85 86 # Handles 87 # ------- 88 89 type Two is Type: 90 Two{a: Array, b: Array} 91 92 def two.join(l: Two, r: Two) -> Two: 93 Two{la, lb} = l 94 Two{ra, rb} = r 95 Two{Array.join(U32, la, ra), Array.join(U32, lb, rb)} 96 97 def two.rejoin(a2: Array, b2: Array, r: Two) -> Two: 98 Two{a, b} = r 99 Two{Array.join(U32, a, a2), Array.join(U32, b, b2)} 100 101 # Pass 1: rows 102 # ------------ 103 104 # n + 1 zeros from at. 105 def zero.go(+n: Nat, +at: U32, t: Array) -> Array: 106 match n: 107 case 0n: 108 t[at] <- 0 109 case 1n+p: 110 zero.go(p, (at + 1 : U32), t[at] <- 0) 111 112 def bump.at(+i: U32, r: Array & U32) -> Array: 113 (t, +v) = r 114 t[i] <- (v + 1 : U32) 115 116 # One more at t[i]. 117 def bump(+i: U32, t: Array) -> Array: 118 bump.at(i, t[i]) 119 120 # r holds the keys and the key x read at j: one more in x's bucket of the 121 # row at row, then the same for the n keys after j. k1 is K - 1. 122 def count.go(+n: Nat, +j: U32, +row: U32, +k1: U32, t: Array, r: Array & U32) -> Two: 123 match n: 124 case 0n: 125 (ks, +x) = r 126 Two{ks, bump((row + U32.min(x, k1) : U32), t)} 127 case 1n+p: 128 (ks, +x) = r 129 +q = (j + 1 : U32) 130 count.go(p, q, row, k1, bump((row + U32.min(x, k1) : U32), t), ks[q]) 131 132 # Block i (2^e keys from i * 2^e): its row, at i * 2^kb, zeroed, then its 133 # keys counted into it. k is K (at least 1). 134 def rows.leaf(+i: U32, +e: Nat, +kb: Nat, +k: U32, ks: Array, t: Array) -> Two: 135 +row = U32.shln(i, kb) 136 +at = U32.shln(i, e) 137 count.go(U32.to_nat((U32.shln(1, e) - 1 : U32)), at, row, (k - 1 : U32), 138 zero.go(U32.to_nat((k - 1 : U32)), row, t), ks[at]) 139 140 # Pass 1's forks: d levels down to the blocks under block index i. fk and 141 # ft are two handles each to the keys and the table. Block i reads only 142 # its keys and reads and writes only its row (entries i * 2^kb to 143 # i * 2^kb + K - 1): no two lanes touch one slot. 144 def rows(+d: Nat, +i: U32, +e: Nat, +kb: Nat, +k: U32, fk: Array & Array, ft: Array & Array) -> Two: 145 match d: 146 case 0n: 147 (ks, ks2) = fk 148 (t, t2) = ft 149 two.rejoin(ks2, t2, rows.leaf(i, e, kb, k, ks, t)) 150 case 1n+p: 151 (k0, k1) = fk 152 (t0, t1) = ft 153 +j = (i * 2 : U32) 154 l r = rows(p, j, e, kb, k, Array.fork(U32, k0), Array.fork(U32, t0)) 155 rows(p, (j + 1 : U32), e, kb, k, Array.fork(U32, k1), Array.fork(U32, t1)) 156 two.join(l, r) 157 158 # Pass 2: columns 159 # --------------- 160 161 # From 2^12 rows up (see small, below), the table's 2^d rows are cut into 162 # 2^r groups of 2^m (r = d / 2): the column passes step over the rows of 163 # a group, then over the groups. Below, one pass steps over all the rows: 164 # on the GPU we measured, the split costs more than it saves there. 165 166 # d is below 12: the table is small, and its rows are not split. 167 def small(+d: Nat) -> Bool: 168 U32.is_lt(U32.from_nat(d), 12) 169 170 # r holds the table and the entry x read at at: adds x and the n entries 171 # after it, each step further on, to s. 172 def col.go(+n: Nat, +at: U32, +step: U32, +s: U32, r: Array & U32) -> Array & U32: 173 match n: 174 case 0n: 175 (t, +x) = r 176 (t, (s + x : U32)) 177 case 1n+p: 178 (t, +x) = r 179 +q = (at + step : U32) 180 col.go(p, q, step, (s + x : U32), t[q]) 181 182 # As col.go, but each entry becomes the sum through it (in place). 183 def ins.go(+n: Nat, +at: U32, +step: U32, +s: U32, r: Array & U32) -> Array: 184 match n: 185 case 0n: 186 (t, +x) = r 187 t[at] <- (s + x : U32) 188 case 1n+p: 189 (t, +x) = r 190 +v = (s + x : U32) 191 +q = (at + step : U32) 192 ins.go(p, q, step, v, (t[at] <- v)[q]) 193 194 def grp.put(+at: U32, r: Array & U32) -> Array: 195 (t, +v) = r 196 t[at] <- v 197 198 # Column c of a group (rows from top, nm + 1 of them, step apart), when 199 # c < K: sc, each entry becomes the group's sum through it; else the sum 200 # goes into the last row (at last). 201 def grp.one(lt: Bool, +sc: Bool, +c: U32, +top: U32, +last: U32, +step: U32, +nm: Nat, t: Array) -> Array: 202 match lt: 203 case True{}: 204 match sc: 205 case True{}: 206 ins.go(nm, (top + c : U32), step, 0, t[(top + c : U32)]) 207 case False{}: 208 grp.put((last + c : U32), col.go(nm, (top + c : U32), step, 0, t[(top + c : U32)])) 209 case False{}: 210 t 211 212 # Columns c to c + n of a group. 213 def grp.go(+n: Nat, +sc: Bool, +c: U32, +k: U32, +top: U32, +last: U32, +step: U32, +nm: Nat, t: Array) -> Array: 214 match n: 215 case 0n: 216 grp.one(U32.is_lt(c, k), sc, c, top, last, step, nm, t) 217 case 1n+q: 218 grp.go(q, sc, (c + 1 : U32), k, top, last, step, nm, grp.one(U32.is_lt(c, k), sc, c, top, last, step, nm, t)) 219 220 # Lane i: rows group i / 2^gc (2^m rows of 2^kb), columns group i % 2^gc 221 # (2^w columns). 222 def grp.leaf(+sc: Bool, +i: U32, +gc: Nat, +w: Nat, +m: Nat, +kb: Nat, +k: U32, t: Array) -> Array: 223 +g = U32.shrn(i, gc) 224 +j = (i - U32.shln(g, gc) : U32) 225 +top = U32.shln(g, Nat.add(m, kb)) 226 +last = (top + U32.shln((U32.shln(1, m) - 1 : U32), kb) : U32) 227 grp.go(U32.to_nat((U32.shln(1, w) - 1 : U32)), sc, U32.shln(j, w), k, top, last, U32.shln(1, kb), 228 U32.to_nat((U32.shln(1, m) - 1 : U32)), t) 229 230 # Pass 2a's forks (sort.bend's too, with sc): d levels down to the lanes 231 # under lane i, with two handles to the table. Lane i reads and writes 232 # only its group's rows in its columns: no two lanes touch one slot. 233 def grp(+d: Nat, +sc: Bool, +i: U32, +gc: Nat, +w: Nat, +m: Nat, +kb: Nat, +k: U32, ft: Array & Array) -> Array: 234 match d: 235 case 0n: 236 (t, t2) = ft 237 Array.join(U32, grp.leaf(sc, i, gc, w, m, kb, k, t), t2) 238 case 1n+p: 239 (t0, t1) = ft 240 +j = (i * 2 : U32) 241 l r = grp(p, sc, j, gc, w, m, kb, k, Array.fork(U32, t0)) 242 grp(p, sc, (j + 1 : U32), gc, w, m, kb, k, Array.fork(U32, t1)) 243 Array.join(U32, l, r) 244 245 def col.put(+c: U32, cs: Array, r: Array & U32) -> Two: 246 (t, +v) = r 247 Two{t, cs[c] <- v} 248 249 # Column c's count into cs[c]: the sum of its nb + 1 entries from o + c, 250 # step apart (the groups' last rows), when c < K; else 0. 251 def col.one(lt: Bool, +c: U32, +o: U32, +step: U32, +nb: Nat, p: Two) -> Two: 252 match lt: 253 case True{}: 254 Two{t, cs} = p 255 col.put(c, cs, col.go(nb, (o + c : U32), step, 0, t[(o + c : U32)])) 256 case False{}: 257 Two{t, cs} = p 258 Two{t, cs[c] <- 0} 259 260 # Columns c to c + n. 261 def cols.go(+n: Nat, +c: U32, +k: U32, +o: U32, +step: U32, +nb: Nat, p: Two) -> Two: 262 match n: 263 case 0n: 264 col.one(U32.is_lt(c, k), c, o, step, nb, p) 265 case 1n+q: 266 cols.go(q, (c + 1 : U32), k, o, step, nb, col.one(U32.is_lt(c, k), c, o, step, nb, p)) 267 268 # Pass 2b's forks: d levels down to groups of 2^w columns under group index 269 # j. ft and fc are two handles each to the table and the counts. Group j 270 # reads only its columns of the table and writes only its counts (entries 271 # j * 2^w to j * 2^w + 2^w - 1): no two lanes write one slot. 272 def cols(+d: Nat, +j: U32, +w: Nat, +k: U32, +o: U32, +step: U32, +nb: Nat, ft: Array & Array, fc: Array & Array) -> Two: 273 match d: 274 case 0n: 275 (t, t2) = ft 276 (cs, cs2) = fc 277 two.rejoin(t2, cs2, cols.go(U32.to_nat((U32.shln(1, w) - 1 : U32)), U32.shln(j, w), k, o, step, nb, Two{t, cs})) 278 case 1n+p: 279 (t0, t1) = ft 280 (c0, c1) = fc 281 +i = (j * 2 : U32) 282 l r = cols(p, i, w, k, o, step, nb, Array.fork(U32, t0), Array.fork(U32, c0)) 283 cols(p, (i + 1 : U32), w, k, o, step, nb, Array.fork(U32, t1), Array.fork(U32, c1)) 284 two.join(l, r) 285 286 # Driver 287 # ------ 288 289 def count.fin(keys: Array, r: Two) -> Array & Bufs: 290 Two{table, counts} = r 291 (keys, Bufs{counts, table}) 292 293 # Pass 2b over the 2^kb columns, 2^min(d, kb) groups: each sums its 294 # columns over the 2^r groups' last rows. 295 def count.cols(+r: Nat, +m: Nat, +g: Nat, +k: U32, +kb: Nat, keys: Array, counts: Array, table: Array) -> Array & Bufs: 296 count.fin(keys, cols(g, 0, Nat.sub(kb, g), k, U32.shln((U32.shln(1, m) - 1 : U32), kb), U32.shln(1, Nat.add(m, kb)), 297 U32.to_nat((U32.shln(1, r) - 1 : U32)), Array.fork(U32, table), Array.fork(U32, fit.eq(kb, Array.size(U32, counts))))) 298 299 # Pass 2a over 2^r groups of 2^m rows by 2^min(d, kb) groups of columns, 300 # then 2b; below 2^12 rows, 2b alone, over 2^d groups of one row. 301 def count.grp(lo: Bool, +d: Nat, +k: U32, +kb: Nat, counts: Array, p: Two) -> Array & Bufs: 302 match lo: 303 case True{}: 304 Two{keys, table} = p 305 count.cols(d, 0n, Nat.min(d, kb), k, kb, keys, counts, table) 306 case False{}: 307 Two{keys, table} = p 308 +g = Nat.min(d, kb) 309 +r = Scan.half(d) 310 +m = Nat.sub(d, r) 311 count.cols(r, m, g, k, kb, keys, counts, grp(Nat.add(r, g), False{}, 0, g, Nat.sub(kb, g), m, kb, k, Array.fork(U32, table))) 312 313 def count.at(+b: Nat, +k: U32, bufs: Bufs, r: Array & U32) -> Array & Bufs: 314 Bufs{counts, table} = bufs 315 (keys, +n) = r 316 +kk = U32.max(k, 1) 317 +kb = kbits(kk) 318 +lg = U32.log2(n) 319 +d = Nat.min(b, lg) 320 count.grp(small(d), d, kk, kb, counts, rows(d, 0, Nat.sub(lg, d), kb, kk, Array.fork(U32, keys), 321 Array.fork(U32, fit.ge(Nat.add(d, kb), Array.size(U32, table))))) 322 323 # API 324 # --- 325 326 # The histogram of keys into k buckets at fork depth b. Returns the keys, 327 # unchanged, and the buffers, whose counts field holds the counts. 328 def count(+b: Nat, +k: U32, keys: Array, bufs: Bufs) -> Array & Bufs: 329 count.at(b, k, bufs, Array.size(U32, keys))