Support dims in sort and sortperm, on top of AK

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...)

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