Repository navigation
Look up the tuning hooks behind an inference barrier, so that loading a backend doesn't invalidate AK's operations - #189
Merged
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
AK's package extensions set the tuning of their backends (
reduce_tuning(::MetalBackend, ::Type)etc.), so that it ships with AK rather than with the back ends. The back ends precompile AK's operations, e.g. Metal.jl's and CUDA.jl's reductions through GPUArrays, and that code calls the hooks directly, so it depends on their generic methods. Loading AK's extension for the backend then invalidates it. With the back ends defining their KernelAbstractions methods insrc(JuliaGPU/Metal.jl#1118, JuliaGPU/CUDA.jl#3493) and JuliaGPU/KernelAbstractions.jl#937, which does the same for KernelAbstractions' own tuning, AK'sreduce_tuningandscan_tuningare the largest invalidation trees when loading Metal, andreduce_tuningthe largest one when loading CUDA.This PR looks the tunings up through
tuning(hook, backend, T), which calls the hook behindBase.inferencebarrierand asserts the result type that each hook declares withtuning_type. Every internal call goes through it, and the hooks and extensions are unchanged. Tunings can already differ per device, so no caller relied on their values being constants.Unique method instances invalidated by
using <backend>, KernelAbstractions, and the time a fresh process takes to load the back end and run its first broadcast,getindex,sum,sum(; dims),accumulate,sort,anyand 4-d broadcast (Julia 1.13.2, with the PRs above):On Metal, the first broadcast takes 0.05 s instead of 0.72 s, as it no longer recompiles, and the first
sum0.13 s instead of 0.22 s. AK has no AMDGPU extension, and oneAPI.jl doesn't precompile AK's operations, so neither changes.The values no longer constant-fold, so where a tuning selects between launches, such as
reduce_then_scanand the sub-group width of scans, the host code for both is compiled. That makes the firstaccumulate0.05–0.1 s slower (0.90 → 1.00 s on the M1). Each lookup also takes a dynamic dispatch, which is small next to the launches that follow it.The tests pass on CUDA, AMDGPU and Metal.