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: 48ffc8d	Previous: ff31ed6	Ratio
amdgpu/synchronization/context/device	547.5 ns	565 ns	0.97
amdgpu/synchronization/stream/blocking	237.5 ns	237.5 ns	1
amdgpu/synchronization/stream/nonblocking	315 ns	325 ns	0.97
applications/bitonic_sort	1107943 ns	1128048.75 ns	0.98
applications/convolution	105011.5 ns	103556.5 ns	1.01
applications/floyd_warshall	8877722.5 ns	8865514.25 ns	1.00
applications/histogram	812966.5 ns	814134.25 ns	1.00
applications/prefix_sum	231338.25 ns	231193.25 ns	1.00
array/accumulate/Float32/1d	73703.5 ns	72801.25 ns	1.01
array/accumulate/Float32/dims=1	249106 ns	266143.75 ns	0.94
array/accumulate/Float32/dims=1L	80526.25 ns	80588.5 ns	1.00
array/accumulate/Float32/dims=2	70106 ns	70851 ns	0.99
array/accumulate/Float32/dims=2L	2609661.75 ns	2755514.75 ns	0.95
array/accumulate/Int64/1d	64583.5 ns	77608.5 ns	0.83
array/accumulate/Int64/dims=1	240913.5 ns	240655.75 ns	1.00
array/accumulate/Int64/dims=1L	83788.75 ns	84236.25 ns	0.99
array/accumulate/Int64/dims=2	84278.5 ns	83808.75 ns	1.01
array/accumulate/Int64/dims=2L	2892033.25 ns	2890936.5 ns	1.00
array/broadcast	72176 ns	72886.25 ns	0.99
array/construct	2242.5 ns	2305 ns	0.97
array/copy	37008 ns	37828 ns	0.98
array/copyto!/cpu_to_gpu	110631.5 ns	110896.75 ns	1.00
array/copyto!/gpu_to_cpu	119876.5 ns	119586.75 ns	1.00
array/copyto!/gpu_to_gpu	50543.25 ns	59243.5 ns	0.85
array/iteration/findall/bool	131414.5 ns	134899.5 ns	0.97
array/iteration/findall/int	147374.75 ns	150132.25 ns	0.98
array/iteration/findfirst/bool	180877.75 ns	183332.75 ns	0.99
array/iteration/findfirst/int	158912 ns	161764.75 ns	0.98
array/iteration/findmin/1d	105209 ns	102961.5 ns	1.02
array/iteration/findmin/2d	106061.5 ns	106979.25 ns	0.99
array/iteration/logical	237243.25 ns	240716 ns	0.99
array/iteration/scalar	292269 ns	298229.5 ns	0.98
array/permutedims/2d	70986.25 ns	71661.25 ns	0.99
array/permutedims/3d	58720.75 ns	71008.5 ns	0.83
array/permutedims/4d	72926 ns	73561 ns	0.99
array/random/rand/Float32	44753 ns	45928.25 ns	0.97
array/random/rand/Int64	53575.75 ns	54225.75 ns	0.99
array/random/rand!/Float32	64128.5 ns	64713.25 ns	0.99
array/random/rand!/Int64	72006 ns	72498.5 ns	0.99
array/random/randn/Float32	65231 ns	67480.75 ns	0.97
array/random/randn!/Float32	72058.5 ns	80928.5 ns	0.89
array/reductions/mapreduce/Float32/1d	93294 ns	83573.75 ns	1.12
array/reductions/mapreduce/Float32/dims=1	85683.75 ns	90116.25 ns	0.95
array/reductions/mapreduce/Float32/dims=1L	831186.75 ns	840622 ns	0.99
array/reductions/mapreduce/Float32/dims=2	82071.25 ns	82486 ns	0.99
array/reductions/mapreduce/Float32/dims=2L	139109.5 ns	138304.25 ns	1.01
array/reductions/mapreduce/Int64/1d	93813.75 ns	93766.5 ns	1.00
array/reductions/mapreduce/Int64/dims=1	84331 ns	80898.75 ns	1.04
array/reductions/mapreduce/Int64/dims=1L	844012 ns	844369.75 ns	1.00
array/reductions/mapreduce/Int64/dims=2	82063.5 ns	81251 ns	1.01
array/reductions/mapreduce/Int64/dims=2L	138374.5 ns	140062.25 ns	0.99
array/reductions/reduce/Float32/1d	94423.75 ns	83631.25 ns	1.13
array/reductions/reduce/Float32/dims=1	84496.25 ns	84958.75 ns	0.99
array/reductions/reduce/Float32/dims=1L	833429.25 ns	840804.75 ns	0.99
array/reductions/reduce/Float32/dims=2	81608.5 ns	82443.75 ns	0.99
array/reductions/reduce/Float32/dims=2L	138146.75 ns	138131.75 ns	1.00
array/reductions/reduce/Int64/1d	94356.5 ns	93916.25 ns	1.00
array/reductions/reduce/Int64/dims=1	85461.25 ns	81016.25 ns	1.05
array/reductions/reduce/Int64/dims=1L	839614.5 ns	844852.25 ns	0.99
array/reductions/reduce/Int64/dims=2	80778.75 ns	81323.75 ns	0.99
array/reductions/reduce/Int64/dims=2L	138719.5 ns	138824.5 ns	1.00
array/reverse/1d	36243 ns	44495.75 ns	0.81
array/reverse/1dL	74078.5 ns	56768.25 ns	1.30
array/reverse/1dL_inplace	78011.25 ns	78978.5 ns	0.99
array/reverse/1d_inplace	59238.25 ns	59793.5 ns	0.99
array/reverse/2d	48990.75 ns	49733.25 ns	0.99
array/reverse/2dL	81626 ns	82093.75 ns	0.99
array/reverse/2dL_inplace	88931.25 ns	89941.25 ns	0.99
array/reverse/2d_inplace	61136 ns	58473.25 ns	1.05
array/sorting/1d	332654.75 ns	335384.75 ns	0.99
gemm/tiled	1911867 ns	1893312 ns	1.01
gemm/tiled_unbounded	1939354.75 ns	1921740 ns	1.01
integration/byval/reference	38940 ns	40971 ns	0.95
integration/byval/slices=1	40941 ns	40891 ns	1.00
integration/byval/slices=2	151733 ns	158863 ns	0.96
integration/byval/slices=3	234224 ns	237383 ns	0.99
integration/volumerhs	4895858 ns	4890440 ns	1.00
kernel/indexing	55578.25 ns	56608.25 ns	0.98
kernel/indexing_checked	41070.5 ns	57785.75 ns	0.71
kernel/launch	1317.5 ns	1332.5 ns	0.99
kernel/rand	97279 ns	97788.75 ns	0.99
latency/import	1710497467 ns	1755878492 ns	0.97
latency/precompile	39735404641 ns	39748477386 ns	1.00
latency/ttfp	2320818700 ns	2321457000 ns	1.00
stencil/diffusion3d	1620615.25 ns	1618492.75 ns	1.00
stencil/diffusion3d_checked	1654065.75 ns	1658485.75 ns	1.00

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.

[luraess]

Closing as addressed by JuliaGPU/AcceleratedKernels.jl#117 and #1076

[luraess]

Coming back to the idea of merging this upon suggestion from @maleadt as stopgap and switch over to the fast implementation in a couple of weeks when it get wired up in GPUArrays.

x-ref JuliaGPU/AcceleratedKernels.jl#117 (comment)

[luraess]

UPDATE:

AcceleratedKernels has since added native dims support in JuliaGPU/AcceleratedKernels.jl#117, which is faster (numbers in #1076), but it is not released yet and the plan is to wire it up through GPUArrays. As suggested in JuliaGPU/AcceleratedKernels.jl#117 (comment), this lands as a stopgap using a tag-and-sort approach, and the switch-over will only need to change src/kernels/sorting.jl, since the tests already match #1076.

[luraess]

cscs-ci run

[luraess]

cscs-ci run

[luraess]

cscs-ci run

[luraess]

Unrelated to this PR but after rebasing on main we are still getting Int128 miscompile for 1.12.7 (On CDNA at least) https://gitlab.com/cscs-ci/ci-testing/webhook-ci/mirrors/3266148763123078/6677263153381401/-/jobs/16581746464#L1259

[luraess]

cscs-ci run