Switch to CompilerCaching.jl.

CUDACore/Project.toml

@@ -54,7 +54,7 @@ ChainRulesCore = "1"
 EnzymeCore = "0.8.2"
 ExprTools = "0.1"
 GPUArrays = "11.5.4"
-GPUCompiler = "1.19"
+GPUCompiler = "2"
 GPUToolbox = "3"
 KernelAbstractions = "0.9.38"
 LLVM = "9.6"
@@ -78,3 +78,6 @@ julia = "1.10"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+
+[sources]
+GPUCompiler = {url = "https://github.com/JuliaGPU/GPUCompiler.jl", rev = "tb/compilercaching"}

CUDACore/src/compiler/compilation.jl

@@ -166,15 +166,21 @@ end
 
 ## compiler implementation (cache, configure, compile, and link)
 
-# cache of compilation caches, per context
-const _compiler_caches = Dict{CuContext, Dict{Any, CuFunction}}();
-function compiler_cache(ctx::CuContext)
-    cache = get(_compiler_caches, ctx, nothing)
-    if cache === nothing
-        cache = Dict{Any, CuFunction}()
-        _compiler_caches[ctx] = cache
-    end
-    return cache
+# GPUCompiler 2.0 caching: back-ends attach a mutable results struct to each cached
+# `CodeInstance` (on Julia 1.11+ this is Julia's integrated code cache, which also persists
+# artifacts through package precompilation; on 1.10 it's a session-local store). We keep
+# session-portable artifacts (the cubin `image` and entry-point `entry`) separate from the
+# session-local `CuFunction` handles, which are context-specific and must not be serialized
+# into a package image.
+mutable struct CUDACompilerResults
+    # session-portable artifacts (safe to persist across sessions)
+    image::Union{Nothing,Vector{UInt8}}
+    entry::Union{Nothing,String}
+
+    # session-local kernel handles, linear-scanned by context; usually holds a single entry
+    kernels::Vector{Tuple{CuContext,CuFunction}}
+
+    CUDACompilerResults() = new(nothing, nothing, Tuple{CuContext,CuFunction}[])
 end
 
 # cache of compiler configurations, per device (but additionally configurable via kwargs)
@@ -497,12 +503,28 @@ function compile(@nospecialize(job::CompilerJob))
     return (image, entry=LLVM.name(meta.entry))
 end
 
-# link into an executable kernel
-function link(@nospecialize(job::CompilerJob), compiled)
-    # load as an executable kernel object
-    ctx = context()
-    mod = CuModule(compiled.image)
-    CuFunction(mod, compiled.entry)
+# link a compiled image into a session-local `CuFunction` on the active context
+function link_kernel(@nospecialize(job::CompilerJob), image::Vector{UInt8}, entry::String)
+    # load as an executable kernel object on the current context
+    mod = CuModule(image)
+    CuFunction(mod, entry)
+end
+
+# look up the cached compilation artifacts for `job`, running the compiler on a miss.
+#
+# Storage is managed by `GPUCompiler.cached_results`: Julia's integrated code cache on 1.11+
+# (which also persists artifacts through precompilation), or a session-local store on 1.10.
+# `image === nothing` identifies a freshly-created `CUDACompilerResults` that hasn't been
+# compiled yet; the `compile_hook` check additionally forces the compile path so that
+# reflection consumers (`@device_code_*`) observe the compilation even on a cache hit.
+function compile_or_lookup(@nospecialize(job::CompilerJob))::CUDACompilerResults
+    res = GPUCompiler.cached_results(CUDACompilerResults, job)
+    if res.image === nothing || GPUCompiler.compile_hook[] !== nothing
+        compiled = compile(job)
+        res.image = compiled.image
+        res.entry = compiled.entry
+    end
+    return res
 end
 
 

CUDACore/src/compiler/execution.jl

@@ -449,11 +449,32 @@ function cufunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT}
     cuda = active_state()
 
     Base.@lock cufunction_lock begin
-        # compile the function
-        cache = compiler_cache(cuda.context)
+        # look up (or generate) the compilation artifacts for this function
         source = methodinstance(F, tt)
         config = compiler_config(cuda.device; kwargs...)::CUDACompilerConfig
-        fun = GPUCompiler.cached_compilation(cache, source, config, compile, link)
+        job = CompilerJob(source, config)
+        res = compile_or_lookup(job)::CUDACompilerResults
+
+        # resolve the `CuFunction` for the active context. this is a session-local handle,
+        # so it lives in the results struct's linear cache rather than being persisted;
+        # the scan is almost always over a single entry (one `===` comparison).
+        ctx = cuda.context
+        fun = nothing
+        for (cached_ctx, cached_fun) in res.kernels
+            if cached_ctx === ctx
+                fun = cached_fun
+                break
+            end
+        end
+        if fun === nothing
+            fun = link_kernel(job, res.image::Vector{UInt8}, res.entry::String)
+            # don't cache session-local handles while generating output: 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, fun))
+            end
+        end
 
         # 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

CUDACore/src/device/runtime.jl

@@ -4,9 +4,6 @@
 
 import Base.Sys: WORD_SIZE
 
-# reset the runtime cache from global scope, so that any change triggers recompilation
-GPUCompiler.reset_runtime()
-
 # load or build the runtime for the most likely compilation jobs
 function precompile_runtime()
     f = ()->return

CUDACore/src/precompile.jl

@@ -42,13 +42,14 @@ if :NVPTX in LLVM.backends()
 end
 
 # kernel launch infrastructure
-precompile(Tuple{typeof(cufunction), typeof(identity), Type{Tuple{Nothing}}})
-precompile(Tuple{typeof(link), CompilerJob, NamedTuple{(:image, :entry), Tuple{Vector{UInt8}, String}}})
+let CUDACompilerJob = CompilerJob{PTXCompilerTarget, CUDACompilerParams}
+    precompile(Tuple{typeof(cufunction), typeof(identity), Type{Tuple{Nothing}}})
+    precompile(Tuple{typeof(link_kernel), CUDACompilerJob, Vector{UInt8}, String})
 
-# GPUCompiler compilation pipeline (specialized for CUDACore's compile/link)
-precompile(Tuple{typeof(GPUCompiler.actual_compilation),
-    Dict{Any, CuFunction}, Core.MethodInstance, UInt64,
-    CUDACompilerConfig, typeof(compile), typeof(link)})
+    # GPUCompiler 2.0 caching pipeline (specialized for CUDACore's results struct)
+    precompile(Tuple{typeof(compile_or_lookup), CUDACompilerJob})
+    precompile(Tuple{typeof(GPUCompiler.cached_results), Type{CUDACompilerResults}, CUDACompilerJob})
+end
 
 # scalar reference (used by cuBLAS for alpha/beta parameters)
 precompile(Tuple{Type{CuRefValue{Float32}}, Float32})

CUDATools/Project.toml

@@ -28,7 +28,7 @@ CUDACore = "=6.2.0"
 CUDA_Compiler_jll = "0.3, 0.4"
 CUPTI = "=6.2.0"
 Crayons = "4"
-GPUCompiler = "1.4"
+GPUCompiler = "2"
 LLVM = "9.3.1"
 NVML = "=6.2.0"
 NVTX = "1"

CUDATools/src/reflection.jl

@@ -76,7 +76,7 @@ function code_sass(io::IO, job::CompilerJob; raw::Bool=false)
     end
 
     compiled = CUDACore.compile(job)
-    CUPTI.@enable! cfg CUDACore.link(job, compiled)
+    CUPTI.@enable! cfg CUDACore.link_kernel(job, compiled.image, compiled.entry)
 
     return
 end