Support dims in sort, sort!, sortperm and sortperm!

[luraess]

Fixes #1030.

src/kernels/sorting.jl forwarded everything to AcceleratedKernels, which has no dims argument yet, so every dims entry point failed, not always loudly:

call	before
sort!(A; dims=2)	MethodError from AK._sort_impl! (the issue)
sort(A; dims=2)	Scalar indexing is disallowed
sortperm(A; dims=2)	MethodError from AK._sortperm_impl!
sortperm!(ix, A; dims=2)	MethodError from AK._sortperm_impl!

sort needs its own method because Base.sort(A; dims) does not route through Base.sort!: it permutes and calls the internal CPU sort_chunks!, falling off the GPU.

Approach

AK tracks dims in JuliaGPU/AcceleratedKernels.jl#59 and JuliaGPU/GPUArrays.jl#608 is blocked on it, so this is a thin layer over AK.sort! meant to be deleted once AK grows dims (not trying to revive removed in #688).

Calling AK.sort! on a view per slice works but serialises into one tiny kernel launch per slice. Instead each element is tagged with the index of its slice and the array is sorted once, ordered lexicographically by (slice, element); slices come out grouped and internally sorted, then get scattered back. Tagging and scatter are plain broadcasts, so no new kernels.

RX 7900 XTX, Float32, ROCm 6.4.4, whole sort!(A; dims) call:

size	dims	slices	per-slice loop	this PR
(100, 100)	1	100	1.22 ms	0.14 ms
(1024, 1024)	1	1024	29.06 ms	0.92 ms
(1024, 1024)	2	1024	30.95 ms	0.81 ms
(8192, 128)	2	8192	115.79 ms	0.82 ms
(128, 8192)	1	8192	114.38 ms	0.76 ms

The cost is a global O(N log²N) sort where per-slice would be O(n log²n), plus the tag array and AK's temporary (~4× the footprint for Float64). A segmented sort upstream fixes both and is the intended replacement.

Note that on a matrix without dims, sort! still sorts flat and sortperm returns a flat vector, where Base throws UndefKeywordError. This is pre-existing, and changing it would be breaking.

[simeonschaub]

Should we perhaps contribute this to AK instead? Otherwise, every backend will have to reimplement its own version of this

[luraess]

Should we perhaps contribute this to AK instead? Otherwise, every backend will have to reimplement its own version of this

We could yeah - this was just a tmp solution until AK lands their implementation we may then want to use. But, we could go ahead and propose this to AK directly.

[github-actions]

AMDGPU.jl Benchmarks

Details

Benchmark suite	Current: bc87309	Previous: 6873c07	Ratio
amdgpu/synchronization/context/device	567.5 ns	562.5 ns	1.01
amdgpu/synchronization/stream/blocking	232.5 ns	235 ns	0.99
amdgpu/synchronization/stream/nonblocking	310 ns	307.5 ns	1.01
array/accumulate/Float32/1d	77961 ns	80728.75 ns	0.97
array/accumulate/Float32/dims=1	269931 ns	283941.25 ns	0.95
array/accumulate/Float32/dims=1L	81121.25 ns	98314 ns	0.83
array/accumulate/Float32/dims=2	81031 ns	94271.25 ns	0.86
array/accumulate/Float32/dims=2L	2616645.25 ns	2756887.75 ns	0.95
array/accumulate/Int64/1d	79608.75 ns	84288.5 ns	0.94
array/accumulate/Int64/dims=1	244625.75 ns	243910.75 ns	1.00
array/accumulate/Int64/dims=1L	84351.25 ns	105756.5 ns	0.80
array/accumulate/Int64/dims=2	86986.25 ns	96848.75 ns	0.90
array/accumulate/Int64/dims=2L	2889469 ns	2896974.5 ns	1.00
array/broadcast	74271 ns	73488.5 ns	1.01
array/construct	2217.5 ns	2497.5 ns	0.89
array/copy	37463 ns	36940.5 ns	1.01
array/copyto!/cpu_to_gpu	111751.5 ns	111344 ns	1.00
array/copyto!/gpu_to_cpu	112044 ns	111624 ns	1.00
array/copyto!/gpu_to_gpu	59665.75 ns	58636 ns	1.02
array/iteration/findall/bool	142712 ns	148087 ns	0.96
array/iteration/findall/int	153284.5 ns	150307 ns	1.02
array/iteration/findfirst/bool	146029.5 ns	146372 ns	1.00
array/iteration/findfirst/int	146094.25 ns	147247 ns	0.99
array/iteration/findmin/1d	111301.5 ns	133386.75 ns	0.83
array/iteration/findmin/2d	114219 ns	124199.25 ns	0.92
array/iteration/logical	251411 ns	251731 ns	1.00
array/iteration/scalar	306571.75 ns	295624 ns	1.04
array/permutedims/2d	73203.5 ns	72068.5 ns	1.02
array/permutedims/3d	72661 ns	71211 ns	1.02
array/permutedims/4d	75288.5 ns	74226 ns	1.01
array/random/rand/Float32	46195.75 ns	44905.75 ns	1.03
array/random/rand/Int64	55740.5 ns	54425.75 ns	1.02
array/random/rand!/Float32	60653.25 ns	65320.75 ns	0.93
array/random/rand!/Int64	74218.5 ns	73206 ns	1.01
array/random/randn/Float32	76723.5 ns	80716 ns	0.95
array/random/randn!/Float32	81843.5 ns	72986 ns	1.12
array/reductions/mapreduce/Float32/1d	102104 ns	111976.5 ns	0.91
array/reductions/mapreduce/Float32/dims=1	89476.25 ns	94548.75 ns	0.95
array/reductions/mapreduce/Float32/dims=1L	836291.25 ns	831646.25 ns	1.01
array/reductions/mapreduce/Float32/dims=2	91993.75 ns	98293.75 ns	0.94
array/reductions/mapreduce/Float32/dims=2L	144339.5 ns	144357 ns	1.00
array/reductions/mapreduce/Int64/1d	102864 ns	108913.75 ns	0.94
array/reductions/mapreduce/Int64/dims=1	91703.75 ns	94923.75 ns	0.97
array/reductions/mapreduce/Int64/dims=1L	832068.75 ns	833828.75 ns	1.00
array/reductions/mapreduce/Int64/dims=2	92863.75 ns	97976.25 ns	0.95
array/reductions/mapreduce/Int64/dims=2L	145364.25 ns	143437 ns	1.01
array/reductions/reduce/Float32/1d	102019 ns	111806.5 ns	0.91
array/reductions/reduce/Float32/dims=1	91866.25 ns	94403.75 ns	0.97
array/reductions/reduce/Float32/dims=1L	834751.25 ns	834681.25 ns	1.00
array/reductions/reduce/Float32/dims=2	89563.75 ns	99041.5 ns	0.90
array/reductions/reduce/Float32/dims=2L	144709.5 ns	143764.25 ns	1.01
array/reductions/reduce/Int64/1d	103231.5 ns	111681.25 ns	0.92
array/reductions/reduce/Int64/dims=1	92013.75 ns	94491.5 ns	0.97
array/reductions/reduce/Int64/dims=1L	832298.75 ns	832031 ns	1.00
array/reductions/reduce/Int64/dims=2	90918.5 ns	97973.75 ns	0.93
array/reductions/reduce/Int64/dims=2L	145637 ns	144357 ns	1.01
array/reverse/1d	46820.75 ns	46065.5 ns	1.02
array/reverse/1dL	76726 ns	76008.5 ns	1.01
array/reverse/1dL_inplace	73651 ns	80596 ns	0.91
array/reverse/1d_inplace	62213.25 ns	52315.5 ns	1.19
array/reverse/2d	51468.25 ns	50530.5 ns	1.02
array/reverse/2dL	88868.5 ns	96663.75 ns	0.92
array/reverse/2dL_inplace	93163.75 ns	91923.75 ns	1.01
array/reverse/2d_inplace	56255.75 ns	63498.25 ns	0.89
array/sorting/1d	338587 ns	335899.5 ns	1.01
integration/byval/reference	39951 ns	39670 ns	1.01
integration/byval/slices=1	40661 ns	39880 ns	1.02
integration/byval/slices=2	157052 ns	140532 ns	1.12
integration/byval/slices=3	237713 ns	238064 ns	1.00
integration/volumerhs	5004257 ns	5020670 ns	1.00
kernel/indexing	58660.75 ns	57780.75 ns	1.02
kernel/indexing_checked	60253.5 ns	59608.25 ns	1.01
kernel/launch	1467.5 ns	1460.25 ns	1.00
kernel/rand	105646.5 ns	111641.5 ns	0.95
latency/import	1674671930 ns	1936446792 ns	0.86
latency/precompile	38379952448 ns	38047934530 ns	1.01
latency/ttfp	5736476824 ns	6106370781 ns	0.94

This comment was automatically generated by workflow using github-action-benchmark.

[luraess]

I am leaning towards merging this and possibly removing a few bits once AK lands support for dims.