diff --git a/CUDACore/Project.toml b/CUDACore/Project.toml index ff56e1aded..d7605597d2 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 = "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"} diff --git a/CUDACore/src/compiler/compilation.jl b/CUDACore/src/compiler/compilation.jl index 07a3d5514e..8091aa12bc 100644 --- a/CUDACore/src/compiler/compilation.jl +++ b/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 diff --git a/CUDACore/src/compiler/execution.jl b/CUDACore/src/compiler/execution.jl index 9f0ee318be..704e812a9c 100644 --- a/CUDACore/src/compiler/execution.jl +++ b/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} diff --git a/CUDACore/src/precompile.jl b/CUDACore/src/precompile.jl index 6297bf7217..44c56fe203 100644 --- a/CUDACore/src/precompile.jl +++ b/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}) diff --git a/CUDATools/src/reflection.jl b/CUDATools/src/reflection.jl index 874603358e..af667b11f0 100644 --- a/CUDATools/src/reflection.jl +++ b/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 diff --git a/test/core/codegen.jl b/test/core/codegen.jl index 9ab5b5aa68..993568b9fe 100644 --- a/test/core/codegen.jl +++ b/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