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).
48	import Base
49	import ./scan.bend as Scan
50	import ./histogram.bend as Hist
51	
52	# Buffers
53	# -------
54	
55	type 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.
59	def 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.
63	def 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	
71	type Three is Type:
72	  Three{a: Array<U32>, b: Array<U32>, c: Array<U32>}
73	
74	def 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	
79	def 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.
89	def 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	
99	def 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.
105	def 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.
115	def 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.
125	def 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).
136	def 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	
155	def 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.
161	def 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].
173	def 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	
182	def 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.
188	def 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.
197	def 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).
205	def 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.
215	def 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.
234	def 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	
244	def 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].
249	def 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.
258	def 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.
268	def 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	
285	def 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.
291	def 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	
295	def 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.
301	def 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).
312	def 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.
321	def 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	
340	def 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.
345	def 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.
350	def 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.
365	def 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.
370	def 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.
381	def 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.
398	def 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.
413	type Moved<-T: Data> is Type:
414	  Moved{perm: Array<U32>, src: Array<T>, dst: Array<T>}
415	
416	type Sd<-T: Data> is Type:
417	  Sd{src: Array<T>, dst: Array<T>}
418	
419	def 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].
424	def 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	
429	def 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.
436	def 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	
446	def 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	
451	def 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.
460	def 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	
478	def 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	
485	def 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	
489	def 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	
495	def 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.
507	def 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.
514	def 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))