From c2a4c1ffec6969cd598b20b2245d5e337a9f8691 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Wed, 8 Jul 2026 04:29:23 +0530 Subject: [PATCH 1/5] fix(accumulate): avoid divergent global read on non-CUDA backends MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _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 --- src/accumulate/accumulate_1d_gpu.jl | 37 +++++++++++++++++------------ 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/src/accumulate/accumulate_1d_gpu.jl b/src/accumulate/accumulate_1d_gpu.jl index 31fee7e..f4153d9 100644 --- a/src/accumulate/accumulate_1d_gpu.jl +++ b/src/accumulate/accumulate_1d_gpu.jl @@ -22,6 +22,8 @@ end len = length(v) @uniform block_size = @groupsize()[1] temp = @localmem eltype(v) (0x2 * block_size + conflict_free_offset(0x2 * block_size),) + s_lastval = @localmem eltype(v) (1,) # holds v[len]; captured at load time so + # the inclusive shift needn't re-read v # NOTE: for many index calculations in this library, computation using zero-indexing leads to # fewer operations (also code is transpiled to CUDA / ROCm / oneAPI / Metal code which do zero @@ -42,16 +44,21 @@ end bank_offset_a = conflict_free_offset(ai) bank_offset_b = conflict_free_offset(bi) - if block_offset + ai < len - temp[ai + bank_offset_a + 0x1] = v[block_offset + ai + 0x1] - else - temp[ai + bank_offset_a + 0x1] = neutral + # Load both elements, keeping them in registers. Re-reading global `v` later + # from inside the single-thread inclusive-shift branch is a divergent global + # read that faults or hangs on some backends, so everything the shift needs is + # captured here: each thread's own last element (my_bi_val) and, in a shared + # slot, the array's final element v[len] (written by whichever thread owns it). + my_ai_val = (block_offset + ai < len) ? v[block_offset + ai + 0x1] : neutral + temp[ai + bank_offset_a + 0x1] = my_ai_val + my_bi_val = (block_offset + bi < len) ? v[block_offset + bi + 0x1] : neutral + temp[bi + bank_offset_b + 0x1] = my_bi_val + + if block_offset + ai == len - 0x1 + s_lastval[0x1] = my_ai_val end - - if block_offset + bi < len - temp[bi + bank_offset_b + 0x1] = v[block_offset + bi + 0x1] - else - temp[bi + bank_offset_b + 0x1] = neutral + if block_offset + bi == len - 0x1 + s_lastval[0x1] = my_bi_val end # Build block reduction down @@ -115,13 +122,13 @@ end end temp[bi - 0x1 + conflict_free_offset(bi - 0x1) + 0x1] = t2 - # ...and accumulate the last value too + # ...and accumulate the last value too, without re-reading global v. A + # full non-last block's last element is this thread's own my_bi_val + # (== v[(iblock+1)*block_size*2]); the final block instead needs the + # array's last element v[len], captured in s_lastval at load time. if bi == 0x2 * block_size - 0x1 - if iblock < num_blocks - 0x1 - temp[bi + bank_offset_b + 0x1] = op(t2, v[(iblock + 0x1) * block_size * 0x2]) - else - temp[bi + bank_offset_b + 0x1] = op(t2, v[len]) - end + last_val = iblock < num_blocks - 0x1 ? my_bi_val : s_lastval[0x1] + temp[bi + bank_offset_b + 0x1] = op(t2, last_val) end end From baca23d100c3aa87bc6bba262a5fae31b0473b70 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Wed, 8 Jul 2026 04:29:23 +0530 Subject: [PATCH 2/5] perf(sort): optimize GPU radix sort (portable, no sub-group intrinsics) 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 --- Project.toml | 2 + src/sort/radix_sort.jl | 310 ++++++++++++++++++++++++++++++----------- src/sort/sort.jl | 39 ++++-- 3 files changed, 257 insertions(+), 94 deletions(-) diff --git a/Project.toml b/Project.toml index b5d6038..4471644 100644 --- a/Project.toml +++ b/Project.toml @@ -5,6 +5,7 @@ version = "0.4.3" [deps] ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197" +Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" @@ -18,6 +19,7 @@ AcceleratedKernelsoneAPIExt = "oneAPI" [compat] ArgCheck = "2" +Atomix = "0.1, 1" GPUArraysCore = "0.2.0" KernelAbstractions = "0.9.34, 0.10" Markdown = "1" diff --git a/src/sort/radix_sort.jl b/src/sort/radix_sort.jl index 263c2d3..a56a759 100644 --- a/src/sort/radix_sort.jl +++ b/src/sort/radix_sort.jl @@ -1,24 +1,34 @@ # LSD radix sort (8-bit, 256 buckets). Stable, GPU-only. # -# Algorithm per pass (P = 4 passes for 32-bit, 8 for 64-bit): -# 1. _radix_hist! — each block counts its contiguous chunk into -# hist[digit * num_blocks + block]; all 256 threads run -# in parallel, each owning one bucket and counting via a -# broadcast-read loop over precomputed shared-mem digits. -# 2. accumulate! — exclusive prefix-sum over the flat hist array gives -# scatter offsets: hist[k*B+b] = # elements with -# digit < k (all blocks) + # elements with digit k in -# blocks 0..b-1 -# 3. _radix_scatter! — each thread stores its digit into shared mem; then -# counts preceding same-digit entries via a broadcast -# read loop (stable, no atomics); all threads scatter -# in parallel using hist[k*B+b] + intra-block rank. +# Backend-portable: no dependency on sub-group / warp intrinsics. Where a +# backend reports shared-memory atomics support, faster atomic-based histogram +# and scatter kernels are used; otherwise a scan-based path runs anywhere. +# +# Algorithm (per non-trivial pass over byte k): +# 1. histogram — count, per block, how many elements have each byte-digit. +# Layout: hist[k * B + b] = count of digit k in block b; B = num_blocks. +# Shared-atomic where supported (_radix_hist_atomic!), portable +# per-bucket scan otherwise (_radix_hist!). +# 2. accumulate! — exclusive prefix sum over hist → global per-(digit, block) offsets. +# 3. scatter — stable scatter to the offsets. Chunked O(32)-rank where shared-memory +# atomics are available (_radix_scatter_chunked!), O(block_size)-rank +# broadcast scan otherwise (_radix_scatter!). +# +# Key optimizations: +# • Fused min/max range (_rs_key_range): one reduction over the sort keys instead of +# separate minimum + maximum. +# • Skip-pass via min/max keys: if min and max share the whole byte-suffix from byte k +# up, every element does too, so the whole pass is skipped (e.g. UInt32 in [0, 255] +# sorts in a single pass). # # Supported element types: UInt32/64, Int32/64, Float32/64. -# For custom lt/by → falls back to merge_sort!. +# Custom lt/by → falls back to merge_sort!. + +import Atomix -const _RS_BITS = UInt32(8) -const _RS_SIZE = UInt32(256) # 2^_RS_BITS +const _RS_BITS = UInt32(8) +const _RS_SIZE = UInt32(256) # 2^_RS_BITS +const _RS_CHUNK = 32 # chunked-scatter chunk width (smaller = cheaper rank; 32 best measured) # ─── Make any supported scalar type sortable as an unsigned integer ─────────── @@ -44,58 +54,83 @@ end ((rev ? ~_to_sort_key(x) : _to_sort_key(x)) >> shift) & (_RS_SIZE - 0x1) -# ─── Phase 1: block-level histogram (parallel, one thread per bucket) ───────── -# Elements are staged in s_elem then each element's digit is precomputed once -# into s_digit (avoids calling _to_sort_key repeatedly — for Float64 it involves -# a 64-bit multiply). Each thread then tallies only its own bucket by scanning -# s_digit: all threads read the same s_digit[jj] each iteration → hardware -# broadcast with zero bank conflicts, no atomics, one global write per thread. +# ─── Phase 1: per-pass histogram — generic scan (all backends) ─────────────── +# hist[k * num_blocks + b] = count of elements with digit k in block b. +# +# Each thread loads its element's digit into s_digit, then scans s_digit for each +# of its assigned buckets (bucket t, t+NI, t+2*NI, …). Uses no shared-memory +# atomics, so it is the portable fallback that runs on every backend. @kernel inbounds=true cpu=false unsafe_indices=true function _radix_hist!( hist, @Const(v), shift::UInt32, rev::Bool, ) - @uniform N = @groupsize()[1] @uniform NI = Int(@groupsize()[1]) - s_elem = @localmem eltype(v) (N,) - s_digit = @localmem UInt32 (N,) + s_digit = @localmem UInt32 (NI,) iblock = Int(@index(Group, Linear)) - 1 ithread = Int(@index(Local, Linear)) - 1 len = Int(length(v)) - num_blocks = Int(length(hist)) ÷ 256 + num_blocks = Int(length(hist)) ÷ Int(_RS_SIZE) + # 0xffffffff doesn't match any valid bucket (0–255); OOB elements are neutral. i = iblock * NI + ithread - if i < len - s_elem[ithread + 1] = v[i + 1] - end + s_digit[ithread + 1] = UInt32(i < len ? _rs_digit(v[i + 1], shift, rev) : 0xffffffff) @synchronize() - block_len = min(NI, len - iblock * NI) - if ithread < block_len - s_digit[ithread + 1] = _rs_digit(s_elem[ithread + 1], shift, rev) + # Thread t handles buckets t, t+NI, t+2*NI, … (covers all 256 when NI ≤ 256). + bucket = ithread + while bucket < Int(_RS_SIZE) + cnt = UInt32(0) + for jj in 1:NI + cnt += UInt32(s_digit[jj] == UInt32(bucket)) + end + hist[bucket * num_blocks + iblock + 1] = cnt + bucket += NI end - @synchronize() +end + + +# ─── Phase 1b: histogram — shared-atomic (backends with shared-mem atomics) ── +# Same per-block 256-bin output, but O(1)/element via a shared-memory atomic +# increment instead of the O(256)/element per-bucket scan above. Selected by +# the driver only where the backend reports atomics support. + +@kernel inbounds=true cpu=false unsafe_indices=true function _radix_hist_atomic!( + hist, @Const(v), shift::UInt32, rev::Bool, +) + @uniform NI = Int(@groupsize()[1]) + s_hist = @localmem UInt32 (Int(_RS_SIZE),) + + iblock = Int(@index(Group, Linear)) - 1 + ithread = Int(@index(Local, Linear)) - 1 + len = Int(length(v)) + num_blocks = Int(length(hist)) ÷ Int(_RS_SIZE) j = ithread - while j < 256 - my_bucket = UInt32(j) - count = UInt32(0) - for jj in 1:block_len - count += UInt32(s_digit[jj] == my_bucket) - end - hist[j * num_blocks + iblock + 1] = count + while j < Int(_RS_SIZE) + s_hist[j + 1] = UInt32(0) j += NI end + @synchronize() + + i = iblock * NI + ithread + if i < len + d = Int(_rs_digit(v[i + 1], shift, rev)) + Atomix.@atomic s_hist[d + 1] += UInt32(1) + end + @synchronize() + + bucket = ithread + while bucket < Int(_RS_SIZE) + hist[bucket * num_blocks + iblock + 1] = s_hist[bucket + 1] + bucket += NI + end end -# ─── Phase 3: stable scatter (parallel rank via broadcast read) ─────────────── -# Each thread stores its digit in s_digit and reads all 256 block offsets from -# the prefix-summed hist into s_gbase — both before the single @synchronize. -# After the barrier each thread counts preceding elements with the same digit via -# a broadcast-read loop (all threads read the same s_digit[jj] each iteration → -# hardware broadcast, zero bank conflicts). Scatter is then one shared-mem read -# (s_gbase) plus one global write (v_out). +# ─── Phase 3: scatter — broadcast-read rank (O(N) per thread) ─────────────── +# `hist` is already the exclusive-prefix-summed per-block offsets; +# `hist[k * num_blocks + b]` = global start for bucket k, block b (1-indexed). @kernel inbounds=true cpu=false unsafe_indices=true function _radix_scatter!( v_out, @Const(v_in), @Const(hist), shift::UInt32, rev::Bool, @@ -111,8 +146,6 @@ end len = Int(length(v_in)) num_blocks = Int(length(hist)) ÷ 256 - # Coissue the v_in load and the strided hist reads; both hit global memory - # simultaneously while the GPU pipelines them. i = iblock * NI + ithread if i < len s_elem[ithread + 1] = v_in[i + 1] @@ -139,6 +172,80 @@ end end +# ─── Phase 3b: scatter — chunked stable rank (O(chunk) per thread) ─────────── +# The broadcast scatter's rank is O(block_size)/element. Split the block into +# 32-wide chunks: per-chunk digit counts (built with shared-memory atomics — +# order doesn't matter for counts) give each chunk's stable base via a cross-chunk +# exclusive prefix, and each element only scans its own chunk (≤32) for the +# intra-chunk part. Rank = cross-chunk-base[digit] + intra-chunk same-digit +# count — still fully stable, but O(32) instead of O(block_size). + +@kernel inbounds=true cpu=false unsafe_indices=true function _radix_scatter_chunked!( + v_out, @Const(v_in), @Const(hist), shift::UInt32, rev::Bool, +) + @uniform N = @groupsize()[1] + @uniform NI = Int(@groupsize()[1]) + @uniform NCH = Int(@groupsize()[1]) ÷ _RS_CHUNK # number of chunks + s_elem = @localmem eltype(v_in) (N,) + s_digit = @localmem UInt32 (N,) + s_gbase = @localmem UInt32 (256,) + s_chist = @localmem UInt32 (256 * NCH,) # per-chunk digit counts → bases + + iblock = Int(@index(Group, Linear)) - 1 + ithread = Int(@index(Local, Linear)) - 1 + len = Int(length(v_in)) + num_blocks = Int(length(hist)) ÷ 256 + + i = iblock * NI + ithread + if i < len + s_elem[ithread + 1] = v_in[i + 1] + end + j = ithread + while j < 256 + s_gbase[j + 1] = hist[j * num_blocks + iblock + 1] + j += NI + end + j = ithread + while j < 256 * NCH + s_chist[j + 1] = UInt32(0) + j += NI + end + @synchronize() + + mychunk = ithread ÷ _RS_CHUNK + my_digit = UInt32(i < len ? _rs_digit(s_elem[ithread + 1], shift, rev) : 0) + s_digit[ithread + 1] = my_digit + if i < len + Atomix.@atomic s_chist[mychunk * 256 + Int(my_digit) + 1] += UInt32(1) + end + @synchronize() + + # cross-chunk exclusive prefix per digit (thread d owns digit d, d+NI, …) + d = ithread + while d < 256 + acc = UInt32(0) + for c in 0:NCH-1 + cnt = s_chist[c * 256 + d + 1] + s_chist[c * 256 + d + 1] = acc + acc += cnt + end + d += NI + end + @synchronize() + + if i < len + chunk_start = mychunk * _RS_CHUNK + cnt = UInt32(0) + for jj in chunk_start:(ithread - 1) + cnt += UInt32(s_digit[jj + 1] == my_digit) + end + rank = s_chist[mychunk * 256 + Int(my_digit) + 1] + cnt + gpos = Int(s_gbase[my_digit + 1]) + Int(rank) + v_out[gpos + 1] = s_elem[ithread + 1] + end +end + + # ─── Implementation ────────────────────────────────────────────────────────── _rs_supported(::Type{T}) where T = @@ -146,22 +253,38 @@ _rs_supported(::Type{T}) where T = T === UInt64 || T === Int64 || T === Float64 -""" - _radix_sort!( - v::AbstractArray, backend::Backend=get_backend(v); - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, +# Return (min_sort_key, max_sort_key) as UInt64, accounting for descending order. +# Used to detect passes where all elements share the same byte-digit (trivial pass). +# +# Single fused reduction: map each element to its (order-preserving) unsigned sort +# key and reduce to (min_key, max_key) in one pass, instead of two separate full +# reductions over the array (minimum + maximum). ~halves the range-finding cost. +function _rs_key_range(v::AbstractArray{T}, descending::Bool) where T + K = typeof(_to_sort_key(zero(T))) # UInt32 for 32-bit types, UInt64 for 64-bit + ident = (typemax(K), typemin(K)) # identity for (min, max) over keys + min_k, max_k = mapreduce( + x -> (k = _to_sort_key(x); (k, k)), + (a, b) -> (min(a[1], b[1]), max(a[2], b[2])), + v; + init=ident, + neutral=ident, ) + if descending + # rev=true flips all bits: digit = (~key >> shift) & 0xFF + # Maximum value → minimum key after negation; swap accordingly. + UInt64(~max_k), UInt64(~min_k) + else + UInt64(min_k), UInt64(max_k) + end +end -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. +""" + _radix_sort!(v, backend; lt, by, rev, order, block_size, temp) -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. +In-place GPU LSD radix sort (8-bit, 256 buckets per pass). Supported types: +`UInt32`, `Int32`, `Float32`, `UInt64`, `Int64`, `Float64`. Falls back to +[`merge_sort!`](@ref) for any other type or when `lt`/`by` are non-default. """ function _radix_sort!( v::AbstractArray{T}, backend::Backend=get_backend(v); @@ -173,7 +296,6 @@ function _radix_sort!( temp::Union{Nothing, AbstractArray}=nothing, ) where T - # Fall back for unsupported types or custom comparators if !_rs_supported(T) || lt !== isless || by !== identity return merge_sort!(v, backend; lt, by, rev, order, block_size, temp) end @@ -186,9 +308,17 @@ function _radix_sort!( descending = (rev === true) || order === Base.Order.Reverse num_blocks = cld(n, block_size) - # hist[k * num_blocks + b] = count of elements with digit k in block b + n_passes = sizeof(T) * 8 ÷ Int(_RS_BITS) # 4 for 32-bit, 8 for 64-bit + + # Single histogram buffer; no need to zero before each pass — _radix_hist! + # zero-initializes its own shared-memory histogram and writes directly here. hist = similar(v, UInt32, Int(_RS_SIZE) * num_blocks) + # Reusable scratch for accumulate!'s per-block prefixes (ScanPrefixes uses a + # 256-thread, 2-elems-per-thread grid → 512 elements per block), so the + # exclusive prefix sum does not re-allocate on every pass. + acc_temp = similar(v, UInt32, cld(length(hist), 512)) + p1 = v p2 = if !isnothing(temp) @argcheck length(temp) >= n && eltype(temp) === T @@ -197,32 +327,54 @@ function _radix_sort!( similar(v) end - n_passes = sizeof(T) * 8 ÷ Int(_RS_BITS) # 4 for 32-bit, 8 for 64-bit - hist_kern! = _radix_hist!(backend, block_size) - scat_kern! = _radix_scatter!(backend, block_size) - ndrange = (block_size * num_blocks,) + ndrange = (block_size * num_blocks,) - for pass in 0:n_passes - 1 - shift = UInt32(pass) * _RS_BITS + # Compute (min, max) sort-key to detect passes where all elements share the + # same digit → skip the full hist+scan+scatter for that byte position. + min_key, max_key = _rs_key_range(p1, descending) + + # The fast histogram (O(1)/element atomic counting) and the chunked scatter + # (per-chunk sub-histograms) both use shared-memory atomics; select them + # where the backend reports atomics support and fall back to the portable + # scan/broadcast kernels otherwise. The chunked scatter additionally needs + # block_size to be a multiple of its 32-wide chunk. + has_atomics = KernelAbstractions.supports_atomics(backend) + hist_kern! = has_atomics ? + _radix_hist_atomic!(backend, block_size) : + _radix_hist!(backend, block_size) + scat_kern! = (has_atomics && block_size % _RS_CHUNK == 0) ? + _radix_scatter_chunked!(backend, block_size) : + _radix_scatter!(backend, block_size) - fill!(hist, 0x0) - hist_kern!(hist, p1, shift, descending; ndrange) - KernelAbstractions.synchronize(backend) + n_actual = 0 - accumulate!(+, hist, backend; init=UInt32(0), inclusive=false) - KernelAbstractions.synchronize(backend) + for pass in 0:n_passes - 1 + shift = UInt64(pass) * UInt64(_RS_BITS) - scat_kern!(p2, p1, hist, shift, descending; ndrange) - KernelAbstractions.synchronize(backend) + # Trivial pass: all elements share the same byte k AND all higher bytes + # are also identical (so no element can have a different byte k). + # Sufficient condition: (min_key >> shift) == (max_key >> shift). + # Checking just the single byte is WRONG when higher bytes differ. + (min_key >> shift) == (max_key >> shift) && continue + + shift32 = UInt32(shift) + # The three kernels run in order on a single backend stream, so no host + # synchronization is needed between them. + hist_kern!(hist, p1, shift32, descending; ndrange) + accumulate!(+, hist, backend; init=UInt32(0), inclusive=false, temp=acc_temp) + scat_kern!(p2, p1, hist, shift32, descending; ndrange) p1, p2 = p2, p1 + n_actual += 1 end - # After an even number of passes p1 === v (result already in v). - # After an odd number, p1 is the temp buffer; copy back. - if isodd(n_passes) + # p1 holds the result; copy back only if it's in the temp buffer. + if isodd(n_actual) copyto!(v, p1) end + # Block once so the sort is complete on return; the passes only enqueue work. + KernelAbstractions.synchronize(backend) + v end diff --git a/src/sort/sort.jl b/src/sort/sort.jl index dc20d41..29ac7d6 100644 --- a/src/sort/sort.jl +++ b/src/sort/sort.jl @@ -22,7 +22,11 @@ end """ RadixSort() -Use GPU radix sort for `sort!` and `sort`. This algorithm does not support `sortperm!`. +Use GPU radix sort for `sort!` and `sort`. A stable LSD radix sort (8-bit digits) supporting +32- and 64-bit integers and floats; other element types or custom `lt`/`by` fall back to +[`merge_sort!`](@ref). It is typically faster than [`MergeSort`](@ref) for these numeric types, +and skips passes automatically for small-range or structured data. Defaults to `block_size=256`. +This algorithm does not support `sortperm!`. """ struct RadixSort <: SortAlgorithm end @@ -55,7 +59,7 @@ struct SampleSort <: SortAlgorithm end alg::Union{Nothing, SortAlgorithm}=nothing, # GPU settings - block_size::Int=256, + block_size::Union{Nothing, Int}=nothing, # Temporary buffer, same size as `v` temp::Union{Nothing, AbstractArray}=nothing, @@ -76,8 +80,9 @@ 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. When `block_size` is left as +`nothing` (the default), each GPU algorithm picks its own tuned value (256 for merge sort, 512 for +radix sort); pass an explicit `block_size` to override. ## Algorithm choice By default, `sort!` uses [`sample_sort!`](@ref) on CPU backends and [`merge_sort!`](@ref) on GPU @@ -131,8 +136,8 @@ function _sort_impl!( alg::Union{Nothing, SortAlgorithm}=nothing, - # GPU settings - block_size::Int=256, + # GPU settings; nothing => each GPU algorithm picks its own tuned default + block_size::Union{Nothing, Int}=nothing, # Temporary buffer, same size as `v` temp::Union{Nothing, AbstractArray}=nothing, @@ -143,14 +148,17 @@ function _sort_impl!( merge_sort!( v, backend; lt, by, rev, order, - block_size, + block_size=isnothing(block_size) ? 256 : block_size, temp, ) elseif alg isa RadixSort + # The scatter's per-element rank cost scales with block_size, so + # overly large blocks hurt; 256 is the portable sweet spot across + # backends. _radix_sort!( v, backend; lt, by, rev, order, - block_size, + block_size=isnothing(block_size) ? 256 : block_size, temp, ) else @@ -189,7 +197,7 @@ end alg::Union{Nothing, SortAlgorithm}=nothing, # GPU settings - block_size::Int=256, + block_size::Union{Nothing, Int}=nothing, # Temporary buffer, same size as `v` temp::Union{Nothing, AbstractArray}=nothing, @@ -228,7 +236,7 @@ end alg::Union{Nothing, SortAlgorithm}=nothing, # GPU settings - block_size::Int=256, + block_size::Union{Nothing, Int}=nothing, # Temporary buffer, same size as `v` temp::Union{Nothing, AbstractArray}=nothing, @@ -272,20 +280,21 @@ function _sortperm_impl!( alg::Union{Nothing, SortAlgorithm}=nothing, - # GPU settings - block_size::Int=256, + # GPU settings; nothing => merge sort's tuned default (sortperm is merge-only) + block_size::Union{Nothing, Int}=nothing, # Temporary buffer, same size as `v` temp::Union{Nothing, AbstractArray}=nothing, ) if use_gpu_algorithm(backend, prefer_threads) alg = isnothing(alg) ? MergeSort() : alg + bs = isnothing(block_size) ? 256 : block_size if alg isa MergeSort if alg.lowmem merge_sortperm_lowmem!( ix, v, backend; lt, by, rev, order, - block_size, + block_size=bs, temp, ) else @@ -296,7 +305,7 @@ function _sortperm_impl!( merge_sortperm!( ix, v, backend; lt, by, rev, order, - block_size, + block_size=bs, temp_ix=temp, # old `temp` was the index buffer; maps directly to temp_ix ) end @@ -340,7 +349,7 @@ end alg::Union{Nothing, SortAlgorithm}=nothing, # GPU settings - block_size::Int=256, + block_size::Union{Nothing, Int}=nothing, # Temporary buffer, same size as `v` temp::Union{Nothing, AbstractArray}=nothing, From e090e0c00c3f08e6222afef91a201f5369ddfc07 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Wed, 22 Jul 2026 12:42:03 +0530 Subject: [PATCH 3/5] fix(accumulate): keep scan barriers uniform so POCL/SPIR-V doesn't deadlock --- src/accumulate/accumulate_1d_gpu.jl | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/src/accumulate/accumulate_1d_gpu.jl b/src/accumulate/accumulate_1d_gpu.jl index f4153d9..e8eb212 100644 --- a/src/accumulate/accumulate_1d_gpu.jl +++ b/src/accumulate/accumulate_1d_gpu.jl @@ -109,14 +109,21 @@ end # For exclusive ScanPrefixes scans, the local exclusive output is already correct. # DecoupledLookback still shifts non-first blocks because _accumulate_previous! - # expects each block's last value to be globally inclusive. - if inclusive || (iblock != 0x0 && !isnothing(flags)) - # To compute an inclusive scan, shift elements left... - @synchronize() - t1 = temp[ai + bank_offset_a + 0x1] - t2 = temp[bi + bank_offset_b + 0x1] - @synchronize() + # expects each block's last value to be globally inclusive. The condition is + # uniform across the group, but it is a runtime value: the @synchronize() barriers + # of the shift must stay OUTSIDE the branch. A barrier inside a conditional the + # compiler cannot prove uniform deadlocks POCL/SPIR-V (its work-item loop does not + # see the barrier reached on every path), so only the shared-memory writes are + # guarded while the barriers always execute. + do_shift = inclusive || (iblock != 0x0 && !isnothing(flags)) + + # To compute an inclusive scan, shift elements left... + @synchronize() + t1 = temp[ai + bank_offset_a + 0x1] + t2 = temp[bi + bank_offset_b + 0x1] + @synchronize() + if do_shift if ai > 0x0 temp[ai - 0x1 + conflict_free_offset(ai - 0x1) + 0x1] = t1 end From 976390b06bbe2980c8807d3ba20ee84353f094c9 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Wed, 22 Jul 2026 15:06:15 +0530 Subject: [PATCH 4/5] perf(sort): process multiple items per thread in the radix histogram and scatter --- src/sort/radix_sort.jl | 138 +++++++++++++++++++++++++++++------------ 1 file changed, 97 insertions(+), 41 deletions(-) diff --git a/src/sort/radix_sort.jl b/src/sort/radix_sort.jl index a56a759..c450a15 100644 --- a/src/sort/radix_sort.jl +++ b/src/sort/radix_sort.jl @@ -17,6 +17,9 @@ # Key optimizations: # • Fused min/max range (_rs_key_range): one reduction over the sort keys instead of # separate minimum + maximum. +# • Items-per-thread tiling: the fast kernels process items_per_thread elements +# per thread, shrinking the per-(digit, block) histogram — and the exclusive +# scan over it — by the same factor. # • Skip-pass via min/max keys: if min and max share the whole byte-suffix from byte k # up, every element does too, so the whole pass is skipped (e.g. UInt32 in [0, 255] # sorts in a single pass). @@ -61,8 +64,10 @@ end # of its assigned buckets (bucket t, t+NI, t+2*NI, …). Uses no shared-memory # atomics, so it is the portable fallback that runs on every backend. +# The trailing Val is always Val(1) on this path; accepted only so the driver can +# call every histogram/scatter kernel with one uniform signature. @kernel inbounds=true cpu=false unsafe_indices=true function _radix_hist!( - hist, @Const(v), shift::UInt32, rev::Bool, + hist, @Const(v), shift::UInt32, rev::Bool, ::Val, ) @uniform NI = Int(@groupsize()[1]) s_digit = @localmem UInt32 (NI,) @@ -92,12 +97,15 @@ end # ─── Phase 1b: histogram — shared-atomic (backends with shared-mem atomics) ── # Same per-block 256-bin output, but O(1)/element via a shared-memory atomic -# increment instead of the O(256)/element per-bucket scan above. Selected by -# the driver only where the backend reports atomics support. +# increment instead of the O(256)/element per-bucket scan above. Each thread +# processes ITEMS elements (block-strided, so global reads stay coalesced) and a +# block covers block_size*ITEMS elements — shrinking the per-(digit, block) +# histogram, and the exclusive scan over it, by ITEMS×. Selected by the driver +# only where the backend reports atomics support. @kernel inbounds=true cpu=false unsafe_indices=true function _radix_hist_atomic!( - hist, @Const(v), shift::UInt32, rev::Bool, -) + hist, @Const(v), shift::UInt32, rev::Bool, ::Val{ITEMS}, +) where ITEMS @uniform NI = Int(@groupsize()[1]) s_hist = @localmem UInt32 (Int(_RS_SIZE),) @@ -105,6 +113,7 @@ end ithread = Int(@index(Local, Linear)) - 1 len = Int(length(v)) num_blocks = Int(length(hist)) ÷ Int(_RS_SIZE) + base = iblock * NI * ITEMS j = ithread while j < Int(_RS_SIZE) @@ -113,10 +122,14 @@ end end @synchronize() - i = iblock * NI + ithread - if i < len - d = Int(_rs_digit(v[i + 1], shift, rev)) - Atomix.@atomic s_hist[d + 1] += UInt32(1) + m = 0 + while m < ITEMS + i = base + ithread + m * NI + if i < len + d = Int(_rs_digit(v[i + 1], shift, rev)) + Atomix.@atomic s_hist[d + 1] += UInt32(1) + end + m += 1 end @synchronize() @@ -132,8 +145,10 @@ end # `hist` is already the exclusive-prefix-summed per-block offsets; # `hist[k * num_blocks + b]` = global start for bucket k, block b (1-indexed). +# The trailing Val is always Val(1) on this path; accepted only so the driver can +# call every histogram/scatter kernel with one uniform signature. @kernel inbounds=true cpu=false unsafe_indices=true function _radix_scatter!( - v_out, @Const(v_in), @Const(hist), shift::UInt32, rev::Bool, + v_out, @Const(v_in), @Const(hist), shift::UInt32, rev::Bool, ::Val, ) @uniform N = @groupsize()[1] @uniform NI = Int(@groupsize()[1]) @@ -173,21 +188,25 @@ end # ─── Phase 3b: scatter — chunked stable rank (O(chunk) per thread) ─────────── -# The broadcast scatter's rank is O(block_size)/element. Split the block into +# The broadcast scatter's rank is O(block_size)/element. Split the tile into # 32-wide chunks: per-chunk digit counts (built with shared-memory atomics — # order doesn't matter for counts) give each chunk's stable base via a cross-chunk # exclusive prefix, and each element only scans its own chunk (≤32) for the # intra-chunk part. Rank = cross-chunk-base[digit] + intra-chunk same-digit # count — still fully stable, but O(32) instead of O(block_size). +# +# Each thread processes ITEMS tile positions (block-strided, coalesced loads); +# the tile covers block_size*ITEMS elements and must match the histogram's +# block coverage so the per-(digit, block) offsets line up. @kernel inbounds=true cpu=false unsafe_indices=true function _radix_scatter_chunked!( - v_out, @Const(v_in), @Const(hist), shift::UInt32, rev::Bool, -) - @uniform N = @groupsize()[1] - @uniform NI = Int(@groupsize()[1]) - @uniform NCH = Int(@groupsize()[1]) ÷ _RS_CHUNK # number of chunks - s_elem = @localmem eltype(v_in) (N,) - s_digit = @localmem UInt32 (N,) + v_out, @Const(v_in), @Const(hist), shift::UInt32, rev::Bool, ::Val{ITEMS}, +) where ITEMS + @uniform NI = Int(@groupsize()[1]) + @uniform TILE = Int(@groupsize()[1]) * ITEMS + @uniform NCH = (Int(@groupsize()[1]) * ITEMS) ÷ _RS_CHUNK # number of chunks + s_elem = @localmem eltype(v_in) (TILE,) + s_digit = @localmem UInt32 (TILE,) s_gbase = @localmem UInt32 (256,) s_chist = @localmem UInt32 (256 * NCH,) # per-chunk digit counts → bases @@ -195,10 +214,23 @@ end ithread = Int(@index(Local, Linear)) - 1 len = Int(length(v_in)) num_blocks = Int(length(hist)) ÷ 256 - - i = iblock * NI + ithread - if i < len - s_elem[ithread + 1] = v_in[i + 1] + base = iblock * TILE + + # Load the tile: keys + digits, in tile-position order (== input order, which + # the stable rank below relies on). 0xffffffff marks out-of-range positions; + # it matches no valid digit, so they count and scatter nothing. + m = 0 + while m < ITEMS + p = ithread + m * NI + i = base + p + if i < len + k = v_in[i + 1] + s_elem[p + 1] = k + s_digit[p + 1] = _rs_digit(k, shift, rev) + else + s_digit[p + 1] = 0xffffffff + end + m += 1 end j = ithread while j < 256 @@ -212,11 +244,14 @@ end end @synchronize() - mychunk = ithread ÷ _RS_CHUNK - my_digit = UInt32(i < len ? _rs_digit(s_elem[ithread + 1], shift, rev) : 0) - s_digit[ithread + 1] = my_digit - if i < len - Atomix.@atomic s_chist[mychunk * 256 + Int(my_digit) + 1] += UInt32(1) + m = 0 + while m < ITEMS + p = ithread + m * NI + d = s_digit[p + 1] + if d != 0xffffffff + Atomix.@atomic s_chist[(p ÷ _RS_CHUNK) * 256 + Int(d) + 1] += UInt32(1) + end + m += 1 end @synchronize() @@ -233,15 +268,23 @@ end end @synchronize() - if i < len - chunk_start = mychunk * _RS_CHUNK - cnt = UInt32(0) - for jj in chunk_start:(ithread - 1) - cnt += UInt32(s_digit[jj + 1] == my_digit) + m = 0 + while m < ITEMS + p = ithread + m * NI + d = s_digit[p + 1] + if d != 0xffffffff + chunk_start = (p ÷ _RS_CHUNK) * _RS_CHUNK + cnt = UInt32(0) + q = chunk_start + while q < p + cnt += UInt32(s_digit[q + 1] == d) + q += 1 + end + rank = s_chist[(p ÷ _RS_CHUNK) * 256 + Int(d) + 1] + cnt + gpos = Int(s_gbase[Int(d) + 1]) + Int(rank) + v_out[gpos + 1] = s_elem[p + 1] end - rank = s_chist[mychunk * 256 + Int(my_digit) + 1] + cnt - gpos = Int(s_gbase[my_digit + 1]) + Int(rank) - v_out[gpos + 1] = s_elem[ithread + 1] + m += 1 end end @@ -293,6 +336,7 @@ function _radix_sort!( rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Forward, block_size::Int=256, + items_per_thread::Int=2, temp::Union{Nothing, AbstractArray}=nothing, ) where T @@ -304,10 +348,21 @@ function _radix_sort!( n <= 1 && return v @argcheck ispow2(block_size) && block_size >= 1 + @argcheck items_per_thread >= 1 descending = (rev === true) || order === Base.Order.Reverse - num_blocks = cld(n, block_size) + # Processing several items per thread shrinks the per-(digit, block) histogram + # (and its exclusive scan) by the same factor. Only the fast atomic kernels + # support it; the portable scan/broadcast fallbacks stay at one item per + # thread. The default of 2 keeps the chunked scatter's shared memory within + # every backend's budget (Metal caps threadgroup memory at 32 KiB); discrete + # NVIDIA/AMD GPUs are measurably faster still at 4. + has_atomics = KernelAbstractions.supports_atomics(backend) + use_fast = has_atomics && block_size % _RS_CHUNK == 0 + items = use_fast ? items_per_thread : 1 + + num_blocks = cld(n, block_size * items) n_passes = sizeof(T) * 8 ÷ Int(_RS_BITS) # 4 for 32-bit, 8 for 64-bit # Single histogram buffer; no need to zero before each pass — _radix_hist! @@ -337,12 +392,13 @@ function _radix_sort!( # (per-chunk sub-histograms) both use shared-memory atomics; select them # where the backend reports atomics support and fall back to the portable # scan/broadcast kernels otherwise. The chunked scatter additionally needs - # block_size to be a multiple of its 32-wide chunk. - has_atomics = KernelAbstractions.supports_atomics(backend) + # block_size to be a multiple of its 32-wide chunk. Both fast kernels must + # agree on the tile size (block_size*items), so they share the same Val. + vitems = Val(items) hist_kern! = has_atomics ? _radix_hist_atomic!(backend, block_size) : _radix_hist!(backend, block_size) - scat_kern! = (has_atomics && block_size % _RS_CHUNK == 0) ? + scat_kern! = use_fast ? _radix_scatter_chunked!(backend, block_size) : _radix_scatter!(backend, block_size) @@ -360,9 +416,9 @@ function _radix_sort!( shift32 = UInt32(shift) # The three kernels run in order on a single backend stream, so no host # synchronization is needed between them. - hist_kern!(hist, p1, shift32, descending; ndrange) + hist_kern!(hist, p1, shift32, descending, vitems; ndrange) accumulate!(+, hist, backend; init=UInt32(0), inclusive=false, temp=acc_temp) - scat_kern!(p2, p1, hist, shift32, descending; ndrange) + scat_kern!(p2, p1, hist, shift32, descending, vitems; ndrange) p1, p2 = p2, p1 n_actual += 1 From 908ff5aba0c35c87adaa5dca97a5baf7577c4387 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Wed, 22 Jul 2026 15:50:41 +0530 Subject: [PATCH 5/5] perf(sort): sort tile-sized arrays in a single kernel launch Co-Authored-By: Claude Opus 4.8 --- src/sort/radix_sort.jl | 153 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 152 insertions(+), 1 deletion(-) diff --git a/src/sort/radix_sort.jl b/src/sort/radix_sort.jl index c450a15..3d17e20 100644 --- a/src/sort/radix_sort.jl +++ b/src/sort/radix_sort.jl @@ -20,6 +20,8 @@ # • Items-per-thread tiling: the fast kernels process items_per_thread elements # per thread, shrinking the per-(digit, block) histogram — and the exclusive # scan over it — by the same factor. +# • Single-block fast path: arrays that fit one tile (2 * block_size elements) +# are sorted entirely in shared memory by one kernel launch. # • Skip-pass via min/max keys: if min and max share the whole byte-suffix from byte k # up, every element does too, so the whole pass is skipped (e.g. UInt32 in [0, 255] # sorts in a single pass). @@ -289,6 +291,140 @@ end end + +# ─── Whole-array single-block sort (small arrays) ──────────────────────────── +# For arrays that fit in one tile (2 * block_size elements) the entire multi-pass +# sort runs inside a single work-group: keys ping-pong between two shared-memory +# buffers, with the same chunked stable rank as the scatter above, and global +# memory is touched exactly twice (load + store). One kernel launch replaces the +# hist/scan/scatter pipeline's 3 launches per pass, which dominate at this size. +# +# The pass loop has a compile-time trip count (NPASS from the element type) and +# the source/destination buffers are chosen with a data select, never a branch, +# so every @synchronize below executes unconditionally on all backends. + +@kernel inbounds=true cpu=false unsafe_indices=true function _radix_sort_block!( + v, rev::Bool, ::Val{NPASS}, +) where NPASS + @uniform NI = Int(@groupsize()[1]) + @uniform TILE = Int(@groupsize()[1]) * 2 + @uniform NCH = (Int(@groupsize()[1]) * 2) ÷ _RS_CHUNK + s_a = @localmem eltype(v) (TILE,) + s_b = @localmem eltype(v) (TILE,) + s_digit = @localmem UInt32 (TILE,) + s_chist = @localmem UInt32 (256 * NCH,) + s_loff = @localmem UInt32 (256,) # per-digit start offsets in the tile + + it = Int(@index(Local, Linear)) - 1 + n = Int(length(v)) + + m = 0 + while m < 2 + p = it + m * NI + if p < n + s_a[p + 1] = v[p + 1] + end + m += 1 + end + @synchronize() + + pass = 0 + while pass < NPASS + sh = UInt32(pass) * _RS_BITS + # Uniform data select between the ping-pong buffers; keeps the barriers + # below outside any conditional. + src = iseven(pass) ? s_a : s_b + dst = iseven(pass) ? s_b : s_a + + j = it + while j < 256 * NCH + s_chist[j + 1] = UInt32(0) + j += NI + end + j = it + while j < 256 + s_loff[j + 1] = UInt32(0) + j += NI + end + @synchronize() + + m = 0 + while m < 2 + p = it + m * NI + if p < n + d = _rs_digit(src[p + 1], sh, rev) + s_digit[p + 1] = d + Atomix.@atomic s_chist[(p ÷ _RS_CHUNK) * 256 + Int(d) + 1] += UInt32(1) + end + m += 1 + end + @synchronize() + + # cross-chunk exclusive prefix per digit; digit totals land in s_loff + d = it + while d < 256 + acc = UInt32(0) + c = 0 + while c < NCH + cnt = s_chist[c * 256 + d + 1] + s_chist[c * 256 + d + 1] = acc + acc += cnt + c += 1 + end + s_loff[d + 1] = acc + d += NI + end + @synchronize() + + # exclusive scan of the 256 digit totals (serial by thread 0: 256 adds, + # negligible next to the tile's rank work) + if it == 0 + run = UInt32(0) + dd = 0 + while dd < 256 + t = s_loff[dd + 1] + s_loff[dd + 1] = run + run += t + dd += 1 + end + end + @synchronize() + + # stable rank (cross-chunk base + intra-chunk same-digit count) -> place + m = 0 + while m < 2 + p = it + m * NI + if p < n + d = s_digit[p + 1] + chunk_start = (p ÷ _RS_CHUNK) * _RS_CHUNK + cnt = UInt32(0) + q = chunk_start + while q < p + cnt += UInt32(s_digit[q + 1] == d) + q += 1 + end + rank = s_chist[(p ÷ _RS_CHUNK) * 256 + Int(d) + 1] + cnt + dst[Int(s_loff[Int(d) + 1]) + Int(rank) + 1] = src[p + 1] + end + m += 1 + end + @synchronize() + + pass += 1 + end + + res = iseven(NPASS) ? s_a : s_b + m = 0 + while m < 2 + p = it + m * NI + if p < n + v[p + 1] = res[p + 1] + end + m += 1 + end +end + + # ─── Implementation ────────────────────────────────────────────────────────── _rs_supported(::Type{T}) where T = @@ -362,8 +498,23 @@ function _radix_sort!( use_fast = has_atomics && block_size % _RS_CHUNK == 0 items = use_fast ? items_per_thread : 1 + n_passes = sizeof(T) * 8 ÷ Int(_RS_BITS) # 4 for 32-bit, 8 for 64-bit + + # Whole-array single-block fast path: the full sort runs in shared memory in + # one launch, skipping the hist/scan/scatter pipeline and its allocations. + # The kernel's shared footprint scales with the tile, so it is only taken up + # to block_size 256 (worst case, 64-bit keys at tile 512: ~27 KiB), which + # fits every backend's budget — Metal caps threadgroup memory at 32 KiB and + # CUDA static shared memory at 48 KiB; block_size 512 with 64-bit keys would + # need ~53 KiB and fail to launch. + if use_fast && block_size <= 256 && n <= 2 * block_size + _radix_sort_block!(backend, block_size)( + v, descending, Val(n_passes); ndrange=block_size) + KernelAbstractions.synchronize(backend) + return v + end + num_blocks = cld(n, block_size * items) - n_passes = sizeof(T) * 8 ÷ Int(_RS_BITS) # 4 for 32-bit, 8 for 64-bit # Single histogram buffer; no need to zero before each pass — _radix_hist! # zero-initializes its own shared-memory histogram and writes directly here.