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

docs/src/api/sort.md

@@ -14,13 +14,14 @@ AcceleratedKernels.sortperm
 ```
 
 Algorithm choice is available on `sort!` / `sort` / `sortperm!` / `sortperm` with `alg=AK.MergeSort()`,
-`alg=AK.MergeSort(lowmem=true)`, `alg=AK.RadixSort()`, or `alg=AK.SampleSort()`, depending on the
-backend and operation.
+`alg=AK.MergeSort(lowmem=true)`, `alg=AK.RadixSort()`, `alg=AK.BitonicSort()`, or
+`alg=AK.SampleSort()`, depending on the backend and operation.
 
 Function signatures:
 ```@docs
 AcceleratedKernels.MergeSort
 AcceleratedKernels.RadixSort
+AcceleratedKernels.BitonicSort
 AcceleratedKernels.SampleSort
 ```
 
@@ -42,8 +43,14 @@ AK.sort!(A; dims=1)             # each column sorted
 ix = AK.sortperm(A; dims=2)     # A[ix] has each row sorted
 ```
 
-On GPU backends `dims` uses merge sort (`RadixSort()` does not support it); on CPU backends each
-slice is sorted with `Base.sort!`.
+On GPU backends `dims` uses merge sort by default (`RadixSort()` does not support it); on CPU
+backends each slice is sorted with `Base.sort!`. `BitonicSort()` is unstable and supports `sort!`
+and `sort`, including `dims`, but not `sortperm!` or `sortperm`. It is fastest for small arrays
+and short slices, slower than `MergeSort`/`RadixSort` for large whole-array sorts:
+```julia
+A = ROCArray(rand(Float32, 64, 100_000))
+AK.sort!(A; dims=1, alg=AK.BitonicSort())
+```
 
 As GPU memory is more expensive, all functions in AcceleratedKernels.jl expose any temporary arrays they will use (the `temp` argument); you can supply your own buffers to make the algorithms not allocate additional GPU storage, e.g.:
 ```julia

src/sort/bitonic_sort.jl

@@ -0,0 +1,155 @@
+# Bitonic sort with every comparator ascending in `ord`. At each level `kk = 2, 4, ..., N`,
+# the first step pairs zero-based `i` with `i ⊻ (kk - 1)`, reflecting the second ascending run.
+# The remaining strides `kk/4, ..., 1` merge the two runs without sorting either descending.
+#
+# For other lengths, imagine padding with values ordered after all real elements. Ascending
+# comparators leave that padding at the end, so we skip pairs whose partner is past the length.
+# No padding storage or sentinel value is needed.
+#
+# Strides below the tile span run in local memory; larger strides each need a global pass.
+# Both kernels use slices.jl layouts, with one flat slice for whole-array sorts.
+
+# Elements `i` and `partner` (zero-based) of comparator `p` at level `kklog`, stride `jlog`
+@inline function bitonic_pair(p, kklog, jlog)
+    j = 1 << jlog
+    i = ((p >> jlog) << (jlog + 1)) | (p & (j - 1))
+    partner = jlog == kklog - 1 ? i ⊻ (2j - 1) : i + j
+    i, partner
+end
+
+@inline function compare_exchange!(elems, ord, i, partner)
+    a = elems[i + 1]
+    b = elems[partner + 1]
+    if Base.Order.lt(ord, b, a)
+        elems[i + 1] = b
+        elems[partner + 1] = a
+    end
+end
+
+
+# Run the local strides of levels `first_level:last_level`. A tile holds part of one long slice
+# (`SPAN == CAP`) or `CAP ÷ SPAN` short slices. Unused entries are never read.
+@kernel cpu=false inbounds=true unsafe_indices=true function bitonic_tile!(
+    vec, ord, layout, tiles_per_slice::Int, first_level::Int, last_level::Int, ::Val{CAP}, ::Val{SPAN},
+) where {CAP, SPAN}
+    @uniform BS = Int(@groupsize()[1])
+    tile = @localmem eltype(vec) (CAP,)
+
+    iblock = Int(@index(Group, Linear)) - 1
+    ithread = Int(@index(Local, Linear)) - 1
+    islice, itile = slice_block(layout, iblock, tiles_per_slice)
+    islice *= CAP ÷ SPAN
+    nslices = slice_count(layout)
+    spanlog = trailing_zeros(SPAN)
+    base = itile * SPAN
+    len = layout.len - base
+
+    # Packing arithmetic folds away when SPAN == CAP.
+    pos = ithread
+    while pos < CAP
+        s = SPAN == CAP ? islice : islice + (pos >> spanlog)
+        off = SPAN == CAP ? pos : pos & (SPAN - 1)
+        if off < len && s < nslices
+            tile[pos + 1] = slice(vec, layout, s)[base + off + 1]
+        end
+        pos += BS
+    end
+    @synchronize()
+
+    kklog = first_level
+    while kklog <= last_level
+        jlog = min(kklog - 1, spanlog - 1)
+        while jlog >= 0
+            p = ithread
+            while p < CAP ÷ 2
+                offset = SPAN == CAP ? 0 : (p >> (spanlog - 1)) * SPAN
+                i, partner = bitonic_pair(SPAN == CAP ? p : p & (SPAN ÷ 2 - 1), kklog, jlog)
+                if partner < len && (SPAN == CAP || islice + offset ÷ SPAN < nslices)
+                    compare_exchange!(tile, ord, offset + i, offset + partner)
+                end
+                p += BS
+            end
+            @synchronize()
+            jlog -= 1
+        end
+        kklog += 1
+    end
+
+    pos = ithread
+    while pos < CAP
+        s = SPAN == CAP ? islice : islice + (pos >> spanlog)
+        off = SPAN == CAP ? pos : pos & (SPAN - 1)
+        if off < len && s < nslices
+            slice(vec, layout, s)[base + off + 1] = tile[pos + 1]
+        end
+        pos += BS
+    end
+end
+
+
+# One global step at stride `2^jlog >= CAP`, one thread per comparator
+@kernel cpu=false inbounds=true function bitonic_global!(
+    vec, ord, layout, pairs_per_slice, kklog::Int, jlog::Int,
+)
+    islice, p = slice_block(layout, Int(@index(Global, Linear)) - 1, pairs_per_slice)
+    elems = slice(vec, layout, islice)
+    i, partner = bitonic_pair(p, kklog, jlog)
+    partner < layout.len && compare_exchange!(elems, ord, i, partner)
+end
+
+
+"""
+    bitonic_defaults(backend::Backend)
+
+Default `block_size` and `items_per_thread` for [`BitonicSort`](@ref) on `backend`.
+"""
+bitonic_defaults(::Backend) = (block_size=256, items_per_thread=8)
+
+
+function bitonic_sort!(
+    v::AbstractArray, backend::Backend=get_backend(v);
+
+    lt=isless,
+    by=identity,
+    rev::Union{Nothing, Bool}=nothing,
+    order::Base.Order.Ordering=Base.Order.Forward,
+
+    block_size::Int=256,
+    items_per_thread::Int=8,
+    dims::Union{Colon, Integer}=Colon(),
+)
+    @argcheck block_size > 0 && ispow2(block_size)
+    @argcheck items_per_thread > 0 && ispow2(items_per_thread)
+    @argcheck block_size <= typemax(Int) ÷ items_per_thread
+    layout = slice_layout(v, dims)
+    len = layout.len
+    (isempty(v) || len <= 1) && return v
+    ord = Base.Order.ord(lt, by, rev, order)
+
+    # Pack short slices into one tile; split long slices across tiles.
+    tile_size = block_size * items_per_thread
+    npow = nextpow(2, len)
+    nslices = slice_count(layout)
+    span = min(npow, tile_size)
+    cap = span < block_size ? span * nextpow(2, min(nslices, tile_size ÷ span)) : span
+    tiles_per_slice = cld(len, span)
+    nblocks = cld(nslices, cap ÷ span) * tiles_per_slice
+    tile_threads = min(block_size, cap)
+    tile! = bitonic_tile!(backend, tile_threads)
+    tile!(v, ord, layout, tiles_per_slice, 1, trailing_zeros(span), Val(cap), Val(span);
+          ndrange=nblocks * tile_threads)
+
+    if npow > span
+        global! = bitonic_global!(backend, block_size)
+        pairs = npow ÷ 2
+        for kklog in (trailing_zeros(span) + 1):trailing_zeros(npow)
+            for jlog in (kklog - 1):-1:trailing_zeros(span)
+                global!(v, ord, layout, pairs, kklog, jlog; ndrange=nslices * pairs)
+            end
+            tile!(v, ord, layout, tiles_per_slice, kklog, kklog, Val(cap), Val(span);
+                  ndrange=nblocks * tile_threads)
+        end
+    end
+
+    v
+end

src/sort/sort.jl

@@ -5,6 +5,7 @@ include("merge_sort_by_key.jl")
 include("merge_sortperm.jl")
 include("cpu_sample_sort.jl")
 include("radix_sort.jl")
+include("bitonic_sort.jl")
 
 
 # Available sorting algorithms
@@ -34,6 +35,25 @@ end
 
 _radix_defaults(::Backend) = (block_size=256, items_per_thread=2)
 
+"""
+    BitonicSort(; block_size=nothing, items_per_thread=nothing)
+
+Use GPU bitonic sort for `sort!` and `sort`, whole-array or along `dims`. Supports GPU-compatible
+element types and `lt`/`by`/`rev`/`order`; `by` is evaluated at each comparison. The sort is
+unstable and does not support `sortperm!` or `sortperm`.
+
+Fastest for small arrays and short slices, slower than `MergeSort`/`RadixSort` for large
+whole-array sorts. Tiles of up to `block_size * items_per_thread` elements sort in local memory;
+larger inputs need global passes.
+
+Both settings must be positive powers of two and default to `bitonic_defaults(backend)`.
+For `block_size`, the `sort!` keyword takes precedence over the backend default.
+"""
+Base.@kwdef struct BitonicSort <: SortAlgorithm
+    block_size::Union{Nothing, Int} = nothing
+    items_per_thread::Union{Nothing, Int} = nothing
+end
+
 """
     SampleSort()
 
@@ -77,8 +97,8 @@ arguments are the same as for `Base.sort`.
 
 By default (`dims=:`) the whole array is sorted as one vector, whatever its shape. Pass an integer
 `dims` to sort each 1D slice along that dimension independently, like `Base.sort!(A; dims)`. On GPU
-backends the slices are sorted with the same merge sort as the whole-array case, each slice its own
-sequence of merges; `RadixSort` does not support `dims`. On CPU backends each slice is sorted with
+backends this uses merge sort by default, or [`BitonicSort`](@ref) when requested; `RadixSort`
+does not support `dims`. On CPU backends each slice is sorted with
 `Base.sort!` (also with `alg=SampleSort()`, and without using `temp`), slices spread over the tasks.
 
 ## CPU
@@ -93,15 +113,17 @@ faster if it is a more compute-heavy operation to hide memory latency - that inc
 - Less cache-predictable data movement, e.g. `sortperm`.
 
 ## GPU
-GPU settings: `block_size` sets the number of threads per block. For `RadixSort`, fields on the
-algorithm take precedence over this keyword, then backend defaults. `items_per_thread` is set on
-`RadixSort` and defaults to 2.
+GPU settings: `block_size` sets the number of threads per block. For `RadixSort` and `BitonicSort`,
+fields on the algorithm take precedence over this keyword, then backend defaults.
+`items_per_thread` is set on the algorithm and defaults to 2 for `RadixSort`, 8 for `BitonicSort`.
 
 ## Algorithm choice
 By default, `sort!` uses sample sort on CPU backends and merge sort on GPU
 backends. Pass `alg=SampleSort()` for the CPU path, `alg=MergeSort()` for the GPU merge-sort path,
-or `alg=RadixSort()` to opt into GPU radix sorting. `RadixSort()` supports 32-bit and 64-bit
-integers and floats with default `lt`/`by`.
+`alg=RadixSort()` to opt into GPU radix sorting, or `alg=BitonicSort()` for the GPU sorting
+network. `RadixSort()` supports 32-bit and 64-bit integers and floats with default `lt`/`by`.
+`BitonicSort()` is unstable: fastest for small arrays and short slices, slower than
+`MergeSort`/`RadixSort` for large whole-array sorts.
 
 For both CPU and GPU backends, the `temp` argument can be used to reuse a temporary buffer of the
 same size as `v` to store the sorted output.
@@ -184,6 +206,14 @@ function _sort_impl!(
                 items_per_thread=radix_items,
                 temp,
             )
+        elseif alg isa BitonicSort
+            defaults = bitonic_defaults(backend)
+            bitonic_sort!(
+                v, backend;
+                lt, by, rev, order, dims,
+                block_size=something(alg.block_size, block_size, defaults.block_size),
+                items_per_thread=something(alg.items_per_thread, defaults.items_per_thread),
+            )
         else
             throw(ArgumentError("$(typeof(alg)) is not supported by sort! on GPU backends"))
         end
@@ -289,7 +319,7 @@ along `dims`. The permutation is stable in both cases.
 ## Algorithm choice
 By default, `sortperm!` uses sample sort on CPU backends and merge sort on GPU
 backends. Pass `alg=MergeSort(lowmem=true)` to use the lower-memory GPU permutation path.
-`RadixSort()` does not provide a permutation path.
+`RadixSort()` and `BitonicSort()` do not provide a permutation path.
 """
 function sortperm!(
     ix::AbstractArray,
@@ -361,6 +391,8 @@ function _sortperm_impl!(
             end
         elseif alg isa RadixSort
             throw(ArgumentError("RadixSort does not support sortperm"))
+        elseif alg isa BitonicSort
+            throw(ArgumentError("BitonicSort does not support sortperm"))
         else
             throw(ArgumentError("$(typeof(alg)) is not supported by sortperm! on GPU backends"))
         end