Skip to content

Specialize the sorts on by and lt - #193

Open
maleadt wants to merge 1 commit into
mainfrom
tb/sort-specialize-by
Open

maleadt wants to merge 1 commit into
mainfrom
tb/sort-specialize-by

Conversation

@maleadt

@maleadt maleadt commented Oct 28, 2026

Copy link
Copy Markdown
Member

Follow-up to #191, for the sorts: a warm sort! with by, and any warm sortperm!, still runs inference on every call. Since GPUArrays 12 that includes Base's sort! and sortperm of GPU arrays.

julia> f(v) = sort!(v; by=abs);   # v::CuVector{Float32}, 1000 elements

julia> @time invokelatest(f, v)
  0.000022 seconds (214 allocations: 8.250 KiB, 0.32% compilation time: 100% of which was recompilation)

sort!, sortperm! and their setup helpers only pass by and lt on, so Julia compiles them for by::Function. In that code, the Base.promote_op(by, eltype(v)) that sizes MergeSort's key buffers can't fold. sortperm! evaluates it even with the default by=identity. This PR names the types of by and lt with where parameters in the sort methods (including sort_by_key!, which shares the ordering code), the same change #191 makes for the predicates.

Afterwards a warm call compiles nothing (the @time above: 135 allocations, 4.2 KiB, no compilation). On an RTX 5080, with 10,000 Float32s called by dynamic dispatch (median of five BenchmarkTools runs):

call main PR
sort!(x; by=abs) 44.5 µs, 335 allocs 41.0 µs, 256 allocs
sortperm!(ix, x; by=abs) 46.7 µs, 335 allocs 44.0 µs, 277 allocs
sortperm!(ix, x; lt=>) 30.2 µs, 311 allocs 27.9 µs, 249 allocs
sort!(X; dims=1, by=abs) 27.2 µs, 216 allocs 23.6 µs, 141 allocs
sort!(x) 36.6 µs, 196 allocs 36.4 µs, 185 allocs

The kernels are unchanged: they were already compiled for the run-time types of by and the ordering.

The reductions' dims=: default (Colon <: Function) causes no inference at run time: warm sum and reduce compile nothing on CPU and CUDA. Neither do the searchsorted*! functions, so they're left as they are.

The new tests check that a warm call of sort! and sortperm! by dynamic dispatch compiles nothing. They fail on main and pass with this PR on CPU threads, the KA host backend, and CUDA (Julia 1.12), and on CPU with Julia 1.10. They use the compile-time helper from #191, added to test/setup.jl as warm_compile_time (not compile_time, which ParallelTestRunner assigns as a global in each test module); #191 can switch to it once both are merged.

`sort!`, `sortperm!` and their setup helpers only pass `by` and `lt`
on, so Julia compiles them for `by::Function`. There,
`Base.promote_op(by, eltype(v))`, which sizes MergeSort's key buffers,
can't fold and runs inference on every call (`sortperm!` even with the
default `by=identity`). That also affects Base's `sort!` and `sortperm`
of GPU arrays, which go through AK since GPUArrays 12.

Name their types with `where` parameters, as the reductions do for
`f` and `op`, and likewise in `sort_by_key!`, which shares the ordering
code. A warm `sort!(x; by=abs)` of a 10,000-element `CuArray` then
compiles nothing and drops from 335 to 256 host allocations.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant