From faa01c093871ecf7b5098ca365684e253b637610 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 30 Jun 2026 14:48:24 +0200 Subject: [PATCH] refactor(sort): route algorithm selection through alg= keyword Make sort!/sort/sortperm!/sortperm the documented surface for choosing a sort implementation via alg=. Remove the docs entries for the algorithm-specific helpers (merge_sort!, sample_sort!, etc.) so alg= is the blessed selection API; the helpers remain available as unexported internals. Update the binary-search example to use the alg keyword. --- docs/src/api/binarysearch.md | 2 +- docs/src/api/sort.md | 16 ---------------- src/sort/radix_sort.jl | 4 ++-- src/sort/sort.jl | 13 ++++++------- 4 files changed, 9 insertions(+), 26 deletions(-) diff --git a/docs/src/api/binarysearch.md b/docs/src/api/binarysearch.md index 898da5a0..6f5ef485 100644 --- a/docs/src/api/binarysearch.md +++ b/docs/src/api/binarysearch.md @@ -13,7 +13,7 @@ using Metal # Sorted array v = MtlArray(rand(Float32, 100_000)) -AK.merge_sort!(v) +AK.sort!(v; alg=AK.MergeSort()) # Elements `x` to place within `v` at indices `ix` x = MtlArray(rand(Float32, 10_000)) diff --git a/docs/src/api/sort.md b/docs/src/api/sort.md index 9c846c9d..904e7658 100644 --- a/docs/src/api/sort.md +++ b/docs/src/api/sort.md @@ -13,12 +13,6 @@ AcceleratedKernels.sortperm! AcceleratedKernels.sortperm ``` -Specific implementations that the interfaces above forward to: -- `sample_sort!` - multithreaded CPU sample sort, deferring to Base.sort! on independent slices. -- `merge_sort!` (in-place), `merge_sort` (out-of-place) - sort arbitrary objects with custom comparisons. -- `merge_sort_by_key!`, `merge_sort_by_key` - sort a vector of keys along with a "payload", a vector of corresponding values. -- `merge_sortperm!`, `merge_sortperm`, `merge_sortperm_lowmem!`, `merge_sortperm_lowmem` - compute a sorting index permutation. - Algorithm choice is available on `sort!` / `sort` / `sortperm!` / `sortperm` with `alg=AK.MergeSort()`, `alg=AK.MergeSort(lowmem=true)`, `alg=AK.RadixSort()`, or `alg=AK.SampleSort()`, depending on the backend and operation. @@ -28,16 +22,6 @@ Function signatures: AcceleratedKernels.MergeSort AcceleratedKernels.RadixSort AcceleratedKernels.SampleSort -AcceleratedKernels.sample_sort! -AcceleratedKernels.sample_sortperm! -AcceleratedKernels.merge_sort! -AcceleratedKernels.merge_sort -AcceleratedKernels.merge_sort_by_key! -AcceleratedKernels.merge_sort_by_key -AcceleratedKernels.merge_sortperm! -AcceleratedKernels.merge_sortperm -AcceleratedKernels.merge_sortperm_lowmem! -AcceleratedKernels.merge_sortperm_lowmem ``` Example: diff --git a/src/sort/radix_sort.jl b/src/sort/radix_sort.jl index 263c2d39..268ec8fd 100644 --- a/src/sort/radix_sort.jl +++ b/src/sort/radix_sort.jl @@ -15,7 +15,7 @@ # in parallel using hist[k*B+b] + intra-block rank. # # Supported element types: UInt32/64, Int32/64, Float32/64. -# For custom lt/by → falls back to merge_sort!. +# For custom lt/by -> falls back to merge sort. const _RS_BITS = UInt32(8) const _RS_SIZE = UInt32(256) # 2^_RS_BITS @@ -158,7 +158,7 @@ _rs_supported(::Type{T}) where T = Sort `v` in-place using a GPU LSD radix sort (8-bit, 256 buckets per pass). Supported element types: `UInt32`, `Int32`, `Float32`, `UInt64`, `Int64`, `Float64`. -Falls back to [`merge_sort!`](@ref) for any other type or when `lt`/`by` are provided. +Falls back to merge sort for any other type or when `lt`/`by` are provided. The temporary buffer `temp` (same type and size as `v`) can be passed to avoid allocating internally. `block_size` must be a power of 2. diff --git a/src/sort/sort.jl b/src/sort/sort.jl index dc20d41c..aa54bdfd 100644 --- a/src/sort/sort.jl +++ b/src/sort/sort.jl @@ -66,7 +66,7 @@ arguments are the same as for `Base.sort`. ## CPU CPU settings: use at most `max_tasks` threads to sort the array such that at least `min_elems` -elements are sorted by each thread. A parallel [`sample_sort!`](@ref) is used, processing +elements are sorted by each thread. A parallel sample sort is used, processing independent slices of the array and deferring to `Base.sort!` for the final local sorts. Note that the Base Julia `sort!` is mainly memory-bound, so multithreaded sorting only becomes @@ -76,15 +76,14 @@ faster if it is a more compute-heavy operation to hide memory latency - that inc - Less cache-predictable data movement, e.g. `sortperm`. ## GPU -GPU settings: use `block_size` threads per block to sort the array. A parallel [`merge_sort!`](@ref) -is used. +GPU settings: use `block_size` threads per block to sort the array. A parallel merge sort is used. ## Algorithm choice -By default, `sort!` uses [`sample_sort!`](@ref) on CPU backends and [`merge_sort!`](@ref) on GPU +By default, `sort!` uses sample sort on CPU backends and merge sort on GPU backends. Pass `alg=SampleSort()` for the CPU path, `alg=MergeSort()` for the GPU merge-sort path, or `alg=RadixSort()` to opt into GPU radix sorting. `RadixSort()` supports 32-bit and 64-bit integers and floats; unsupported element types or custom `lt`/`by` settings fall back to -[`merge_sort!`](@ref). +merge sort. For both CPU and GPU backends, the `temp` argument can be used to reuse a temporary buffer of the same size as `v` to store the sorted output. @@ -239,8 +238,8 @@ Save into `ix` the index permutation of `v` such that `v[ix]` is sorted. The `lt [`sort!`](@ref) with custom by-index comparators. ## Algorithm choice -By default, `sortperm!` uses [`sample_sortperm!`](@ref) on CPU backends and [`merge_sortperm!`](@ref) -on GPU backends. Pass `alg=MergeSort(lowmem=true)` to use the lower-memory GPU permutation path. +By default, `sortperm!` uses sample sort on CPU backends and merge sort on GPU +backends. Pass `alg=MergeSort(lowmem=true)` to use the lower-memory GPU permutation path. `RadixSort()` does not provide a permutation path. """ function sortperm!(