From aaed884d2f85ccc87d44034d1653282e404e61c7 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Mon, 24 Aug 2026 15:17:36 +0530 Subject: [PATCH 1/2] Add dims support to sort, sort!, sortperm and sortperm! Sort or permute each 1D slice along an integer `dims` independently, matching Base. The default `dims=:` keeps the existing flat behaviour. There is no batched sort kernel, so each element is tagged with its slice and the whole array is sorted once by (slice, value), then scattered back into place. sortperm carries the original index as the payload and uses it to break ties, keeping the permutation stable. This reuses the backend's tuned sort and runs unchanged on CPU and every GPU backend. --- src/sort/sort.jl | 146 +++++++++++++++++++++++++++++++++++++++++++ test/generic/sort.jl | 97 ++++++++++++++++++++++++++++ 2 files changed, 243 insertions(+) diff --git a/src/sort/sort.jl b/src/sort/sort.jl index 8a7d3c99..05e2ac6d 100644 --- a/src/sort/sort.jl +++ b/src/sort/sort.jl @@ -61,6 +61,9 @@ struct SampleSort <: SortAlgorithm end # Algorithm choice alg::Union{Nothing, SortAlgorithm}=nothing, + # Sort each slice along this dimension; `:` sorts the whole array flat + dims::Union{Colon, Integer}=Colon(), + # GPU settings block_size::Union{Nothing, Int}=nothing, @@ -71,6 +74,10 @@ 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`. +With the default `dims=:` the whole array is sorted as one flat vector. Pass an integer `dims` to +sort each 1D slice along that dimension independently, matching `Base.sort(A; dims)`. The `dims` +form always uses a comparison sort and so ignores `alg`. + ## CPU CPU settings: use at most `max_tasks` threads to sort the array such that at least `min_elems` elements are sorted by each thread. A parallel sample sort is used, processing @@ -138,12 +145,24 @@ function _sort_impl!( alg::Union{Nothing, SortAlgorithm}=nothing, + # Sort each 1D slice along this dimension; `:` sorts the whole array as one flat vector + dims::Union{Colon, Integer}=Colon(), + # GPU settings; nothing => each GPU algorithm picks its own tuned default block_size::Union{Nothing, Int}=nothing, # Temporary buffer, same size as `v` temp::Union{Nothing, AbstractArray}=nothing, ) + if !(dims isa Colon) + return _sort_dims!( + v, backend, Int(dims); + lt, by, rev, order, + max_tasks, min_elems, prefer_threads, + block_size, + ) + end + if use_gpu_algorithm(backend, prefer_threads) alg = isnothing(alg) ? MergeSort() : alg if alg isa MergeSort @@ -189,6 +208,54 @@ function _sort_impl!( end +# Sort each slice along `dim` on its own, like Base.sort(A; dims). +# We have no batched sort kernel, so tag every element with its slice, sort the whole array +# once by (slice, value), then scatter each element back to its place. Works on any backend. +function _sort_dims!( + v::AbstractArray{T, N}, backend::Backend, dim::Int; + lt, by, rev, order, + max_tasks, min_elems, prefer_threads, + block_size, +) where {T, N} + 1 <= dim <= N || throw(ArgumentError("dimension $dim is not 1 ≤ dims ≤ $N")) + slice_len = size(v, dim) + (length(v) <= 1 || slice_len <= 1) && return v # every slice is a singleton + + len = length(v) + bs = isnothing(block_size) ? 256 : block_size + + # a slice is picked by the other dims, so collapse `dim` to 1 + other_size = ntuple(d -> d == dim ? 1 : size(v, d), N) + slice_lin = LinearIndices(other_size) + slice_car = CartesianIndices(other_size) + elem_car = CartesianIndices(v) + + keys = similar(v, Tuple{Int, T}, len) + foreachindex(v, backend; max_tasks, min_elems, prefer_threads, block_size=bs) do i + ci = elem_car[i] + proj = ntuple(d -> d == dim ? 1 : ci[d], N) + sid = slice_lin[CartesianIndex(proj)] - 1 + @inbounds keys[i] = (sid, v[i]) + end + + # slices in order, values within a slice by the user's ordering + o = Base.Order.ord(lt, by, rev, order) + comp = (a, b) -> a[1] != b[1] ? a[1] < b[1] : Base.Order.lt(o, a[2], b[2]) + _sort_impl!(keys, backend; lt=comp, max_tasks, min_elems, prefer_threads, block_size) + + # sorted keys are grouped by slice, so slot k lands at row (k-1)%slice_len of its slice + foreachindex(keys, backend; max_tasks, min_elems, prefer_threads, block_size=bs) do k + s = (k - 1) ÷ slice_len + r = (k - 1) % slice_len + oc = slice_car[s + 1] + dst = ntuple(d -> d == dim ? r + 1 : oc[d], N) + @inbounds v[CartesianIndex(dst)] = keys[k][2] + end + + v +end + + """ sort( v::AbstractArray, backend::Backend=get_backend(v); @@ -244,6 +311,9 @@ end # Algorithm choice alg::Union{Nothing, SortAlgorithm}=nothing, + # Permute each slice along this dimension; `:` permutes the whole array flat + dims::Union{Colon, Integer}=Colon(), + # GPU settings block_size::Union{Nothing, Int}=nothing, @@ -255,6 +325,11 @@ Save into `ix` the index permutation of `v` such that `v[ix]` is sorted. The `lt `order` arguments are the same as for `Base.sortperm`. The same algorithms are used as for [`sort!`](@ref) with custom by-index comparators. +With the default `dims=:` the whole array is permuted as one flat vector. Pass an integer `dims` to +permute each 1D slice along that dimension independently, matching `Base.sortperm(A; dims)`; then +`ix` holds global linear indices and must have the same axes as `v`. The `dims` form always uses a +comparison sort and so ignores `alg`. + ## 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. @@ -289,12 +364,24 @@ function _sortperm_impl!( alg::Union{Nothing, SortAlgorithm}=nothing, + # Permute each slice along this dimension; `:` permutes the whole array as one flat vector + dims::Union{Colon, Integer}=Colon(), + # GPU settings; nothing => merge sort's tuned default (sortperm is merge-only) block_size::Union{Nothing, Int}=nothing, # Temporary buffer, same size as `v` temp::Union{Nothing, AbstractArray}=nothing, ) + if !(dims isa Colon) + return _sortperm_dims!( + ix, v, backend, Int(dims); + lt, by, rev, order, + max_tasks, min_elems, prefer_threads, + block_size, + ) + end + if use_gpu_algorithm(backend, prefer_threads) alg = isnothing(alg) ? MergeSort() : alg bs = isnothing(block_size) ? 256 : block_size @@ -340,6 +427,65 @@ function _sortperm_impl!( end +# Same idea as _sort_dims!, but produce the permutation instead of sorting in place. +# We tag with (slice, value, index), sort, and write the original index into `ix`. The index +# also breaks ties, which keeps the permutation stable like Base.sortperm(A; dims). +function _sortperm_dims!( + ix::AbstractArray, v::AbstractArray{T, N}, backend::Backend, dim::Int; + lt, by, rev, order, + max_tasks, min_elems, prefer_threads, + block_size, +) where {T, N} + 1 <= dim <= N || throw(ArgumentError("dimension $dim is not 1 ≤ dims ≤ $N")) + axes(ix) == axes(v) || throw(ArgumentError("index array must have the same axes as the input")) + slice_len = size(v, dim) + len = length(v) + bs = isnothing(block_size) ? 256 : block_size + + if len <= 1 || slice_len <= 1 # each slice holds one element + foreachindex(v, backend; max_tasks, min_elems, prefer_threads, block_size=bs) do i + @inbounds ix[i] = i + end + return ix + end + + # a slice is picked by the other dims, so collapse `dim` to 1 + other_size = ntuple(d -> d == dim ? 1 : size(v, d), N) + slice_lin = LinearIndices(other_size) + slice_car = CartesianIndices(other_size) + elem_car = CartesianIndices(v) + + keys = similar(v, Tuple{Int, T, Int}, len) + foreachindex(v, backend; max_tasks, min_elems, prefer_threads, block_size=bs) do i + ci = elem_car[i] + proj = ntuple(d -> d == dim ? 1 : ci[d], N) + sid = slice_lin[CartesianIndex(proj)] - 1 + @inbounds keys[i] = (sid, v[i], i) + end + + # slices in order, then values, then index to break ties so the result stays stable + o = Base.Order.ord(lt, by, rev, order) + comp = (a, b) -> begin + a[1] != b[1] && return a[1] < b[1] + Base.Order.lt(o, a[2], b[2]) && return true + Base.Order.lt(o, b[2], a[2]) && return false + return a[3] < b[3] + end + _sort_impl!(keys, backend; lt=comp, max_tasks, min_elems, prefer_threads, block_size) + + # sorted keys are grouped by slice, so slot k lands at row (k-1)%slice_len of its slice + foreachindex(keys, backend; max_tasks, min_elems, prefer_threads, block_size=bs) do k + s = (k - 1) ÷ slice_len + r = (k - 1) % slice_len + oc = slice_car[s + 1] + dst = ntuple(d -> d == dim ? r + 1 : oc[d], N) + @inbounds ix[CartesianIndex(dst)] = keys[k][3] + end + + ix +end + + """ sortperm( v::AbstractArray, diff --git a/test/generic/sort.jl b/test/generic/sort.jl index d8b941de..39ffdd8d 100644 --- a/test/generic/sort.jl +++ b/test/generic/sort.jl @@ -831,4 +831,101 @@ end @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.RadixSort()) end end + + +@testset "sort_dims" begin + Random.seed!(0) + + # Fuzzy correctness against Base.sort(A; dims) for 2D and 3D arrays + for _ in 1:200 + nd = rand(2:3) + sz = ntuple(_ -> rand(1:15), nd) + for T in valid_backend_eltypes(BACKEND, (Int32, Float32)) + A_h = rand(T, sz...) + A = array_from_host(A_h) + for dim in 1:nd, rev in (false, true) + @test Array(AK.sort(A; prefer_threads, dims=dim, rev=rev)) == + sort(A_h; dims=dim, rev=rev) + end + end + end + + # by= and order= act on the values within each slice + A_h = rand(Float32, 17, 23) + A = array_from_host(A_h) + @test Array(AK.sort(A; prefer_threads, dims=1, by=x->-x)) == sort(A_h; dims=1, by=x->-x) + @test Array(AK.sort(A; prefer_threads, dims=2, + order=Base.Order.Reverse)) == sort(A_h; dims=2, order=Base.Order.Reverse) + + # In-place sorts each slice, leaves the array otherwise intact + A_h = rand(Int32, 40, 31) + A = array_from_host(A_h) + AK.sort!(A; prefer_threads, dims=2) + @test Array(A) == sort(A_h; dims=2) + + # dims=1 on a vector is a full sort + v_h = rand(Int32, 5000) + v = array_from_host(v_h) + @test Array(AK.sort(v; prefer_threads, dims=1)) == sort(v_h) + + # Singleton slice dimension is a no-op + A_h = rand(Float32, 1, 64) + A = array_from_host(A_h) + @test Array(AK.sort(A; prefer_threads, dims=1)) == A_h + + # Out-of-range dimension errors + A = array_from_host(rand(Float32, 8, 8)) + @test_throws ArgumentError AK.sort(A; prefer_threads, dims=3) + @test_throws ArgumentError AK.sort(A; prefer_threads, dims=0) +end + + +@testset "sortperm_dims" begin + Random.seed!(0) + + # Fuzzy correctness against Base.sortperm(A; dims); small integer ranges give many ties, so + # matching Base's index array exactly also checks that the permutation is stable + for _ in 1:200 + nd = rand(2:3) + sz = ntuple(_ -> rand(1:15), nd) + for T in valid_backend_eltypes(BACKEND, (Int32, Float32)) + A_h = T <: Integer ? rand(T(0):T(4), sz...) : rand(T, sz...) + A = array_from_host(A_h) + for dim in 1:nd, rev in (false, true) + ix = Array(AK.sortperm(A; prefer_threads, dims=dim, rev=rev)) + @test ix == sortperm(A_h; dims=dim, rev=rev) + @test A_h[ix] == sort(A_h; dims=dim, rev=rev) + end + end + end + + # by= and order= act on the values within each slice + A_h = rand(Float32, 17, 23) + A = array_from_host(A_h) + @test Array(AK.sortperm(A; prefer_threads, dims=1, by=x->-x)) == sortperm(A_h; dims=1, by=x->-x) + @test Array(AK.sortperm(A; prefer_threads, dims=2, + order=Base.Order.Reverse)) == sortperm(A_h; dims=2, order=Base.Order.Reverse) + + # In-place fills ix with the same global linear indices as Base + A_h = rand(Int32(0):Int32(5), 40, 31) + A = array_from_host(A_h) + ix = array_from_host(zeros(Int, 40, 31)) + AK.sortperm!(ix, A; prefer_threads, dims=2) + @test Array(ix) == sortperm(A_h; dims=2) + + # dims=1 on a vector is a full sortperm + v_h = rand(Int32(0):Int32(9), 5000) + v = array_from_host(v_h) + @test Array(AK.sortperm(v; prefer_threads, dims=1)) == sortperm(v_h) + + # Singleton slice dimension yields the identity index array + A_h = rand(Float32, 1, 64) + A = array_from_host(A_h) + @test Array(AK.sortperm(A; prefer_threads, dims=1)) == reshape(1:64, 1, 64) + + # Out-of-range dimension errors + A = array_from_host(rand(Float32, 8, 8)) + @test_throws ArgumentError AK.sortperm(A; prefer_threads, dims=3) + @test_throws ArgumentError AK.sortperm(A; prefer_threads, dims=0) +end end From 927d2dc5781614ba411a625b5e2e785f2f1f4551 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 16 Sep 2026 12:32:31 +0200 Subject: [PATCH 2/2] Sort along dims with slice-aware merge sort kernels Replace the tag-and-sort implementation of `dims` (tuple keys, a global sort of the whole array and a scatter back) with merge sort kernels that know about slices. A `SliceLayout` describes the 1D slices of an array along a dimension as `len` elements `stride` apart; each block maps its linear index to a slice and a tile within it, and accesses the slice through a `SliceView`. Slices up to twice the block size are sorted in local memory in one launch, larger ones get per-slice global merge passes. The whole-array sort is a one-slice `FlatLayout` whose accessor returns the array itself, so that path is unchanged. `merge_sort!`, `merge_sort_by_key!`, `merge_sortperm!` and `merge_sortperm_lowmem!` take `dims` directly; `sortperm` keeps linear indices as the values of a key/value sort, which is stable. On CPU each slice is sorted with `Base.sort!` on a view, slices spread over tasks, and `sortperm` sorts linear indices with a `Perm` ordering like Base. `RadixSort` with `dims` throws instead of silently switching algorithms, and the low-level sortperm entry points now default to `lt=isless`. Sorting 1024x1024 Float32 along a dimension on an RTX 5080 drops from 0.51 to 0.16 ms (CUDA.jl's bitonic sort: 0.31 ms), sortperm from 0.66 to 0.19 ms; the M1 and an Intel iGPU see 3-7x. The dims tests now also run on the CPU path and cover tile boundaries, global merge passes, ties, NaNs, strided views and empty slices. --- docs/src/api/sort.md | 12 +++ src/sort/merge_sort.jl | 89 ++++++++++------- src/sort/merge_sort_by_key.jl | 76 +++++++++----- src/sort/merge_sortperm.jl | 55 ++++++++--- src/sort/slices.jl | 71 ++++++++++++++ src/sort/sort.jl | 180 ++++++++-------------------------- test/generic/sort.jl | 130 +++++++++++++++++++----- 7 files changed, 373 insertions(+), 240 deletions(-) create mode 100644 src/sort/slices.jl diff --git a/docs/src/api/sort.md b/docs/src/api/sort.md index 904e7658..20af5a85 100644 --- a/docs/src/api/sort.md +++ b/docs/src/api/sort.md @@ -33,6 +33,18 @@ v = ROCArray(rand(Int32, 100_000)) AK.sort!(v) ``` +Multidimensional arrays are sorted as one flat vector by default; pass `dims` to sort each 1D slice +along that dimension independently, like `Base.sort!(A; dims)`. `sortperm` along `dims` returns +linear indices into the array, so `A[ix]` is sorted along `dims`: +```julia +A = ROCArray(rand(Float32, 1000, 1000)) +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!`. + 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 v = ROCArray(rand(Float32, 100_000)) diff --git a/src/sort/merge_sort.jl b/src/sort/merge_sort.jl index b81cdbf9..72a04dc3 100644 --- a/src/sort/merge_sort.jl +++ b/src/sort/merge_sort.jl @@ -1,26 +1,29 @@ -@kernel inbounds=true cpu=false unsafe_indices=true function _merge_sort_block!(vec, comp) +@kernel inbounds=true cpu=false unsafe_indices=true function _merge_sort_block!( + vec, comp, layout, blocks_per_slice, +) @uniform N = @groupsize()[1] s_buf = @localmem eltype(vec) (N * 0x2,) T = eltype(vec) I = typeof(N) - len = length(vec) - # NOTE: for many index calculations in this library, computation using zero-indexing leads to - # fewer operations (also code is transpiled to CUDA / ROCm / oneAPI / Metal code which do zero - # indexing). Internal calculations will be done using zero indexing except when actually - # accessing memory. As with C, the lower bound is inclusive, the upper bound exclusive. + # Use zero-based indices internally and half-open search bounds. # Group (block) and local (thread) indices iblock = @index(Group, Linear) - 0x1 ithread = @index(Local, Linear) - 0x1 + # Each block sorts a tile of one slice + islice, iblock = slice_block(layout, iblock, blocks_per_slice) + elems = slice(vec, layout, islice) + len = layout.len + i = ithread + iblock * N * 0x2 - i < len && (s_buf[ithread + 0x1] = vec[i + 0x1]) + i < len && (s_buf[ithread + 0x1] = elems[i + 0x1]) i = ithread + N + iblock * N * 0x2 - i < len && (s_buf[ithread + N + 0x1] = vec[i + 0x1]) + i < len && (s_buf[ithread + N + 0x1] = elems[i + 0x1]) @synchronize() @@ -68,28 +71,30 @@ end i = ithread + iblock * N * 0x2 - i < len && (vec[i + 0x1] = s_buf[ithread + 0x1]) + i < len && (elems[i + 0x1] = s_buf[ithread + 0x1]) i = ithread + N + iblock * N * 0x2 - i < len && (vec[i + 0x1] = s_buf[ithread + N + 0x1]) + i < len && (elems[i + 0x1] = s_buf[ithread + N + 0x1]) end @kernel inbounds=true cpu=false unsafe_indices=true function _merge_sort_global!( - @Const(vec_in), vec_out, comp, half_size_group, + @Const(vec_in), vec_out, comp, half_size_group, layout, blocks_per_slice, ) - len = length(vec_in) N = @groupsize()[1] - # NOTE: for many index calculations in this library, computation using zero-indexing leads to - # fewer operations (also code is transpiled to CUDA / ROCm / oneAPI / Metal code which do zero - # indexing). Internal calculations will be done using zero indexing except when actually - # accessing memory. As with C, the lower bound is inclusive, the upper bound exclusive. + # Use zero-based indices internally and half-open search bounds. # Group (block) and local (thread) indices iblock = @index(Group, Linear) - 0x1 ithread = @index(Local, Linear) - 0x1 + # Merges never cross slice boundaries + islice, iblock = slice_block(layout, iblock, blocks_per_slice) + slice_in = slice(vec_in, layout, islice) + slice_out = slice(vec_out, layout, islice) + len = layout.len + idx = ithread + iblock * N size_group = half_size_group * 0x2 gid = idx ÷ half_size_group @@ -101,14 +106,14 @@ end if lo >= len # Incomplete left half, nothing to swap on the right, simply copy elements to be sorted # in next iteration - pos_in < len && (vec_out[pos_in + 0x1] = vec_in[pos_in + 0x1]) + pos_in < len && (slice_out[pos_in + 0x1] = slice_in[pos_in + 0x1]) else hi = (gid + 0x1) * size_group hi > len && (hi = len) - pos_out = pos_in + _lower_bound_s0(vec_in, vec_in[pos_in + 0x1], lo, hi, comp) - lo - vec_out[pos_out + 0x1] = vec_in[pos_in + 0x1] + pos_out = pos_in + _lower_bound_s0(slice_in, slice_in[pos_in + 0x1], lo, hi, comp) - lo + slice_out[pos_out + 0x1] = slice_in[pos_in + 0x1] # Right half pos_in = gid * size_group + half_size_group + idx % half_size_group @@ -116,8 +121,8 @@ end if pos_in < len lo = gid * size_group hi = lo + half_size_group - pos_out = pos_in - half_size_group + _upper_bound_s0(vec_in, vec_in[pos_in + 0x1], lo, hi, comp) - lo - vec_out[pos_out + 0x1] = vec_in[pos_in + 0x1] + pos_out = pos_in - half_size_group + _upper_bound_s0(slice_in, slice_in[pos_in + 0x1], lo, hi, comp) - lo + slice_out[pos_out + 0x1] = slice_in[pos_in + 0x1] end end end @@ -134,6 +139,9 @@ end block_size::Int=256, temp::Union{Nothing, AbstractArray}=nothing, + + # Sort each 1D slice along this dimension; `:` sorts the whole array as one vector + dims::Union{Colon, Integer}=Colon(), ) """ function merge_sort!( @@ -146,40 +154,43 @@ function merge_sort!( block_size::Int=256, temp::Union{Nothing, AbstractArray}=nothing, + dims::Union{Colon, Integer}=Colon(), ) # Simple sanity checks @argcheck block_size > 0 + layout = slice_layout(v, dims) + ord = Base.Order.ord(lt, by, rev, order) + if !isnothing(temp) + @argcheck length(temp) == length(v) + @argcheck eltype(temp) === eltype(v) + end + (isempty(v) || layout.len <= 1) && return v - # Hoist `by` transform: broadcast it once to produce a key array, then sort - # (key, value) pairs. Without hoisting, by(elem) fires inside every binary-search - # comparison in the O(n log²n) merge hot-path — once per element per merge step. + # Compute keys once instead of evaluating `by` in every comparison. if by !== identity keys = by.(v) merge_sort_by_key!( keys, v, backend; - lt, rev, order, block_size, + lt, rev, order, block_size, dims, temp_values=temp, # temp was for v swap buffer; maps to temp_values here ) return v end - if !isnothing(temp) - @argcheck length(temp) == length(v) - @argcheck eltype(temp) === eltype(v) - end - # Construct comparator - ord = Base.Order.ord(lt, by, rev, order) comp = (x, y) -> Base.Order.lt(ord, x, y) - # Block level - blocks = (length(v) + block_size * 2 - 1) ÷ (block_size * 2) - _merge_sort_block!(backend, block_size)(v, comp, ndrange=(block_size * blocks,)) + # Block level: each block sorts a tile of one slice in local memory + len = layout.len + blocks = (len + block_size * 2 - 1) ÷ (block_size * 2) + _merge_sort_block!(backend, block_size)( + v, comp, layout, blocks, + ndrange=(block_size * blocks * slice_count(layout),), + ) - # Global level + # Global level: merge the sorted tiles of each slice, doubling the run length every pass half_size_group = Int32(block_size * 2) size_group = half_size_group * 2 - len = length(v) if len > half_size_group p1 = v p2 = isnothing(temp) ? similar(v) : temp @@ -189,7 +200,10 @@ function merge_sort!( niter = 0 while len > half_size_group blocks = ((len + half_size_group - 1) ÷ half_size_group + 1) ÷ 2 * (half_size_group ÷ block_size) - kernel!(p1, p2, comp, half_size_group, ndrange=(block_size * blocks,)) + kernel!( + p1, p2, comp, half_size_group, layout, blocks, + ndrange=(block_size * blocks * slice_count(layout),), + ) half_size_group = half_size_group << 1; size_group = size_group << 1; @@ -218,6 +232,7 @@ end block_size::Int=256, temp::Union{Nothing, AbstractArray}=nothing, + dims::Union{Colon, Integer}=Colon(), ) """ function merge_sort( diff --git a/src/sort/merge_sort_by_key.jl b/src/sort/merge_sort_by_key.jl index f6de5f3c..6d10b47b 100644 --- a/src/sort/merge_sort_by_key.jl +++ b/src/sort/merge_sort_by_key.jl @@ -1,31 +1,35 @@ -@kernel inbounds=true cpu=false unsafe_indices=true function _merge_sort_by_key_block!(keys, values, comp) +@kernel inbounds=true cpu=false unsafe_indices=true function _merge_sort_by_key_block!( + keys, values, comp, layout, blocks_per_slice, +) @uniform N = @groupsize()[1] s_keys = @localmem eltype(keys) (N * 0x2,) s_values = @localmem eltype(values) (N * 0x2,) I = typeof(N) - len = length(keys) - # NOTE: for many index calculations in this library, computation using zero-indexing leads to - # fewer operations (also code is transpiled to CUDA / ROCm / oneAPI / Metal code which do zero - # indexing). Internal calculations will be done using zero indexing except when actually - # accessing memory. As with C, the lower bound is inclusive, the upper bound exclusive. + # Use zero-based indices internally and half-open search bounds. # Group (block) and local (thread) indices iblock = @index(Group, Linear) - 0x1 ithread = @index(Local, Linear) - 0x1 + # Each block sorts a tile of one slice + islice, iblock = slice_block(layout, iblock, blocks_per_slice) + slice_keys = slice(keys, layout, islice) + slice_values = slice(values, layout, islice) + len = layout.len + i = ithread + iblock * N * 0x2 if i < len - s_keys[ithread + 0x1] = keys[i + 0x1] - s_values[ithread + 0x1] = values[i + 0x1] + s_keys[ithread + 0x1] = slice_keys[i + 0x1] + s_values[ithread + 0x1] = slice_values[i + 0x1] end i = ithread + N + iblock * N * 0x2 if i < len - s_keys[ithread + N + 0x1] = keys[i + 0x1] - s_values[ithread + N + 0x1] = values[i + 0x1] + s_keys[ithread + N + 0x1] = slice_keys[i + 0x1] + s_values[ithread + N + 0x1] = slice_values[i + 0x1] end @synchronize() @@ -85,14 +89,14 @@ i = ithread + iblock * N * 0x2 if i < len - keys[i + 0x1] = s_keys[ithread + 0x1] - values[i + 0x1] = s_values[ithread + 0x1] + slice_keys[i + 0x1] = s_keys[ithread + 0x1] + slice_values[i + 0x1] = s_values[ithread + 0x1] end i = ithread + N + iblock * N * 0x2 if i < len - keys[i + 0x1] = s_keys[ithread + N + 0x1] - values[i + 0x1] = s_values[ithread + N + 0x1] + slice_keys[i + 0x1] = s_keys[ithread + N + 0x1] + slice_values[i + 0x1] = s_values[ithread + N + 0x1] end end @@ -100,21 +104,25 @@ end @kernel inbounds=true cpu=false unsafe_indices=true function _merge_sort_by_key_global!( @Const(keys_in), keys_out, @Const(values_in), values_out, - comp, half_size_group, + comp, half_size_group, layout, blocks_per_slice, ) - len = length(keys_in) N = @groupsize()[1] - # NOTE: for many index calculations in this library, computation using zero-indexing leads to - # fewer operations (also code is transpiled to CUDA / ROCm / oneAPI / Metal code which do zero - # indexing). Internal calculations will be done using zero indexing except when actually - # accessing memory. As with C, the lower bound is inclusive, the upper bound exclusive. + # Use zero-based indices internally and half-open search bounds. # Group (block) and local (thread) indices iblock = @index(Group, Linear) - 0x1 ithread = @index(Local, Linear) - 0x1 + # Merges never cross slice boundaries + islice, iblock = slice_block(layout, iblock, blocks_per_slice) + keys_in = slice(keys_in, layout, islice) + keys_out = slice(keys_out, layout, islice) + values_in = slice(values_in, layout, islice) + values_out = slice(values_out, layout, islice) + len = layout.len + idx = ithread + iblock * N size_group = half_size_group * 0x2 gid = idx ÷ half_size_group @@ -167,6 +175,9 @@ end block_size::Int=256, temp_keys::Union{Nothing, AbstractArray}=nothing, temp_values::Union{Nothing, AbstractArray}=nothing, + + # Sort each 1D slice along this dimension; `:` sorts the whole array as one vector + dims::Union{Colon, Integer}=Colon(), ) """ function merge_sort_by_key!( @@ -182,10 +193,15 @@ function merge_sort_by_key!( block_size::Int=256, temp_keys::Union{Nothing, AbstractArray}=nothing, temp_values::Union{Nothing, AbstractArray}=nothing, + dims::Union{Colon, Integer}=Colon(), ) # Simple sanity checks @argcheck block_size > 0 @argcheck length(keys) == length(values) + layout = slice_layout(keys, dims) + if !(dims isa Colon) + @argcheck axes(keys) == axes(values) + end if !isnothing(temp_keys) @argcheck length(temp_keys) == length(keys) @argcheck eltype(temp_keys) === eltype(keys) @@ -197,16 +213,20 @@ function merge_sort_by_key!( # Construct comparator ord = Base.Order.ord(lt, by, rev, order) + (isempty(keys) || layout.len <= 1) && return keys, values comp = (x, y) -> Base.Order.lt(ord, x, y) - # Block level - blocks = (length(keys) + block_size * 2 - 1) ÷ (block_size * 2) - _merge_sort_by_key_block!(backend, block_size)(keys, values, comp, ndrange=(block_size * blocks,)) + # Block level: each block sorts a tile of one slice in local memory + len = layout.len + blocks = (len + block_size * 2 - 1) ÷ (block_size * 2) + _merge_sort_by_key_block!(backend, block_size)( + keys, values, comp, layout, blocks, + ndrange=(block_size * blocks * slice_count(layout),), + ) - # Global level + # Global level: merge the sorted tiles of each slice, doubling the run length every pass half_size_group = Int32(block_size * 2) size_group = half_size_group * 2 - len = length(keys) if len > half_size_group pk1 = keys pk2 = isnothing(temp_keys) ? similar(keys) : temp_keys @@ -219,7 +239,10 @@ function merge_sort_by_key!( niter = 0 while len > half_size_group blocks = ((len + half_size_group - 1) ÷ half_size_group + 1) ÷ 2 * (half_size_group ÷ block_size) - kernel!(pk1, pk2, pv1, pv2, comp, half_size_group, ndrange=(block_size * blocks,)) + kernel!( + pk1, pk2, pv1, pv2, comp, half_size_group, layout, blocks, + ndrange=(block_size * blocks * slice_count(layout),), + ) half_size_group = half_size_group << 1; size_group = size_group << 1; @@ -253,6 +276,7 @@ end block_size::Int=256, temp_keys::Union{Nothing, AbstractArray}=nothing, temp_values::Union{Nothing, AbstractArray}=nothing, + dims::Union{Colon, Integer}=Colon(), ) """ function merge_sort_by_key( diff --git a/src/sort/merge_sortperm.jl b/src/sort/merge_sortperm.jl index 6b97061b..44b50f46 100644 --- a/src/sort/merge_sortperm.jl +++ b/src/sort/merge_sortperm.jl @@ -4,7 +4,7 @@ v::AbstractArray, backend::Backend=get_backend(v); - lt=(<), + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, @@ -13,6 +13,9 @@ block_size::Int=256, temp_ix::Union{Nothing, AbstractArray}=nothing, temp_v::Union{Nothing, AbstractArray}=nothing, + + # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector + dims::Union{Colon, Integer}=Colon(), ) """ function merge_sortperm!( @@ -20,7 +23,7 @@ function merge_sortperm!( v::AbstractArray, backend::Backend=get_backend(v); - lt=(<), + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, @@ -29,10 +32,14 @@ function merge_sortperm!( block_size::Int=256, temp_ix::Union{Nothing, AbstractArray}=nothing, temp_v::Union{Nothing, AbstractArray}=nothing, + dims::Union{Colon, Integer}=Colon(), ) # Simple sanity checks @argcheck block_size > 0 @argcheck length(ix) == length(v) + dims isa Colon || @argcheck axes(ix) == axes(v) + layout = slice_layout(v, dims) + Base.Order.ord(lt, by, rev, order) # validate the ordering keywords before touching ix if !isnothing(temp_ix) @argcheck length(temp_ix) == length(ix) @argcheck eltype(temp_ix) === eltype(ix) @@ -43,15 +50,16 @@ function merge_sortperm!( @argcheck eltype(temp_v) === eltype(v) end - # Initialise indices that will be sorted by the keys in v + # Initialise the linear indices that will be sorted by the keys in v foreachindex(ix, block_size=block_size) do i @inbounds ix[i] = i end + (isempty(v) || layout.len <= 1) && return ix keys = inplace ? v : copy(v) merge_sort_by_key!( keys, ix, backend; - lt, by, rev, order, block_size, + lt, by, rev, order, block_size, dims, temp_keys=temp_v, temp_values=temp_ix, ) @@ -63,7 +71,7 @@ end merge_sortperm( v::AbstractArray, backend::Backend=get_backend(v); - lt=(<), + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, @@ -72,6 +80,9 @@ end block_size::Int=256, temp_ix::Union{Nothing, AbstractArray}=nothing, temp_v::Union{Nothing, AbstractArray}=nothing, + + # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector + dims::Union{Colon, Integer}=Colon(), ) """ function merge_sortperm( @@ -92,13 +103,16 @@ end v::AbstractArray, backend::Backend=get_backend(v); - lt=(<), + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, block_size::Int=256, temp::Union{Nothing, AbstractArray}=nothing, + + # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector + dims::Union{Colon, Integer}=Colon(), ) """ function merge_sortperm_lowmem!( @@ -106,39 +120,46 @@ function merge_sortperm_lowmem!( v::AbstractArray, backend::Backend=get_backend(v); - lt=(<), + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, block_size::Int=256, temp::Union{Nothing, AbstractArray}=nothing, + dims::Union{Colon, Integer}=Colon(), ) # Simple sanity checks @argcheck block_size > 0 @argcheck length(ix) == length(v) + dims isa Colon || @argcheck axes(ix) == axes(v) + layout = slice_layout(v, dims) + ord = Base.Order.ord(lt, by, rev, order) if !isnothing(temp) @argcheck length(temp) == length(ix) @argcheck eltype(temp) === eltype(ix) end - # Initialise indices that will be sorted by the keys in v + # Initialise the linear indices that will be sorted by the keys in v foreachindex(ix, block_size=block_size) do i @inbounds ix[i] = i end + (isempty(ix) || layout.len <= 1) && return ix # Construct custom comparator indexing into global array v - ord = Base.Order.ord(lt, by, rev, order) comp = (ix, iy) -> Base.Order.lt(ord, v[ix], v[iy]) # Block level - blocks = (length(ix) + block_size * 2 - 1) ÷ (block_size * 2) - _merge_sort_block!(backend, block_size)(ix, comp, ndrange=(block_size * blocks,)) + len = layout.len + blocks = (len + block_size * 2 - 1) ÷ (block_size * 2) + _merge_sort_block!(backend, block_size)( + ix, comp, layout, blocks, + ndrange=(block_size * blocks * slice_count(layout),), + ) # Global level half_size_group = Int32(block_size * 2) size_group = half_size_group * 2 - len = length(ix) if len > half_size_group p1 = ix p2 = isnothing(temp) ? similar(ix) : temp @@ -148,7 +169,10 @@ function merge_sortperm_lowmem!( niter = 0 while len > half_size_group blocks = ((len + half_size_group - 1) ÷ half_size_group + 1) ÷ 2 * (half_size_group ÷ block_size) - kernel!(p1, p2, comp, half_size_group, ndrange=(block_size * blocks,)) + kernel!( + p1, p2, comp, half_size_group, layout, blocks, + ndrange=(block_size * blocks * slice_count(layout),), + ) half_size_group = half_size_group << 1; size_group = size_group << 1; @@ -170,13 +194,16 @@ end merge_sortperm_lowmem( v::AbstractArray, backend::Backend=get_backend(v); - lt=(<), + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, block_size::Int=256, temp::Union{Nothing, AbstractArray}=nothing, + + # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector + dims::Union{Colon, Integer}=Colon(), ) """ function merge_sortperm_lowmem( diff --git a/src/sort/slices.jl b/src/sort/slices.jl new file mode 100644 index 00000000..8cebb781 --- /dev/null +++ b/src/sort/slices.jl @@ -0,0 +1,71 @@ +# Sorting along a dimension works on the 1D slices of an array: for `dims=d`, slice `i` (zero-based) +# holds `len = size(v, d)` elements spaced `stride = prod(size(v)[1:d-1])` apart. The kernels only +# see slices through `slice(v, layout, i)`; the whole-array sort is a flat layout with one slice, +# for which `slice` returns the array itself, keeping its indexing untouched. +struct FlatLayout + len::Int +end + +struct SliceLayout + len::Int # elements per slice + stride::Int # distance between consecutive elements of a slice + count::Int # number of slices +end + +slice_layout(v::AbstractArray, ::Colon) = FlatLayout(length(v)) + +function slice_layout(v::AbstractArray, dims::Integer) + 1 <= dims <= ndims(v) || throw(ArgumentError("dimension out of range")) + Base.require_one_based_indexing(v) + len = size(v, dims) + stride = 1 + for d in 1:dims - 1 + stride *= size(v, d) + end + SliceLayout(len, stride, len == 0 ? 0 : length(v) ÷ len) +end + +slice_count(::FlatLayout) = 1 +slice_count(layout::SliceLayout) = layout.count + + +# One slice of `parent`, indexed like a vector (one-based) +struct SliceView{A <: AbstractArray} + parent::A + offset::Int # zero-based linear offset of the first element + stride::Int +end + +Base.@propagate_inbounds Base.getindex(s::SliceView, i) = + s.parent[s.offset + (i - 1) * s.stride + 1] +Base.@propagate_inbounds Base.setindex!(s::SliceView, x, i) = + s.parent[s.offset + (i - 1) * s.stride + 1] = x +Base.eltype(::Type{SliceView{A}}) where A = eltype(A) + +# Split a kernel's linear block index into (slice, block within the slice) +@inline slice_block(::FlatLayout, iblock, blocks_per_slice) = (0, iblock) +@inline slice_block(::SliceLayout, iblock, blocks_per_slice) = divrem(iblock, blocks_per_slice) + +@inline slice(v, ::FlatLayout, i) = v + +@inline function slice(v, layout::SliceLayout, i) + # `i` is an index straight from the kernel, so it may be unsigned on some backends + i = Int(i) + offset = (i % layout.stride) + (i ÷ layout.stride) * layout.stride * layout.len + SliceView(v, offset, layout.stride) +end + + +# CPU: run `f(slice)` on a view of each slice of `v` along `dims`, slices split across tasks +function foreach_slice(f, v::AbstractArray{T, N}, dims::Integer; max_tasks, min_elems) where {T, N} + 1 <= dims <= N || throw(ArgumentError("dimension out of range")) + others = CartesianIndices(ntuple(d -> d == dims ? Base.OneTo(1) : axes(v, d), N)) + min_slices = cld(min_elems, max(size(v, dims), 1)) + task_partition(length(others), max_tasks, min_slices) do irange + for i in irange + other = others[i] + f(view(v, ntuple(d -> d == dims ? Colon() : other[d], N)...)) + end + end + v +end diff --git a/src/sort/sort.jl b/src/sort/sort.jl index 05e2ac6d..75b96e17 100644 --- a/src/sort/sort.jl +++ b/src/sort/sort.jl @@ -1,4 +1,5 @@ include("utils.jl") +include("slices.jl") include("merge_sort.jl") include("merge_sort_by_key.jl") include("merge_sortperm.jl") @@ -61,7 +62,7 @@ struct SampleSort <: SortAlgorithm end # Algorithm choice alg::Union{Nothing, SortAlgorithm}=nothing, - # Sort each slice along this dimension; `:` sorts the whole array flat + # Sort each 1D slice along this dimension; `:` sorts the whole array as one vector dims::Union{Colon, Integer}=Colon(), # GPU settings @@ -74,9 +75,11 @@ 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`. -With the default `dims=:` the whole array is sorted as one flat vector. Pass an integer `dims` to -sort each 1D slice along that dimension independently, matching `Base.sort(A; dims)`. The `dims` -form always uses a comparison sort and so ignores `alg`. +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 +`Base.sort!` (also with `alg=SampleSort()`, and without using `temp`), slices spread over the tasks. ## CPU CPU settings: use at most `max_tasks` threads to sort the array such that at least `min_elems` @@ -145,7 +148,7 @@ function _sort_impl!( alg::Union{Nothing, SortAlgorithm}=nothing, - # Sort each 1D slice along this dimension; `:` sorts the whole array as one flat vector + # Sort each 1D slice along this dimension; `:` sorts the whole array as one vector dims::Union{Colon, Integer}=Colon(), # GPU settings; nothing => each GPU algorithm picks its own tuned default @@ -154,15 +157,6 @@ function _sort_impl!( # Temporary buffer, same size as `v` temp::Union{Nothing, AbstractArray}=nothing, ) - if !(dims isa Colon) - return _sort_dims!( - v, backend, Int(dims); - lt, by, rev, order, - max_tasks, min_elems, prefer_threads, - block_size, - ) - end - if use_gpu_algorithm(backend, prefer_threads) alg = isnothing(alg) ? MergeSort() : alg if alg isa MergeSort @@ -170,9 +164,10 @@ function _sort_impl!( v, backend; lt, by, rev, order, block_size=isnothing(block_size) ? 256 : block_size, - temp, + temp, dims, ) elseif alg isa RadixSort + dims isa Colon || throw(ArgumentError("RadixSort does not support sorting along `dims`")) _rs_supported(eltype(v)) || throw(ArgumentError("RadixSort is not supported for eltype \"$(eltype(v))\"")) ordering = Base.Order.ord(lt, by, rev, order) ordering === Base.Order.Forward || ordering === Base.Order.Reverse || @@ -194,7 +189,9 @@ function _sort_impl!( end else alg = isnothing(alg) ? SampleSort() : alg - if alg isa SampleSort + if !(alg isa SampleSort) + throw(ArgumentError("$(typeof(alg)) is not supported by sort! on CPU backends")) + elseif dims isa Colon sample_sort!( v; lt, by, rev, order, @@ -202,56 +199,12 @@ function _sort_impl!( temp, ) else - throw(ArgumentError("$(typeof(alg)) is not supported by sort! on CPU backends")) + ord = Base.Order.ord(lt, by, rev, order) + foreach_slice(v, dims; max_tasks, min_elems) do slice + Base.sort!(slice; order=ord) + end end end -end - - -# Sort each slice along `dim` on its own, like Base.sort(A; dims). -# We have no batched sort kernel, so tag every element with its slice, sort the whole array -# once by (slice, value), then scatter each element back to its place. Works on any backend. -function _sort_dims!( - v::AbstractArray{T, N}, backend::Backend, dim::Int; - lt, by, rev, order, - max_tasks, min_elems, prefer_threads, - block_size, -) where {T, N} - 1 <= dim <= N || throw(ArgumentError("dimension $dim is not 1 ≤ dims ≤ $N")) - slice_len = size(v, dim) - (length(v) <= 1 || slice_len <= 1) && return v # every slice is a singleton - - len = length(v) - bs = isnothing(block_size) ? 256 : block_size - - # a slice is picked by the other dims, so collapse `dim` to 1 - other_size = ntuple(d -> d == dim ? 1 : size(v, d), N) - slice_lin = LinearIndices(other_size) - slice_car = CartesianIndices(other_size) - elem_car = CartesianIndices(v) - - keys = similar(v, Tuple{Int, T}, len) - foreachindex(v, backend; max_tasks, min_elems, prefer_threads, block_size=bs) do i - ci = elem_car[i] - proj = ntuple(d -> d == dim ? 1 : ci[d], N) - sid = slice_lin[CartesianIndex(proj)] - 1 - @inbounds keys[i] = (sid, v[i]) - end - - # slices in order, values within a slice by the user's ordering - o = Base.Order.ord(lt, by, rev, order) - comp = (a, b) -> a[1] != b[1] ? a[1] < b[1] : Base.Order.lt(o, a[2], b[2]) - _sort_impl!(keys, backend; lt=comp, max_tasks, min_elems, prefer_threads, block_size) - - # sorted keys are grouped by slice, so slot k lands at row (k-1)%slice_len of its slice - foreachindex(keys, backend; max_tasks, min_elems, prefer_threads, block_size=bs) do k - s = (k - 1) ÷ slice_len - r = (k - 1) % slice_len - oc = slice_car[s + 1] - dst = ntuple(d -> d == dim ? r + 1 : oc[d], N) - @inbounds v[CartesianIndex(dst)] = keys[k][2] - end - v end @@ -272,6 +225,9 @@ end # Algorithm choice alg::Union{Nothing, SortAlgorithm}=nothing, + # Sort each 1D slice along this dimension; `:` sorts the whole array as one vector + dims::Union{Colon, Integer}=Colon(), + # GPU settings block_size::Union{Nothing, Int}=nothing, @@ -311,7 +267,7 @@ end # Algorithm choice alg::Union{Nothing, SortAlgorithm}=nothing, - # Permute each slice along this dimension; `:` permutes the whole array flat + # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector dims::Union{Colon, Integer}=Colon(), # GPU settings @@ -325,10 +281,10 @@ Save into `ix` the index permutation of `v` such that `v[ix]` is sorted. The `lt `order` arguments are the same as for `Base.sortperm`. The same algorithms are used as for [`sort!`](@ref) with custom by-index comparators. -With the default `dims=:` the whole array is permuted as one flat vector. Pass an integer `dims` to -permute each 1D slice along that dimension independently, matching `Base.sortperm(A; dims)`; then -`ix` holds global linear indices and must have the same axes as `v`. The `dims` form always uses a -comparison sort and so ignores `alg`. +By default (`dims=:`) the whole array is permuted as one vector. Pass an integer `dims` to permute +each 1D slice along that dimension independently, like `Base.sortperm!(ix, A; dims)`: `ix` must +then have the same axes as `v` and receives linear indices into `v`, so that `v[ix]` is sorted +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 @@ -364,7 +320,7 @@ function _sortperm_impl!( alg::Union{Nothing, SortAlgorithm}=nothing, - # Permute each slice along this dimension; `:` permutes the whole array as one flat vector + # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector dims::Union{Colon, Integer}=Colon(), # GPU settings; nothing => merge sort's tuned default (sortperm is merge-only) @@ -374,12 +330,9 @@ function _sortperm_impl!( temp::Union{Nothing, AbstractArray}=nothing, ) if !(dims isa Colon) - return _sortperm_dims!( - ix, v, backend, Int(dims); - lt, by, rev, order, - max_tasks, min_elems, prefer_threads, - block_size, - ) + 1 <= dims <= ndims(v) || throw(ArgumentError("dimension out of range")) + axes(ix) == axes(v) || + throw(ArgumentError("index array must have the same axes as the input, $(axes(ix)) != $(axes(v))")) end if use_gpu_algorithm(backend, prefer_threads) @@ -391,7 +344,7 @@ function _sortperm_impl!( ix, v, backend; lt, by, rev, order, block_size=bs, - temp, + temp, dims, ) else # merge_sortperm! copies keys alongside indices in shared memory so comparisons @@ -403,6 +356,7 @@ function _sortperm_impl!( lt, by, rev, order, block_size=bs, temp_ix=temp, # old `temp` was the index buffer; maps directly to temp_ix + dims, ) end elseif alg isa RadixSort @@ -412,7 +366,9 @@ function _sortperm_impl!( end else alg = isnothing(alg) ? SampleSort() : alg - if alg isa SampleSort + if !(alg isa SampleSort) + throw(ArgumentError("$(typeof(alg)) is not supported by sortperm! on CPU backends")) + elseif dims isa Colon sample_sortperm!( ix, v; lt, by, rev, order, @@ -421,67 +377,14 @@ function _sortperm_impl!( temp, ) else - throw(ArgumentError("$(typeof(alg)) is not supported by sortperm! on CPU backends")) - end - end -end - - -# Same idea as _sort_dims!, but produce the permutation instead of sorting in place. -# We tag with (slice, value, index), sort, and write the original index into `ix`. The index -# also breaks ties, which keeps the permutation stable like Base.sortperm(A; dims). -function _sortperm_dims!( - ix::AbstractArray, v::AbstractArray{T, N}, backend::Backend, dim::Int; - lt, by, rev, order, - max_tasks, min_elems, prefer_threads, - block_size, -) where {T, N} - 1 <= dim <= N || throw(ArgumentError("dimension $dim is not 1 ≤ dims ≤ $N")) - axes(ix) == axes(v) || throw(ArgumentError("index array must have the same axes as the input")) - slice_len = size(v, dim) - len = length(v) - bs = isnothing(block_size) ? 256 : block_size - - if len <= 1 || slice_len <= 1 # each slice holds one element - foreachindex(v, backend; max_tasks, min_elems, prefer_threads, block_size=bs) do i - @inbounds ix[i] = i + # Like Base: sort the linear indices of each slice by the values they point to + ord = Base.Order.Perm(Base.Order.ord(lt, by, rev, order), vec(v)) + copyto!(ix, LinearIndices(v)) + foreach_slice(ix, dims; max_tasks, min_elems) do slice + Base.sort!(slice; order=ord) + end end - return ix end - - # a slice is picked by the other dims, so collapse `dim` to 1 - other_size = ntuple(d -> d == dim ? 1 : size(v, d), N) - slice_lin = LinearIndices(other_size) - slice_car = CartesianIndices(other_size) - elem_car = CartesianIndices(v) - - keys = similar(v, Tuple{Int, T, Int}, len) - foreachindex(v, backend; max_tasks, min_elems, prefer_threads, block_size=bs) do i - ci = elem_car[i] - proj = ntuple(d -> d == dim ? 1 : ci[d], N) - sid = slice_lin[CartesianIndex(proj)] - 1 - @inbounds keys[i] = (sid, v[i], i) - end - - # slices in order, then values, then index to break ties so the result stays stable - o = Base.Order.ord(lt, by, rev, order) - comp = (a, b) -> begin - a[1] != b[1] && return a[1] < b[1] - Base.Order.lt(o, a[2], b[2]) && return true - Base.Order.lt(o, b[2], a[2]) && return false - return a[3] < b[3] - end - _sort_impl!(keys, backend; lt=comp, max_tasks, min_elems, prefer_threads, block_size) - - # sorted keys are grouped by slice, so slot k lands at row (k-1)%slice_len of its slice - foreachindex(keys, backend; max_tasks, min_elems, prefer_threads, block_size=bs) do k - s = (k - 1) ÷ slice_len - r = (k - 1) % slice_len - oc = slice_car[s + 1] - dst = ntuple(d -> d == dim ? r + 1 : oc[d], N) - @inbounds ix[CartesianIndex(dst)] = keys[k][3] - end - ix end @@ -503,6 +406,9 @@ end # Algorithm choice alg::Union{Nothing, SortAlgorithm}=nothing, + # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector + dims::Union{Colon, Integer}=Colon(), + # GPU settings block_size::Union{Nothing, Int}=nothing, diff --git a/test/generic/sort.jl b/test/generic/sort.jl index 39ffdd8d..5910950e 100644 --- a/test/generic/sort.jl +++ b/test/generic/sort.jl @@ -832,30 +832,52 @@ end end end +end + @testset "sort_dims" begin Random.seed!(0) # Fuzzy correctness against Base.sort(A; dims) for 2D and 3D arrays - for _ in 1:200 + for _ in 1:100 nd = rand(2:3) sz = ntuple(_ -> rand(1:15), nd) - for T in valid_backend_eltypes(BACKEND, (Int32, Float32)) + for T in valid_backend_eltypes(BACKEND, (Int32, Float32, Float64)) A_h = rand(T, sz...) A = array_from_host(A_h) for dim in 1:nd, rev in (false, true) - @test Array(AK.sort(A; prefer_threads, dims=dim, rev=rev)) == - sort(A_h; dims=dim, rev=rev) + @test Array(AK.sort(A; prefer_threads, dims=dim, rev)) == sort(A_h; dims=dim, rev) end end end - # by= and order= act on the values within each slice - A_h = rand(Float32, 17, 23) + # Slice lengths around the block tile (2 * block_size) and beyond it, so that both the + # block-level sort and the global merge passes are exercised, with duplicate-heavy data + for len in (511, 512, 513, 1023, 1024, 1025, 2049, 10_000), nslices in (1, 3) + A_h = rand(Int32(0):Int32(7), len, nslices) + @test Array(AK.sort(array_from_host(A_h); prefer_threads, dims=1)) == sort(A_h; dims=1) + @test Array(AK.sort(array_from_host(A_h); prefer_threads, dims=1, rev=true)) == sort(A_h; dims=1, rev=true) + A_h = rand(Float32, nslices, len) + @test Array(AK.sort(array_from_host(A_h); prefer_threads, dims=2)) == sort(A_h; dims=2) + end + + # by, lt, order and temp act on the values within each slice + A_h = rand(Float32, 300, 700) A = array_from_host(A_h) @test Array(AK.sort(A; prefer_threads, dims=1, by=x->-x)) == sort(A_h; dims=1, by=x->-x) - @test Array(AK.sort(A; prefer_threads, dims=2, - order=Base.Order.Reverse)) == sort(A_h; dims=2, order=Base.Order.Reverse) + @test Array(AK.sort(A; prefer_threads, dims=2, lt=(>))) == sort(A_h; dims=2, lt=(>)) + @test Array(AK.sort(A; prefer_threads, dims=2, order=Base.Order.Reverse)) == sort(A_h; dims=2, order=Base.Order.Reverse) + @test Array(AK.sort(A; prefer_threads, dims=2, temp=similar(A))) == sort(A_h; dims=2) + if !prefer_threads + @test Array(AK.sort(A; prefer_threads, dims=1, block_size=64)) == sort(A_h; dims=1) + @test_throws ArgumentError AK.sort(A; prefer_threads, dims=1, alg=AK.RadixSort()) + end + + # NaNs, infinities and signed zeros order like Base + A_h = Float32[NaN 1 -0.0; 0.0 -Inf NaN; 2 NaN Inf] + A = array_from_host(A_h) + @test isequal(Array(AK.sort(A; prefer_threads, dims=1)), sort(A_h; dims=1)) + @test isequal(Array(AK.sort(A; prefer_threads, dims=2, rev=true)), sort(A_h; dims=2, rev=true)) # In-place sorts each slice, leaves the array otherwise intact A_h = rand(Int32, 40, 31) @@ -868,10 +890,18 @@ end v = array_from_host(v_h) @test Array(AK.sort(v; prefer_threads, dims=1)) == sort(v_h) - # Singleton slice dimension is a no-op - A_h = rand(Float32, 1, 64) + # 4D arrays + A_h = rand(Int32(0):Int32(3), 3, 4, 5, 6) A = array_from_host(A_h) - @test Array(AK.sort(A; prefer_threads, dims=1)) == A_h + for dim in 1:4 + @test Array(AK.sort(A; prefer_threads, dims=dim)) == sort(A_h; dims=dim) + end + + # Empty and singleton slices + for sz in ((0, 5), (5, 0), (1, 64), (64, 1)), dim in 1:2 + A_h = rand(Float32, sz...) + @test Array(AK.sort(array_from_host(A_h); prefer_threads, dims=dim)) == sort!(copy(A_h); dims=dim) + end # Out-of-range dimension errors A = array_from_host(rand(Float32, 8, 8)) @@ -885,26 +915,39 @@ end # Fuzzy correctness against Base.sortperm(A; dims); small integer ranges give many ties, so # matching Base's index array exactly also checks that the permutation is stable - for _ in 1:200 + for _ in 1:100 nd = rand(2:3) sz = ntuple(_ -> rand(1:15), nd) - for T in valid_backend_eltypes(BACKEND, (Int32, Float32)) + for T in valid_backend_eltypes(BACKEND, (Int32, Float32, Float64)) A_h = T <: Integer ? rand(T(0):T(4), sz...) : rand(T, sz...) A = array_from_host(A_h) for dim in 1:nd, rev in (false, true) - ix = Array(AK.sortperm(A; prefer_threads, dims=dim, rev=rev)) - @test ix == sortperm(A_h; dims=dim, rev=rev) - @test A_h[ix] == sort(A_h; dims=dim, rev=rev) + ix = Array(AK.sortperm(A; prefer_threads, dims=dim, rev)) + @test ix == sortperm(A_h; dims=dim, rev) + @test A_h[ix] == sort(A_h; dims=dim, rev) end end end - # by= and order= act on the values within each slice - A_h = rand(Float32, 17, 23) + # Ties across block tiles and global merge passes must stay stable + for len in (511, 512, 513, 1025, 2049, 10_000), nslices in (1, 3) + A_h = rand(Int32(0):Int32(3), len, nslices) + @test Array(AK.sortperm(array_from_host(A_h); prefer_threads, dims=1)) == sortperm(A_h; dims=1) + @test Array(AK.sortperm(array_from_host(A_h); prefer_threads, dims=1, rev=true)) == sortperm(A_h; dims=1, rev=true) + A_h = rand(Int32(0):Int32(3), nslices, len) + @test Array(AK.sortperm(array_from_host(A_h); prefer_threads, dims=2)) == sortperm(A_h; dims=2) + end + + # by, order, temp and the low-memory GPU path + A_h = rand(Float32, 300, 700) A = array_from_host(A_h) @test Array(AK.sortperm(A; prefer_threads, dims=1, by=x->-x)) == sortperm(A_h; dims=1, by=x->-x) - @test Array(AK.sortperm(A; prefer_threads, dims=2, - order=Base.Order.Reverse)) == sortperm(A_h; dims=2, order=Base.Order.Reverse) + @test Array(AK.sortperm(A; prefer_threads, dims=2, order=Base.Order.Reverse)) == sortperm(A_h; dims=2, order=Base.Order.Reverse) + @test Array(AK.sortperm(A; prefer_threads, dims=2, temp=similar(A, Int))) == sortperm(A_h; dims=2) + if !prefer_threads + @test Array(AK.sortperm(A; prefer_threads, dims=2, alg=AK.MergeSort(lowmem=true))) == sortperm(A_h; dims=2) + @test Array(AK.sortperm(A; prefer_threads, dims=1, alg=AK.MergeSort(lowmem=true), block_size=64)) == sortperm(A_h; dims=1) + end # In-place fills ix with the same global linear indices as Base A_h = rand(Int32(0):Int32(5), 40, 31) @@ -918,14 +961,49 @@ end v = array_from_host(v_h) @test Array(AK.sortperm(v; prefer_threads, dims=1)) == sortperm(v_h) - # Singleton slice dimension yields the identity index array - A_h = rand(Float32, 1, 64) - A = array_from_host(A_h) - @test Array(AK.sortperm(A; prefer_threads, dims=1)) == reshape(1:64, 1, 64) + # Empty and singleton slices + for sz in ((0, 5), (5, 0), (1, 64), (64, 1)), dim in 1:2 + A_h = rand(Float32, sz...) + @test Array(AK.sortperm(array_from_host(A_h); prefer_threads, dims=dim)) == sortperm(A_h; dims=dim) + end + + # Slice offsets refer to the view's linear indexing, not its parent's strides. + A_h = rand(Int32(0):Int32(3), 6, 1030) + A = array_from_host(A_h) + V_h = view(A_h, 1:2:6, 1:2:1030) + V = view(A, 1:2:6, 1:2:1030) + @test Array(AK.sortperm(V; prefer_threads, dims=2)) == sortperm(V_h; dims=2) + AK.sort!(V; prefer_threads, dims=2) + sort!(V_h; dims=2) + @test Array(A) == A_h + + # Floating-point ordering also applies to the low-level entry points. + A_h = Float32[NaN 1 -0.0; 0.0 -Inf NaN; 2 NaN Inf] + A = array_from_host(A_h) + for dim in 1:2, rev in (false, true) + expected = sortperm(A_h; dims=dim, rev) + @test Array(AK.sortperm(A; prefer_threads, dims=dim, rev)) == expected + if !prefer_threads + @test Array(AK.merge_sortperm(A; dims=dim, rev)) == expected + @test Array(AK.merge_sortperm_lowmem(A; dims=dim, rev)) == expected + end + end - # Out-of-range dimension errors + # Invalid dimensions must not overwrite the output. + ix = array_from_host(fill(-1, size(A_h))) + @test_throws ArgumentError AK.sortperm!(ix, A; prefer_threads, dims=3) + @test Array(ix) == fill(-1, size(A_h)) + if !prefer_threads + @test_throws ArgumentError AK.merge_sortperm!(ix, A; dims=3) + @test Array(ix) == fill(-1, size(A_h)) + @test_throws ArgumentError AK.merge_sortperm_lowmem!(ix, A; dims=3) + @test Array(ix) == fill(-1, size(A_h)) + @test_throws ArgumentError AK.merge_sort_by_key!(copy(A), similar(ix, length(ix)); dims=1) + end + + # Out-of-range dimension and mismatched index array errors A = array_from_host(rand(Float32, 8, 8)) @test_throws ArgumentError AK.sortperm(A; prefer_threads, dims=3) @test_throws ArgumentError AK.sortperm(A; prefer_threads, dims=0) -end + @test_throws ArgumentError AK.sortperm!(array_from_host(zeros(Int, 64)), A; prefer_threads, dims=1) end