diff --git a/src/kernels/sorting.jl b/src/kernels/sorting.jl index be55db12a..55c239c15 100644 --- a/src/kernels/sorting.jl +++ b/src/kernels/sorting.jl @@ -1,3 +1,12 @@ Base.sort!(x::AnyROCArray; kwargs...) = (AK.sort!(x; kwargs...); return x) + +# Base's out-of-place `sort(x; dims)` would sort chunks on the CPU, tripping scalar indexing. +Base.sort(x::AnyROCArray; kwargs...) = sort!(copy(x); kwargs...) + 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...) + +# AcceleratedKernels wants the index array shaped like `x` when `dims` is given. +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...) diff --git a/test/hip_rocarray/sorting.jl b/test/hip_rocarray/sorting.jl index 37fed8c97..0fb0af6ab 100644 --- a/test/hip_rocarray/sorting.jl +++ b/test/hip_rocarray/sorting.jl @@ -26,3 +26,49 @@ 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=(>)), + (; 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 Array(p) == sortperm(x; dims, kwargs...) + @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, which must be shaped like `x`. + ix = ROCArray(reshape(collect(1:4), 2, 2)) + @test Array(sortperm!(ix, x; dims=2)) == [3 1; 2 4] + @test_throws ArgumentError sortperm!(ROCArray(collect(1:4)), x; dims=2) + + # 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