From 4b2f306e07eab0c776fe3db74ebde4f9f566bbf3 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 27 Oct 2026 14:19:12 +0100 Subject: [PATCH] Look up the tuning hooks behind an inference barrier Backend packages precompile AK's operations, e.g. Metal.jl's reductions through GPUArrays, and that code depended on the generic methods of the tuning hooks, so loading AK's extension for the backend invalidated it. Look the tunings up through a `tuning` helper that hides the call from inference. Tunings may already differ per device, so callers never relied on their values being constants. --- src/accumulate/accumulate.jl | 2 +- src/accumulate/accumulate_1d_gpu.jl | 2 +- src/accumulate/accumulate_nd.jl | 2 +- src/accumulate/tuning.jl | 3 ++- src/algorithms.jl | 8 ++++++++ src/block.jl | 2 +- src/findall.jl | 3 ++- src/predicates.jl | 3 ++- src/reduce/mapreduce_nd.jl | 2 +- src/reduce/reduce.jl | 4 ++-- src/reduce/shapes.jl | 4 ++-- src/reduce/tuning.jl | 3 ++- src/sort/tuning.jl | 3 ++- 13 files changed, 27 insertions(+), 14 deletions(-) diff --git a/src/accumulate/accumulate.jl b/src/accumulate/accumulate.jl index 86d39a98..ba5c9665 100644 --- a/src/accumulate/accumulate.jl +++ b/src/accumulate/accumulate.jl @@ -150,7 +150,7 @@ function scan_setup( (; prefixes=scratch(typeof(seed), prefixes_length(n, a))) elseif a isa SliceScan # The chunk totals of slices split into chunks (none for slices in one chunk) - m = slice_scratch(sz, Int(dims), a, scan_tuning(backend, A), typeof(seed)) + m = slice_scratch(sz, Int(dims), a, tuning(scan_tuning, backend, A), typeof(seed)) (; totals=scratch(typeof(seed), m)) else (;) diff --git a/src/accumulate/accumulate_1d_gpu.jl b/src/accumulate/accumulate_1d_gpu.jl index 54fff954..913d25c7 100644 --- a/src/accumulate/accumulate_1d_gpu.jl +++ b/src/accumulate/accumulate_1d_gpu.jl @@ -156,7 +156,7 @@ function accumulate_1d_gpu!( scan_totals!() = accumulate_1d_gpu!(op, totals, backend, alg; init=NoInit(), seed, inclusive=true, prefixes=view(prefixes, num_blocks+1:length(prefixes))) - if scan_tuning(backend, eltype(v)).reduce_then_scan + if tuning(scan_tuning, backend, eltype(v)).reduce_then_scan # Reduce the tiles, then scan them from their prefix: this reads the source twice and # writes the result once tile_aggregates_kernel!(backend, block_size)(op, totals, src, load, start, seed, diff --git a/src/accumulate/accumulate_nd.jl b/src/accumulate/accumulate_nd.jl index a5a51c03..d0576590 100644 --- a/src/accumulate/accumulate_nd.jl +++ b/src/accumulate/accumulate_nd.jl @@ -9,7 +9,7 @@ function accumulate_nd!( totals::Union{Nothing, AbstractVector}=nothing, ) where {OP} n_inner, len, n_outer = slice_extents(size(v), dims) - t = scan_tuning(backend, eltype(v)) + t = tuning(scan_tuning, backend, eltype(v)) slice_scan!(op, v, backend, alg, t; src, load, start=scan_start(v, init, seed), seed, inclusive, n_inner, len, n_outer, totals) return v diff --git a/src/accumulate/tuning.jl b/src/accumulate/tuning.jl index 60e01052..19bead24 100644 --- a/src/accumulate/tuning.jl +++ b/src/accumulate/tuning.jl @@ -93,6 +93,7 @@ The scan tuning for element type `T` (the running-value type the scan computes i conventions. """ scan_tuning(::Backend, ::Type) = ScanTuning() +tuning_type(::typeof(scan_tuning)) = ScanTuning # CPU devices (the host backend's, and OpenCL's CPU devices) implement sub-group operations with # barriers of the whole block, and keep a scanned tile in their caches for adding its prefix cpu_scan_tuning() = ScanTuning(reduce_then_scan=false, subgroups=false) @@ -109,7 +110,7 @@ or `CPUThreads.Partitioned` with every field set, or throws an `ArgumentError`. """ function resolve_scan(alg::Algorithm, backend::Backend, ::Type{T}, dims, ::Type{S}=T) where {T, S} - t = scan_tuning(backend, T) + t = tuning(scan_tuning, backend, T) a = alg isa Auto ? select_scan(backend, dims) : alg a = fill_defaults(a, t, S) check_scan(a, backend, S, dims) diff --git a/src/algorithms.jl b/src/algorithms.jl index 3d28f255..2b7b2f4e 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -150,6 +150,14 @@ runs_kernels(::Backend) = true runs_kernels(::HostBackend) = nameof(HostBackend) !== :CPU +# Tunings: values chosen per device by the `*_tuning(backend, T)` hooks, to which AK's package +# extensions add their backends' methods. They are looked up with `tuning(hook, backend, T)`, +# hidden from inference: backend packages precompile AK's operations, and code that depends on +# a hook's generic method is invalidated when the extension that adds the backend's method is +# loaded. Each hook declares the type it returns with `tuning_type`. +tuning(hook, backend, T) = Base.inferencebarrier(hook)(backend, T)::tuning_type(hook) + + # Checks shared by the algorithm families # Algorithm names in error messages diff --git a/src/block.jl b/src/block.jl index 6627f0ef..b86117d6 100644 --- a/src/block.jl +++ b/src/block.jl @@ -23,7 +23,7 @@ end # sub-group scans the sub-groups' totals. @static if isdefined(KernelAbstractions, :KernelInterface) function scan_subgroup_width(backend, ::Type{T}, block_size) where {T} - scan_tuning(backend, T).subgroups || return Val(0) + tuning(scan_tuning, backend, T).subgroups || return Val(0) KI.supports_linear_subgroups(backend) && KI.supports_shuffle(backend, T) || return Val(0) W = KI.sub_group_size(backend) return ispow2(W) && W <= block_size <= W * W ? Val(W) : Val(0) diff --git a/src/findall.jl b/src/findall.jl index a3c99dbf..d81b0334 100644 --- a/src/findall.jl +++ b/src/findall.jl @@ -31,6 +31,7 @@ The `findall` tuning for element type `T` (the input's) on `backend`'s current d [`sort_tuning`](@ref) for the conventions. """ findall_tuning(::Backend, ::Type) = FindallTuning() +tuning_type(::typeof(findall_tuning)) = FindallTuning """ resolve_findall(alg, backend, T) -> Algorithm @@ -39,7 +40,7 @@ Resolve `alg` for `findall` over elements of type `T` on `backend`: `ScanScatter `CPUThreads.Partitioned` with every field set, or an `ArgumentError`. """ function resolve_findall(alg::Algorithm, backend::Backend, ::Type{T}) where {T} - t = findall_tuning(backend, T) + t = tuning(findall_tuning, backend, T) a = alg isa Auto ? (runs_threads(backend) ? CPUThreads.Partitioned() : ScanScatter()) : alg a = fill_defaults(a, t, T) check_findall(a, backend) diff --git a/src/predicates.jl b/src/predicates.jl index 06a96ee0..cebd6c6b 100644 --- a/src/predicates.jl +++ b/src/predicates.jl @@ -47,6 +47,7 @@ The `any`/`all` tuning for element type `T` (the input's) on `backend`'s current [`sort_tuning`](@ref) for the conventions. """ predicate_tuning(::Backend, ::Type) = PredicateTuning() +tuning_type(::typeof(predicate_tuning)) = PredicateTuning """ resolve_predicate(alg, backend, T) -> Algorithm @@ -56,7 +57,7 @@ Resolve `alg` for `any`/`all` over elements of type `T` on `backend`: `Concurren `ArgumentError`. """ function resolve_predicate(alg::Algorithm, backend::Backend, ::Type{T}) where {T} - t = predicate_tuning(backend, T) + t = tuning(predicate_tuning, backend, T) a = alg isa Auto ? select_predicate(backend, t) : alg a = fill_defaults(a, t, backend) check_predicate(a, backend) diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index c346ae0c..93997591 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/src/reduce/mapreduce_nd.jl @@ -65,7 +65,7 @@ function mapreduce_nd!( Base.inferencebarrier(mapreduce_1d_gpu)( f, op, linear_source(src, launch.buffer, launch.offset), backend; init, seed, block_size=alg.block_size, items_per_thread=alg.items_per_thread, - max_blocks=reduce_tuning(backend, A).target_blocks, + max_blocks=tuning(reduce_tuning, backend, A).target_blocks, partials=get(bufs, :partials, nothing), switch_below=0, commutative=!shape.ordered, dst) return dst diff --git a/src/reduce/reduce.jl b/src/reduce/reduce.jl index 15210c83..e20742b7 100644 --- a/src/reduce/reduce.jl +++ b/src/reduce/reduce.jl @@ -193,7 +193,7 @@ function mapreduce_setup(f::F, op::OP, src, backend, init, neutral, acctype, dim # (the buffer's name does not depend on the length, so that results infer) sizes = if a isa BlockReduce && seed !== nothing n = mapreduce_1d_host_partials(length(src), a.block_size, a.items_per_thread, - reduce_tuning(backend, A).target_blocks, + tuning(reduce_tuning, backend, A).target_blocks, a.switch_below) (; partials=scratch(typeof(seed), n)) else @@ -409,7 +409,7 @@ function mapreduce_whole(f::F, op::OP, src, backend, alg, ::Type{A}; init, seed, f, op, src, backend; init, seed, block_size=alg.block_size, items_per_thread=alg.items_per_thread, - max_blocks=reduce_tuning(backend, A).target_blocks, + max_blocks=tuning(reduce_tuning, backend, A).target_blocks, partials, switch_below=alg.switch_below, commutative, ) else diff --git a/src/reduce/shapes.jl b/src/reduce/shapes.jl index 634e3ed8..7f456b29 100644 --- a/src/reduce/shapes.jl +++ b/src/reduce/shapes.jl @@ -255,7 +255,7 @@ function select_launch(src, backend, alg::BlockReduce, ::Type{A}, rdims, dst_siz if dst_size == 1 && (isnothing(strided) || nr == 1 && rstrides[1] == 1) return launch(Whole(!commutative)) end - shape = reduce_shape(dst_size, len, segs, alg.block_size, reduce_tuning(backend, A), A) + shape = reduce_shape(dst_size, len, segs, alg.block_size, tuning(reduce_tuning, backend, A), A) commutative || (shape = ordered_shape(shape, dst_size, len, alg.block_size)) return launch(shape) end @@ -270,7 +270,7 @@ function launch_scratch(launch::DimsLaunch, src, backend, alg, ::Type{A}, seed, if shape isa Whole len = Base.prod(size(src)[d] for d in rdims; init=1) n = mapreduce_1d_partials(len, alg.block_size, alg.items_per_thread, - reduce_tuning(backend, A).target_blocks) + tuning(reduce_tuning, backend, A).target_blocks) return n == 0 ? (;) : (; partials=scratch(typeof(seed), n)) end groups = split_groups(shape) diff --git a/src/reduce/tuning.jl b/src/reduce/tuning.jl index 0680f608..353714d9 100644 --- a/src/reduce/tuning.jl +++ b/src/reduce/tuning.jl @@ -73,6 +73,7 @@ The reduction tuning for accumulator type `T` on `backend`'s current device; see [`sort_tuning`](@ref) for the conventions. """ reduce_tuning(::Backend, ::Type) = ReduceTuning() +tuning_type(::typeof(reduce_tuning)) = ReduceTuning # Kernels on CPU devices (the host backend's, and OpenCL's CPU devices) reduce short runs faster # with a thread per output, and long ones faster with groups of threads: with PoCL on a Ryzen # 9950X, summing an 8×10^6 matrix along `dims=1` takes half the time, and a 4096×4096 one along @@ -105,7 +106,7 @@ function resolve_reduce(alg::Algorithm, backend::Backend, ::Type{T}, dims) where throw(ArgumentError("BlockReduce: `items_per_thread` and `switch_below` only apply to " * "whole-array reductions (`dims=:`)")) end - t = reduce_tuning(backend, T) + t = tuning(reduce_tuning, backend, T) a = alg isa Auto ? select_reduce(backend) : alg a = fill_defaults(a, t, T) check_reduce(a, backend, T, dims) diff --git a/src/sort/tuning.jl b/src/sort/tuning.jl index 352758a5..a27e946e 100644 --- a/src/sort/tuning.jl +++ b/src/sort/tuning.jl @@ -116,6 +116,7 @@ task would launch on). AK defines this generic method; AK's package extensions a per backend type, which may choose different values per device. """ sort_tuning(::Backend, ::Type) = SortTuning() +tuning_type(::typeof(sort_tuning)) = SortTuning # The operation an algorithm is resolved for, for error messages @@ -140,7 +141,7 @@ Base.@constprop :aggressive function resolve_sort( ) T = eltype(v) layout = slice_layout(v, dims) - t = sort_tuning(backend, T) + t = tuning(sort_tuning, backend, T) a = alg isa Auto ? select_sort(alg, backend, t, T, layout, ord; perm, pairs) : alg a = fill_defaults(a, t, T, merge_eltypes) check_sort(a, backend, T, layout, ord; perm, pairs, merge_eltypes)