sort.bend

bend-parallel@0.1.0.0 · 515 lines · 53 definitions

1# sort.bend: a parallel stable counting sort by U32 key, with no atomics,
2# and a parallel gather to reorder payloads by its permutation.
3#
4#   import ./sort.bend as Sort
5#   Sort.by_key(8n, 256, keys, Sort.Bufs.new())   # (keys, bufs)
6#   Sort.Bufs{perm, counts, offs, table} = bufs
7#   # perm[j]: the input index of the key the sort puts at j
8#   # counts[c], offs[c]: bucket c's size and first slot
9#   Sort.gather(~U32, 8n, perm, keys, [0 : U32^d])   # Moved{perm, keys, sorted}
10#
11# (As in scan.bend: take a call's result apart in a def's parameter.)
12#
13# k is the bucket count K (0 counts as 1), any U32; a key at or past K
14# sorts as K - 1. The sort is stable: keys of one bucket keep their input
15# order, so the result equals ref.bend's perm, offsets and histogram bit
16# for bit, at every fork depth b. counts and offs hold 2^ceil(log2 K)
17# entries: past K, counts are 0 and offs are n. A call under `!` runs its
18# forks on the GPU.
19#
20# Bufs carries the buffers from call to call (the table of 2^b rows of
21# 2^ceil(log2 K) counts, the counts, the offsets and the permutation). A
22# call keeps a buffer that fits and allocates otherwise; under `!` an
23# allocation fills on one lane, so size them once with Bufs.sized outside
24# the `!` (see histogram.bend).
25#
26# Passes:
27#   1. per block of keys, in parallel: the block's counts per bucket into
28#      its row of the table (histogram.bend's pass 1);
29#   2. from 2^12 rows up, the table's rows are cut into 2^(b/2) groups
30#      of rows (see histogram.bend): a) per (group of rows, group of
31#      buckets), in parallel: each column becomes its running sum within
32#      the group;
33#      b) per group of buckets, in parallel: over the groups' last rows,
34#      each bucket's total into counts and offs, and each group's base
35#      (the counts of the groups above it) into its last row; then offs
36#      becomes its exclusive scan (scan.bend), the buckets' first slots;
37#      c) per (group of rows, group of buckets) again: each table entry
38#      becomes its bucket's first slot plus the counts of the rows above
39#      it (base plus the running sum of the row before it): where the
40#      block's keys of that bucket start. Below 2^12 rows, one group
41#      of all the rows, and no a): b) sums each column, c) walks it down;
42#   3. per block, in parallel: each key, in order, takes the next slot of
43#      its bucket in the block's row, and perm[slot] becomes its index.
44#
45# Sharing: as in scan.bend, lanes share arrays through Base's @unsafe
46# Array.fork and Array.join; no def here is @unsafe itself, and none may
47# be evaluated in the checker (call them from an IO main).
48import Base
49import bend-parallel@0.1.0.0/scan.bend as Scan
50import bend-parallel@0.1.0.0/histogram.bend as Hist
51
52# Buffers
53# -------
54
55type Bufs is Type:
56  Bufs{perm: Array<U32>, counts: Array<U32>, offs: Array<U32>, table: Array<U32>}
57
58# Empty buffers: the first call allocates what it needs.
59def Bufs.new() -> Bufs:
60  Bufs{[0 : U32^0n], [0 : U32^0n], [0 : U32^0n], [0 : U32^0n]}
61
62# Buffers for 2^lg keys at fork depth b and k buckets, allocated now.
63def Bufs.sized(+b: Nat, +k: U32, +lg: Nat) -> Bufs:
64  +kb = Hist.kbits(U32.max(k, 1))
65  Bufs{Array.new(U32, lg, 0), Array.new(U32, kb, 0), Array.new(U32, kb, 0),
66    Array.new(U32, Nat.add(Nat.min(b, lg), kb), 0)}
67
68# Handles
69# -------
70
71type Three is Type:
72  Three{a: Array<U32>, b: Array<U32>, c: Array<U32>}
73
74def three.join(l: Three, r: Three) -> Three:
75  Three{la, lb, lc} = l
76  Three{ra, rb, rc} = r
77  Three{Array.join(U32, la, ra), Array.join(U32, lb, rb), Array.join(U32, lc, rc)}
78
79def three.rejoin(a2: Array<U32>, b2: Array<U32>, c2: Array<U32>, r: Three) -> Three:
80  Three{a, b, c} = r
81  Three{Array.join(U32, a, a2), Array.join(U32, b, b2), Array.join(U32, c, c2)}
82
83# Pass 2b: over the groups, column totals and group bases
84# --------------------------------------------------------
85
86# r holds the table and the entry x read at at: it becomes s, and the n
87# entries after it (each step further on) follow, from s + x. Returns the
88# table and the sum past the last.
89def xs.go(+n: Nat, +at: U32, +step: U32, +s: U32, r: Array<U32> & U32) -> Array<U32> & U32:
90  match n:
91    case 0n:
92      (t, +x) = r
93      (t[at] <- s, (s + x : U32))
94    case 1n+p:
95      (t, +x) = r
96      +q = (at + step : U32)
97      xs.go(p, q, step, (s + x : U32), (t[at] <- s)[q])
98
99def sum.put(+c: U32, cs: Array<U32>, os: Array<U32>, r: Array<U32> & U32) -> Three:
100  (t, +v) = r
101  Three{t, cs[c] <- v, os[c] <- v}
102
103# The nb + 1 entries from at, step apart: their sum; when wr, each also
104# becomes the sum of those before it.
105def sum.col(+wr: Bool, +at: U32, +step: U32, +nb: Nat, t: Array<U32>) -> Array<U32> & U32:
106  match wr:
107    case True{}:
108      xs.go(nb, at, step, 0, t[at])
109    case False{}:
110      Hist.col.go(nb, at, step, 0, t[at])
111
112# Column c, when c < K: over its nb + 1 entries from o + c, step apart
113# (the groups' last rows, or all the rows of a small table), the total
114# into cs[c] and os[c], by sum.col. Else 0 into both.
115def sum.one(lt: Bool, +wr: Bool, +c: U32, +o: U32, +step: U32, +nb: Nat, p: Three) -> Three:
116  match lt:
117    case True{}:
118      Three{t, cs, os} = p
119      sum.put(c, cs, os, sum.col(wr, (o + c : U32), step, nb, t))
120    case False{}:
121      Three{t, cs, os} = p
122      Three{t, cs[c] <- 0, os[c] <- 0}
123
124# Columns c to c + n.
125def sum.go(+n: Nat, +wr: Bool, +c: U32, +k: U32, +o: U32, +step: U32, +nb: Nat, p: Three) -> Three:
126  match n:
127    case 0n:
128      sum.one(U32.is_lt(c, k), wr, c, o, step, nb, p)
129    case 1n+q:
130      sum.go(q, wr, (c + 1 : U32), k, o, step, nb, sum.one(U32.is_lt(c, k), wr, c, o, step, nb, p))
131
132# Pass 2b's forks: d levels down to groups of 2^w columns under group j,
133# with two handles each to the table, the counts and the offsets. Group j
134# reads (and when wr writes) only its columns of the table and writes
135# only its own entries of counts and offs (j * 2^w to j * 2^w + 2^w - 1).
136def sums(+d: Nat, +wr: Bool, +j: U32, +w: Nat, +k: U32, +o: U32, +step: U32, +nb: Nat, ft: Array<U32> & Array<U32>, fc: Array<U32> & Array<U32>, fo: Array<U32> & Array<U32>) -> Three:
137  match d:
138    case 0n:
139      (t, t2) = ft
140      (cs, cs2) = fc
141      (os, os2) = fo
142      three.rejoin(t2, cs2, os2, sum.go(U32.to_nat((U32.shln(1, w) - 1 : U32)), wr, U32.shln(j, w), k, o, step, nb, Three{t, cs, os}))
143    case 1n+p:
144      (t0, t1) = ft
145      (c0, c1) = fc
146      (o0, o1) = fo
147      +i = (j * 2 : U32)
148      l r = sums(p, wr, i, w, k, o, step, nb, Array.fork(U32, t0), Array.fork(U32, c0), Array.fork(U32, o0))
149        sums(p, wr, (i + 1 : U32), w, k, o, step, nb, Array.fork(U32, t1), Array.fork(U32, c1), Array.fork(U32, o1))
150      three.join(l, r)
151
152# Pass 2c: each table entry becomes its block's first slot
153# --------------------------------------------------------
154
155def base.fin(+q: U32, +s: U32, t: Array<U32>) -> Array<U32>:
156  t[q] <- s
157
158# r holds the table and x, the entry a row (step) above at, which holds
159# the group's running sum through that row: at becomes s + x, and so up
160# for the n rows above; the group's first row, the next above, becomes s.
161def base.go(+n: Nat, +at: U32, +step: U32, +s: U32, r: Array<U32> & U32) -> Array<U32>:
162  match n:
163    case 0n:
164      (t, +x) = r
165      base.fin((at - step : U32), s, t[at] <- (s + x : U32))
166    case 1n+p:
167      (t, +x) = r
168      +q = (at - step : U32)
169      base.go(p, q, step, s, (t[at] <- (s + x : U32))[(q - step : U32)])
170
171# A group's column from its last row at (nm + 1 rows): r holds the table
172# and the group's base b, which the last row holds; o is offs[c].
173def base.at(+at: U32, +step: U32, +nm: Nat, +o: U32, r: Array<U32> & U32) -> Array<U32>:
174  match nm:
175    case 0n:
176      (t, +b) = r
177      t[at] <- (o + b : U32)
178    case 1n+p:
179      (t, +b) = r
180      base.go(p, at, step, (o + b : U32), t[(at - step : U32)])
181
182def base.os(+c: U32, +last: U32, +step: U32, +nm: Nat, t: Array<U32>, r: Array<U32> & U32) -> Hist.Two:
183  (os, +o) = r
184  Hist.Two{base.at((last + c : U32), step, nm, o, t[(last + c : U32)]), os}
185
186# Column c of a group, when c < K: from its first slot offs[c] plus the
187# group's base, which its last row holds.
188def base.one(lt: Bool, +c: U32, +last: U32, +step: U32, +nm: Nat, p: Hist.Two) -> Hist.Two:
189  match lt:
190    case True{}:
191      Hist.Two{t, os} = p
192      base.os(c, last, step, nm, t, os[c])
193    case False{}:
194      p
195
196# Columns c to c + n of a group.
197def base.cols(+n: Nat, +c: U32, +k: U32, +last: U32, +step: U32, +nm: Nat, p: Hist.Two) -> Hist.Two:
198  match n:
199    case 0n:
200      base.one(U32.is_lt(c, k), c, last, step, nm, p)
201    case 1n+q:
202      base.cols(q, (c + 1 : U32), k, last, step, nm, base.one(U32.is_lt(c, k), c, last, step, nm, p))
203
204# Lane i: rows group i / 2^gc, columns group i % 2^gc (as Hist.grp.leaf).
205def base.leaf(+i: U32, +gc: Nat, +w: Nat, +m: Nat, +kb: Nat, +k: U32, p: Hist.Two) -> Hist.Two:
206  +g = U32.shrn(i, gc)
207  +j = (i - U32.shln(g, gc) : U32)
208  +last = (U32.shln(g, Nat.add(m, kb)) + U32.shln((U32.shln(1, m) - 1 : U32), kb) : U32)
209  base.cols(U32.to_nat((U32.shln(1, w) - 1 : U32)), U32.shln(j, w), k, last, U32.shln(1, kb),
210    U32.to_nat((U32.shln(1, m) - 1 : U32)), p)
211
212# Pass 2c's forks, shaped as pass 2a's (Hist.grp), with two handles each
213# to the table and the offsets. Lane i reads the offsets and reads and
214# writes only its group's rows in its columns.
215def bases(+d: Nat, +i: U32, +gc: Nat, +w: Nat, +m: Nat, +kb: Nat, +k: U32, ft: Array<U32> & Array<U32>, fo: Array<U32> & Array<U32>) -> Hist.Two:
216  match d:
217    case 0n:
218      (t, t2) = ft
219      (os, os2) = fo
220      Hist.two.rejoin(t2, os2, base.leaf(i, gc, w, m, kb, k, Hist.Two{t, os}))
221    case 1n+p:
222      (t0, t1) = ft
223      (o0, o1) = fo
224      +j = (i * 2 : U32)
225      l r = bases(p, j, gc, w, m, kb, k, Array.fork(U32, t0), Array.fork(U32, o0))
226        bases(p, (j + 1 : U32), gc, w, m, kb, k, Array.fork(U32, t1), Array.fork(U32, o1))
227      Hist.two.join(l, r)
228
229# Pass 2c for a small table (see Hist.small): each column in one walk
230# ------------------------------------------------------------------
231
232# r holds the table and the entry x read at at: it becomes s, and the n
233# entries after it (each step further on) follow, from s + x.
234def fwd.go(+n: Nat, +at: U32, +step: U32, +s: U32, r: Array<U32> & U32) -> Array<U32>:
235  match n:
236    case 0n:
237      (t, +x) = r
238      t[at] <- s
239    case 1n+p:
240      (t, +x) = r
241      +q = (at + step : U32)
242      fwd.go(p, q, step, (s + x : U32), (t[at] <- s)[q])
243
244def fwd.at(+c: U32, +kb: Nat, +nb: Nat, t: Array<U32>, r: Array<U32> & U32) -> Hist.Two:
245  (os, +s) = r
246  Hist.Two{fwd.go(nb, c, U32.shln(1, kb), s, t[c]), os}
247
248# Column c, when c < K: its entries from its first slot offs[c].
249def fwd.one(lt: Bool, +c: U32, +kb: Nat, +nb: Nat, p: Hist.Two) -> Hist.Two:
250  match lt:
251    case True{}:
252      Hist.Two{t, os} = p
253      fwd.at(c, kb, nb, t, os[c])
254    case False{}:
255      p
256
257# Columns c to c + n.
258def fwd.cols(+n: Nat, +c: U32, +k: U32, +kb: Nat, +nb: Nat, p: Hist.Two) -> Hist.Two:
259  match n:
260    case 0n:
261      fwd.one(U32.is_lt(c, k), c, kb, nb, p)
262    case 1n+q:
263      fwd.cols(q, (c + 1 : U32), k, kb, nb, fwd.one(U32.is_lt(c, k), c, kb, nb, p))
264
265# The forks, shaped as pass 2b's, with two handles each to the table and
266# the offsets. Group j reads only its own offsets and reads and writes
267# only its columns of the table.
268def fwds(+d: Nat, +j: U32, +w: Nat, +k: U32, +kb: Nat, +nb: Nat, ft: Array<U32> & Array<U32>, fo: Array<U32> & Array<U32>) -> Hist.Two:
269  match d:
270    case 0n:
271      (t, t2) = ft
272      (os, os2) = fo
273      Hist.two.rejoin(t2, os2, fwd.cols(U32.to_nat((U32.shln(1, w) - 1 : U32)), U32.shln(j, w), k, kb, nb, Hist.Two{t, os}))
274    case 1n+p:
275      (t0, t1) = ft
276      (o0, o1) = fo
277      +i = (j * 2 : U32)
278      l r = fwds(p, i, w, k, kb, nb, Array.fork(U32, t0), Array.fork(U32, o0))
279        fwds(p, (i + 1 : U32), w, k, kb, nb, Array.fork(U32, t1), Array.fork(U32, o1))
280      Hist.two.join(l, r)
281
282# Pass 3: scatter
283# ---------------
284
285def put.at(+j: U32, +slot: U32, pm: Array<U32>, r: Array<U32> & U32) -> Hist.Two:
286  (t, +at) = r
287  Hist.Two{t[slot] <- (at + 1 : U32), pm[at] <- j}
288
289# Key j, whose bucket's next slot is t[slot]: perm[that slot] <- j, and
290# the slot moves on.
291def put.one(+j: U32, +slot: U32, p: Hist.Two) -> Hist.Two:
292  Hist.Two{t, pm} = p
293  put.at(j, slot, pm, t[slot])
294
295def put.fin(ks: Array<U32>, p: Hist.Two) -> Three:
296  Hist.Two{t, pm} = p
297  Three{ks, t, pm}
298
299# r holds the keys and the key x read at j: key j goes to the next slot
300# of its bucket in the row at row, then the n keys after j. k1 is K - 1.
301def put.go(+n: Nat, +j: U32, +row: U32, +k1: U32, p: Hist.Two, r: Array<U32> & U32) -> Three:
302  match n:
303    case 0n:
304      (ks, +x) = r
305      put.fin(ks, put.one(j, (row + U32.min(x, k1) : U32), p))
306    case 1n+q:
307      (ks, +x) = r
308      +i = (j + 1 : U32)
309      put.go(q, i, row, k1, put.one(j, (row + U32.min(x, k1) : U32), p), ks[i])
310
311# Block i (2^e keys from i * 2^e, its row at i * 2^kb).
312def puts.leaf(+i: U32, +e: Nat, +kb: Nat, +k: U32, ks: Array<U32>, t: Array<U32>, pm: Array<U32>) -> Three:
313  +at = U32.shln(i, e)
314  put.go(U32.to_nat((U32.shln(1, e) - 1 : U32)), at, U32.shln(i, kb), (k - 1 : U32), Hist.Two{t, pm}, ks[at])
315
316# Pass 3's forks: d levels down to the blocks under block i, with two
317# handles each to the keys, the table and perm. Block i reads only its
318# keys, reads and writes only its row of the table, and writes perm only
319# at the slots its row hands it: pass 2 gave each (block, bucket) pair its
320# own run of slots, so no two lanes write one slot.
321def puts(+d: Nat, +i: U32, +e: Nat, +kb: Nat, +k: U32, fk: Array<U32> & Array<U32>, ft: Array<U32> & Array<U32>, fp: Array<U32> & Array<U32>) -> Three:
322  match d:
323    case 0n:
324      (ks, ks2) = fk
325      (t, t2) = ft
326      (pm, pm2) = fp
327      three.rejoin(ks2, t2, pm2, puts.leaf(i, e, kb, k, ks, t, pm))
328    case 1n+p:
329      (k0, k1) = fk
330      (t0, t1) = ft
331      (p0, p1) = fp
332      +j = (i * 2 : U32)
333      l r = puts(p, j, e, kb, k, Array.fork(U32, k0), Array.fork(U32, t0), Array.fork(U32, p0))
334        puts(p, (j + 1 : U32), e, kb, k, Array.fork(U32, k1), Array.fork(U32, t1), Array.fork(U32, p1))
335      three.join(l, r)
336
337# Driver
338# ------
339
340def run.fin(cs: Array<U32>, os: Array<U32>, r: Three) -> Array<U32> & Bufs:
341  Three{ks, t, pm} = r
342  (ks, Bufs{pm, cs, os, t})
343
344# Pass 3.
345def run.puts(+d: Nat, +e: Nat, +kb: Nat, +k: U32, ks: Array<U32>, pm: Array<U32>, cs: Array<U32>, r: Hist.Two) -> Array<U32> & Bufs:
346  Hist.Two{t, os} = r
347  run.fin(cs, os, puts(d, 0, e, kb, k, Array.fork(U32, ks), Array.fork(U32, t), Array.fork(U32, pm)))
348
349# Pass 2c: by fwds when lo (a small table), else by bases.
350def run.bases(lo: Bool, +d: Nat, +e: Nat, +kb: Nat, +k: U32, ks: Array<U32>, pm: Array<U32>, cs: Array<U32>, t: Array<U32>, r: Array<U32> & U32) -> Array<U32> & Bufs:
351  match lo:
352    case True{}:
353      (os, +n) = r
354      +g = Nat.min(d, kb)
355      run.puts(d, e, kb, k, ks, pm, cs, fwds(g, 0, Nat.sub(kb, g), k, kb,
356        U32.to_nat((U32.shln(1, d) - 1 : U32)), Array.fork(U32, t), Array.fork(U32, os)))
357    case False{}:
358      (os, +n) = r
359      +g = Nat.min(d, kb)
360      +h = Scan.half(d)
361      run.puts(d, e, kb, k, ks, pm, cs, bases(Nat.add(h, g), 0, g, Nat.sub(kb, g), Nat.sub(d, h), kb, k,
362        Array.fork(U32, t), Array.fork(U32, os)))
363
364# Pass 2b's scan: offs becomes its exclusive scan.
365def run.scan(lo: Bool, +d: Nat, +e: Nat, +kb: Nat, +k: U32, ks: Array<U32>, pm: Array<U32>, r: Three) -> Array<U32> & Bufs:
366  Three{t, cs, os} = r
367  run.bases(lo, d, e, kb, k, ks, pm, cs, t, Scan.exclusive(d, os))
368
369# Pass 2b, after 2a (t): over the 2^h groups' last rows.
370def run.tots(+d: Nat, +e: Nat, +kb: Nat, +k: U32, ks: Array<U32>, pm: Array<U32>, cs: Array<U32>, os: Array<U32>, t: Array<U32>) -> Array<U32> & Bufs:
371  +g = Nat.min(d, kb)
372  +h = Scan.half(d)
373  +m = Nat.sub(d, h)
374  run.scan(False{}, d, e, kb, k, ks, pm, sums(g, True{}, 0, Nat.sub(kb, g), k, U32.shln((U32.shln(1, m) - 1 : U32), kb),
375    U32.shln(1, Nat.add(m, kb)), U32.to_nat((U32.shln(1, h) - 1 : U32)), Array.fork(U32, t),
376    Array.fork(U32, Hist.fit.eq(kb, Array.size(U32, cs))),
377    Array.fork(U32, Hist.fit.eq(kb, Array.size(U32, os)))))
378
379# Pass 2a: each group's running column sums (Hist.grp), then 2b. For a
380# small table (lo), no 2a: 2b sums each column over all the rows.
381def run.sums(lo: Bool, +d: Nat, +e: Nat, +kb: Nat, +k: U32, pm: Array<U32>, cs: Array<U32>, os: Array<U32>, r: Hist.Two) -> Array<U32> & Bufs:
382  match lo:
383    case True{}:
384      Hist.Two{ks, t} = r
385      +g = Nat.min(d, kb)
386      run.scan(True{}, d, e, kb, k, ks, pm, sums(g, False{}, 0, Nat.sub(kb, g), k, 0, U32.shln(1, kb),
387        U32.to_nat((U32.shln(1, d) - 1 : U32)), Array.fork(U32, t),
388        Array.fork(U32, Hist.fit.eq(kb, Array.size(U32, cs))),
389        Array.fork(U32, Hist.fit.eq(kb, Array.size(U32, os)))))
390    case False{}:
391      Hist.Two{ks, t} = r
392      +g = Nat.min(d, kb)
393      +h = Scan.half(d)
394      run.tots(d, e, kb, k, ks, pm, cs, os, Hist.grp(Nat.add(h, g), True{}, 0, g, Nat.sub(kb, g), Nat.sub(d, h), kb, k,
395        Array.fork(U32, t)))
396
397# Pass 1.
398def run(+b: Nat, +k: U32, bufs: Bufs, r: Array<U32> & U32) -> Array<U32> & Bufs:
399  Bufs{pm, cs, os, t} = bufs
400  (ks, +n) = r
401  +kk = U32.max(k, 1)
402  +kb = Hist.kbits(kk)
403  +lg = U32.log2(n)
404  +d = Nat.min(b, lg)
405  +e = Nat.sub(lg, d)
406  run.sums(Hist.small(d), d, e, kb, kk, Hist.fit.eq(lg, Array.size(U32, pm)), cs, os, Hist.rows(d, 0, e, kb, kk,
407    Array.fork(U32, ks), Array.fork(U32, Hist.fit.ge(Nat.add(d, kb), Array.size(U32, t)))))
408
409# Gather
410# ------
411
412# A gather's arrays: the permutation, the source and the destination.
413type Moved<-T: Data> is Type:
414  Moved{perm: Array<U32>, src: Array<T>, dst: Array<T>}
415
416type Sd<-T: Data> is Type:
417  Sd{src: Array<T>, dst: Array<T>}
418
419def gat.at(~T: Data, +j: U32, dst: Array<T>, r: Array<T> & T) -> Sd<T>:
420  (src, v) = r
421  Sd{src, Array.set(T, dst, j, v)}
422
423# dst[j] <- src[x].
424def gat.one(~T: Data, +j: U32, +x: U32, p: Sd<T>) -> Sd<T>:
425  match p:
426    case Sd{src, dst}:
427      gat.at(~T, j, dst, Array.get(T, src, x))
428
429def gat.fin(~T: Data, pm: Array<U32>, p: Sd<T>) -> Moved<T>:
430  match p:
431    case Sd{src, dst}:
432      Moved{pm, src, dst}
433
434# r holds perm and its entry x at j: dst[j] <- src[x], then the n entries
435# after j.
436def gat.go(~T: Data, +n: Nat, +j: U32, p: Sd<T>, r: Array<U32> & U32) -> Moved<T>:
437  match n:
438    case 0n:
439      (pm, +x) = r
440      gat.fin(~T, pm, gat.one(~T, j, x, p))
441    case 1n+q:
442      (pm, +x) = r
443      +i = (j + 1 : U32)
444      gat.go(~T, q, i, gat.one(~T, j, x, p), pm[i])
445
446def moved.join(~T: Data, l: Moved<T>, r: Moved<T>) -> Moved<T>:
447  match l r:
448    case Moved{lp, ls, ld} Moved{rp, rs, rd}:
449      Moved{Array.join(U32, lp, rp), Array.join(T, ls, rs), Array.join(T, ld, rd)}
450
451def moved.rejoin(~T: Data, p2: Array<U32>, s2: Array<T>, d2: Array<T>, r: Moved<T>) -> Moved<T>:
452  match r:
453    case Moved{p, s, d}:
454      Moved{Array.join(U32, p, p2), Array.join(T, s, s2), Array.join(T, d, d2)}
455
456# The gather's forks: d levels down to blocks of 2^e slots under block i,
457# with two handles each to perm, src and dst. Block i reads only its
458# entries of perm, reads src anywhere (no lane writes src) and writes only
459# its slots of dst.
460def gats(~T: Data, +d: Nat, +i: U32, +e: Nat, fp: Array<U32> & Array<U32>, fs: Array<T> & Array<T>, fd: Array<T> & Array<T>) -> Moved<T>:
461  match d:
462    case 0n:
463      (pm, pm2) = fp
464      (src, src2) = fs
465      (dst, dst2) = fd
466      +at = U32.shln(i, e)
467      moved.rejoin(~T, pm2, src2, dst2, gat.go(~T, U32.to_nat((U32.shln(1, e) - 1 : U32)), at,
468        Sd{src, dst}, pm[at]))
469    case 1n+p:
470      (p0, p1) = fp
471      (s0, s1) = fs
472      (d0, d1) = fd
473      +j = (i * 2 : U32)
474      l r = gats(~T, p, j, e, Array.fork(U32, p0), Array.fork(T, s0), Array.fork(T, d0))
475        gats(~T, p, (j + 1 : U32), e, Array.fork(U32, p1), Array.fork(T, s1), Array.fork(T, d1))
476      moved.join(~T, l, r)
477
478def gather.pick(~T: Data, ok: Bool, +lg: Nat, +v: T, src: Array<T>, dst: Array<T>) -> Sd<T>:
479  match ok:
480    case True{}:
481      Sd{src, dst}
482    case False{}:
483      Sd{src, Array.new(T, lg, v)}
484
485def gather.dst(~T: Data, ok: Bool, +lg: Nat, dst: Array<T>, r: Array<T> & T) -> Sd<T>:
486  (src, v) = r
487  gather.pick(~T, ok, lg, v, src, dst)
488
489def gather.run(~T: Data, +b: Nat, +lg: Nat, pm: Array<U32>, p: Sd<T>) -> Moved<T>:
490  match p:
491    case Sd{src, dst}:
492      +d = Nat.min(b, lg)
493      gats(~T, d, 0, Nat.sub(lg, d), Array.fork(U32, pm), Array.fork(T, src), Array.fork(T, dst))
494
495def gather.at(~T: Data, +b: Nat, src: Array<T>, r: Array<U32> & U32, s: Array<T> & U32) -> Moved<T>:
496  (pm, +n) = r
497  (dst, +m) = s
498  +lg = U32.log2(n)
499  gather.run(~T, b, lg, pm, gather.dst(~T, U32.is_eq(m, n), lg, dst, Array.get(T, src, 0)))
500
501# API
502# ---
503
504# The stable counting sort of keys into k buckets at fork depth b. Returns
505# the keys, unchanged, and the buffers: perm, counts and offs hold the
506# result.
507def by_key(+b: Nat, +k: U32, keys: Array<U32>, bufs: Bufs) -> Array<U32> & Bufs:
508  run(b, k, bufs, Array.size(U32, keys))
509
510# dst[j] <- src[perm[j]] for every slot j, at fork depth b: with by_key's
511# perm, src in sorted order. src may be any size (indices wrap); dst is
512# reused when it has perm's size, else replaced by a new array. Returns
513# the three arrays.
514def gather(~T: Data, +b: Nat, perm: Array<U32>, src: Array<T>, dst: Array<T>) -> Moved<T>:
515  gather.at(~T, b, src, Array.size(U32, perm), Array.size(T, dst))