Repository navigation
Share one kernel pair and driver among the GPU merge sorts - #180
Merged
Merged
Conversation
The merge sorts had two copies of the tile kernel and the merge-pass kernel (keys only, and keys with values) and three of the host ping-pong driver (sort!, sort_by_key! with the normal sortperm!, and the lowmem sortperm!), differing only in whether a values array moves along with the keys. The kernels now take `values` as an array or `nothing`, with `getvalue`/`setvalue!` doing nothing for `nothing`, so the keys-only specialisation compiles to the same code as before. One driver, `merge_sort_slices!`, serves every entry point; lowmem sortperm is the keys-only sort of the indices under a comparator that reads `v`. The three files are merged into merge_sort.jl. Performance is unchanged within noise. Min of up to 500 runs over two processes each, before and after, RTX 5080: sort! Float32 1e7 4778 -> 4783 us, Int64 1e7 5702 -> 5685 us, 16-byte struct 1e7 9148 -> 9148 us, sort_by_key! 1e7 6822 -> 6819 us, sortperm! 1e7 7009 -> 7045 us, lowmem 10239 -> 10251 us, sort! dims=2 on 1000x10000 3387 -> 3377 us; 1e5 cases within +-1%, 1e3 within +-3% (host overhead dominates there, and runs of the same code vary as much). First-call compile times are unchanged. The AMDGPU iGPU and PoCL show the same.
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.
The GPU merge sort was written out more than once: two copies of the tile kernel and the merge-pass kernel (keys only, and keys with values), and three of the host driver that runs the tile pass, the merge passes and the copy-back after an odd number of passes (
sort!;sort_by_key!andsortperm!; thelowmemsortperm!). They differed only in whether a values array moves along with the keys, so every fix or optimization had to be made two or three times.Now there is one kernel pair and one driver. The kernels take
valuesas an array ornothing; loads and stores ofnothingvalues do nothing and the local-memory buffer for them is not allocated, which is decided at compile time, so the keys-only kernels compile to what they did before. Thelowmemsortperm!is the keys-only sort of the index array under a comparator that readsv, as before. The algorithm, scratch sizes and workspace layout are unchanged. The three files become one, about 120 lines shorter.Performance is unchanged within noise (min of up to 500 runs, two processes each, RTX 5080):
sort!of 10^7Float32s 4778 → 4783 µs,Int64s 5702 → 5685 µs, 16-byte structs 9148 → 9148 µs,sort_by_key!6822 → 6819 µs,sortperm!7009 → 7045 µs, lowmem 10239 → 10251 µs,sort!alongdims=2of 1000×10000 3387 → 3377 µs; smaller sizes within the run-to-run variation. The AMD iGPU and PoCL show the same, and first-call compile times don't change.Tested on CPU threads, PoCL (also with AK's kernels on the host backend), CUDA (Julia 1.12 and 1.10) and AMDGPU.