Skip to content

Add dims support to sort, sort!, sortperm and sortperm! - #117

Open
shreyas-omkar wants to merge 1 commit into
JuliaGPU:mainfrom
shreyas-omkar:sh/sort-dims
Open

Add dims support to sort, sort!, sortperm and sortperm!#117
shreyas-omkar wants to merge 1 commit into
JuliaGPU:mainfrom
shreyas-omkar:sh/sort-dims

Conversation

@shreyas-omkar

@shreyas-omkar shreyas-omkar commented Aug 24, 2026

Copy link
Copy Markdown
Member

Adds a dims keyword to sort/sort!/sortperm/sortperm! for portable per-slice N-D sorting across CPU and every GPU backend, using the same tag-and-lexicographic-sort strategy as AMDGPU.jl #1033

Sort or permute each 1D slice along an integer `dims` independently,
matching Base. The default `dims=:` keeps the existing flat behaviour.

There is no batched sort kernel, so each element is tagged with its slice
and the whole array is sorted once by (slice, value), then scattered back
into place. sortperm carries the original index as the payload and uses it
to break ties, keeping the permutation stable. This reuses the backend's
tuned sort and runs unchanged on CPU and every GPU backend.
@shreyas-omkar shreyas-omkar changed the title Add dims support to sort and sort! Add dims support to sort, sort!, sortperm and sortperm! Aug 24, 2026
@shreyas-omkar
shreyas-omkar marked this pull request as draft August 24, 2026 09:49
@shreyas-omkar

Copy link
Copy Markdown
Member Author

@christiangnrd Please take a look. :)

@shreyas-omkar
shreyas-omkar marked this pull request as ready for review August 26, 2026 06:51
@shreyas-omkar

Copy link
Copy Markdown
Member Author

NVIDIA RTX 5080 (CUDA 13.3.0)

1D sort!(v) Float32

n AK ms CUDA ms AK/CUDA
262,144 0.154 0.451 2.93x
1,048,576 0.393 1.631 4.15x
4,194,304 1.660 7.952 4.79x
16,777,216 9.287 41.792 4.50x
67,108,864 42.302 234.772 5.55x

1D sort!(v) Int32

n AK ms CUDA ms AK/CUDA
262,144 0.153 0.485 3.17x
1,048,576 0.351 1.579 4.50x
4,194,304 1.456 7.775 5.34x
16,777,216 8.613 41.032 4.76x
67,108,864 39.077 231.374 5.92x

1D sortperm(v) Float32

n AK ms CUDA ms AK/CUDA
262,144 0.227 1.130 4.97x
1,048,576 0.476 4.535 9.52x
4,194,304 2.646 49.288 18.63x
16,777,216 11.958 234.213 19.59x

N-D sort(A; dims) Float32, 1,048,576 elements

size dims AK ms CUDA ms AK/CUDA
1024 x 1024 1 0.649 0.349 0.54x
1024 x 1024 2 0.671 0.469 0.70x
256 x 4096 1 0.641 0.435 0.68x
256 x 4096 2 0.682 0.706 1.04x
4096 x 256 1 0.659 0.494 0.75x
4096 x 256 2 0.663 0.543 0.82x

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant