Switch to CompilerCaching.jl

Project.toml

@@ -31,9 +31,9 @@ SPIRVIntrinsics = {path = "lib/intrinsics"}
 [compat]
 Adapt = "4"
 GPUArrays = "11.2.1"
-GPUCompiler = "1.7.1"
+GPUCompiler = "2"
 KernelAbstractions = "0.9.38"
-LLVM = "9.1"
+LLVM = "9.6"
 LinearAlgebra = "1"
 OpenCL_jll = "=2024.10.24"
 Preferences = "1"

src/OpenCL.jl

@@ -37,6 +37,7 @@ include("compiler/capabilities.jl")
 include("compiler/compilation.jl")
 include("compiler/execution.jl")
 include("compiler/reflection.jl")
+include("compiler/precompile.jl")
 
 # integrations and specialized functionality
 include("util.jl")

src/compiler/compilation.jl

@@ -12,6 +12,32 @@ end
 const OpenCLCompilerConfig = CompilerConfig{SPIRVCompilerTarget, OpenCLCompilerParams}
 const OpenCLCompilerJob = CompilerJob{SPIRVCompilerTarget,OpenCLCompilerParams}
 
+"""
+    OpenCLResults
+
+Cached compilation results for an OpenCL kernel job, managed by
+`GPUCompiler.cached_results`. Fields are populated through the compile pipeline:
+`obj` (SPIR-V bytes) + `entry` + `device_rng` after codegen, and `kernels` after the
+session-local link onto an OpenCL context. The first three are session-portable
+(cached through precompilation, except when GPUCompiler marks the job
+session-dependent and wipes its entries before image serialization); `kernels` is
+session-local and never populated during precompilation. `obj === nothing`
+identifies a job that has not been compiled yet.
+
+`kernels` is a small linear cache of `(cl.Context, cl.Kernel)` pairs. The cache partition
+already covers everything that affects codegen via `GPUCompiler.cache_owner`, so the only
+runtime-visible dimension left is the OpenCL context that owns the linked `cl.Kernel`.
+A linear scan with `===` is fastest in the common case (n=1) and stays cheap for the
+rare workload that bounces between a handful of contexts on the same device.
+"""
+mutable struct OpenCLResults
+    obj::Union{Nothing, Vector{UInt8}}                   # SPIR-V binary
+    entry::Union{Nothing, String}
+    device_rng::Bool
+    kernels::Vector{Tuple{cl.Context, cl.Kernel}}        # session-local; linear-scanned
+    OpenCLResults() = new(nothing, nothing, false, Tuple{cl.Context, cl.Kernel}[])
+end
+
 GPUCompiler.runtime_module(::CompilerJob{<:Any,OpenCLCompilerParams}) = OpenCL
 
 GPUCompiler.method_table_view(job::OpenCLCompilerJob) =
@@ -125,18 +151,7 @@ function GPUCompiler.finish_linked_module!(@nospecialize(job::OpenCLCompilerJob)
     return
 end
 
-## compiler implementation (cache, configure, compile, and link)
-
-# cache of compilation caches, per context
-const _compiler_caches = Dict{cl.Context, Dict{Any, Any}}()
-function compiler_cache(ctx::cl.Context)
-    cache = get(_compiler_caches, ctx, nothing)
-    if cache === nothing
-        cache = Dict{Any, Any}()
-        _compiler_caches[ctx] = cache
-    end
-    return cache
-end
+## compiler implementation (configure, compile, and link)
 
 # cache of compiler configurations, per device (but additionally configurable via kwargs)
 const _toolchain = Ref{Any}()
@@ -157,34 +172,35 @@ end
 const SPIRV_VERSION = v"1.4"
 
 @noinline function _compiler_config(dev, backend; kernel=true, name=nothing, always_inline=false,
-                                     sub_group_size::Union{Nothing,Int}=_sub_group_size(dev), kwargs...)
+                                     sub_group_size::Union{Nothing,Int}=_sub_group_size(dev),
+                                     extensions::AbstractVector{<:AbstractString}=String[], kwargs...)
     supports_fp16 = "cl_khr_fp16" in dev.extensions
     supports_fp64 = "cl_khr_fp64" in dev.extensions
 
     if sub_group_size !== nothing && !("cl_intel_required_subgroup_size" in dev.extensions)
         error("Device does not support cl_intel_required_subgroup_size")
     end
 
+    spirv_ext = join(("+$ext" for ext in extensions), ",")
+
     # create GPUCompiler objects
     target = SPIRVCompilerTarget(; version=SPIRV_VERSION, supports_fp16, supports_fp64,
-                                   validate=true, kwargs...)
+                                   validate=true, extensions=spirv_ext, kwargs...)
     params = OpenCLCompilerParams(; sub_group_size, features=device_features(dev),
                                     program_backend=backend)
     CompilerConfig(target, params; kernel, name, always_inline)
 end
 
-# compile to executable machine code
+# run inference + LLVM codegen + SPIR-V emission. returns `(obj, entry, device_rng)`,
+# all session-portable so they survive precompilation when stored on a cached `CodeInstance`.
 const compilations = Threads.Atomic{Int}(0)
-function compile(@nospecialize(job::CompilerJob))
-    compilations[] += 1
+function compile_to_obj(@nospecialize(job::CompilerJob))
+    Threads.atomic_add!(compilations, 1)
 
-    # TODO: this creates a context; cache those.
-    obj, meta = JuliaContext() do ctx
+    JuliaContext() do ctx
         obj, meta = GPUCompiler.compile(:obj, job)
-
         entry = LLVM.name(meta.entry)
         device_rng = StringAttribute("julia.opencl.rng", "") in collect(function_attributes(meta.entry))
-
         (; obj, entry, device_rng)
     end
 end
@@ -289,9 +305,9 @@ function dump_artifacts(artifacts::Pair{String}...)
     return paths
 end
 
-# link into an executable kernel
-function link(@nospecialize(job::CompilerJob), compiled)
-    spirv_bitcode = compiled.obj
+# link the SPIR-V bytes into a session-local `cl.Kernel` on the active context.
+function link_kernel(@nospecialize(job::CompilerJob), obj::Vector{UInt8}, entry::String)
+    spirv_bitcode = obj
     clc_source = nothing
     build_options = ""
 
@@ -342,5 +358,5 @@ function link(@nospecialize(job::CompilerJob), compiled)
         msg *= "\nIf you think this is a bug, please file an issue and attach $(join(files, " and "))"
         error(msg)
     end
-    (; kernel=cl.Kernel(prog, compiled.entry), compiled.device_rng)
+    return cl.Kernel(prog, entry)
 end

src/compiler/execution.jl

@@ -180,27 +180,57 @@ end
 const clfunction_lock = ReentrantLock()
 
 function clfunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT}
-    ctx = cl.context()
-    dev = cl.device()
-
     Base.@lock clfunction_lock begin
-        # compile the function
-        cache = compiler_cache(ctx)
+        config = compiler_config(cl.device(); kwargs...)::OpenCLCompilerConfig
         source = methodinstance(F, tt)
-        config = compiler_config(dev; kwargs...)::OpenCLCompilerConfig
-        linked = GPUCompiler.cached_compilation(cache, source, config, compile, link)
-
-        # create a callable object that captures the function instance. we don't need to think
-        # about world age here, as GPUCompiler already does and will return a different object
-        h = hash(linked.kernel, hash(f, hash(tt)))
-        kernel = get(_kernel_instances, h, nothing)
+        job = CompilerJob(source, config)
+
+        res = compile_or_lookup(job)::OpenCLResults
+
+        # Resolve the cl.Kernel for the active context. Linear scan over the
+        # session-local cache; almost always n=1, so this is one `===` compare.
+        ctx = cl.context()
+        kernel = nothing
+        @inbounds for (cached_ctx, cached_kernel) in res.kernels
+            if cached_ctx === ctx
+                kernel = cached_kernel
+                break
+            end
+        end
         if kernel === nothing
-            # create the kernel state object
-            kernel = HostKernel{F,tt}(f, linked.kernel, linked.device_rng)
-            _kernel_instances[h] = kernel
+            kernel = link_kernel(job, res.obj::Vector{UInt8}, res.entry::String)
+            # Don't cache session-local kernel handles while precompiling: the
+            # results struct is serialized into the package image along with its
+            # CodeInstance, and the handles would come back dangling.
+            if ccall(:jl_generating_output, Cint, ()) != 1
+                push!(res.kernels, (ctx, kernel))
+            end
         end
-        return kernel::HostKernel{F,tt}
+
+        h = hash(kernel, hash(f, hash(tt)))
+        get!(_kernel_instances, h) do
+            HostKernel{F,tt}(f, kernel, res.device_rng)
+        end::HostKernel{F,tt}
+    end
+end
+
+# Look up cached compile artifacts for `job`, compiling on miss. Storage is managed
+# by `GPUCompiler.cached_results` (Julia's integrated code cache on 1.11+, which also
+# persists artifacts through precompilation; a session-local store on 1.10).
+#
+# `obj === nothing` identifies an `OpenCLResults` that hasn't been compiled yet. The
+# `compile_hook` check additionally forces the compile path so reflection-style
+# consumers (`@device_code_*`) observe the compilation even on a cache hit.
+function compile_or_lookup(@nospecialize(job::CompilerJob))::OpenCLResults
+    res = GPUCompiler.cached_results(OpenCLResults, job)
+    if res === nothing || res.obj === nothing || GPUCompiler.compile_hook[] !== nothing
+        compiled = compile_to_obj(job)
+        res = @something res GPUCompiler.cached_results(OpenCLResults, job)
+        res.obj = compiled.obj
+        res.entry = compiled.entry
+        res.device_rng = compiled.device_rng
     end
+    return res
 end
 
 # cache of kernel instances

src/compiler/precompile.jl

@@ -0,0 +1,6 @@
+# Cache plumbing reached on every kernel launch. These specializations need no OpenCL device
+# and keep the first launch from paying for GPUCompiler's generic results machinery itself.
+let Job = OpenCLCompilerJob
+    precompile(Tuple{typeof(compile_or_lookup), Job})
+    precompile(Tuple{typeof(GPUCompiler.cached_results), Type{OpenCLResults}, Job})
+end

src/device/runtime.jl

@@ -1,6 +1,3 @@
-# reset the runtime cache from global scope, so that any change triggers recompilation
-GPUCompiler.reset_runtime()
-
 signal_exception() = return
 
 malloc(sz) = C_NULL

test/Project.toml

@@ -1,8 +1,10 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+CompilerCaching = "9db33cc3-5358-4881-8759-fa4194144afd"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 IOCapture = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"

test/execution.jl

@@ -167,3 +167,45 @@ end
         @test occursin("54321", ir) && !occursin("12345", ir)
     end
 end
+
+@testset "compilation cache" begin
+    mod = @eval module $(gensym())
+        @noinline child() = return
+        kernel() = child()
+    end
+
+    count() = OpenCL.compilations[]
+    launch() = @opencl mod.kernel()
+
+    # the initial launch compiles
+    n = count()
+    Base.invokelatest(launch)
+    @test count() == n+1
+
+    # a second launch hits the cache
+    Base.invokelatest(launch)
+    @test count() == n+1
+
+    # jobs differing only in codegen-level settings get their own artifacts...
+    OpenCL.clfunction(mod.kernel, Tuple{}; name="custom")
+    @test count() == n+2
+    # ... which are cached as well
+    OpenCL.clfunction(mod.kernel, Tuple{}; name="custom")
+    @test count() == n+2
+
+    # reflection observes already-compiled kernels (by forcing recompilation,
+    # which must leave the cached entry in a usable state)
+    @test !isempty(sprint(io->(@device_code_llvm io=io Base.invokelatest(launch))))
+    n = count()
+    Base.invokelatest(launch)
+    @test count() == n
+
+    # redefining the kernel recompiles...
+    @eval mod kernel() = (child(); child())
+    Base.invokelatest(launch)
+    @test count() == n+1
+    # ... as does redefining a callee
+    @eval mod @noinline child() = nothing
+    Base.invokelatest(launch)
+    @test count() == n+2
+end