histogram.bend

bend-parallel@0.1.0.0 · 329 lines · 34 definitions

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).
40import Base
41import bend-parallel@0.1.0.0/scan.bend as Scan
42
43# Buffers
44# -------
45
46type Bufs is Type:
47  Bufs{counts: Array<U32>, table: Array<U32>}
48
49# Empty buffers: the first call allocates what it needs.
50def Bufs.new() -> Bufs:
51  Bufs{[0 : U32^0n], [0 : U32^0n]}
52
53def 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.
61def 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.
65def 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
69def fit.pick(ok: Bool, +d: Nat, t: Array<U32>) -> Array<U32>:
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.
77def fit.eq(+d: Nat, r: Array<U32> & U32) -> Array<U32>:
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.
82def fit.ge(+d: Nat, r: Array<U32> & U32) -> Array<U32>:
83  (t, +n) = r
84  fit.pick(U32.is_ge(n, U32.shln(1, d)), d, t)
85
86# Handles
87# -------
88
89type Two is Type:
90  Two{a: Array<U32>, b: Array<U32>}
91
92def 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
97def two.rejoin(a2: Array<U32>, b2: Array<U32>, 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.
105def zero.go(+n: Nat, +at: U32, t: Array<U32>) -> Array<U32>:
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
112def bump.at(+i: U32, r: Array<U32> & U32) -> Array<U32>:
113  (t, +v) = r
114  t[i] <- (v + 1 : U32)
115
116# One more at t[i].
117def bump(+i: U32, t: Array<U32>) -> Array<U32>:
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.
122def count.go(+n: Nat, +j: U32, +row: U32, +k1: U32, t: Array<U32>, r: Array<U32> & 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).
134def rows.leaf(+i: U32, +e: Nat, +kb: Nat, +k: U32, ks: Array<U32>, t: Array<U32>) -> 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.
144def rows(+d: Nat, +i: U32, +e: Nat, +kb: Nat, +k: U32, fk: Array<U32> & Array<U32>, ft: Array<U32> & Array<U32>) -> 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.
167def 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.
172def col.go(+n: Nat, +at: U32, +step: U32, +s: U32, r: Array<U32> & U32) -> Array<U32> & 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).
183def ins.go(+n: Nat, +at: U32, +step: U32, +s: U32, r: Array<U32> & U32) -> Array<U32>:
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
194def grp.put(+at: U32, r: Array<U32> & U32) -> Array<U32>:
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).
201def grp.one(lt: Bool, +sc: Bool, +c: U32, +top: U32, +last: U32, +step: U32, +nm: Nat, t: Array<U32>) -> Array<U32>:
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.
213def grp.go(+n: Nat, +sc: Bool, +c: U32, +k: U32, +top: U32, +last: U32, +step: U32, +nm: Nat, t: Array<U32>) -> Array<U32>:
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).
222def grp.leaf(+sc: Bool, +i: U32, +gc: Nat, +w: Nat, +m: Nat, +kb: Nat, +k: U32, t: Array<U32>) -> Array<U32>:
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.
233def grp(+d: Nat, +sc: Bool, +i: U32, +gc: Nat, +w: Nat, +m: Nat, +kb: Nat, +k: U32, ft: Array<U32> & Array<U32>) -> Array<U32>:
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
245def col.put(+c: U32, cs: Array<U32>, r: Array<U32> & 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.
251def 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.
261def 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.
272def cols(+d: Nat, +j: U32, +w: Nat, +k: U32, +o: U32, +step: U32, +nb: Nat, ft: Array<U32> & Array<U32>, fc: Array<U32> & Array<U32>) -> 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
289def count.fin(keys: Array<U32>, r: Two) -> Array<U32> & 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.
295def count.cols(+r: Nat, +m: Nat, +g: Nat, +k: U32, +kb: Nat, keys: Array<U32>, counts: Array<U32>, table: Array<U32>) -> Array<U32> & 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.
301def count.grp(lo: Bool, +d: Nat, +k: U32, +kb: Nat, counts: Array<U32>, p: Two) -> Array<U32> & 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
313def count.at(+b: Nat, +k: U32, bufs: Bufs, r: Array<U32> & U32) -> Array<U32> & 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.
328def count(+b: Nat, +k: U32, keys: Array<U32>, bufs: Bufs) -> Array<U32> & Bufs:
329  count.at(b, k, bufs, Array.size(U32, keys))