sort.bend relies on unsafe/foreign

raw source on the hub · import 0xeef3a8486d410ea6f25709aac90dbf40/sort.bend as Sort

sort.bend: a parallel stable counting sort by U32 key, with no atomics, and a parallel gather to reorder payloads by its permutation.

import ./sort.bend as Sort Sort.by_key(8n, 256, keys, Sort.Bufs.new()) # (keys, bufs) Sort.Bufs{perm, counts, offs, table} = bufs # perm[j]: the input index of the key the sort puts at j # counts[c], offs[c]: bucket c's size and first slot Sort.gather(~U32, 8n, perm, keys, [0 : U32^d]) # Moved{perm, keys, sorted}

(As in scan.bend: take a call's result apart in a def's parameter.)

k is the bucket count K (0 counts as 1), any U32; a key at or past K sorts as K - 1. The sort is stable: keys of one bucket keep their input order, so the result equals ref.bend's perm, offsets and histogram bit for bit, at every fork depth b. counts and offs hold 2^ceil(log2 K) entries: past K, counts are 0 and offs are n. A call under ! runs its forks on the GPU.

Bufs carries the buffers from call to call (the table of 2^b rows of 2^ceil(log2 K) counts, the counts, the offsets and the permutation). A call keeps a buffer that fits and allocates otherwise; under ! an allocation fills on one lane, so size them once with Bufs.sized outside the ! (see histogram.bend).

Passes: 1. per block of keys, in parallel: the block's counts per bucket into its row of the table (histogram.bend's pass 1); 2. from 2^12 rows up, the table's rows are cut into 2^(b/2) groups of rows (see histogram.bend): a) per (group of rows, group of buckets), in parallel: each column becomes its running sum within the group; b) per group of buckets, in parallel: over the groups' last rows, each bucket's total into counts and offs, and each group's base (the counts of the groups above it) into its last row; then offs becomes its exclusive scan (scan.bend), the buckets' first slots; c) per (group of rows, group of buckets) again: each table entry becomes its bucket's first slot plus the counts of the rows above it (base plus the running sum of the row before it): where the block's keys of that bucket start. Below 2^12 rows, one group of all the rows, and no a): b) sums each column, c) walks it down; 3. per block, in parallel: each key, in order, takes the next slot of its bucket in the block's row, and perm[slot] becomes its index.

Sharing: as in scan.bend, lanes share arrays through Base's @unsafe Array.fork and Array.join; no def here is @unsafe itself, and none may be evaluated in the checker (call them from an IO main).

3 imports

import Base
import ./scan.bend as Scan
import ./histogram.bend as Hist

Types

type Bufs source · line 55 · raw

Type

Bufs

@perm:Array<U32> -> @counts:Array<U32> -> @offs:Array<U32> -> @table:Array<U32> -> Bufs

type Three source · line 71 · raw

Type

Three

@a:Array<U32> -> @b:Array<U32> -> @c:Array<U32> -> Three

type Moved source · line 413 · raw

@-T:Data -> Type

A gather's arrays: the permutation, the source and the destination.

Moved

@-T:Data -> @perm:Array<U32> -> @src:Array<T> -> @dst:Array<T> -> Moved<T>

type Sd source · line 416 · raw

@-T:Data -> Type

Sd

@-T:Data -> @src:Array<T> -> @dst:Array<T> -> Sd<T>

Definitions

def Bufs.new source · line 59 · raw

Bufs

Empty buffers: the first call allocates what it needs.

def Bufs.sized source · line 63 · raw

@+b:Nat -> @+k:U32 -> @+lg:Nat -> Bufs

Buffers for 2^lg keys at fork depth b and k buckets, allocated now.

def three.join source · line 74 · raw

@l:Three -> @r:Three -> Three

def three.rejoin source · line 79 · raw

@a2:Array<U32> -> @b2:Array<U32> -> @c2:Array<U32> -> @r:Three -> Three

def xs.go source · line 89 · raw

@+n:Nat -> @+at:U32 -> @+step:U32 -> @+s:U32 -> @r:Pair(Array<U32>, U32) -> Pair(Array<U32>, U32)

r holds the table and the entry x read at at: it becomes s, and the n entries after it (each step further on) follow, from s + x. Returns the table and the sum past the last.

def sum.put source · line 99 · raw

@+c:U32 -> @cs:Array<U32> -> @os:Array<U32> -> @r:Pair(Array<U32>, U32) -> Three

def sum.col source · line 105 · raw

@+wr:Bool -> @+at:U32 -> @+step:U32 -> @+nb:Nat -> @t:Array<U32> -> Pair(Array<U32>, U32)

The nb + 1 entries from at, step apart: their sum; when wr, each also becomes the sum of those before it.

def sum.one source · line 115 · raw

@lt:Bool -> @+wr:Bool -> @+c:U32 -> @+o:U32 -> @+step:U32 -> @+nb:Nat -> @p:Three -> Three

Column c, when c < K: over its nb + 1 entries from o + c, step apart (the groups' last rows, or all the rows of a small table), the total into cs[c] and os[c], by sum.col. Else 0 into both.

def sum.go source · line 125 · raw

@+n:Nat -> @+wr:Bool -> @+c:U32 -> @+k:U32 -> @+o:U32 -> @+step:U32 -> @+nb:Nat -> @p:Three -> Three

Columns c to c + n.

def sums source · line 136 · raw

@+d:Nat -> @+wr:Bool -> @+j:U32 -> @+w:Nat -> @+k:U32 -> @+o:U32 -> @+step:U32 -> @+nb:Nat -> @ft:Pair(Array<U32>, Array<U32>) -> @fc:Pair(Array<U32>, Array<U32>) -> @fo:Pair(Array<U32>, Array<U32>) -> Three

Pass 2b's forks: d levels down to groups of 2^w columns under group j, with two handles each to the table, the counts and the offsets. Group j reads (and when wr writes) only its columns of the table and writes only its own entries of counts and offs (j * 2^w to j * 2^w + 2^w - 1).

def base.fin source · line 155 · raw

@+q:U32 -> @+s:U32 -> @t:Array<U32> -> Array<U32>

def base.go source · line 161 · raw

@+n:Nat -> @+at:U32 -> @+step:U32 -> @+s:U32 -> @r:Pair(Array<U32>, U32) -> Array<U32>

r holds the table and x, the entry a row (step) above at, which holds the group's running sum through that row: at becomes s + x, and so up for the n rows above; the group's first row, the next above, becomes s.

def base.at source · line 173 · raw

@+at:U32 -> @+step:U32 -> @+nm:Nat -> @+o:U32 -> @r:Pair(Array<U32>, U32) -> Array<U32>

A group's column from its last row at (nm + 1 rows): r holds the table and the group's base b, which the last row holds; o is offs[c].

def base.os source · line 182 · raw

@+c:U32 -> @+last:U32 -> @+step:U32 -> @+nm:Nat -> @t:Array<U32> -> @r:Pair(Array<U32>, U32) -> 0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two

def base.one source · line 188 · raw

@lt:Bool -> @+c:U32 -> @+last:U32 -> @+step:U32 -> @+nm:Nat -> @p:0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two -> 0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two

Column c of a group, when c < K: from its first slot offs[c] plus the group's base, which its last row holds.

def base.cols source · line 197 · raw

@+n:Nat -> @+c:U32 -> @+k:U32 -> @+last:U32 -> @+step:U32 -> @+nm:Nat -> @p:0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two -> 0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two

Columns c to c + n of a group.

def base.leaf source · line 205 · raw

@+i:U32 -> @+gc:Nat -> @+w:Nat -> @+m:Nat -> @+kb:Nat -> @+k:U32 -> @p:0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two -> 0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two

Lane i: rows group i / 2^gc, columns group i % 2^gc (as Hist.grp.leaf).

def bases source · line 215 · raw

@+d:Nat -> @+i:U32 -> @+gc:Nat -> @+w:Nat -> @+m:Nat -> @+kb:Nat -> @+k:U32 -> @ft:Pair(Array<U32>, Array<U32>) -> @fo:Pair(Array<U32>, Array<U32>) -> 0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two

Pass 2c's forks, shaped as pass 2a's (Hist.grp), with two handles each to the table and the offsets. Lane i reads the offsets and reads and writes only its group's rows in its columns.

def fwd.go source · line 234 · raw

@+n:Nat -> @+at:U32 -> @+step:U32 -> @+s:U32 -> @r:Pair(Array<U32>, U32) -> Array<U32>

r holds the table and the entry x read at at: it becomes s, and the n entries after it (each step further on) follow, from s + x.

def fwd.at source · line 244 · raw

@+c:U32 -> @+kb:Nat -> @+nb:Nat -> @t:Array<U32> -> @r:Pair(Array<U32>, U32) -> 0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two

def fwd.one source · line 249 · raw

@lt:Bool -> @+c:U32 -> @+kb:Nat -> @+nb:Nat -> @p:0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two -> 0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two

Column c, when c < K: its entries from its first slot offs[c].

def fwd.cols source · line 258 · raw

@+n:Nat -> @+c:U32 -> @+k:U32 -> @+kb:Nat -> @+nb:Nat -> @p:0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two -> 0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two

Columns c to c + n.

def fwds source · line 268 · raw

@+d:Nat -> @+j:U32 -> @+w:Nat -> @+k:U32 -> @+kb:Nat -> @+nb:Nat -> @ft:Pair(Array<U32>, Array<U32>) -> @fo:Pair(Array<U32>, Array<U32>) -> 0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two

The forks, shaped as pass 2b's, with two handles each to the table and the offsets. Group j reads only its own offsets and reads and writes only its columns of the table.

def put.at source · line 285 · raw

@+j:U32 -> @+slot:U32 -> @pm:Array<U32> -> @r:Pair(Array<U32>, U32) -> 0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two

def put.one source · line 291 · raw

@+j:U32 -> @+slot:U32 -> @p:0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two -> 0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two

Key j, whose bucket's next slot is t[slot]: perm[that slot] <- j, and the slot moves on.

def put.fin source · line 295 · raw

@ks:Array<U32> -> @p:0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two -> Three

def put.go source · line 301 · raw

@+n:Nat -> @+j:U32 -> @+row:U32 -> @+k1:U32 -> @p:0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two -> @r:Pair(Array<U32>, U32) -> Three

r holds the keys and the key x read at j: key j goes to the next slot of its bucket in the row at row, then the n keys after j. k1 is K - 1.

def puts.leaf source · line 312 · raw

@+i:U32 -> @+e:Nat -> @+kb:Nat -> @+k:U32 -> @ks:Array<U32> -> @t:Array<U32> -> @pm:Array<U32> -> Three

Block i (2^e keys from i * 2^e, its row at i * 2^kb).

def puts source · line 321 · raw

@+d:Nat -> @+i:U32 -> @+e:Nat -> @+kb:Nat -> @+k:U32 -> @fk:Pair(Array<U32>, Array<U32>) -> @ft:Pair(Array<U32>, Array<U32>) -> @fp:Pair(Array<U32>, Array<U32>) -> Three

Pass 3's forks: d levels down to the blocks under block i, with two handles each to the keys, the table and perm. Block i reads only its keys, reads and writes only its row of the table, and writes perm only at the slots its row hands it: pass 2 gave each (block, bucket) pair its own run of slots, so no two lanes write one slot.

def run.fin source · line 340 · raw

@cs:Array<U32> -> @os:Array<U32> -> @r:Three -> Pair(Array<U32>, Bufs)

def run.puts source · line 345 · raw

@+d:Nat -> @+e:Nat -> @+kb:Nat -> @+k:U32 -> @ks:Array<U32> -> @pm:Array<U32> -> @cs:Array<U32> -> @r:0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two -> Pair(Array<U32>, Bufs)

Pass 3.

def run.bases source · line 350 · raw

@lo:Bool -> @+d:Nat -> @+e:Nat -> @+kb:Nat -> @+k:U32 -> @ks:Array<U32> -> @pm:Array<U32> -> @cs:Array<U32> -> @t:Array<U32> -> @r:Pair(Array<U32>, U32) -> Pair(Array<U32>, Bufs)

Pass 2c: by fwds when lo (a small table), else by bases.

def run.scan source · line 365 · raw

@lo:Bool -> @+d:Nat -> @+e:Nat -> @+kb:Nat -> @+k:U32 -> @ks:Array<U32> -> @pm:Array<U32> -> @r:Three -> Pair(Array<U32>, Bufs)

Pass 2b's scan: offs becomes its exclusive scan.

def run.tots source · line 370 · raw

@+d:Nat -> @+e:Nat -> @+kb:Nat -> @+k:U32 -> @ks:Array<U32> -> @pm:Array<U32> -> @cs:Array<U32> -> @os:Array<U32> -> @t:Array<U32> -> Pair(Array<U32>, Bufs)

Pass 2b, after 2a (t): over the 2^h groups' last rows.

def run.sums source · line 381 · raw

@lo:Bool -> @+d:Nat -> @+e:Nat -> @+kb:Nat -> @+k:U32 -> @pm:Array<U32> -> @cs:Array<U32> -> @os:Array<U32> -> @r:0xeef3a8486d410ea6f25709aac90dbf40/histogram.Two -> Pair(Array<U32>, Bufs)

Pass 2a: each group's running column sums (Hist.grp), then 2b. For a small table (lo), no 2a: 2b sums each column over all the rows.

def run source · line 398 · raw

@+b:Nat -> @+k:U32 -> @bufs:Bufs -> @r:Pair(Array<U32>, U32) -> Pair(Array<U32>, Bufs)

Pass 1.

def by_key source · line 507 · raw

@+b:Nat -> @+k:U32 -> @keys:Array<U32> -> @bufs:Bufs -> Pair(Array<U32>, Bufs)

The stable counting sort of keys into k buckets at fork depth b. Returns the keys, unchanged, and the buffers: perm, counts and offs hold the result.

Templates

template gat.at source · line 419 · raw

@-T:Data -> @+j:U32 -> @dst:Array<T> -> @r:Pair(Array<T>, T) -> Sd<T>

template gat.one source · line 424 · raw

@-T:Data -> @+j:U32 -> @+x:U32 -> @p:Sd<T> -> Sd<T>

dst[j] <- src[x].

template gat.fin source · line 429 · raw

@-T:Data -> @pm:Array<U32> -> @p:Sd<T> -> Moved<T>

template gat.go source · line 436 · raw

@-T:Data -> @+n:Nat -> @+j:U32 -> @p:Sd<T> -> @r:Pair(Array<U32>, U32) -> Moved<T>

r holds perm and its entry x at j: dst[j] <- src[x], then the n entries after j.

template moved.join source · line 446 · raw

@-T:Data -> @l:Moved<T> -> @r:Moved<T> -> Moved<T>

template moved.rejoin source · line 451 · raw

@-T:Data -> @p2:Array<U32> -> @s2:Array<T> -> @d2:Array<T> -> @r:Moved<T> -> Moved<T>

template gats source · line 460 · raw

@-T:Data -> @+d:Nat -> @+i:U32 -> @+e:Nat -> @fp:Pair(Array<U32>, Array<U32>) -> @fs:Pair(Array<T>, Array<T>) -> @fd:Pair(Array<T>, Array<T>) -> Moved<T>

The gather's forks: d levels down to blocks of 2^e slots under block i, with two handles each to perm, src and dst. Block i reads only its entries of perm, reads src anywhere (no lane writes src) and writes only its slots of dst.

template gather.pick source · line 478 · raw

@-T:Data -> @ok:Bool -> @+lg:Nat -> @+v:T -> @src:Array<T> -> @dst:Array<T> -> Sd<T>

template gather.dst source · line 485 · raw

@-T:Data -> @ok:Bool -> @+lg:Nat -> @dst:Array<T> -> @r:Pair(Array<T>, T) -> Sd<T>

template gather.run source · line 489 · raw

@-T:Data -> @+b:Nat -> @+lg:Nat -> @pm:Array<U32> -> @p:Sd<T> -> Moved<T>

template gather.at source · line 495 · raw

@-T:Data -> @+b:Nat -> @src:Array<T> -> @r:Pair(Array<U32>, U32) -> @s:Pair(Array<T>, U32) -> Moved<T>

template gather source · line 514 · raw

@-T:Data -> @+b:Nat -> @perm:Array<U32> -> @src:Array<T> -> @dst:Array<T> -> Moved<T>

dst[j] <- src[perm[j]] for every slot j, at fork depth b: with by_key's perm, src in sorted order. src may be any size (indices wrap); dst is reused when it has perm's size, else replaced by a new array. Returns the three arrays.