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: 4c32359	Previous: 16f5974	Ratio
amdgpu/synchronization/context/device	555 ns	755 ns	0.74
amdgpu/synchronization/stream/blocking	227.5 ns	317.5 ns	0.72
amdgpu/synchronization/stream/nonblocking	312.5 ns	432.25 ns	0.72
array/accumulate/Float32/1d	71426 ns	97092.25 ns	0.74
array/accumulate/Float32/dims=1	284614 ns	299718.5 ns	0.95
array/accumulate/Float32/dims=1L	79863.75 ns	121446.75 ns	0.66
array/accumulate/Float32/dims=2	71576 ns	130896.75 ns	0.55
array/accumulate/Float32/dims=2L	2749303.5 ns	2806321.5 ns	0.98
array/accumulate/Int64/1d	78591.25 ns	113217 ns	0.69
array/accumulate/Int64/dims=1	245510.75 ns	269376.25 ns	0.91
array/accumulate/Int64/dims=1L	84301.25 ns	144501.75 ns	0.58
array/accumulate/Int64/dims=2	84711.25 ns	127879.25 ns	0.66
array/accumulate/Int64/dims=2L	2898550.75 ns	2948585.75 ns	0.98
array/broadcast	67691 ns	64054.75 ns	1.06
array/construct	2305.25 ns	3220 ns	0.72
array/copy	37515.5 ns	39815 ns	0.94
array/copyto!/cpu_to_gpu	111291.5 ns	94539.5 ns	1.18
array/copyto!/gpu_to_cpu	110874 ns	95039.5 ns	1.17
array/copyto!/gpu_to_gpu	59465.75 ns	42115 ns	1.41
array/iteration/findall/bool	138209.5 ns	198036.5 ns	0.70
array/iteration/findall/int	151994.5 ns	207411.5 ns	0.73
array/iteration/findfirst/bool	145454.75 ns	167311.75 ns	0.87
array/iteration/findfirst/int	145934.5 ns	155156.75 ns	0.94
array/iteration/findmin/1d	111496.5 ns	179444.25 ns	0.62
array/iteration/findmin/2d	109474 ns	161089.25 ns	0.68
array/iteration/logical	249033.5 ns	317088.25 ns	0.79
array/iteration/scalar	304476.75 ns	343653.25 ns	0.89
array/permutedims/2d	72483.75 ns	82694.75 ns	0.88
array/permutedims/3d	71778.5 ns	82232.25 ns	0.87
array/permutedims/4d	74183.5 ns	84369.75 ns	0.88
array/random/rand/Float32	45743.25 ns	53677.5 ns	0.85
array/random/rand/Int64	53815.75 ns	66274.75 ns	0.81
array/random/rand!/Float32	66113.5 ns	49389.75 ns	1.34
array/random/rand!/Int64	66646 ns	59654.75 ns	1.12
array/random/randn/Float32	77511 ns	93112.25 ns	0.83
array/random/randn!/Float32	81508.75 ns	69269.75 ns	1.18
array/reductions/mapreduce/Float32/1d	98861.5 ns	141249.25 ns	0.70
array/reductions/mapreduce/Float32/dims=1	85443.75 ns	105462 ns	0.81
array/reductions/mapreduce/Float32/dims=1L	834564.5 ns	862533.5 ns	0.97
array/reductions/mapreduce/Float32/dims=2	85686.25 ns	108059.5 ns	0.79
array/reductions/mapreduce/Float32/dims=2L	144304.5 ns	155446.75 ns	0.93
array/reductions/mapreduce/Int64/1d	98456.5 ns	140932 ns	0.70
array/reductions/mapreduce/Int64/dims=1	84351 ns	104764.5 ns	0.81
array/reductions/mapreduce/Int64/dims=1L	831494.25 ns	860696 ns	0.97
array/reductions/mapreduce/Int64/dims=2	85366.25 ns	107757 ns	0.79
array/reductions/mapreduce/Int64/dims=2L	145964.5 ns	163496.75 ns	0.89
array/reductions/reduce/Float32/1d	98741.5 ns	145184.25 ns	0.68
array/reductions/reduce/Float32/dims=1	83461.25 ns	105042 ns	0.79
array/reductions/reduce/Float32/dims=1L	836079.25 ns	856898.25 ns	0.98
array/reductions/reduce/Float32/dims=2	85388.75 ns	106659.5 ns	0.80
array/reductions/reduce/Float32/dims=2L	144847 ns	161129.25 ns	0.90
array/reductions/reduce/Int64/1d	98649 ns	141302 ns	0.70
array/reductions/reduce/Int64/dims=1	83266.25 ns	105144.5 ns	0.79
array/reductions/reduce/Int64/dims=1L	834216.75 ns	851721 ns	0.98
array/reductions/reduce/Int64/dims=2	84761.25 ns	107839.5 ns	0.79
array/reductions/reduce/Int64/dims=2L	145289.5 ns	162989.25 ns	0.89
array/reverse/1d	44310.75 ns	53687.25 ns	0.83
array/reverse/1dL	73771 ns	78502.25 ns	0.94
array/reverse/1dL_inplace	61953.5 ns	62952.25 ns	0.98
array/reverse/1d_inplace	40115.5 ns	45130 ns	0.89
array/reverse/2d	51193 ns	55499.75 ns	0.92
array/reverse/2dL	82478.75 ns	111822 ns	0.74
array/reverse/2dL_inplace	92616.25 ns	74057 ns	1.25
array/reverse/2d_inplace	56725.75 ns	49312.25 ns	1.15
array/sorting/1d	334804.75 ns	359238.75 ns	0.93
integration/byval/reference	39440 ns	41860 ns	0.94
integration/byval/slices=1	40430 ns	42700 ns	0.95
integration/byval/slices=2	146812 ns	133939 ns	1.10
integration/byval/slices=3	239714 ns	239880 ns	1.00
integration/volumerhs	5006840 ns	4916200 ns	1.02
kernel/indexing	58333.25 ns	44594.75 ns	1.31
kernel/indexing_checked	59811 ns	47782.25 ns	1.25
kernel/launch	1380 ns	1785 ns	0.77
kernel/rand	99936.25 ns	110472 ns	0.90
latency/import	1715435625 ns	2242598207 ns	0.76
latency/precompile	39849130654 ns	52102215931 ns	0.76
latency/ttfp	2341714089 ns	3070234127 ns	0.76

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.