Make sortperm! resilient to type mismatches. (#2051)

Author: maleadt

src/sorting.jl

@@ -977,15 +977,14 @@ function Base.partialsort(c::AnyCuArray, k::Union{Integer, OrdinalRange}; kwargs
     return partialsort!(copy(c), k; kwargs...)
 end
 
-function Base.sortperm!(I::AnyCuArray{T}, c::AnyCuArray; initialized=false, kwargs...) where T
-    if length(I) != length(c)
-        throw(ArgumentError("index vector must have the same length/indices as the source vector"))
-    end
+function Base.sortperm!(ix::AnyCuArray{T}, A::AnyCuArray; initialized=false, kwargs...) where T
+    axes(ix) == axes(A) || throw(ArgumentError("index array must have the same size/axes as the source array, $(axes(ix)) != $(axes(A))"))
+
     if !initialized
-        I .= one(T):T(length(I))
+        ix .= LinearIndices(A)
     end
-    bitonic_sort!((c, I); kwargs...)
-    return I
+    bitonic_sort!((A, ix); kwargs...)
+    return ix
 end
 
 function Base.sortperm(c::AnyCuArray; kwargs...)

test/base/sorting.jl

@@ -243,15 +243,15 @@ function check_sortperm!(i, T, N; kwargs...)
     I = CuArray(i)
     a = rand(T, N)
     c = CuArray(a)
-    CUDA.@sync sortperm!(I, c; kwargs...)
+    sortperm!(I, c; kwargs...)
     return Array(I) == sortperm!(i, a; kwargs...)
 end
 
 
 function check_sortperm(T, N; kwargs...)
     a = rand(T, N)
     c = CuArray(a)
-    I = CUDA.@sync sortperm(c; kwargs...)
+    I = sortperm(c; kwargs...)
     return Array(I) == sortperm(a; kwargs...)
 end
 
@@ -397,6 +397,8 @@ end
     @test check_sortperm!(collect(Int32(1):Int32(1000000)), Float32, 1000000; initialized=false)
     # expected error case
     @test_throws ArgumentError sortperm!(CuArray(1:3), CuArray(1:4))
+    # mismatched types (JuliaGPU/CUDA.jl#2046)
+    @test check_sortperm!(collect(UInt64(1):UInt64(1000000)), Int64, 1000000)
 end
 
 end