Use relocatable IR linked at load time

CUDACore/Project.toml

@@ -54,7 +54,7 @@ ChainRulesCore = "1"
 EnzymeCore = "0.8.2"
 ExprTools = "0.1"
 GPUArrays = "11.5.4"
-GPUCompiler = "2.1"
+GPUCompiler = "2.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/relocations"}

CUDACore/src/compiler/compilation.jl

@@ -21,6 +21,9 @@ const AnyCUDAJob = CompilerJob{PTXCompilerTarget, <:AbstractCUDACompilerParams}
 const minreq = (; ptx=v"8.0", sm=sm"50")
 
 GPUCompiler.runtime_module(@nospecialize(job::AnyCUDAJob)) = CUDACore
+# keep host references symbolic and emit them as patchable globals; `link_kernel`
+# patches every site after loading the module
+GPUCompiler.relocation_lowering(@nospecialize(job::AnyCUDAJob)) = :patch
 
 # filter out functions from libdevice and cudadevrt
 GPUCompiler.isintrinsic(@nospecialize(job::AnyCUDAJob), fn::String) =
@@ -189,18 +192,21 @@ end
 # 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
+# conditionally session-portable artifacts (the cubin `image`, entry point, and relocations)
+# 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}
+    relocations::GPUCompiler.Relocations
 
     # session-local kernel handles, linear-scanned by context; usually holds a single entry
-    kernels::Vector{Tuple{CuContext,CuFunction}}
+    kernels::Vector{Tuple{CuContext,CuFunction,Vector{Any}}}
 
-    CUDACompilerResults() = new(nothing, nothing, Tuple{CuContext,CuFunction}[])
+    CUDACompilerResults() = new(nothing, nothing, GPUCompiler.Relocations(),
+                                Tuple{CuContext,CuFunction,Vector{Any}}[])
 end
 
 # cache of compiler configurations, per device (but additionally configurable via kwargs)
@@ -539,14 +545,26 @@ function compile(@nospecialize(job::CompilerJob))
         rm(ptxas_output)
     end
 
-    return (image, entry=LLVM.name(meta.entry))
+    return (image, entry=LLVM.name(meta.entry), relocations=meta.relocations)
 end
 
 # link a compiled image into a session-local `CuFunction` on the active context
-function link_kernel(@nospecialize(job::CompilerJob), image::Vector{UInt8}, entry::String)
+function link_kernel(image::Vector{UInt8}, entry::String,
+                     relocs::GPUCompiler.Relocations)
     # load as an executable kernel object on the current context
     mod = CuModule(image)
-    CuFunction(mod, entry)
+    relocations = GPUCompiler.resolved_relocations(relocs)
+    for (site, value) in relocations.sites
+        ptr_ref = Ref{CuPtr{Cvoid}}()
+        size_ref = Ref{Csize_t}()
+        cuModuleGetGlobal_v2(ptr_ref, size_ref, mod, site.name)
+        site.offset + sizeof(UInt) <= size_ref[] ||
+            error("Relocation '$(site.name)+$(site.offset)' is outside its " *
+                  "$(size_ref[])-byte global")
+        value_ref = Ref(value)
+        cuMemcpyHtoD_v2(ptr_ref[] + site.offset, value_ref, sizeof(UInt))
+    end
+    return CuFunction(mod, entry), relocations.roots
 end
 
 # look up the cached compilation artifacts for `job`, running the compiler on a miss.
@@ -560,9 +578,16 @@ function compile_or_lookup(@nospecialize(job::CompilerJob))::CUDACompilerResults
     res = GPUCompiler.cached_results(CUDACompilerResults, job)
     if res === nothing || res.image === nothing || GPUCompiler.compile_hook[] !== nothing
         compiled = compile(job)
-        res = @something res GPUCompiler.cached_results(CUDACompilerResults, job)
+        if GPUCompiler.supports_relocatable_ir()
+            res = @something res GPUCompiler.cached_results(CUDACompilerResults, job)
+        else
+            # The CodeInstance may be serialized, but generated code from a Julia runtime
+            # without a reliable GV table is session-bound.
+            res = CUDACompilerResults()
+        end
         res.image = compiled.image
         res.entry = compiled.entry
+        res.relocations = compiled.relocations
     end
     return res
 end

CUDACore/src/compiler/execution.jl

@@ -474,6 +474,7 @@ struct HostKernel{F,TT} <: AbstractKernel{F,TT}
     f::F
     fun::CuFunction
     state::KernelState
+    roots::Vector{Any}
 end
 
 @doc (@doc AbstractKernel) HostKernel
@@ -580,19 +581,22 @@ function cufunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT}
         # the scan is almost always over a single entry (one `===` comparison).
         ctx = cuda.context
         fun = nothing
-        for (cached_ctx, cached_fun) in res.kernels
+        roots = Any[]
+        for (cached_ctx, cached_fun, cached_roots) in res.kernels
             if cached_ctx === ctx
                 fun = cached_fun
+                roots = cached_roots
                 break
             end
         end
         if fun === nothing
-            fun = link_kernel(job, res.image::Vector{UInt8}, res.entry::String)
+            fun, roots = link_kernel(res.image::Vector{UInt8}, res.entry::String,
+                                     res.relocations)
             # 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))
+                push!(res.kernels, (ctx, fun, roots))
             end
         end
 
@@ -604,7 +608,7 @@ function cufunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT}
             # create the kernel state object
             state = KernelState(create_exceptions!(fun.mod), UInt32(0))
 
-            kernel = HostKernel{F,tt}(f, fun, state)
+            kernel = HostKernel{F,tt}(f, fun, state, roots)
             _kernel_instances[key] = kernel
         end
         return kernel::HostKernel{F,tt}

CUDACore/src/precompile.jl

@@ -45,7 +45,7 @@ end
 # kernel launch infrastructure
 let CUDACompilerJob = CompilerJob{PTXCompilerTarget, CUDACompilerParams}
     precompile(Tuple{typeof(cufunction), typeof(identity), Type{Tuple{Nothing}}})
-    precompile(Tuple{typeof(link_kernel), CUDACompilerJob, Vector{UInt8}, String})
+    precompile(Tuple{typeof(link_kernel), Vector{UInt8}, String, GPUCompiler.Relocations})
 
     # GPUCompiler 2.0 caching pipeline (specialized for CUDACore's results struct)
     precompile(Tuple{typeof(compile_or_lookup), CUDACompilerJob})

CUDATools/src/reflection.jl

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

test/core/codegen.jl

@@ -255,6 +255,46 @@ end
 
 ############################################################################################
 
+@testset "relocation patching" begin
+    runtime_relocs = @eval module $(gensym())
+        using ..CUDACore
+        const relocated_type_tag = CUDACore.GPUCompiler.Runtime.type_tag
+        relocation_kernel(out) = (out[1] = relocated_type_tag(Val(:float32)); return)
+    end
+
+    out = CUDA.zeros(UInt, 1)
+    @cuda threads=1 runtime_relocs.relocation_kernel(out)
+    expected = UInt(unsafe_load(cglobal(:jl_float32_type, Ptr{UInt})))
+    @test Array(out)[] == expected
+
+    value_relocs = @eval module $(gensym())
+        const relocated_symbol = Symbol("value#global")
+        relocated_symbol_kernel(out, name) =
+            (out[1] = UInt(name === relocated_symbol); return)
+    end
+
+    out = CUDA.zeros(UInt, 1)
+    @cuda threads=1 value_relocs.relocated_symbol_kernel(out, Symbol("value#global"))
+    @test Array(out)[] == 1
+
+    interior_relocs = @eval module $(gensym())
+        @noinline produce(cond::Bool, value::Int32) = cond ? value : 1.5
+        function interior_relocation_kernel(out, cond::Bool, value::Int32)
+            x = produce(cond, value)
+            out[1] = UInt(x isa Float64)
+            return
+        end
+    end
+
+    out = CUDA.zeros(UInt, 1)
+    kernel = @cuda launch=false interior_relocs.interior_relocation_kernel(out, false, Int32(1))
+    kernel(out, false, Int32(1); threads=1)
+    @test Array(out)[] == 1
+    @test any(root -> root === Float64, kernel.roots)
+end
+
+############################################################################################
+
 @testset "SASS" begin
 
 @testset "basic reflection" begin