From ce702f4ced2e88180b2c442cec6e8bcc53b9caf1 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Wed, 16 Sep 2026 08:13:11 +0530 Subject: [PATCH 1/3] feat(sort): add BitonicSort algorithm Add `BitonicSort` as an opt-in GPU sort (`alg=BitonicSort()`), a portable sorting network on KernelAbstractions primitives. Small inputs sort entirely in shared memory; larger inputs use global compare-exchange for long strides and shared-memory batches for short ones. Supports the same 32/64-bit eltypes as RadixSort with forward or reverse ordering. Co-Authored-By: Claude Opus 4.8 --- src/sort/bitonic_sort.jl | 186 +++++++++++++++++++++++++++++++++++++++ src/sort/sort.jl | 38 ++++++++ test/generic/sort.jl | 90 +++++++++++++++++++ 3 files changed, 314 insertions(+) create mode 100644 src/sort/bitonic_sort.jl diff --git a/src/sort/bitonic_sort.jl b/src/sort/bitonic_sort.jl new file mode 100644 index 00000000..64f1d0d3 --- /dev/null +++ b/src/sort/bitonic_sort.jl @@ -0,0 +1,186 @@ +# Portable GPU bitonic sort on KernelAbstractions primitives. A sorting network has no +# data-dependent branching, so it maps well onto GPUs; inputs that fit a single workgroup sort +# entirely in shared memory. Non-power-of-two lengths are padded with a sentinel (`typemax` +# ascending, `typemin` descending) so it sorts to the tail. + +# Radix's eltypes: comparable, with a `typemax`/`typemin` sentinel. +_bs_supported(::Type{T}) where T = + T === UInt32 || T === Int32 || T === Float32 || + T === UInt64 || T === Int64 || T === Float64 + +# Single-workgroup shared-memory budget in bytes; overridable per backend. Matches the radix sort. +_bs_shmem_bytes(::Backend) = 32 * 1024 +@inline _prevpow2(x::Int) = 1 << (8 * sizeof(Int) - leading_zeros(x) - 1) + + +# Sort one CAP-sized window from scratch in shared memory; the window's global base fixes each +# comparator's direction. Used for a whole small array and to build the windows of the large path. +@kernel cpu=false inbounds=true function _bitonic_full!( + w, descending::Bool, ::Val{CAP}, ::Val{BS}, ::Val{IPT}, +) where {CAP, BS, IPT} + tile = @localmem eltype(w) (CAP,) + blk = @index(Group, Linear) + t = Int(@index(Local, Linear)) - 1 + base = Int(blk - 0x1) * CAP + + m = 0 + while m < IPT + pos = t + m * BS + @inbounds tile[pos + 1] = w[base + pos + 1] + m += 1 + end + @synchronize() + + kklog = 1 + while (1 << kklog) <= CAP + jlog = kklog - 1 + while jlog >= 0 + j = 1 << jlog + p = t + while p < (CAP >> 1) + i = ((p >> jlog) << (jlog + 1)) | (p & (j - 1)) + partner = i + j + asc = ((base + i) & (1 << kklog)) == 0 + a = @inbounds tile[i + 1] + b = @inbounds tile[partner + 1] + if (a > b) == (asc != descending) + @inbounds tile[i + 1] = b + @inbounds tile[partner + 1] = a + end + p += BS + end + @synchronize() + jlog -= 1 + end + kklog += 1 + end + + m = 0 + while m < IPT + pos = t + m * BS + @inbounds w[base + pos + 1] = tile[pos + 1] + m += 1 + end +end + + +# Finish the short strides (`j_start` down to 0, all < CAP) of merge level `kklog` for one window +# in shared memory; every comparator's pair stays inside the window. +@kernel cpu=false inbounds=true function _bitonic_block!( + w, kklog::Int, j_start::Int, descending::Bool, + ::Val{CAP}, ::Val{BS}, ::Val{IPT}, +) where {CAP, BS, IPT} + tile = @localmem eltype(w) (CAP,) + blk = @index(Group, Linear) + t = Int(@index(Local, Linear)) - 1 + base = Int(blk - 0x1) * CAP + kk = 1 << kklog + + m = 0 + while m < IPT + pos = t + m * BS + @inbounds tile[pos + 1] = w[base + pos + 1] + m += 1 + end + @synchronize() + + jlog = j_start + while jlog >= 0 + j = 1 << jlog + p = t + while p < (CAP >> 1) + i = ((p >> jlog) << (jlog + 1)) | (p & (j - 1)) + partner = i + j + asc = ((base + i) & kk) == 0 + a = @inbounds tile[i + 1] + b = @inbounds tile[partner + 1] + if (a > b) == (asc != descending) + @inbounds tile[i + 1] = b + @inbounds tile[partner + 1] = a + end + p += BS + end + @synchronize() + jlog -= 1 + end + + m = 0 + while m < IPT + pos = t + m * BS + @inbounds w[base + pos + 1] = tile[pos + 1] + m += 1 + end +end + + +# One global compare-exchange for a large-stride merge level, launched one thread per pair +# (`npow ÷ 2`): every lane works, no divergent guard, and consecutive threads stay coalesced. +@kernel cpu=false inbounds=true function _bitonic_global!( + w, kklog::Int, jlog::Int, descending::Bool, +) + p = Int(@index(Global, Linear)) - 1 + j = 1 << jlog + i = ((p >> jlog) << (jlog + 1)) | (p & (j - 1)) + partner = i + j + asc = (i & (1 << kklog)) == 0 + a = @inbounds w[i + 1] + b = @inbounds w[partner + 1] + if (a > b) == (asc != descending) + @inbounds w[i + 1] = b + @inbounds w[partner + 1] = a + end +end + + +function _bitonic_sort!( + v::AbstractVector{T}, backend::Backend; + descending::Bool, + block_size::Union{Nothing, Int}=nothing, +) where T + len = length(v) + len <= 1 && return v + + npow = nextpow(2, len) + budget = _prevpow2(_bs_shmem_bytes(backend) ÷ sizeof(T)) + cap = min(npow, budget) + bs = isnothing(block_size) ? min(256, cap) : min(_prevpow2(block_size), cap) + ipt = cap ÷ bs + + # Pad to a power of two; the sentinel sinks to the tail so real values fill 1:len. + pad = descending ? typemin(T) : typemax(T) + if npow == len + w = v + else + w = similar(v, npow) + copyto!(w, 1, v, 1, len) + fill!(view(w, (len + 1):npow), pad) + end + + if npow <= cap + _bitonic_full!(backend, bs)(w, descending, Val(cap), Val(bs), Val(ipt); ndrange = bs) + else + nlog = trailing_zeros(npow) + caplog = trailing_zeros(cap) + nblocks = npow ÷ cap + full = _bitonic_full!(backend, bs) + block = _bitonic_block!(backend, bs) + global_step = _bitonic_global!(backend, 256) + + # Build sorted windows, then merge them: global passes for j >= cap, one shared-memory + # batch for the remaining j < cap of each merge level. + full(w, descending, Val(cap), Val(bs), Val(ipt); ndrange = nblocks * bs) + for kklog in (caplog + 1):nlog + jlog = kklog - 1 + while jlog >= caplog + global_step(w, kklog, jlog, descending; ndrange = npow ÷ 2) + jlog -= 1 + end + block(w, kklog, caplog - 1, descending, Val(cap), Val(bs), Val(ipt); + ndrange = nblocks * bs) + end + end + KernelAbstractions.synchronize(backend) + + npow == len || copyto!(v, 1, w, 1, len) + v +end diff --git a/src/sort/sort.jl b/src/sort/sort.jl index 75b96e17..deb8dcf4 100644 --- a/src/sort/sort.jl +++ b/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,31 @@ end _radix_defaults(::Backend) = (block_size=256, items_per_thread=2) +""" + BitonicSort(; block_size=nothing) + +Use a GPU bitonic sort for `sort!` and `sort`. Supports `UInt32`, `Int32`, `Float32`, `UInt64`, +`Int64`, and `Float64` with forward or reverse ordering. This algorithm does not support +`sortperm!`. + +A bitonic sorting network has no data-dependent control flow, so every comparator runs +unconditionally — a good fit for GPUs. For an input that fits a single workgroup's shared memory +the whole sort runs in one kernel with no global-memory round-trips, which is faster than the other +GPU algorithms for small arrays. Above that size the network spills to global memory and its +`O(n·log²n)` comparisons cost more than the `O(n)`-pass [`RadixSort`](@ref) or `MergeSort`, so it +grows slower with `n`: + +- Small arrays (up to roughly a single workgroup, a few thousand elements): fastest of the three. +- Medium arrays: comparable to `MergeSort`. +- Large arrays: slower than `RadixSort` (measured ~1.1–1.8× on an RTX 5080 and RX 9060 XT), because + the number of global compare-exchange passes grows with `log²n`. + +Non-power-of-two lengths are padded up internally; `block_size` sets the threads per workgroup. +""" +Base.@kwdef struct BitonicSort <: SortAlgorithm + block_size::Union{Nothing, Int} = nothing +end + """ SampleSort() @@ -184,6 +210,18 @@ function _sort_impl!( items_per_thread=radix_items, temp, ) + elseif alg isa BitonicSort + v isa AbstractVector || + throw(ArgumentError("BitonicSort only supports vectors (got a $(ndims(v))-dimensional array)")) + _bs_supported(eltype(v)) || throw(ArgumentError("BitonicSort is not supported for eltype \"$(eltype(v))\"")) + ordering = Base.Order.ord(lt, by, rev, order) + ordering === Base.Order.Forward || ordering === Base.Order.Reverse || + throw(ArgumentError("BitonicSort only supports forward or reverse ordering")) + _bitonic_sort!( + v, backend; + descending=ordering === Base.Order.Reverse, + block_size=isnothing(alg.block_size) ? block_size : alg.block_size, + ) else throw(ArgumentError("$(typeof(alg)) is not supported by sort! on GPU backends")) end diff --git a/test/generic/sort.jl b/test/generic/sort.jl index 5910950e..576a959f 100644 --- a/test/generic/sort.jl +++ b/test/generic/sort.jl @@ -1007,3 +1007,93 @@ end @test_throws ArgumentError AK.sortperm(A; prefer_threads, dims=0) @test_throws ArgumentError AK.sortperm!(array_from_host(zeros(Int, 64)), A; prefer_threads, dims=1) end + + +@testset "bitonic_sort_alg" begin + if !prefer_threads + Random.seed!(0) + + # ── Correctness: fuzzy testing across supported types ───────────────── + for T in valid_backend_eltypes(BACKEND, + (UInt32, Int32, Float32, UInt64, Int64, Float64)) + for _ in 1:100 + n = rand(1:100_000) + v_h = rand(T, n) + v = array_from_host(v_h) + AK.sort!(v; prefer_threads, alg=AK.BitonicSort()) + @test Array(v) == sort(v_h) + end + end + + # ── rev=true ────────────────────────────────────────────────────────── + for T in valid_backend_eltypes(BACKEND, + (UInt32, Int32, Float32, UInt64, Int64, Float64)) + n = 10_000 + v_h = rand(T, n) + v = array_from_host(v_h) + AK.sort!(v; prefer_threads, alg=AK.BitonicSort(), rev=true) + @test Array(v) == sort(v_h; rev=true) + end + + # ── Sizes around the single-block / multi-block boundary ────────────── + # Non-power-of-two lengths exercise the sentinel padding; large ones the + # global-memory merge passes. + for n in (2, 3, 7, 8, 255, 256, 257, 1023, 1024, 1025, + 8191, 8192, 8193, 100_000, 1_000_000) + v_h = rand(Float32, n) + v = array_from_host(v_h) + AK.sort!(v; prefer_threads, alg=AK.BitonicSort()) + @test Array(v) == sort(v_h) + end + + # ── Adversarial patterns ────────────────────────────────────────────── + for n in (1000, 8192, 65536) + for v_h in (fill(2.5f0, n), Float32.(1:n), Float32.(n:-1:1), + Float32.(rand(0:1, n)), Float32.(rand(0:3, n))) + v = array_from_host(v_h) + AK.sort!(v; prefer_threads, alg=AK.BitonicSort()) + @test Array(v) == sort(v_h) + end + end + + # ── Ordering composition ────────────────────────────────────────────── + v_h = rand(Int32, 10_000) + for (rev, order) in ((nothing, Base.Order.Reverse), (true, Base.Order.Forward), + (true, Base.Order.Reverse)) + v = array_from_host(v_h) + AK.sort!(v; prefer_threads, alg=AK.BitonicSort(), rev, order) + @test Array(v) == sort(v_h; rev, order) + end + + # ── Block-size tuning ───────────────────────────────────────────────── + v_h = rand(UInt32, 20_000) + for block_size in (64, 128, 256, 512) + v = array_from_host(v_h) + AK.sort!(v; prefer_threads, alg=AK.BitonicSort(; block_size)) + @test Array(v) == sort(v_h) + end + + # ── Edge cases ──────────────────────────────────────────────────────── + @test length(Array(AK.sort!(array_from_host(Int32[]); prefer_threads, alg=AK.BitonicSort()))) == 0 + @test Array(AK.sort!(array_from_host(Int32[42]); prefer_threads, alg=AK.BitonicSort())) == Int32[42] + @test Array(AK.sort!(array_from_host(Int32[2, 1]); prefer_threads, alg=AK.BitonicSort())) == Int32[1, 2] + + # ── Out-of-place: input unchanged ───────────────────────────────────── + v_h = rand(Float32, 10_000) + v = array_from_host(v_h) + w = AK.sort(v; prefer_threads, alg=AK.BitonicSort()) + @test Array(w) == sort(v_h) + @test Array(v) == v_h + + # ── Rejected: N-D array, custom by/lt, unsupported element type ──────── + m = array_from_host(rand(Float32, 8, 8)) + @test_throws ArgumentError AK.sort!(m; prefer_threads, alg=AK.BitonicSort()) + + v = array_from_host(rand(Int32, 16)) + @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.BitonicSort(), by=abs) + @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.BitonicSort(), lt=(>)) + + v = array_from_host(rand(Int16, 16)) + @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.BitonicSort()) + end +end From e10df2ea3bb2c777bd7c2040153293553ce41b86 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Wed, 16 Sep 2026 08:15:21 +0530 Subject: [PATCH 2/3] feat(sort): BitonicSort along dims Add a `dims` path to `sort!`/`sort` (`sort(A; dims=d)`), sorting each 1-D slice independently. On the GPU each slice is sorted by its own workgroup entirely in shared memory, so per-slice sorting of many small slices is fast; each slice must fit the single-workgroup budget. Only BitonicSort implements a `dims` path; the CPU backend falls back to per-slice Base.sort!. Co-Authored-By: Claude Opus 4.8 --- src/sort/bitonic_sort.jl | 83 ++++++++++++++++++++++++++++++++++++++++ src/sort/sort.jl | 43 ++++++++++++++++++++- test/generic/sort.jl | 52 +++++++++++++++++++++++++ 3 files changed, 177 insertions(+), 1 deletion(-) diff --git a/src/sort/bitonic_sort.jl b/src/sort/bitonic_sort.jl index 64f1d0d3..73cc4c0c 100644 --- a/src/sort/bitonic_sort.jl +++ b/src/sort/bitonic_sort.jl @@ -132,6 +132,89 @@ end end +# Sort each length-`slen` slice along one dimension, one workgroup per slice, in shared memory. +# Slices along dim `d` are strided by `inner` (product of the sizes below `d`); element `j` of +# slice `s` (both 0-based) is at `(s % inner) + (s ÷ inner) * inner * slen + j * inner`. +@kernel cpu=false inbounds=true function _bitonic_slice!( + A, slen::Int, inner::Int, descending::Bool, ::Val{CAP}, ::Val{BS}, ::Val{IPT}, +) where {CAP, BS, IPT} + tile = @localmem eltype(A) (CAP,) + s = Int(@index(Group, Linear)) - 1 + t = Int(@index(Local, Linear)) - 1 + sbase = (s % inner) + (s ÷ inner) * inner * slen + pad = descending ? typemin(eltype(A)) : typemax(eltype(A)) + + m = 0 + while m < IPT + j = t + m * BS + @inbounds tile[j + 1] = j < slen ? A[sbase + j * inner + 1] : pad + m += 1 + end + @synchronize() + + kklog = 1 + while (1 << kklog) <= CAP + jlog = kklog - 1 + while jlog >= 0 + jj = 1 << jlog + p = t + while p < (CAP >> 1) + i = ((p >> jlog) << (jlog + 1)) | (p & (jj - 1)) + partner = i + jj + asc = (i & (1 << kklog)) == 0 + a = @inbounds tile[i + 1] + b = @inbounds tile[partner + 1] + if (a > b) == (asc != descending) + @inbounds tile[i + 1] = b + @inbounds tile[partner + 1] = a + end + p += BS + end + @synchronize() + jlog -= 1 + end + kklog += 1 + end + + m = 0 + while m < IPT + j = t + m * BS + j < slen && (@inbounds A[sbase + j * inner + 1] = tile[j + 1]) + m += 1 + end +end + + +# Sort each 1-D slice of `A` along `dim`, matching `Base.sort(A; dims)`. Each slice is sorted by one +# workgroup, so a slice must fit the single-workgroup budget; larger slices are rejected. +function _bitonic_sort_dims!( + A::AbstractArray{T, N}, backend::Backend, dim::Int; + descending::Bool, + block_size::Union{Nothing, Int}=nothing, +) where {T, N} + slen = size(A, dim) + slen <= 1 && return A + + npow = nextpow(2, slen) + budget = _prevpow2(_bs_shmem_bytes(backend) ÷ sizeof(T)) + npow <= budget || throw(ArgumentError( + "BitonicSort along dims=$dim needs each slice ($slen elements) to fit the single-workgroup " * + "shared-memory budget ($budget elements for $T); use a flat sort or a smaller slice")) + + cap = npow + bs = isnothing(block_size) ? min(256, cap) : min(_prevpow2(block_size), cap) + ipt = cap ÷ bs + inner = dim == 1 ? 1 : Base.prod(ntuple(k -> size(A, k), dim - 1)) + nslices = length(A) ÷ slen + + _bitonic_slice!(backend, bs)( + A, slen, inner, descending, Val(cap), Val(bs), Val(ipt); ndrange = nslices * bs, + ) + KernelAbstractions.synchronize(backend) + A +end + + function _bitonic_sort!( v::AbstractVector{T}, backend::Backend; descending::Bool, diff --git a/src/sort/sort.jl b/src/sort/sort.jl index deb8dcf4..f2ac9fc5 100644 --- a/src/sort/sort.jl +++ b/src/sort/sort.jl @@ -55,6 +55,10 @@ grows slower with `n`: the number of global compare-exchange passes grows with `log²n`. Non-power-of-two lengths are padded up internally; `block_size` sets the threads per workgroup. + +`BitonicSort` also implements `sort(A; dims)`: each 1-D slice is sorted by its own workgroup in +shared memory, so per-slice sorting of many small slices is its best case. This requires each slice +to fit the single-workgroup budget. """ Base.@kwdef struct BitonicSort <: SortAlgorithm block_size::Union{Nothing, Int} = nothing @@ -99,7 +103,9 @@ struct SampleSort <: SortAlgorithm end ) Sorts the array `v` in-place using the specified backend. The `lt`, `by`, `rev`, and `order` -arguments are the same as for `Base.sort`. +arguments are the same as for `Base.sort`. Pass an integer `dims` to sort each 1-D slice along that +dimension independently, matching `Base.sort(A; dims)`; this uses [`BitonicSort`](@ref) and requires +each slice to fit a single workgroup's shared memory. 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 @@ -183,6 +189,11 @@ function _sort_impl!( # Temporary buffer, same size as `v` temp::Union{Nothing, AbstractArray}=nothing, ) + if !(dims isa Colon) + return _sort_dims_impl!(v, backend, Int(dims); + lt, by, rev, order, alg, prefer_threads, block_size) + end + if use_gpu_algorithm(backend, prefer_threads) alg = isnothing(alg) ? MergeSort() : alg if alg isa MergeSort @@ -247,6 +258,36 @@ function _sort_impl!( end +# Sort each 1-D slice of `v` along `dim`, matching `Base.sort(A; dims)`. Only `BitonicSort` provides +# a `dims` path (each slice sorts in one workgroup); the CPU path falls back to `Base.sort!`. +function _sort_dims_impl!( + v::AbstractArray{T, N}, backend::Backend, dim::Int; + lt, by, rev, order, alg, prefer_threads, block_size, +) where {T, N} + 1 <= dim <= N || throw(ArgumentError("dimension $dim is out of range 1:$N")) + + if use_gpu_algorithm(backend, prefer_threads) + (isnothing(alg) || alg isa BitonicSort) || + throw(ArgumentError("sort along `dims` is only supported by BitonicSort, got $(typeof(alg))")) + _bs_supported(T) || + throw(ArgumentError("BitonicSort is not supported for eltype \"$T\"")) + ordering = Base.Order.ord(lt, by, rev, order) + ordering === Base.Order.Forward || ordering === Base.Order.Reverse || + throw(ArgumentError("BitonicSort only supports forward or reverse ordering")) + bs_block = (alg isa BitonicSort && !isnothing(alg.block_size)) ? alg.block_size : block_size + _bitonic_sort_dims!(v, backend, dim; descending=ordering === Base.Order.Reverse, + block_size=bs_block) + else + ordering = Base.Order.ord(lt, by, rev, order) + other_dims = Tuple(d for d in 1:N if d != dim) # slices vary only along `dim` + for slice in eachslice(v; dims=other_dims) + sort!(slice; order=ordering) + end + v + end +end + + """ sort( v::AbstractArray, backend::Backend=get_backend(v); diff --git a/test/generic/sort.jl b/test/generic/sort.jl index 576a959f..11853264 100644 --- a/test/generic/sort.jl +++ b/test/generic/sort.jl @@ -1097,3 +1097,55 @@ end @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.BitonicSort()) end end + + +@testset "bitonic_sort_dims" begin + if !prefer_threads + Random.seed!(0) + + # ── Per-slice correctness vs Base.sort(A; dims), 2-D and 3-D ────────── + for T in valid_backend_eltypes(BACKEND, (UInt32, Int32, Float32, UInt64, Int64, Float64)) + for (L, ncols) in ((8, 5), (256, 16), (1024, 4), (100, 50), (2048, 8)) + for dim in (1, 2) + sz = dim == 1 ? (L, ncols) : (ncols, L) + h = rand(T, sz...) + v = array_from_host(h) + AK.sort!(v; prefer_threads, dims=dim, alg=AK.BitonicSort()) + @test Array(v) == sort(h; dims=dim) + + v = array_from_host(h) + AK.sort!(v; prefer_threads, dims=dim, alg=AK.BitonicSort(), rev=true) + @test Array(v) == sort(h; dims=dim, rev=true) + end + end + + h = rand(T, 7, 40, 5) + for dim in (1, 2, 3) + v = array_from_host(h) + AK.sort!(v; prefer_threads, dims=dim, alg=AK.BitonicSort()) + @test Array(v) == sort(h; dims=dim) + end + end + + # ── Out-of-place leaves the input untouched ─────────────────────────── + h = rand(Float32, 256, 10) + v = array_from_host(h) + w = AK.sort(v; prefer_threads, dims=1, alg=AK.BitonicSort()) + @test Array(w) == sort(h; dims=1) + @test Array(v) == h + + # ── The default `dims` path (alg=nothing) also uses bitonic ─────────── + h = rand(Float32, 512, 8) + v = array_from_host(h) + AK.sort!(v; prefer_threads, dims=2) + @test Array(v) == sort(h; dims=2) + + # ── Rejected: slice too large, unsupported eltype, wrong algorithm ───── + @test_throws ArgumentError AK.sort!(array_from_host(rand(Float32, 20_000, 4)); + prefer_threads, dims=1, alg=AK.BitonicSort()) + @test_throws ArgumentError AK.sort!(array_from_host(rand(Int16, 64, 4)); + prefer_threads, dims=1, alg=AK.BitonicSort()) + @test_throws ArgumentError AK.sort!(array_from_host(rand(Float32, 64, 4)); + prefer_threads, dims=1, alg=AK.RadixSort()) + end +end From 3df08bc29fb7261ed844109554c5d3e614d21fac Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 16 Sep 2026 14:09:41 +0200 Subject: [PATCH 3/3] Rework BitonicSort on the slice layout Sort the network ascending in `ord` with the reflected first step of every merge level, and skip comparators whose partner lies past the end: any length sorts without padding, so the sentinel buffer, the eltype whitelist and the Forward/Reverse restriction go away and NaNs order like Base. Every kernel takes the slice layout from #117, so `dims` needs no separate path, slices of any length work, and slices shorter than a workgroup share a tile. Tunables live on `BitonicSort(; block_size, items_per_thread)` with backend defaults from `bitonic_defaults`. --- docs/src/api/sort.md | 15 +- src/sort/bitonic_sort.jl | 350 +++++++++++++-------------------------- src/sort/sort.jl | 103 ++++-------- test/generic/sort.jl | 169 +++++++++++-------- 4 files changed, 255 insertions(+), 382 deletions(-) diff --git a/docs/src/api/sort.md b/docs/src/api/sort.md index 20af5a85..e01d9c6d 100644 --- a/docs/src/api/sort.md +++ b/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 diff --git a/src/sort/bitonic_sort.jl b/src/sort/bitonic_sort.jl index 73cc4c0c..60678bbd 100644 --- a/src/sort/bitonic_sort.jl +++ b/src/sort/bitonic_sort.jl @@ -1,51 +1,71 @@ -# Portable GPU bitonic sort on KernelAbstractions primitives. A sorting network has no -# data-dependent branching, so it maps well onto GPUs; inputs that fit a single workgroup sort -# entirely in shared memory. Non-power-of-two lengths are padded with a sentinel (`typemax` -# ascending, `typemin` descending) so it sorts to the tail. - -# Radix's eltypes: comparable, with a `typemax`/`typemin` sentinel. -_bs_supported(::Type{T}) where T = - T === UInt32 || T === Int32 || T === Float32 || - T === UInt64 || T === Int64 || T === Float64 - -# Single-workgroup shared-memory budget in bytes; overridable per backend. Matches the radix sort. -_bs_shmem_bytes(::Backend) = 32 * 1024 -@inline _prevpow2(x::Int) = 1 << (8 * sizeof(Int) - leading_zeros(x) - 1) +# 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 -# Sort one CAP-sized window from scratch in shared memory; the window's global base fixes each -# comparator's direction. Used for a whole small array and to build the windows of the large path. -@kernel cpu=false inbounds=true function _bitonic_full!( - w, descending::Bool, ::Val{CAP}, ::Val{BS}, ::Val{IPT}, -) where {CAP, BS, IPT} - tile = @localmem eltype(w) (CAP,) - blk = @index(Group, Linear) - t = Int(@index(Local, Linear)) - 1 - base = Int(blk - 0x1) * CAP - m = 0 - while m < IPT - pos = t + m * BS - @inbounds tile[pos + 1] = w[base + pos + 1] - m += 1 +# 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 = 1 - while (1 << kklog) <= CAP - jlog = kklog - 1 + kklog = first_level + while kklog <= last_level + jlog = min(kklog - 1, spanlog - 1) while jlog >= 0 - j = 1 << jlog - p = t - while p < (CAP >> 1) - i = ((p >> jlog) << (jlog + 1)) | (p & (j - 1)) - partner = i + j - asc = ((base + i) & (1 << kklog)) == 0 - a = @inbounds tile[i + 1] - b = @inbounds tile[partner + 1] - if (a > b) == (asc != descending) - @inbounds tile[i + 1] = b - @inbounds tile[partner + 1] = a + 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 @@ -55,215 +75,81 @@ _bs_shmem_bytes(::Backend) = 32 * 1024 kklog += 1 end - m = 0 - while m < IPT - pos = t + m * BS - @inbounds w[base + pos + 1] = tile[pos + 1] - m += 1 - end -end - - -# Finish the short strides (`j_start` down to 0, all < CAP) of merge level `kklog` for one window -# in shared memory; every comparator's pair stays inside the window. -@kernel cpu=false inbounds=true function _bitonic_block!( - w, kklog::Int, j_start::Int, descending::Bool, - ::Val{CAP}, ::Val{BS}, ::Val{IPT}, -) where {CAP, BS, IPT} - tile = @localmem eltype(w) (CAP,) - blk = @index(Group, Linear) - t = Int(@index(Local, Linear)) - 1 - base = Int(blk - 0x1) * CAP - kk = 1 << kklog - - m = 0 - while m < IPT - pos = t + m * BS - @inbounds tile[pos + 1] = w[base + pos + 1] - m += 1 - end - @synchronize() - - jlog = j_start - while jlog >= 0 - j = 1 << jlog - p = t - while p < (CAP >> 1) - i = ((p >> jlog) << (jlog + 1)) | (p & (j - 1)) - partner = i + j - asc = ((base + i) & kk) == 0 - a = @inbounds tile[i + 1] - b = @inbounds tile[partner + 1] - if (a > b) == (asc != descending) - @inbounds tile[i + 1] = b - @inbounds tile[partner + 1] = a - end - p += BS + 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 - @synchronize() - jlog -= 1 - end - - m = 0 - while m < IPT - pos = t + m * BS - @inbounds w[base + pos + 1] = tile[pos + 1] - m += 1 + pos += BS end end -# One global compare-exchange for a large-stride merge level, launched one thread per pair -# (`npow ÷ 2`): every lane works, no divergent guard, and consecutive threads stay coalesced. -@kernel cpu=false inbounds=true function _bitonic_global!( - w, kklog::Int, jlog::Int, descending::Bool, +# 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, ) - p = Int(@index(Global, Linear)) - 1 - j = 1 << jlog - i = ((p >> jlog) << (jlog + 1)) | (p & (j - 1)) - partner = i + j - asc = (i & (1 << kklog)) == 0 - a = @inbounds w[i + 1] - b = @inbounds w[partner + 1] - if (a > b) == (asc != descending) - @inbounds w[i + 1] = b - @inbounds w[partner + 1] = a - end + 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 -# Sort each length-`slen` slice along one dimension, one workgroup per slice, in shared memory. -# Slices along dim `d` are strided by `inner` (product of the sizes below `d`); element `j` of -# slice `s` (both 0-based) is at `(s % inner) + (s ÷ inner) * inner * slen + j * inner`. -@kernel cpu=false inbounds=true function _bitonic_slice!( - A, slen::Int, inner::Int, descending::Bool, ::Val{CAP}, ::Val{BS}, ::Val{IPT}, -) where {CAP, BS, IPT} - tile = @localmem eltype(A) (CAP,) - s = Int(@index(Group, Linear)) - 1 - t = Int(@index(Local, Linear)) - 1 - sbase = (s % inner) + (s ÷ inner) * inner * slen - pad = descending ? typemin(eltype(A)) : typemax(eltype(A)) +""" + bitonic_defaults(backend::Backend) - m = 0 - while m < IPT - j = t + m * BS - @inbounds tile[j + 1] = j < slen ? A[sbase + j * inner + 1] : pad - m += 1 - end - @synchronize() +Default `block_size` and `items_per_thread` for [`BitonicSort`](@ref) on `backend`. +""" +bitonic_defaults(::Backend) = (block_size=256, items_per_thread=8) - kklog = 1 - while (1 << kklog) <= CAP - jlog = kklog - 1 - while jlog >= 0 - jj = 1 << jlog - p = t - while p < (CAP >> 1) - i = ((p >> jlog) << (jlog + 1)) | (p & (jj - 1)) - partner = i + jj - asc = (i & (1 << kklog)) == 0 - a = @inbounds tile[i + 1] - b = @inbounds tile[partner + 1] - if (a > b) == (asc != descending) - @inbounds tile[i + 1] = b - @inbounds tile[partner + 1] = a - end - p += BS - end - @synchronize() - jlog -= 1 - end - kklog += 1 - end - - m = 0 - while m < IPT - j = t + m * BS - j < slen && (@inbounds A[sbase + j * inner + 1] = tile[j + 1]) - m += 1 - end -end +function bitonic_sort!( + v::AbstractArray, backend::Backend=get_backend(v); -# Sort each 1-D slice of `A` along `dim`, matching `Base.sort(A; dims)`. Each slice is sorted by one -# workgroup, so a slice must fit the single-workgroup budget; larger slices are rejected. -function _bitonic_sort_dims!( - A::AbstractArray{T, N}, backend::Backend, dim::Int; - descending::Bool, - block_size::Union{Nothing, Int}=nothing, -) where {T, N} - slen = size(A, dim) - slen <= 1 && return A - - npow = nextpow(2, slen) - budget = _prevpow2(_bs_shmem_bytes(backend) ÷ sizeof(T)) - npow <= budget || throw(ArgumentError( - "BitonicSort along dims=$dim needs each slice ($slen elements) to fit the single-workgroup " * - "shared-memory budget ($budget elements for $T); use a flat sort or a smaller slice")) - - cap = npow - bs = isnothing(block_size) ? min(256, cap) : min(_prevpow2(block_size), cap) - ipt = cap ÷ bs - inner = dim == 1 ? 1 : Base.prod(ntuple(k -> size(A, k), dim - 1)) - nslices = length(A) ÷ slen - - _bitonic_slice!(backend, bs)( - A, slen, inner, descending, Val(cap), Val(bs), Val(ipt); ndrange = nslices * bs, - ) - KernelAbstractions.synchronize(backend) - A -end - - -function _bitonic_sort!( - v::AbstractVector{T}, backend::Backend; - descending::Bool, - block_size::Union{Nothing, Int}=nothing, -) where T - len = length(v) - len <= 1 && return 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) - budget = _prevpow2(_bs_shmem_bytes(backend) ÷ sizeof(T)) - cap = min(npow, budget) - bs = isnothing(block_size) ? min(256, cap) : min(_prevpow2(block_size), cap) - ipt = cap ÷ bs - - # Pad to a power of two; the sentinel sinks to the tail so real values fill 1:len. - pad = descending ? typemin(T) : typemax(T) - if npow == len - w = v - else - w = similar(v, npow) - copyto!(w, 1, v, 1, len) - fill!(view(w, (len + 1):npow), pad) - end - - if npow <= cap - _bitonic_full!(backend, bs)(w, descending, Val(cap), Val(bs), Val(ipt); ndrange = bs) - else - nlog = trailing_zeros(npow) - caplog = trailing_zeros(cap) - nblocks = npow ÷ cap - full = _bitonic_full!(backend, bs) - block = _bitonic_block!(backend, bs) - global_step = _bitonic_global!(backend, 256) - - # Build sorted windows, then merge them: global passes for j >= cap, one shared-memory - # batch for the remaining j < cap of each merge level. - full(w, descending, Val(cap), Val(bs), Val(ipt); ndrange = nblocks * bs) - for kklog in (caplog + 1):nlog - jlog = kklog - 1 - while jlog >= caplog - global_step(w, kklog, jlog, descending; ndrange = npow ÷ 2) - jlog -= 1 + 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 - block(w, kklog, caplog - 1, descending, Val(cap), Val(bs), Val(ipt); - ndrange = nblocks * bs) + tile!(v, ord, layout, tiles_per_slice, kklog, kklog, Val(cap), Val(span); + ndrange=nblocks * tile_threads) end end - KernelAbstractions.synchronize(backend) - npow == len || copyto!(v, 1, w, 1, len) v end diff --git a/src/sort/sort.jl b/src/sort/sort.jl index f2ac9fc5..7876b8d9 100644 --- a/src/sort/sort.jl +++ b/src/sort/sort.jl @@ -36,32 +36,22 @@ end _radix_defaults(::Backend) = (block_size=256, items_per_thread=2) """ - BitonicSort(; block_size=nothing) + BitonicSort(; block_size=nothing, items_per_thread=nothing) -Use a GPU bitonic sort for `sort!` and `sort`. Supports `UInt32`, `Int32`, `Float32`, `UInt64`, -`Int64`, and `Float64` with forward or reverse ordering. This algorithm does not support -`sortperm!`. - -A bitonic sorting network has no data-dependent control flow, so every comparator runs -unconditionally — a good fit for GPUs. For an input that fits a single workgroup's shared memory -the whole sort runs in one kernel with no global-memory round-trips, which is faster than the other -GPU algorithms for small arrays. Above that size the network spills to global memory and its -`O(n·log²n)` comparisons cost more than the `O(n)`-pass [`RadixSort`](@ref) or `MergeSort`, so it -grows slower with `n`: +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`. -- Small arrays (up to roughly a single workgroup, a few thousand elements): fastest of the three. -- Medium arrays: comparable to `MergeSort`. -- Large arrays: slower than `RadixSort` (measured ~1.1–1.8× on an RTX 5080 and RX 9060 XT), because - the number of global compare-exchange passes grows with `log²n`. +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. -Non-power-of-two lengths are padded up internally; `block_size` sets the threads per workgroup. - -`BitonicSort` also implements `sort(A; dims)`: each 1-D slice is sorted by its own workgroup in -shared memory, so per-slice sorting of many small slices is its best case. This requires each slice -to fit the single-workgroup budget. +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 """ @@ -103,14 +93,12 @@ struct SampleSort <: SortAlgorithm end ) Sorts the array `v` in-place using the specified backend. The `lt`, `by`, `rev`, and `order` -arguments are the same as for `Base.sort`. Pass an integer `dims` to sort each 1-D slice along that -dimension independently, matching `Base.sort(A; dims)`; this uses [`BitonicSort`](@ref) and requires -each slice to fit a single workgroup's shared memory. +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 @@ -125,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. @@ -189,11 +179,6 @@ function _sort_impl!( # Temporary buffer, same size as `v` temp::Union{Nothing, AbstractArray}=nothing, ) - if !(dims isa Colon) - return _sort_dims_impl!(v, backend, Int(dims); - lt, by, rev, order, alg, prefer_threads, block_size) - end - if use_gpu_algorithm(backend, prefer_threads) alg = isnothing(alg) ? MergeSort() : alg if alg isa MergeSort @@ -222,16 +207,12 @@ function _sort_impl!( temp, ) elseif alg isa BitonicSort - v isa AbstractVector || - throw(ArgumentError("BitonicSort only supports vectors (got a $(ndims(v))-dimensional array)")) - _bs_supported(eltype(v)) || throw(ArgumentError("BitonicSort is not supported for eltype \"$(eltype(v))\"")) - ordering = Base.Order.ord(lt, by, rev, order) - ordering === Base.Order.Forward || ordering === Base.Order.Reverse || - throw(ArgumentError("BitonicSort only supports forward or reverse ordering")) - _bitonic_sort!( + defaults = bitonic_defaults(backend) + bitonic_sort!( v, backend; - descending=ordering === Base.Order.Reverse, - block_size=isnothing(alg.block_size) ? block_size : alg.block_size, + 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")) @@ -258,36 +239,6 @@ function _sort_impl!( end -# Sort each 1-D slice of `v` along `dim`, matching `Base.sort(A; dims)`. Only `BitonicSort` provides -# a `dims` path (each slice sorts in one workgroup); the CPU path falls back to `Base.sort!`. -function _sort_dims_impl!( - v::AbstractArray{T, N}, backend::Backend, dim::Int; - lt, by, rev, order, alg, prefer_threads, block_size, -) where {T, N} - 1 <= dim <= N || throw(ArgumentError("dimension $dim is out of range 1:$N")) - - if use_gpu_algorithm(backend, prefer_threads) - (isnothing(alg) || alg isa BitonicSort) || - throw(ArgumentError("sort along `dims` is only supported by BitonicSort, got $(typeof(alg))")) - _bs_supported(T) || - throw(ArgumentError("BitonicSort is not supported for eltype \"$T\"")) - ordering = Base.Order.ord(lt, by, rev, order) - ordering === Base.Order.Forward || ordering === Base.Order.Reverse || - throw(ArgumentError("BitonicSort only supports forward or reverse ordering")) - bs_block = (alg isa BitonicSort && !isnothing(alg.block_size)) ? alg.block_size : block_size - _bitonic_sort_dims!(v, backend, dim; descending=ordering === Base.Order.Reverse, - block_size=bs_block) - else - ordering = Base.Order.ord(lt, by, rev, order) - other_dims = Tuple(d for d in 1:N if d != dim) # slices vary only along `dim` - for slice in eachslice(v; dims=other_dims) - sort!(slice; order=ordering) - end - v - end -end - - """ sort( v::AbstractArray, backend::Backend=get_backend(v); @@ -368,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, @@ -440,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 diff --git a/test/generic/sort.jl b/test/generic/sort.jl index 11853264..de7bd002 100644 --- a/test/generic/sort.jl +++ b/test/generic/sort.jl @@ -1012,89 +1012,97 @@ end @testset "bitonic_sort_alg" begin if !prefer_threads Random.seed!(0) + alg = AK.BitonicSort() - # ── Correctness: fuzzy testing across supported types ───────────────── - for T in valid_backend_eltypes(BACKEND, - (UInt32, Int32, Float32, UInt64, Int64, Float64)) - for _ in 1:100 + # Fuzz across element types, including ones RadixSort cannot handle + for T in valid_backend_eltypes(BACKEND, (UInt8, Int16, Int32, UInt32, Float32, Int64, UInt64, Float64)) + for _ in 1:20 n = rand(1:100_000) v_h = rand(T, n) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.BitonicSort()) + AK.sort!(v; prefer_threads, alg) @test Array(v) == sort(v_h) end end - # ── rev=true ────────────────────────────────────────────────────────── - for T in valid_backend_eltypes(BACKEND, - (UInt32, Int32, Float32, UInt64, Int64, Float64)) - n = 10_000 - v_h = rand(T, n) - v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.BitonicSort(), rev=true) - @test Array(v) == sort(v_h; rev=true) - end - - # ── Sizes around the single-block / multi-block boundary ────────────── - # Non-power-of-two lengths exercise the sentinel padding; large ones the - # global-memory merge passes. - for n in (2, 3, 7, 8, 255, 256, 257, 1023, 1024, 1025, - 8191, 8192, 8193, 100_000, 1_000_000) + # Lengths around the tile size (2048 by default) and through several global levels + for n in (1, 2, 3, 7, 8, 255, 256, 257, 1023, 1024, 1025, 2047, 2048, 2049, 4095, 4096, + 4097, 8191, 8192, 8193, 100_000, 1_000_000) v_h = rand(Float32, n) - v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.BitonicSort()) + v = array_from_host(v_h) + AK.sort!(v; prefer_threads, alg) @test Array(v) == sort(v_h) end - # ── Adversarial patterns ────────────────────────────────────────────── + # Adversarial patterns for n in (1000, 8192, 65536) for v_h in (fill(2.5f0, n), Float32.(1:n), Float32.(n:-1:1), Float32.(rand(0:1, n)), Float32.(rand(0:3, n))) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.BitonicSort()) + AK.sort!(v; prefer_threads, alg) @test Array(v) == sort(v_h) end end - # ── Ordering composition ────────────────────────────────────────────── + # NaNs and signed zeros follow `isless`, like Base + v_h = rand(Float32, 5000) + v_h[rand(1:5000, 100)] .= NaN32 + v_h[rand(1:5000, 100)] .= -0.0f0 + v_h[rand(1:5000, 100)] .= 0.0f0 + for rev in (false, true) + v = array_from_host(v_h) + AK.sort!(v; prefer_threads, alg, rev) + @test isequal(Array(v), sort(v_h; rev)) + end + + # lt, by, rev and order v_h = rand(Int32, 10_000) - for (rev, order) in ((nothing, Base.Order.Reverse), (true, Base.Order.Forward), - (true, Base.Order.Reverse)) + for kw in ((rev=true,), (order=Base.Order.Reverse,), (rev=true, order=Base.Order.Reverse), + (lt=(>),), (by=abs,), (by=x -> x % Int32(7), rev=true), (lt=(a, b) -> a % 5 < b % 5,)) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.BitonicSort(), rev, order) - @test Array(v) == sort(v_h; rev, order) + AK.sort!(v; prefer_threads, alg, kw...) + sorted = Array(v) + ord = Base.Order.ord(get(kw, :lt, isless), get(kw, :by, identity), + get(kw, :rev, nothing), get(kw, :order, Base.Order.Forward)) + @test issorted(sorted; order=ord) + @test sort(sorted) == sort(v_h) end - # ── Block-size tuning ───────────────────────────────────────────────── + # Tuning v_h = rand(UInt32, 20_000) - for block_size in (64, 128, 256, 512) + for block_size in (32, 128, 256, 512), items_per_thread in (1, 2, 8, 16) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.BitonicSort(; block_size)) + AK.sort!(v; prefer_threads, alg=AK.BitonicSort(; block_size, items_per_thread)) @test Array(v) == sort(v_h) end + v = array_from_host(v_h) + AK.sort!(v; prefer_threads, alg, block_size=64) + @test Array(v) == sort(v_h) + @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.BitonicSort(block_size=100)) + @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.BitonicSort(items_per_thread=3)) + @test_throws ArgumentError AK.sort!(v; prefer_threads, + alg=AK.BitonicSort(block_size=2, items_per_thread=1 << (Sys.WORD_SIZE - 2))) - # ── Edge cases ──────────────────────────────────────────────────────── - @test length(Array(AK.sort!(array_from_host(Int32[]); prefer_threads, alg=AK.BitonicSort()))) == 0 - @test Array(AK.sort!(array_from_host(Int32[42]); prefer_threads, alg=AK.BitonicSort())) == Int32[42] - @test Array(AK.sort!(array_from_host(Int32[2, 1]); prefer_threads, alg=AK.BitonicSort())) == Int32[1, 2] + # Empty input + @test isempty(Array(AK.sort!(array_from_host(Int32[]); prefer_threads, alg))) - # ── Out-of-place: input unchanged ───────────────────────────────────── + # Matrices are sorted as one flat vector by default + m_h = rand(Float32, 100, 30) + m = array_from_host(m_h) + AK.sort!(m; prefer_threads, alg) + @test vec(Array(m)) == sort(vec(m_h)) + + # Out-of-place: input unchanged v_h = rand(Float32, 10_000) - v = array_from_host(v_h) - w = AK.sort(v; prefer_threads, alg=AK.BitonicSort()) + v = array_from_host(v_h) + w = AK.sort(v; prefer_threads, alg) @test Array(w) == sort(v_h) @test Array(v) == v_h - # ── Rejected: N-D array, custom by/lt, unsupported element type ──────── - m = array_from_host(rand(Float32, 8, 8)) - @test_throws ArgumentError AK.sort!(m; prefer_threads, alg=AK.BitonicSort()) - - v = array_from_host(rand(Int32, 16)) - @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.BitonicSort(), by=abs) - @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.BitonicSort(), lt=(>)) - - v = array_from_host(rand(Int16, 16)) - @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.BitonicSort()) + # No permutation path + @test_throws ArgumentError AK.sortperm(v; prefer_threads, alg) + else + @test_throws ArgumentError AK.sort!(rand(Int32, 16); prefer_threads, alg=AK.BitonicSort()) end end @@ -1102,19 +1110,20 @@ end @testset "bitonic_sort_dims" begin if !prefer_threads Random.seed!(0) + alg = AK.BitonicSort() - # ── Per-slice correctness vs Base.sort(A; dims), 2-D and 3-D ────────── - for T in valid_backend_eltypes(BACKEND, (UInt32, Int32, Float32, UInt64, Int64, Float64)) - for (L, ncols) in ((8, 5), (256, 16), (1024, 4), (100, 50), (2048, 8)) + # Slices that fit one tile and slices that need global passes, 2D and 3D + for T in valid_backend_eltypes(BACKEND, (Int16, Int32, Float32, Int64, Float64)) + for (L, ncols) in ((8, 5), (256, 16), (1024, 4), (100, 50), (2048, 8), (2049, 3), (5000, 3), (20_000, 2)) for dim in (1, 2) sz = dim == 1 ? (L, ncols) : (ncols, L) h = rand(T, sz...) v = array_from_host(h) - AK.sort!(v; prefer_threads, dims=dim, alg=AK.BitonicSort()) + AK.sort!(v; prefer_threads, dims=dim, alg) @test Array(v) == sort(h; dims=dim) v = array_from_host(h) - AK.sort!(v; prefer_threads, dims=dim, alg=AK.BitonicSort(), rev=true) + AK.sort!(v; prefer_threads, dims=dim, alg, rev=true) @test Array(v) == sort(h; dims=dim, rev=true) end end @@ -1122,30 +1131,48 @@ end h = rand(T, 7, 40, 5) for dim in (1, 2, 3) v = array_from_host(h) - AK.sort!(v; prefer_threads, dims=dim, alg=AK.BitonicSort()) + AK.sort!(v; prefer_threads, dims=dim, alg) @test Array(v) == sort(h; dims=dim) end end - # ── Out-of-place leaves the input untouched ─────────────────────────── - h = rand(Float32, 256, 10) + # Packed slices spanning two tiles, with a partial final tile and captured orderings + packed_alg = AK.BitonicSort(block_size=8, items_per_thread=4) + mask = UInt32(0x55) + for len in (3, 4), dim in (1, 2) + h = rand(UInt32, dim == 1 ? (len, 9) : (9, len)) + for kw in ((by=x -> xor(x, mask),), (lt=(a, b) -> xor(a, mask) < xor(b, mask),)) + v = array_from_host(h) + AK.sort!(v; prefer_threads, dims=dim, alg=packed_alg, kw...) + @test Array(v) == sort(h; dims=dim, kw...) + end + end + + # by / lt, NaNs, vectors, empty and singleton slices + h = rand(Float32, 300, 700) + @test Array(AK.sort(array_from_host(h); prefer_threads, dims=1, alg, by=x -> -x)) == sort(h; dims=1, by=x -> -x) + @test Array(AK.sort(array_from_host(h); prefer_threads, dims=2, alg, lt=(>))) == sort(h; dims=2, lt=(>)) + h[rand(1:length(h), 1000)] .= NaN32 + @test isequal(Array(AK.sort(array_from_host(h); prefer_threads, dims=2, alg)), sort(h; dims=2)) + h = rand(Float32, 1000) + @test Array(AK.sort(array_from_host(h); prefer_threads, dims=1, alg)) == sort(h) + @test size(AK.sort(array_from_host(rand(Float32, 0, 5)); prefer_threads, dims=1, alg)) == (0, 5) + @test size(AK.sort(array_from_host(rand(Float32, 5, 0)); prefer_threads, dims=1, alg)) == (5, 0) + h = rand(Float32, 1, 100) + @test Array(AK.sort(array_from_host(h); prefer_threads, dims=1, alg)) == h + + # Tuning applies per slice; out-of-place leaves the input untouched + h = rand(Float32, 3000, 10) + for items_per_thread in (1, 4, 16) + v = array_from_host(h) + AK.sort!(v; prefer_threads, dims=1, alg=AK.BitonicSort(; block_size=128, items_per_thread)) + @test Array(v) == sort(h; dims=1) + end v = array_from_host(h) - w = AK.sort(v; prefer_threads, dims=1, alg=AK.BitonicSort()) + w = AK.sort(v; prefer_threads, dims=1, alg) @test Array(w) == sort(h; dims=1) @test Array(v) == h - # ── The default `dims` path (alg=nothing) also uses bitonic ─────────── - h = rand(Float32, 512, 8) - v = array_from_host(h) - AK.sort!(v; prefer_threads, dims=2) - @test Array(v) == sort(h; dims=2) - - # ── Rejected: slice too large, unsupported eltype, wrong algorithm ───── - @test_throws ArgumentError AK.sort!(array_from_host(rand(Float32, 20_000, 4)); - prefer_threads, dims=1, alg=AK.BitonicSort()) - @test_throws ArgumentError AK.sort!(array_from_host(rand(Int16, 64, 4)); - prefer_threads, dims=1, alg=AK.BitonicSort()) - @test_throws ArgumentError AK.sort!(array_from_host(rand(Float32, 64, 4)); - prefer_threads, dims=1, alg=AK.RadixSort()) + @test_throws ArgumentError AK.sort!(array_from_host(h); prefer_threads, dims=3, alg) end end