diff --git a/src/execution.jl b/src/execution.jl index 0dcb3113..e37e9925 100644 --- a/src/execution.jl +++ b/src/execution.jl @@ -68,3 +68,129 @@ function assign_args!(code, _args) return vars, var_exprs end + + +## kernel lookup cache + +""" + cached_lookup(lookup, V, F, tt, key)::V + +Return the result of `lookup()`, reusing an earlier result for the same kernel signature +(`F`, `tt`) and `key`. This is meant for the functions that back-ends call on every kernel +launch to look up the compiled kernel (e.g. `cufunction`), which otherwise have to resolve +the method instance, build a `CompilerJob` and query the compiler's caches every time: + +```julia +function cufunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT} + fun, state = GPUCompiler.cached_lookup(Tuple{CuFunction,KernelState}, F, tt, + (context(), values(kwargs))) do + cufunction_link(methodinstance(F, tt); kwargs...) + end + return HostKernel{F,tt}(f, fun, state) +end +``` + +An earlier result is reused when it was looked up by the same type of `lookup`, in the same +world age (see `Base.tls_world_age`), and with a key that is `isequal` to `key`. Matching the +world exactly is what keeps this correct when methods are added or redefined: compiled code +can still be valid in a newer world, but the kernel call may now dispatch to another method. +In another world `lookup` runs again, which normally finds the compiled code in the +compiler's cache. + +Everything that affects the result of `lookup` besides the kernel signature and the world +must be part of `key`, e.g. the context or device and the compiler keyword arguments. Use +values with stable `isequal` semantics: mutating an object referenced by an earlier key does +not invalidate the results stored with it. `V` is the type of the result, which should be +concrete and should only contain compiled artifacts or session handles: do not return the +kernel callable `f` itself, as it holds the values captured by the kernel, which need to +come from each launch. The cache keeps neither `lookup` nor `f` alive, but it does keep +the keys and results of a few lookups per signature. + +Results are not stored while a compile hook is active (e.g. by the `@device_code_*` macros), +so that the hook sees every lookup (see [`run_compile_hook`](@ref)), or while generating +output during precompilation, as session handles cannot be serialized. Exceptions thrown by +`lookup` propagate, and nothing is stored. + +Hits don't take a lock. Concurrent misses may each run `lookup`, so it has to be thread-safe +itself, and stored results may be discarded at any point. +""" +@inline function cached_lookup(lookup::L, ::Type{V}, ::Type{F}, ::Type{TT}, + key::K) where {L,V,F,TT,K} + cache = lookup_cache(L, F, TT, K, V) + entry = @atomic :acquire cache.head + if entry !== nothing && entry.world == tls_world_age() && isequal(entry.key, key) && + !compile_hook_active() + return entry.value + end + return cached_lookup_slow(cache, lookup, key) +end + +@public cached_lookup + +# The cache is an immutable list of results, newest first, so that hits don't need a lock. +struct LookupCacheEntry{K,V} + world::UInt + key::K + value::V + next::Union{Nothing,LookupCacheEntry{K,V}} +end + +mutable struct LookupCache{K,V} + @atomic head::Union{Nothing,LookupCacheEntry{K,V}} + LookupCache{K,V}() where {K,V} = new{K,V}(nothing) +end + +# enough for a signature that is looked up for a couple of devices, or in a couple of worlds +const LOOKUP_CACHE_SIZE = 4 + +# A cache per lookup and signature, as a constant in the code of every caller. The generator +# only creates an empty cache; it may run several times for the same signature, creating +# several caches, which only costs some missed hits. Callers that are serialized during +# precompilation serialize the cache too, which is why it is never written to then. +@generated function lookup_cache(::Type{L}, ::Type{F}, ::Type{TT}, ::Type{K}, + ::Type{V}) where {L,F,TT,K,V} + cache = LookupCache{K,V}() + return :($cache) +end + +@noinline compile_hook_active() = active_compile_hook[] !== nothing + +# Don't specialize on the lookup, which is a different closure for every kernel. The world +# is the one `lookup` runs in, as it's called from the same task. +@noinline Base.@nospecializeinfer function cached_lookup_slow(cache::LookupCache{K,V}, + @nospecialize(lookup), + key::K) where {K,V} + world = tls_world_age() + head = @atomic :acquire cache.head + store = !compile_hook_active() && ccall(:jl_generating_output, Cint, ()) != 1 + if store + # the head was checked by the caller + entry = head === nothing ? nothing : head.next + while entry !== nothing + if entry.world == world && isequal(entry.key, key) + return entry.value + end + entry = entry.next + end + end + + value = lookup()::V + if store + # keep the most recent results, except for one with the same world and key that + # another task may have stored meanwhile + kept = LookupCacheEntry{K,V}[] + entry = @atomic :acquire cache.head + while entry !== nothing && length(kept) < LOOKUP_CACHE_SIZE - 1 + if entry.world != world || !isequal(entry.key, key) + push!(kept, entry) + end + entry = entry.next + end + tail = nothing + for entry in Iterators.reverse(kept) + tail = LookupCacheEntry{K,V}(entry.world, entry.key, entry.value, tail) + end + @atomic :release cache.head = LookupCacheEntry{K,V}(world, key, value, tail) + end + return value +end diff --git a/test/execution.jl b/test/execution.jl new file mode 100644 index 00000000..bc73790b --- /dev/null +++ b/test/execution.jl @@ -0,0 +1,256 @@ +@testset "cached_lookup" begin + mod = @eval module $(gensym()) + using GPUCompiler + + # stands in for a back-end's kernel lookup: returns the method that the kernel call + # dispatches to, which is what a compiler job for the kernel would be built from + const lookups = Ref(0) + function lookup_method(f::F, tt::TT=Tuple{}; key=nothing) where {F,TT} + GPUCompiler.cached_lookup(Method, F, tt, key) do + lookups[] += 1 + methodinstance(F, tt).def + end + end + + @noinline helper(x) = x + 1 + kernel(x::Number) = helper(x) + launch(x; key=nothing) = lookup_method(kernel, Tuple{typeof(x)}; key) + + other_lookup(f::F, tt::TT) where {F,TT} = + GPUCompiler.cached_lookup(() -> nothing, Method, F, tt, nothing) + other_launch(x) = other_lookup(kernel, Tuple{typeof(x)}) + end + + lookups() = mod.lookups[] + launch(x; kwargs...) = Base.invokelatest(mod.launch, x; kwargs...) + + m = launch(1) + @test m.sig == Tuple{typeof(mod.kernel), Number} + n = lookups() + @test launch(1) === m + @test lookups() == n + + # in another world the lookup runs again, but finds the same method + @eval mod unrelated() = nothing + @test launch(1) === m + @test lookups() == n + 1 + @test launch(1) === m + @test lookups() == n + 1 + + # redefining a callee doesn't change the method, but it is looked up again + @eval mod @noinline helper(x) = x + 2 + @test launch(1) === m + @test lookups() == n + 2 + + # a more specific method is used, although the code for the other one is still valid + world = Base.get_world_counter() + @eval mod kernel(x::Int) = -x + m′ = launch(1) + @test m′.sig == Tuple{typeof(mod.kernel), Int} + @test lookups() == n + 3 + + # lookups in an older world see the methods of that world, and are still cached + @test Base.invoke_in_world(world, mod.launch, 1) === m + @test launch(1) === m′ + @test lookups() == n + 3 + + # other argument types and keys are looked up separately + @test launch(1.0) === m + @test lookups() == n + 4 + @test launch(1; key=:a) === m′ + @test lookups() == n + 5 + @test launch(1; key=:b) === m′ + @test lookups() == n + 6 + for _ in 1:2, key in (:a, :b) + @test launch(1; key) === m′ + end + @test lookups() == n + 6 + + # keys are compared with `isequal` + @test launch(1; key=NaN) === m′ + @test launch(1; key=NaN) === m′ + @test lookups() == n + 7 + @test launch(1; key=-0.0) === m′ + @test launch(1; key=0.0) === m′ + @test lookups() == n + 9 + + # only the most recent results are kept + keys = 1:GPUCompiler.LOOKUP_CACHE_SIZE + for key in keys + launch(1; key) + end + n′ = lookups() + for key in keys + launch(1; key) + end + @test lookups() == n′ + launch(1; key=0) + launch(1; key=first(keys)) + @test lookups() == n′ + 2 + + # different lookups with the same types don't share results + @test launch(1) === m′ + @test_throws TypeError Base.invokelatest(mod.other_launch, 1) +end + +@testset "cached_lookup: compile hooks" begin + mod = @eval module $(gensym()) + using GPUCompiler + const lookups = Ref(0) + kernel() = nothing + function launch(key) + GPUCompiler.cached_lookup(Int, typeof(kernel), Tuple{}, key) do + lookups[] += 1 + GPUCompiler.run_compile_hook(key) + key + end + end + function hooked_launch(keys) + seen = Int[] + GPUCompiler.with_compile_hook(key -> push!(seen, key)) do + foreach(launch, keys) + end + return seen + end + end + + lookups() = mod.lookups[] + launch(key) = Base.invokelatest(mod.launch, key) + + @test launch(1) == 1 + n = lookups() + + # a hook sees every lookup, also of results that were cached before + @test Base.invokelatest(mod.hooked_launch, [1, 1, 2]) == [1, 1, 2] + @test lookups() == n + 3 + + # and nothing that is looked up while it is active is cached + @test launch(1) == 1 + @test lookups() == n + 3 + @test launch(2) == 2 + @test lookups() == n + 4 +end + +@testset "cached_lookup: lifetime and errors" begin + mod = @eval module $(gensym()) + using GPUCompiler + const lookups = Ref(0) + + mutable struct Callable + x::Int + end + (c::Callable)() = c.x + + # the result only depends on the type of the callable; its values come from the call + function launch(f::F, key=nothing) where {F} + code = GPUCompiler.cached_lookup(Int, F, Tuple{}, key) do + lookups[] += 1 + # capture the callable, which mustn't keep it alive + f.x + 100 + end + return code + f() + end + + # without leaving a reference to the callable on the caller's stack + @noinline function launch_new(finalized) + c = Callable(3) + finalizer(_ -> finalized[] = true, c) + launch(c, :new) + return + end + + const fail = Ref(false) + function failing_launch() + GPUCompiler.cached_lookup(Int, typeof(failing_launch), Tuple{}, nothing) do + lookups[] += 1 + fail[] && error("lookup failed") + 42 + end + end + end + + lookups() = mod.lookups[] + + @test Base.invokelatest(mod.launch, mod.Callable(1)) == 101 + n = lookups() + @test Base.invokelatest(mod.launch, mod.Callable(2)) == 102 + @test lookups() == n + + finalized = Threads.Atomic{Bool}(false) + Base.invokelatest(mod.launch_new, finalized) + # finalizers may run on another thread, or be delayed by a concurrent collection + for _ in 1:10 + GC.gc(true) + yield() + finalized[] && break + sleep(0.1) + end + @test finalized[] + + # failed lookups aren't cached + mod.fail[] = true + n = lookups() + @test_throws ErrorException Base.invokelatest(mod.failing_launch) + @test_throws ErrorException Base.invokelatest(mod.failing_launch) + @test lookups() == n + 2 + mod.fail[] = false + @test Base.invokelatest(mod.failing_launch) == 42 + @test Base.invokelatest(mod.failing_launch) == 42 + @test lookups() == n + 3 +end + +@testset "cached_lookup: concurrency" begin + mod = @eval module $(gensym()) + using GPUCompiler + const lookups = Threads.Atomic{Int}(0) + kernel() = nothing + function launch(key) + GPUCompiler.cached_lookup(Int, typeof(kernel), Tuple{}, key) do + Threads.atomic_add!(lookups, 1) + yield() + key + end + end + end + + # more keys than the cache holds, so that tasks keep replacing each other's results + results = Base.invokelatest() do + tasks = map(1:8 * Threads.nthreads()) do i + Threads.@spawn begin + ok = true + for j in 1:1000 + key = (i + j) % 8 + ok &= mod.launch(key) == key + end + ok + end + end + fetch.(tasks) + end + @test all(results) + @test mod.lookups[] > 0 +end + +precompile_test_harness("cached_lookup: precompilation") do load_path + write(joinpath(load_path, "LookupCacheUser.jl"), """ + module LookupCacheUser + using GPUCompiler + kernel() = nothing + lookup_kernel() = 42 + launch() = GPUCompiler.cached_lookup(lookup_kernel, Int, typeof(kernel), Tuple{}, nothing) + # serializes the code of `launch`, including its cache + launch() + end + """) + Base.compilecache(Base.PkgId("LookupCacheUser"), stderr, stdout) + + @eval let + using LookupCacheUser + # results looked up while precompiling, like session handles, don't survive + cache = GPUCompiler.lookup_cache(typeof(LookupCacheUser.lookup_kernel), + typeof(LookupCacheUser.kernel), Tuple{}, Nothing, Int) + @test (@atomic cache.head) === nothing + @test LookupCacheUser.launch() == 42 + end +end