sortperm fails with dims keyword

[Sleort]

Describe the bug

Since Julia 1.9, Base.sortperm has accepted a dims keyword. This doesn't work with CuArrays.

To reproduce

The Minimal Working Example (MWE) for this bug:

julia> using CUDA

julia> xcpu = rand(3,3)
3×3 Matrix{Float64}:
 0.783023   0.240864   0.854814
 0.0488199  0.0840899  0.600029
 0.190407   0.872661   0.386219

julia> sortperm(xcpu; dims=1)
3×3 Matrix{Int64}:
 2  5  9
 3  4  8
 1  6  7

julia> xgpu = cu(xcpu)
3×3 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
 0.783023   0.240864   0.854814
 0.0488199  0.0840899  0.600029
 0.190407   0.872661   0.386219

julia> sortperm(xgpu; dims=1)
ERROR: MethodError: no method matching bitonic_sort!(::Tuple{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}; dims::Int64)

Closest candidates are:
  bitonic_sort!(::Any; by, lt, rev) got unsupported keyword argument "dims"
   @ CUDA ~/.julia/packages/CUDA/tVtYo/src/sorting.jl:850

Stacktrace:
 [1] kwerr(::NamedTuple{(:dims,), Tuple{Int64}}, ::Function, ::Tuple{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}})
   @ Base ./error.jl:165
 [2] sortperm!(I::CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}, c::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}; initialized::Bool, kwargs::Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:dims,), Tuple{Int64}}})
   @ CUDA ~/.julia/packages/CUDA/tVtYo/src/sorting.jl:984
 [3] sortperm(c::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}; kwargs::Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:dims,), Tuple{Int64}}})
   @ CUDA ~/.julia/packages/CUDA/tVtYo/src/sorting.jl:989
 [4] top-level scope
   @ REPL[5]:1
 [5] top-level scope
   @ ~/.julia/packages/CUDA/tVtYo/src/initialization.jl:185

Manifest.toml

(@v1.9) pkg> st CUDA
Status `~/.julia/environments/v1.9/Project.toml`
  [052768ef] CUDA v4.4.0

[[deps.CUDA]]
deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CUDA_Driver_jll", "CUDA_Runtime_Discovery", "CUDA_Runtime_jll", "ExprTools", "GPUArrays", "GPUCompiler", "KernelAbstractions", "LLVM", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "Preferences", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "SpecialFunctions", "UnsafeAtomicsLLVM"]
git-tree-sha1 = "35160ef0f03b14768abfd68b830f8e3940e8e0dc"
uuid = "052768ef-5323-5732-b1bb-66c8b64840ba"
version = "4.4.0"

[[deps.CUDA_Driver_jll]]
deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg"]
git-tree-sha1 = "498f45593f6ddc0adff64a9310bb6710e851781b"
uuid = "4ee394cb-3365-5eb0-8335-949819d2adfc"
version = "0.5.0+1"

[[deps.CUDA_Runtime_Discovery]]
deps = ["Libdl"]
git-tree-sha1 = "bcc4a23cbbd99c8535a5318455dcf0f2546ec536"
uuid = "1af6417a-86b4-443c-805f-a4643ffb695f"
version = "0.2.2"

[[deps.CUDA_Runtime_jll]]
deps = ["Artifacts", "CUDA_Driver_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"]
git-tree-sha1 = "5248d9c45712e51e27ba9b30eebec65658c6ce29"
uuid = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
version = "0.6.0+0"

[[deps.CUDNN_jll]]
deps = ["Artifacts", "CUDA_Runtime_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"]
git-tree-sha1 = "75923dce4275ead3799b238e10178a68c07dbd3b"
uuid = "62b44479-cb7b-5706-934f-f13b2eb2e645"
version = "8.9.4+0"

Expected behavior

The same behavior of sortperm when using CuArrays as with Arrays.

Version info

Details on Julia:

julia> versioninfo()
Julia Version 1.9.3
Commit bed2cd540a1 (2023-08-24 14:43 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 16 × Intel(R) Core(TM) i7-10875H CPU @ 2.30GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, skylake)
  Threads: 1 on 16 virtual cores

Details on CUDA:

julia> CUDA.versioninfo()
CUDA runtime 12.1, artifact installation
CUDA driver 12.2
NVIDIA driver 535.98.0

CUDA libraries: 
- CUBLAS: 12.1.3
- CURAND: 10.3.2
- CUFFT: 11.0.2
- CUSOLVER: 11.4.5
- CUSPARSE: 12.1.0
- CUPTI: 18.0.0
- NVML: 12.0.0+535.98

Julia packages: 
- CUDA: 4.4.0
- CUDA_Driver_jll: 0.5.0+1
- CUDA_Runtime_jll: 0.6.0+0

Toolchain:
- Julia: 1.9.3
- LLVM: 14.0.6
- PTX ISA support: 3.2, 4.0, 4.1, 4.2, 4.3, 5.0, 6.0, 6.1, 6.3, 6.4, 6.5, 7.0, 7.1, 7.2, 7.3, 7.4, 7.5
- Device capability support: sm_37, sm_50, sm_52, sm_53, sm_60, sm_61, sm_62, sm_70, sm_72, sm_75, sm_80, sm_86

1 device:
  0: NVIDIA GeForce GTX 1650 Ti (sm_75, 3.596 GiB / 4.000 GiB available)

Additional context

I believe this bug is due to bitonic_sort! being used instead of quicksort! within sortperm!, and bitonic_sort! doesn't support a dims keyword (?)

[ngphuoc]

I also hit into this bug. I need this sortperm with dims keyword to define a ChainRule for sort.