Optimize Radix Sort.#97
Draft
shreyas-omkar wants to merge 2 commits into
Draft
Conversation
_accumulate_block!'s inclusive path re-read global `v` from inside a single-thread branch after the scan barriers — a divergent global read that faults or hangs on backends other than CUDA. Capture everything the inclusive shift needs at load time (each thread's own last element, plus the array's final element in a shared slot) so no divergent global read remains. Fixes accumulate!/cumsum and the GPU radix sort (which relies on accumulate!) on non-CUDA backends. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Rework the opt-in GPU LSD radix sort to be fast and portable across backends, with no dependency on sub-group / warp intrinsics: - Shared-atomic histogram (O(1)/element) where the backend reports shared-memory atomics support; portable per-bucket scan otherwise. - Chunked stable scatter rank: O(32)/element instead of O(block_size), via per-chunk sub-histograms plus a cross-chunk exclusive prefix; gated on atomics support, broadcast-scan fallback otherwise. - block_size default 256 (the scatter rank cost scales with block_size, so this is the portable sweet spot). - Fused min/max range in one reduction; skip passes where every element shares a byte; reuse the accumulate! scratch across passes; a single final sync instead of per-kernel syncs. Supports UInt32/64, Int32/64, Float32/64; custom lt/by fall back to merge_sort!. Correct across all types, sizes and directions. On backends without sub-group ballot, radix now beats merge_sort for every supported type (previously slower on several). Adds Atomix as a dependency. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
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
1 participant
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.
CUDA (RTX 5080), 4M elements
AMD (RX 9060 XT), 4M elements