improve interface

- add initialized kwarg
- safe index typing
- argument error if index arr len != val len

Author: xaellison

src/sorting.jl

@@ -853,13 +853,14 @@ function bitonic_sort!(c; by = identity, lt = isless, rev = false) where {T}
         j_final = (1 + k0 - k)
         for j = 1:j_final
 
-            # use Int32 args for indexing --> ~10% faster kernels
-            args1 = (c, map(Int32, (c_len, k, j, j_final))..., by, lt, Val(rev))
+            # use Int32 args for indexing when possible --> ~10% faster kernels
+            I = c_len <= typemax(Int32) ? Int32 : Int
+            args1 = (c, map(I, (c_len, k, j, j_final))..., by, lt, Val(rev))
             kernel1 = @cuda launch = false comparator_small_kernel(args1...)
             config1 = launch_configuration(kernel1.fun, shmem = threads -> bitonic_shmem(c, threads))
             threads1 = prevpow(2, config1.threads)
 
-            args2 = (c, map(Int32, (c_len, k, j))..., by, lt, Val(rev))
+            args2 = (c, map(I, (c_len, k, j))..., by, lt, Val(rev))
             kernel2 = @cuda launch = false comparator_kernel(args2...)
             config2 = launch_configuration(kernel2.fun, shmem = threads -> bitonic_shmem(c, threads))
             threads2 = prevpow(2, config2.threads)
@@ -939,11 +940,17 @@ function Base.partialsort(c::AnyCuArray, k::Union{Integer, OrdinalRange}; kwargs
     return partialsort!(copy(c), k; kwargs...)
 end
 
-function Base.sortperm!(I::AnyCuArray, c::AnyCuArray; kwargs...)
+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
+    if !initialized
+        I .= one(T):T(length(I))
+    end
     bitonic_sort!((c, I); kwargs...)
     return I
 end
 
 function Base.sortperm(c::AnyCuArray; kwargs...)
-    sortperm!(CuArray(1:length(c)), c; kwargs...)
+    sortperm!(CuArray(1:length(c)), c; initialized=true, kwargs...)
 end

test/sorting.jl

@@ -356,9 +356,13 @@ end
     @test check_sortperm(Float64, 1000000; rev=true)
     @test check_sortperm(Float64, 1000000; by=x->abs(x-0.5))
     @test check_sortperm(Float64, 1000000; rev=true, by=x->abs(x-0.5))
-    # check with int32 indices
+    # check with Int32 indices
     @test check_sortperm!(collect(Int32(1):Int32(1000000)), Float32, 1000000)
-
+    # `initialized` kwarg
+    @test check_sortperm!(collect(Int32(1):Int32(1000000)), Float32, 1000000; initialized=true)
+    @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))
 end
 
 end