Repository navigation
Pick bitonic and radix sorts on GPUs where they are faster - #184
Merged
Merged
Conversation
maleadt
added this pull request to stack #178
October 25, 2026 05:33
maleadt
force-pushed
the
tb/sort-tuning
branch
from
October 25, 2026 06:37
38e12ff to
e31b8e0
Compare
A BitonicSort whose tile of `block_size * items_per_thread` elements exceeds the device's local memory failed while compiling the kernel, with a different error on each backend (ptxas exiting with code 255 on CUDA, a compiler crash on Metal, running out of memory on oneAPI), e.g. `BitonicSort(block_size=1024, items_per_thread=16)` for Int64 or the default settings for 64-byte elements. Reject such tiles in `check_sort` with an ArgumentError, against the same 32 KiB budget as RadixSort, and make `Auto` fall back to MergeSort for element types too large for the tuned tile.
MergeSort holds a tile of `2 * block_size` keys, and the values moved along with them, in local memory, without checking its size: sorting elements of 256 bytes with the default block size failed to compile on CUDA (ptxas exiting with code 255), so `Auto` could not sort them. Resolution now gets the element types of the tiles from the setups (the `by` keys and the values of `sort!`, `sortperm!` and `sort_by_key!`), shrinks the default block size to a power of two whose tiles fit the 32 KiB budget, and rejects explicit block sizes that do not fit.
With 256 threads and chunks of 32, RadixSort was far from the best setting on several GPUs, and on an Arc A770 it never beat MergeSort. Measured over a sweep of block size, items per thread and chunk size, `(512, 2, 64)` takes this much less time than `(256, 2, 32)` to sort 1M and 16M random Float32 and Int64 elements: RTX 5080 5-8% RTX 6000 Ada 9-15% M1 6-24% Arc A770 39-44% Radeon gfx1036 19-52% and is within 12% of the fastest setting on each of them. Re-measured on the RTX 5080 on this branch, Float32 and UInt64 at 1M and 16M elements take 8-13% less time (Float32 1M: 217 -> 199 us, 16M: 2.81 -> 2.48 ms; UInt64 1M: 432 -> 385 us, 16M: 6.05 -> 5.50 ms). PoCL prefers other settings at some sizes, but is no slower with these at the lengths where RadixSort beats MergeSort.
`Auto` applied `bitonic_max_len` to whole arrays and slices alike and never chose RadixSort along `dims`. The crossovers differ: BitonicSort keeps beating MergeSort along `dims` for slices several times longer than the whole arrays it wins on, and RadixSort beats MergeSort for long slices of 32-bit elements on most devices, but only where the array is large enough for its passes to pay off at all. `SortTuning` gets `bitonic_max_slice_len` and `radix_min_slice_len`, which apply to the slice length along `dims`; RadixSort also needs `radix_min_len` elements in all. A single slice (a vector, or a matrix with one non-singleton dimension) is sorted like a whole array, which lets RadixSort sort it. The new thresholds default to never, so no tuning changes its choices yet.
Add `sort_tuning` methods to the CUDA, Metal, oneAPI and OpenCL
extensions, from timings of MergeSort against BitonicSort and RadixSort
on random keys of 32- and 64-bit integers and floats, as whole arrays
and along `dims=1` and `dims=2` for totals from 2^12 to 2^24 elements:
bitonic radix 32/64-bit bitonic radix 32/64-bit
whole whole slices slices
CUDA 512 2^22 / 2^23 256 2^19 / never
Metal 512 2^17 / 2^18 1024 2^15 / never
oneAPI 128 2^21 / never 128 2^15 / never
OpenCL (CPU) 128 2^17 / 2^19 1024 512 / 2^18
measured on an RTX 5080 and an RTX 6000 Ada (taking the lower where
they disagree), an M1, an Arc A770 and PoCL on a Ryzen 9 9950X. Each
method states the speedups it was chosen for. OpenCL GPU devices have
not been measured and keep MergeSort.
The docs said GPUs sort with MergeSort until a backend's tuning says otherwise. CUDA, Metal, oneAPI and PoCL now choose BitonicSort for short arrays and slices and RadixSort for long ones, so say that, and that devices of other backends keep MergeSort.
maleadt
force-pushed
the
tb/sort-tuning
branch
from
October 25, 2026 06:47
e31b8e0 to
e2a5063
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Autosorting has thresholds for pickingBitonicSort(short inputs) andRadixSort(long ones, of the element types it supports) overMergeSort, but no backend set them, so every GPU sort was a merge sort, which ranks elements by binary search and is the slowest of the three where the others apply. This sets them from measurements, per backend:Autosort!on an RTX 5080MergeSort)Int32sFloat32sFloat64sInt32s alongdims=1Float32s alongdims=1Sorting along
dimsgets thresholds of its own (bitonic_max_slice_len,radix_min_slice_len), on the slices' length: bitonic sort keeps winning on many short slices up to lengths where it already loses on whole arrays, and radix sort wins on long slices, whereAutonever used it. A single slice sorts like a whole array. The extensions set them for CUDA (measured on an RTX 5080 and an RTX 6000 Ada), Metal (M1), oneAPI (Arc A770) and OpenCL's CPU devices (PoCL); each comment states the crossovers and speedups. Thresholds are conservative where devices disagree: 64-bit radix sort stays off on the A770, where it is sometimes slower, and OpenCL GPU devices and AMD GPUs (only an integrated one was measured) keep merge sort.sortperm!andsort_by_key!keep merge sort, the only GPU algorithm that supports them. Unstable bitonic sort is only picked where ties can't be observed, as before.RadixSort's default settings change from 256 threads with chunks of 32 to 512 threads with chunks of 64, which takes 8-13% less time on the RTX 5080 and up to 44% on the A770 (where radix sort never beat merge sort with the old setting).Two bugs this turned up, fixed with tests:
BitonicSortandMergeSortdidn't check that their tiles fit in local memory, so large settings or element types (e.g.Autoon 256-byte elements) failed in the backend's compiler (ptxas exit codes, Metal compiler crashes) instead of with anArgumentError;Autonow shrinks merge sort's block size to fit, and falls back from bitonic sort.Known trade-off, stated in the extension comments: in small arrays (fewer than 2^16 elements) bitonic sort packs short slices into fewer blocks than merge sort, so some of those take a few microseconds longer (up to 41% on CUDA and 69% on the A770); fixing that is a kernel change for later.
Tested on CPU threads, PoCL (also with AK's kernels on the host backend), CUDA (RTX 5080 on Julia 1.12 and 1.10, RTX 6000 Ada), oneAPI and Metal.