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 new file mode 100644 index 00000000..60678bbd --- /dev/null +++ b/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 diff --git a/src/sort/sort.jl b/src/sort/sort.jl index 75b96e17..7876b8d9 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,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 diff --git a/test/generic/sort.jl b/test/generic/sort.jl index 5910950e..de7bd002 100644 --- a/test/generic/sort.jl +++ b/test/generic/sort.jl @@ -1007,3 +1007,172 @@ 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) + alg = AK.BitonicSort() + + # 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) + @test Array(v) == sort(v_h) + end + end + + # 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) + @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) + @test Array(v) == sort(v_h) + end + end + + # 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 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, 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 + + # Tuning + v_h = rand(UInt32, 20_000) + 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, 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))) + + # Empty input + @test isempty(Array(AK.sort!(array_from_host(Int32[]); prefer_threads, alg))) + + # 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) + @test Array(w) == sort(v_h) + @test Array(v) == v_h + + # 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 + + +@testset "bitonic_sort_dims" begin + if !prefer_threads + Random.seed!(0) + alg = AK.BitonicSort() + + # 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) + @test Array(v) == sort(h; dims=dim) + + v = array_from_host(h) + AK.sort!(v; prefer_threads, dims=dim, alg, 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) + @test Array(v) == sort(h; dims=dim) + end + end + + # 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) + @test Array(w) == sort(h; dims=1) + @test Array(v) == h + + @test_throws ArgumentError AK.sort!(array_from_host(h); prefer_threads, dims=3, alg) + end +end