From bc8730988e76fba416b4ab00690f6464997cac5b Mon Sep 17 00:00:00 2001 From: Ludovic Raess Date: Mon, 17 Aug 2026 16:21:19 +0200 Subject: [PATCH] Fix sorting --- src/kernels/sorting.jl | 96 ++++++++++++++++++++++++++++++++++-- test/hip_rocarray/sorting.jl | 44 +++++++++++++++++ 2 files changed, 137 insertions(+), 3 deletions(-) diff --git a/src/kernels/sorting.jl b/src/kernels/sorting.jl index be55db12a..93e336a0a 100644 --- a/src/kernels/sorting.jl +++ b/src/kernels/sorting.jl @@ -1,3 +1,93 @@ -Base.sort!(x::AnyROCArray; kwargs...) = (AK.sort!(x; kwargs...); return x) -Base.sortperm!(ix::AnyROCArray, x::AnyROCArray; kwargs...) = (AK.sortperm!(ix, x; kwargs...); return ix) -Base.sortperm(x::AnyROCArray; kwargs...) = sortperm!(ROCArray(1:length(x)), x; kwargs...) +Base.sort!(x::AnyROCArray; dims::Union{Nothing, Integer}=nothing, kwargs...) = + isnothing(dims) ? (AK.sort!(x; kwargs...); x) : _sort_dims!(x, dims; kwargs...) + +# Base's out-of-place `sort(x; dims)` sorts chunks with the CPU algorithm, which trips +# scalar indexing - route it through `sort!` instead. +Base.sort(x::AnyROCArray; kwargs...) = sort!(copy(x); kwargs...) + +Base.sortperm!(ix::AnyROCArray, x::AnyROCArray; dims::Union{Nothing, Integer}=nothing, kwargs...) = + isnothing(dims) ? + (AK.sortperm!(ix, x; kwargs...); ix) : + _sortperm_dims!(ix, x, dims; kwargs...) + +Base.sortperm(x::AnyROCArray; dims::Union{Nothing, Integer}=nothing, kwargs...) = + isnothing(dims) ? + sortperm!(ROCArray(1:length(x)), x; kwargs...) : + sortperm!(reshape(ROCArray(1:length(x)), size(x)), x; dims, kwargs...) + +# Sorting along `dims`. AcceleratedKernels has no `dims` argument yet +# (JuliaGPU/AcceleratedKernels.jl#59), and one sort per slice would serialize into too many +# tiny kernels. Instead each element is tagged with the index of its slice and the array is +# sorted once, ordered lexicographically by `(slice, element)`; slices then come out grouped +# and internally sorted, and are scattered back. Drop this once AK supports `dims` directly. + +# Tag element `(i, j - 1)` of the `(sd * n, rest)` view of `x` with its slice index, plus its +# linear index for `sortperm`. Broadcast alongside `x` itself, so the ranges stay lazy. +_slice_kv(::Type{K}, sd, i, j, v) where K = (K(mod1(i, sd) + sd * j), v) +_slice_kvi(::Type{K}, ::Type{I}, sd, n, i, j, v) where {K, I} = + (K(mod1(i, sd) + sd * j), v, I(i + sd * n * j)) + +_slice_value(kv) = kv[2] +_slice_index(kv) = kv[3] + +# `x` seen as `(sd, n, rest)`: `n` elements along `dims`, strided by `sd`, over `rest` groups. +function _slice_layout(x::AnyROCArray, dims::Integer) + 1 ≤ dims ≤ ndims(x) || throw(ArgumentError("dimension out of range")) + n = size(x, dims) + sd = prod(size(x)[1:dims - 1]) + return n, sd, length(x) ÷ (n * sd) +end + +# Smallest key type that can enumerate the slices of `x`. +_slice_keytype(x::AnyROCArray) = length(x) ≤ typemax(Int32) ? Int32 : Int64 + +# `(slice, element)` ordering, ties broken by the user-supplied ordering. +_slice_lt(ord) = (a, b) -> a[1] == b[1] ? Base.Order.lt(ord, a[2], b[2]) : a[1] < b[1] + +# Scatter the sorted - and therefore slice-grouped - `kv` back over `dst`. +function _slice_scatter!(f::F, dst, kv, n, sd, rest) where F + if sd == 1 + reshape(dst, n, rest) .= f.(reshape(kv, n, rest)) + else + reshape(dst, sd, n, rest) .= f.(PermutedDimsArray(reshape(kv, n, sd, rest), (2, 1, 3))) + end + return dst +end + +function _sort_dims!( + x::AnyROCArray, dims::Integer; + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, + order::Base.Order.Ordering=Base.Order.Forward, kwargs..., +) + isempty(x) && return x + n, sd, rest = _slice_layout(x, dims) + K = _slice_keytype(x) + + kv = similar(x, Tuple{K, eltype(x)}, length(x)) + reshape(kv, sd * n, rest) .= _slice_kv.( + K, sd, 1:sd * n, (0:rest - 1)', reshape(x, sd * n, rest)) + + AK.sort!(kv; lt=_slice_lt(Base.Order.ord(lt, by, rev, order)), kwargs...) + + return _slice_scatter!(_slice_value, x, kv, n, sd, rest) +end + +function _sortperm_dims!( + ix::AnyROCArray, x::AnyROCArray, dims::Integer; + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, + order::Base.Order.Ordering=Base.Order.Forward, kwargs..., +) + axes(ix) == axes(x) || throw(DimensionMismatch( + "index array has axes $(axes(ix)), but sorted array has axes $(axes(x))")) + isempty(x) && return ix + n, sd, rest = _slice_layout(x, dims) + K, I = _slice_keytype(x), eltype(ix) + + kv = similar(x, Tuple{K, eltype(x), I}, length(x)) + reshape(kv, sd * n, rest) .= _slice_kvi.( + K, I, sd, n, 1:sd * n, (0:rest - 1)', reshape(x, sd * n, rest)) + + AK.sort!(kv; lt=_slice_lt(Base.Order.ord(lt, by, rev, order)), kwargs...) + + return _slice_scatter!(_slice_index, ix, kv, n, sd, rest) +end diff --git a/test/hip_rocarray/sorting.jl b/test/hip_rocarray/sorting.jl index 37fed8c97..497619f02 100644 --- a/test/hip_rocarray/sorting.jl +++ b/test/hip_rocarray/sorting.jl @@ -26,3 +26,47 @@ using AMDGPU: ROCArray @test y ≈ Array(xd[sortperm(xd; by=k -> 2 * k)]) end end + +@testset "Sorting along dims" begin + for sz in ((1, 1), (4, 7), (33, 65), (100, 100), (7, 5, 3), (1, 64, 1)), T in ( + Float32, Float64, Int32, Int64, + ) + x = rand(T, sz...) + xd = ROCArray(x) + + for dims in 1:length(sz), kwargs in ( + (;), (; rev=true), (; by=k -> 2 * k), (; lt=!isless), + (; order=Base.Order.Reverse), + ) + y = sort(x; dims, kwargs...) + @test y == Array(sort(xd; dims, kwargs...)) + + yd = copy(xd) + @test y == Array(sort!(yd; dims, kwargs...)) + @test y == Array(yd) + + p = sortperm(xd; dims, kwargs...) + @test size(p) == size(x) + @test y == Array(xd)[Array(p)] + end + end + + # `sort` must not mutate its argument. + x = ROCArray(Float32[3 1; 2 4]) + @test Array(sort(x; dims=2)) == Float32[1 3; 2 4] + @test Array(x) == Float32[3 1; 2 4] + + # `sortperm!` writes into the given index array. + ix = ROCArray(reshape(collect(1:4), 2, 2)) + @test Array(sortperm!(ix, x; dims=2)) == [3 1; 2 4] + + # Out-of-range dimensions error like `Base.sort!` does. + @test_throws ArgumentError sort!(x; dims=3) + @test_throws ArgumentError sort!(x; dims=0) + + # Larger-than-a-block slices, exercising the global merge passes. + x = rand(Float32, 1024, 1024) + for dims in 1:2 + @test sort(x; dims) == Array(sort(ROCArray(x); dims)) + end +end