From 03be940536d03cb90936d5149b5cc4931f778ea2 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 27 Oct 2026 14:19:12 +0100 Subject: [PATCH] Look up fold_threshold behind an inference barrier Launches compiled before KernelAbstractions' extension for a backend is loaded, such as those precompiled by AMDGPU.jl and Metal.jl through GPUArrays and AcceleratedKernels, depended on the generic method of fold_threshold, so loading the extension that adds the backend's method invalidated them. Look tuned values up through a `tuning` helper that hides the call from inference, and only for the launches that can be folded: tuned ones of more than three dimensions. --- src/launch.jl | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/launch.jl b/src/launch.jl index cba017925..3dd7180d8 100644 --- a/src/launch.jl +++ b/src/launch.jl @@ -82,10 +82,12 @@ function select_launch(kernel::Kernel, workgroupsize, iterspace) groups = size(blocks(iterspace)) items = size(workitems(iterspace)) tuned = KernelAbstractions.workgroupsize(kernel) <: DynamicSize && workgroupsize === nothing + # only tuned launches of more than three dimensions can be folded + fold_items = tuned && length(groups) > 3 ? tuning(fold_threshold, b) : 0 return select_launch( map(mul_extent, groups, items), tuned ? nothing : items, KI.max_work_group_size(b), KI.max_work_group_dims(b), KI.max_num_groups(b), - narrow_index_type(b); fold_items = fold_threshold(b) + narrow_index_type(b); fold_items ) end @@ -135,10 +137,17 @@ with divisions instead. Where that happens depends on the hardware, so backends measured value keep the default of 64, chosen for NVIDIA GPUs. Set here and in KernelAbstractions' extensions, not by the backends, so that it can be -tuned without releasing them. +tuned without releasing them. Look it up with `tuning`. """ fold_threshold(backend) = 64 +# Values tuned per backend, like `fold_threshold`, are looked up with `tuning(hook, backend)`, +# hidden from inference. The backends' packages precompile launches, and code that depends on +# a hook's generic method is invalidated when the extension that adds the backend's method +# is loaded. +tuning(hook, backend) = Base.inferencebarrier(hook)(backend)::tuning_type(hook) +tuning_type(::typeof(fold_threshold)) = Int + # arithmetic on the extents of an iteration space, which has to fit an `Int` @noinline throw_too_large() = throw(ArgumentError("Iteration space has more than typemax(Int) work-items"))