Add cached_lookup to cache kernel lookups per signature

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

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