Switch to CompilerCaching.jl.

.buildkite/pipeline.yml

@@ -14,6 +14,7 @@ steps:
       julia --project -e '
         println("--- :julia: Instantiating project")
         using Pkg
+        Pkg.add(url="https://github.com/JuliaGPU/GPUCompiler.jl", rev="main")
         Pkg.instantiate()
         Pkg.build()
 
@@ -35,14 +36,18 @@ steps:
     plugins:
       - JuliaCI/julia#v1:
           version: "1.10"
-      - JuliaCI/julia-test#v1:
       - JuliaCI/julia-coverage#v1:
           codecov: true
     agents:
       queue: "rocm"
       rocmgpu: "*"
     if: build.message !~ /\[skip tests\]/
-    command: "julia --project -e 'using Pkg; Pkg.update()'"
+    command: |
+      git clone --depth 1 --branch main https://github.com/JuliaGPU/GPUCompiler.jl GPUCompiler
+      julia --project -e '
+        using Pkg
+        Pkg.develop(path="GPUCompiler")
+        Pkg.test(; coverage=true)'
     timeout_in_minutes: 90
     env:
       JULIA_NUM_THREADS: 4
@@ -132,13 +137,16 @@ steps:
     plugins:
       - JuliaCI/julia#v1:
           version: "1.10"
-      - JuliaCI/julia-test#v1:
-          test_args: "enzyme"
     agents:
       queue: "rocm"
       rocmgpu: "*"
     if: build.message !~ /\[skip tests\]/
-    command: "julia --project -e 'using Pkg; Pkg.update()'"
+    command: |
+      git clone --depth 1 --branch main https://github.com/JuliaGPU/GPUCompiler.jl GPUCompiler
+      julia --project -e '
+        using Pkg
+        Pkg.develop(path="GPUCompiler")
+        Pkg.test(; test_args=["enzyme"])'
     timeout_in_minutes: 45
     env:
       JULIA_NUM_THREADS: 4
@@ -168,10 +176,11 @@ steps:
     plugins:
       - JuliaCI/julia#v1:
           version: "1.10"
-      - JuliaCI/julia-test#v1:
-          run_tests: false
     command: |
+      git clone --depth 1 --branch main https://github.com/JuliaGPU/GPUCompiler.jl GPUCompiler
       julia --project -e '
+        using Pkg
+        Pkg.develop(path="GPUCompiler")
         using AMDGPU
         @assert !AMDGPU.functional()'
     agents:

Project.toml

@@ -49,6 +49,9 @@ AMDGPUEnzymeCoreExt = "EnzymeCore"
 AMDGPUSparseMatricesCSRExt = "SparseMatricesCSR"
 AMDGPUSpecialFunctionsExt = "SpecialFunctions"
 
+[sources]
+GPUCompiler = {url = "https://github.com/JuliaGPU/GPUCompiler.jl", rev = "main"}
+
 [compat]
 AbstractFFTs = "1.0"
 AcceleratedKernels = "0.3.1, 0.4"
@@ -60,7 +63,7 @@ ChainRulesCore = "1"
 EnzymeCore = "0.8"
 ExprTools = "0.1"
 GPUArrays = "11.3.1"
-GPUCompiler = "1"
+GPUCompiler = "2"
 GPUToolbox = "0.1.0, 0.2, 0.3, 1, 2, 3"
 KernelAbstractions = "0.9.2"
 LLD_jll = "15, 16, 17, 18, 19, 20, 21.1"

src/compiler/codegen.jl

@@ -12,15 +12,33 @@ end
 const HIPCompilerConfig = CompilerConfig{GCNCompilerTarget, HIPCompilerParams}
 const HIPCompilerJob = CompilerJob{GCNCompilerTarget, HIPCompilerParams}
 
-const _hip_compiler_cache = Dict{HIP.HIPDevice, Dict{Any, HIP.HIPFunction}}()
+"""
+    HIPResults
+
+Cached compilation results for a HIP kernel job, managed by
+`GPUCompiler.cached_results`. Session-portable artifacts (the lld-linked shared
+object `obj`, the entry-point name `entry`, and the detected `global_hostcalls`)
+are populated after codegen and persist across sessions (e.g. through package
+precompilation). The session-local `functions` are `HIPFunction` handles linked
+onto a specific device; they are device-specific and never populated during
+precompilation. `obj === nothing` identifies a job that has not been compiled yet.
+
+`functions` is a small linear cache of `(HIPDevice, HIPFunction)` pairs, matching the
+old per-device cache semantics; the scan is almost always over a single entry.
+"""
+mutable struct HIPResults
+    # session-portable artifacts
+    obj::Union{Nothing,Vector{UInt8}}       # lld-linked shared object
+    entry::Union{Nothing,String}
+    global_hostcalls::Vector{Symbol}
+    # session-local handles (never populated during precompilation)
+    functions::Vector{Tuple{HIP.HIPDevice,HIP.HIPFunction}}
+    HIPResults() = new(nothing, nothing, Symbol[], Tuple{HIP.HIPDevice,HIP.HIPFunction}[])
+end
 
 # hash(fun, hash(f, hash(tt))) => HIPKernel
 const _kernel_instances = Dict{UInt, Runtime.HIPKernel}()
 
-function compiler_cache(dev::HIP.HIPDevice)
-    get!(() -> Dict{UInt, Any}(), _hip_compiler_cache, dev)
-end
-
 GPUCompiler.runtime_module(@nospecialize(::HIPCompilerJob)) = AMDGPU
 
 GPUCompiler.method_table(@nospecialize(::HIPCompilerJob)) = AMDGPU.method_table
@@ -168,12 +186,33 @@ The following kwargs are supported:
 function hipfunction(f::F, tt::TT = Tuple{}; kwargs...) where {F <: Core.Function, TT}
     Base.@lock hipfunction_lock begin
         dev = AMDGPU.device()
-        cache = compiler_cache(dev)
         config = compiler_config(dev; kwargs...)
 
         source = methodinstance(F, tt)
-        fun = GPUCompiler.cached_compilation(
-            cache, source, config, hipcompile, hiplink)
+        job = CompilerJob(source, config)
+        res = compile_or_lookup(job)
+
+        # Resolve the `HIPFunction` for the active device. 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, matching the old
+        # per-device cache (`==` compare, as `HIPDevice` was the Dict key before).
+        fun = nothing
+        for (cached_dev, cached_fun) in res.functions
+            if cached_dev == dev
+                fun = cached_fun
+                break
+            end
+        end
+        if fun === nothing
+            fun = hiplink(job, res.obj::Vector{UInt8}, res.entry::String,
+                          res.global_hostcalls)
+            # 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.functions, (dev, fun))
+            end
+        end
 
         h = hash(fun, hash(f, hash(tt)))
         kernel = get!(_kernel_instances, h) do
@@ -183,6 +222,25 @@ function hipfunction(f::F, tt::TT = Tuple{}; kwargs...) where {F <: Core.Functio
     end
 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. `obj === nothing` identifies a freshly-created `HIPResults` 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))::HIPResults
+    res = GPUCompiler.cached_results(HIPResults, job)
+    if res === nothing || res.obj === nothing || GPUCompiler.compile_hook[] !== nothing
+        compiled = hipcompile(job)
+        res = @something res GPUCompiler.cached_results(HIPResults, job)
+        res.obj = compiled.obj
+        res.entry = compiled.entry
+        res.global_hostcalls = compiled.global_hostcalls
+    end
+    return res
+end
+
 function create_executable(obj)
     lld = if AMDGPU.lld_artifact
         `$(LLD_jll.lld()) -flavor gnu`
@@ -250,8 +308,8 @@ function hipcompile(@nospecialize(job::CompilerJob))
     (; obj=create_executable(codeunits(obj)), entry, global_hostcalls)
 end
 
-function hiplink(@nospecialize(job::CompilerJob), compiled)
-    (; obj, entry, global_hostcalls) = compiled
+# link a compiled shared object into a session-local `HIPFunction` on the active device.
+function hiplink(@nospecialize(job::CompilerJob), obj, entry, global_hostcalls)
     mod = HIP.HIPModule(obj)
     HIP.HIPFunction(mod, entry, global_hostcalls)
 end

src/device/runtime.jl

@@ -2,9 +2,6 @@ using Core: LLVMPtr
 
 ## GPU runtime library
 
-# reset the runtime cache from global scope, so that any change triggers recompilation
-GPUCompiler.reset_runtime()
-
 @inline @generated kernel_state() = GPUCompiler.kernel_state_value(AMDGPU.KernelState)
 
 @generated function llvm_atomic_cas(ptr::LLVMPtr{T,A}, cmp::T, val::T) where {T, A}

src/precompile.jl

@@ -18,9 +18,9 @@ if :AMDGPU in LLVM.backends()
             end
 
             # Build a device-free compiler config for a baseline GCN target.
-            # `gfx1030` (RDNA2, wavefront 32) is a portable baseline that exercises
-            # the full pipeline; the cached *code* is reused regardless of the
-            # actual device's ISA at runtime (only the kernel binary differs).
+            # `gfx1030` (RDNA2, wavefront 32) is a representative baseline that exercises
+            # the full pipeline. GPU artifacts remain correctly partitioned by target ISA;
+            # the reusable benefit here is precompiled host-side compiler machinery.
             #
             # NOTE: the ISA must be RDNA/CDNA, not pre-RDNA. The `wavefrontsize*`
             # LLVM features only exist on gfx10+, so pairing them with e.g. gfx900
@@ -52,9 +52,9 @@ if :AMDGPU in LLVM.backends()
             # MIs into native compilation, causing LLVM errors. Guard like CUDA.jl.
             @static if VERSION >= v"1.12-"
                 if !instrumented
-                    GPUCompiler.JuliaContext() do ctx
-                        GPUCompiler.compile(:obj, job)
-                    end
+                    # Exercise the same compile-or-lookup path used by kernel launches, and
+                    # attach its artifact to the package-image CI.
+                    Compiler.compile_or_lookup(job)
 
                     # The compile above runs during precompilation, when ROCm
                     # discovery (`__init__`) has NOT run, so `libdevice_libs` is
@@ -72,14 +72,17 @@ end
 
 # Kernel launch infrastructure that the workload above cannot reach, because it
 # requires a live device (mirrors CUDA.jl's explicit precompile directives:
-# `cufunction`, `link`, and `actual_compilation`).
-precompile(Tuple{typeof(Compiler.hipfunction), typeof(identity), Type{Tuple{Nothing}}})
-precompile(Tuple{typeof(GPUCompiler.actual_compilation),
-    Dict{Any, HIP.HIPFunction}, Core.MethodInstance, UInt64,
-    Compiler.HIPCompilerConfig, typeof(Compiler.hipcompile), typeof(Compiler.hiplink)})
-precompile(Tuple{typeof(Compiler.hiplink), Compiler.HIPCompilerJob,
-    NamedTuple{(:obj, :entry, :global_hostcalls),
-        Tuple{Vector{UInt8}, String, Vector{Symbol}}}})
+# `cufunction`, `link_kernel`, `compile_or_lookup`, and `cached_results`).
+let HIPCompilerJob = Compiler.HIPCompilerJob
+    precompile(Tuple{typeof(Compiler.hipfunction), typeof(identity), Type{Tuple{Nothing}}})
+    precompile(Tuple{typeof(Compiler.hiplink), HIPCompilerJob,
+        Vector{UInt8}, String, Vector{Symbol}})
+
+    # GPUCompiler 2.0 caching pipeline (specialized for AMDGPU's results struct)
+    precompile(Tuple{typeof(Compiler.compile_or_lookup), HIPCompilerJob})
+    precompile(Tuple{typeof(GPUCompiler.cached_results),
+        Type{Compiler.HIPResults}, HIPCompilerJob})
+end
 
 # Hot entry points of the bundled ROCm libraries, mirroring CUDA.jl's per-library
 # precompile directives. These compile the (GPU-free) Julia wrappers so the first