From a62bb19fc0ef0aaca5825a93e4abbc1a5616b268 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 3 Jul 2026 08:45:02 +0200 Subject: [PATCH 1/2] Switch to CompilerCaching.jl. --- CUDACore/Project.toml | 5 ++- CUDACore/src/compiler/compilation.jl | 52 ++++++++++++++++++++-------- CUDACore/src/compiler/execution.jl | 27 +++++++++++++-- CUDACore/src/precompile.jl | 13 +++---- CUDATools/Project.toml | 2 +- CUDATools/src/reflection.jl | 2 +- 6 files changed, 74 insertions(+), 27 deletions(-) diff --git a/CUDACore/Project.toml b/CUDACore/Project.toml index 34cebbf1bb..42ac69350d 100644 --- a/CUDACore/Project.toml +++ b/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"} diff --git a/CUDACore/src/compiler/compilation.jl b/CUDACore/src/compiler/compilation.jl index 52dc2af0a3..6e126bb1e1 100644 --- a/CUDACore/src/compiler/compilation.jl +++ b/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 diff --git a/CUDACore/src/compiler/execution.jl b/CUDACore/src/compiler/execution.jl index 61dddc54c6..4711555b33 100644 --- a/CUDACore/src/compiler/execution.jl +++ b/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 diff --git a/CUDACore/src/precompile.jl b/CUDACore/src/precompile.jl index 8817aa74cc..3a61f987c1 100644 --- a/CUDACore/src/precompile.jl +++ b/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}) diff --git a/CUDATools/Project.toml b/CUDATools/Project.toml index 1003b32801..0cfdb9a488 100644 --- a/CUDATools/Project.toml +++ b/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" diff --git a/CUDATools/src/reflection.jl b/CUDATools/src/reflection.jl index 0d5df7cdea..874603358e 100644 --- a/CUDATools/src/reflection.jl +++ b/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 From 625eaec68361f1cc78ce5babee98972ad1be1322 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Mon, 6 Jul 2026 15:09:34 +0200 Subject: [PATCH 2/2] Remove the obsolete reset_runtime call. Co-Authored-By: Claude Fable 5 --- CUDACore/src/device/runtime.jl | 3 --- 1 file changed, 3 deletions(-) diff --git a/CUDACore/src/device/runtime.jl b/CUDACore/src/device/runtime.jl index 6739d27d68..bd041651f6 100644 --- a/CUDACore/src/device/runtime.jl +++ b/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