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)