Look up the tuning hooks behind an inference barrier, so that loading a backend doesn't invalidate AK's operations

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
         (;)

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,

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

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)

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

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)

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)

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)

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

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

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)

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)

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)