Add BitonicSort, a GPU sorting network for small arrays and short slices

[shreyas-omkar]

BitonicSort is a new opt-in algorithm for sort! and sort on GPU backends, whole-array or along dims. Each workgroup sorts a tile of block_size * items_per_thread elements (2048 by default) entirely in local memory, so an array or slice that fits a tile is sorted by a single launch, and slices shorter than a workgroup are packed several per launch. Longer inputs add one global-memory pass per long stride of the network. This is the case merge sort handles worst and radix sort cannot handle at all (RadixSort has no dims), and it is where the new algorithm wins by a wide margin:

import AcceleratedKernels as AK
using CUDA

A = CuArray(rand(Float32, 32, 100_000))
AK.sort!(A; dims=1, alg=AK.BitonicSort())          # each column sorted, one launch
B = AK.sort(A; dims=2, rev=true, alg=AK.BitonicSort())
v = CuArray(rand(Int16, 5000))
AK.sort!(v; alg=AK.BitonicSort(), by=abs)           # any eltype, any lt/by/order
AK.sort!(v; alg=AK.BitonicSort(block_size=128, items_per_thread=4))

The algorithm is not stable and has no sortperm path (an ArgumentError), so it stays opt-in: MergeSort remains the default for both whole-array and dims sorts. Everything else follows Base.sort!: lt, by, rev and order compose the same way, NaNs and signed zeros order like isless, and any element type works. Both tunables must be powers of two; block_size falls back to the sort! keyword and then to bitonic_defaults(backend), which backends can specialise (256 threads, 8 items per thread).

How it works

A bitonic network sorts by a fixed sequence of compare-exchange steps: for each merge level kk = 2, 4, ..., N the strides j = kk/2, ..., 1. Two things keep the kernels simple:

• Every comparator sorts ascending in the requested ordering. The first step of each level pairs i with its mirror image i ⊻ (kk - 1), which turns two ascending runs into a bitonic sequence without ever sorting a run descending. Because no comparator moves a larger element below a smaller one, elements past the end behave as if padded with values that sort last, so pairs whose partner is past the length are simply skipped. Any length sorts in place with no padding buffer, no sentinel, and no restriction on eltype or comparator.
• bitonic_tile! runs all strides below the tile size for a range of levels in local memory, and bitonic_global! runs one long stride over global memory with one thread per pair. A sort is one tile launch, then per larger level its global steps followed by one tile launch.

Along dims, each launch covers every slice: long slices are split across tiles like the flat case, and slices shorter than a workgroup share a tile (SPAN < CAP in the tile kernel), which is what makes many tiny slices fast.

Performance

sort!, Float32, whole call, ms. "Merge" is MergeSort (the default, from #117 for dims); "vendor" is the backend package's own sort! where it has one.

RTX 5080, sort!(A; dims=1), slices x count. CUB DeviceSegmentedSort::SortKeys (what PyTorch uses) is listed for scale; it is kernel time only and needs contiguous segments.

slice len	slices	Bitonic	Merge	CUDA.jl	CUB
8	131072	0.044	1.73	2.44	0.020
32	32768	0.058	0.64	1.33	0.034
128	8192	0.078	0.24	0.59	1.008
512	2048	0.094	0.134	0.256	0.265
1024	1024	0.105	0.155	0.308	0.144
4096	256	0.135	0.205	0.442	0.043
8192	128	0.178	0.235	0.532	0.079

dims=2 (strided slices) is within 10% of dims=1 up to 1024-element slices. The submitted version is faster for medium slices because it compared with > and had no bounds guards; the generic isless comparator alone costs about 10%.

RTX 5080, whole-array sort!:

n	Bitonic	Merge	Radix	CUDA.jl
1024	0.017	0.022	0.056	0.031
16384	0.053	0.041	0.083	0.063
262144	0.141	0.139	0.126	0.380
1048576	0.406	0.393	0.255	1.52
16777216	7.17	9.18	2.91	40.8

Apple M1 (Metal.jl) and Intel Xe iGPU (oneAPI.jl, which forwards sort! to AK), sort!(A; dims=1):

slice len	slices	M1 Bitonic	M1 Merge	Metal.jl	Xe Bitonic	Xe Merge
8	131072	1.6	62.6	946	0.59	108
32	32768	1.9	20.7	237	1.07	36.1
128	8192	2.0	6.8	60.7	1.6	12.7
1024	1024	3.6	4.2	8.7	3.5	6.4
8192	128	6.8	6.9	3.4	5.7	10.0

For whole-array sorts on those two GPUs bitonic is only ahead of merge sort below about 16k elements, and the global passes are slow on the Intel iGPU (16M: 779 ms vs 383 ms merge), which the docstring reflects: use it for small arrays and short slices.

Tuning sweeps over block_size in {128, 256, 512} and items_per_thread in {1, ..., 32} on all three GPUs put the default (256, 8) within 10% of the best setting for every dims workload; 512 x 32 exceeds Metal's 32 KiB local memory for 64-bit elements.

[shreyas-omkar]

Hey @maleadt please take a look. The win is sort(A; dims) with small slices. Each slice is sorted by its own workgroup entirely in shared memory no global memory round-trips, and all slices run in parallel.

[maleadt]

Slightly reworked the implementation on top of your two commits and rebased the branch onto #117, so the dims machinery is shared with the merge sort.

Changes made:

• No padding, no sentinels. The network now sorts ascending in ord throughout, with the first step of each merge level pairing i with i ⊻ (kk - 1) (the "reflected" formulation CUDA.jl also uses) and comparators skipped when the partner is past the end. Any length sorts in place, so the similar(v, npow) buffer, typemax/typemin padding and the eltype/ordering whitelist are gone: any eltype, lt, by and order work, and NaNs order like Base. The previous version put NaNs in the middle ([-0.0, NaN, -1.0, 1.0, NaN, 0.0, ...]) because the sentinel path compared with >.
• One dims path. The kernels take the SliceLayout from Add dims support to sort, sort!, sortperm and sortperm! #117, so _bitonic_slice!, _sort_dims_impl! and the eachslice CPU fallback are gone, slices longer than a tile work (global steps stay within a slice), and the CPU side reuses Add dims support to sort, sort!, sortperm and sortperm! #117's multi-tasked slice loop.
• Two kernels instead of four. bitonic_tile! runs the strides below the tile size for a range of levels, bitonic_global! runs one long stride with a thread per pair. Slices shorter than a workgroup are packed several per tile, which made 8-element slices about 2x faster on the RTX 5080 and 3-8x on Metal and oneAPI.
• Tunables on the struct. BitonicSort(; block_size, items_per_thread) like RadixSort, with bitonic_defaults(::Backend) = (256, 8) chosen from a sweep on three GPUs, instead of the 32 KiB byte budget. Non-power-of-two settings throw instead of being rounded down.
• KernelAbstractions.synchronize calls dropped; sortperm with alg=BitonicSort() throws a clear ArgumentError.
• Tests cover the new generality (more eltypes, NaN/signed zero, lt/by/order, packed partial tiles, empty slices) and pass on CUDA, POCL, Metal, oneAPI and CPU.

Numbers (RTX 5080, Float32, sort!(A; dims=1), ms): 8 x 131072: 0.088 -> 0.044; 32 x 32768: 0.060 -> 0.058; 1024 x 1024: 0.075 -> 0.105. The medium-slice cases got a bit slower because the generic isless comparator and the bounds guards cost ~10% each versus a raw > on padded data, which I think is the right trade for matching Base. Flat sorts are unchanged within noise.