Make GPUCompiler IR relocatable across Julia sessions

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUCompiler"
 uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
-version = "2.1.1"
+version = "2.2.0"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
 
 [workspace]

src/GPUCompiler.jl

@@ -54,6 +54,7 @@ include("mangling.jl")
 
 # compiler interface and implementations
 include("interface.jl")
+include("relocation.jl")
 include("error.jl")
 include("native.jl")
 include("ptx.jl")
@@ -84,13 +85,10 @@ include("precompile.jl")
 
 function __init__()
     STDERR_HAS_COLOR[] = get(stderr, :color, false)
+    empty!(session_results_cache)
 
     @static if !HAS_INTEGRATED_CACHE
-        # session-local results keyed by CodeInstance; entries serialized during
-        # GPUCompiler's own precompilation can never be valid in a later session
-        empty!(legacy_job_results)
-        # ditto for the in-process CodeCaches: CIs deposited by our own precompile
-        # workload carry world ages from the precompilation process
+        # CodeInstances created by GPUCompiler's precompile workload are process-local.
         empty!(GLOBAL_CI_CACHES)
     end
 

src/deprecated.jl

@@ -229,35 +229,6 @@ function CC.findsup(@nospecialize(sig::Type), table::StackedMethodTable)
 end
 
 
-## 1.10 `cached_results`
-#
-# Session-local storage for the per-job results structs; on 1.11+ these live on the
-# `CodeInstance`s of Julia's integrated cache instead (see `interface.jl`). Keep the
-# same identity here by associating results with a foreign `CodeInstance`: unrelated
-# world-age advances can reuse a still-valid CI, while invalidation makes the lookup
-# resolve to a new CI and therefore a new results struct.
-struct LegacyJobResultEntry
-    config::CompilerConfig
-    value::Any
-end
-
-const legacy_job_results = IdDict{CodeInstance,Vector{LegacyJobResultEntry}}()
-const job_results_lock = ReentrantLock()
-
-function job_results(::Type{V}, ci::CodeInstance, config::CompilerConfig) where {V}
-    Base.@lock job_results_lock begin
-        entries = get!(legacy_job_results, ci) do
-            LegacyJobResultEntry[]
-        end
-        for entry in entries
-            entry.config === config && entry.value isa V && return entry.value::V
-        end
-        v = V()
-        push!(entries, LegacyJobResultEntry(config, v))
-        return v
-    end
-end
-
 function job_code_instance(@nospecialize(job::CompilerJob))
     cache = WorldView(get_code_cache(job), job.world, job.world)
     CC.get(cache, job.source, nothing)
@@ -268,18 +239,9 @@ end
 function cached_results(::Type{V}, job::CompilerJob) where {V}
     ci = job_code_instance(job)
     ci === nothing && return nothing
-    return job_results(V, ci, job.config)
+    return session_results(V, ci, job.config)
 end
 
-
-## 1.10 session-dependent results
-#
-# Nothing to wipe: `legacy_job_results` never persists meaningfully across sessions (its
-# CodeInstance keys are session-specific, and our own image's entries are cleared in
-# `__init__`; entries written by a downstream package's workload don't make it into that
-# package's image at all, cf. cross-image mutation loss).
-mark_session_dependent!(@nospecialize(job::CompilerJob)) = nothing
-
 end # !HAS_INTEGRATED_CACHE
 
 

src/driver.jl

@@ -50,6 +50,10 @@ export compile
 
 Compile a `job` to one of the following formats as specified by the `target` argument:
 `:llvm` for LLVM IR, `:asm` for assembly, or `:obj` for object code.
+
+The default [`relocation_lowering`](@ref) strategy resolves Julia-value relocations in the
+`:llvm` result. Other strategies retain relocation metadata for their loader; `:defer`
+consumers resolve it with [`apply_relocations!`](@ref).
 """
 function compile(target::Symbol, @nospecialize(job::CompilerJob))
     if compile_hook[] !== nothing
@@ -58,7 +62,8 @@ function compile(target::Symbol, @nospecialize(job::CompilerJob))
     return compile_unhooked(target, job)
 end
 
-function compile_unhooked(output::Symbol, @nospecialize(job::CompilerJob))
+function compile_unhooked(output::Symbol, @nospecialize(job::CompilerJob);
+                          resolve_relocations::Bool=true)
     if context(; throw_error=false) === nothing
         error("No active LLVM context. Use `JuliaContext()` do-block syntax to create one.")
     end
@@ -73,7 +78,7 @@ function compile_unhooked(output::Symbol, @nospecialize(job::CompilerJob))
 
     ## LLVM IR
 
-    ir, ir_meta = emit_llvm(job)
+    ir, ir_meta = emit_llvm(job; resolve_relocations)
 
     if output == :llvm
         if job.config.strip
@@ -93,7 +98,7 @@ function compile_unhooked(output::Symbol, @nospecialize(job::CompilerJob))
     else
         error("Unknown assembly format $output")
     end
-    asm, asm_meta = emit_asm(job, ir, format)
+    asm, asm_meta = emit_asm(job, ir, ir_meta.relocations, format)
 
     if output == :asm || output == :obj
         return asm, (; asm_meta..., ir_meta..., ir)
@@ -169,7 +174,8 @@ end
 
 const __llvm_initialized = Ref(false)
 
-@locked function emit_llvm(@nospecialize(job::CompilerJob))
+@locked function emit_llvm(@nospecialize(job::CompilerJob);
+                           resolve_relocations::Bool=true)
     if !__llvm_initialized[]
         InitializeAllTargets()
         InitializeAllTargetInfos()
@@ -180,7 +186,7 @@ const __llvm_initialized = Ref(false)
     end
 
     @tracepoint "IR generation" begin
-        ir, compiled, gv_to_value = irgen(job)
+        ir, compiled, relocations = irgen(job)
         if job.config.entry_abi === :specfunc
             entry_fn = compiled[job.source].specfunc
         else
@@ -245,11 +251,9 @@ const __llvm_initialized = Ref(false)
                     dyn_ir, dyn_meta = @invokelatest deferred_codegen(dyn_job, job)
                     dyn_entry_fn = LLVM.name(dyn_meta.entry)
                     merge!(compiled, dyn_meta.compiled)
-                    if haskey(dyn_meta, :gv_to_value)
-                        merge!(gv_to_value, dyn_meta.gv_to_value)
-                    end
                     @assert context(dyn_ir) == context(ir)
-                    link!(ir, dyn_ir)
+                    link_relocatable!(ir, relocations, dyn_ir,
+                                      dyn_meta.relocations)
                     changed = true
                     dyn_entry_fn
                 end
@@ -292,7 +296,7 @@ const __llvm_initialized = Ref(false)
     if job.config.toplevel && job.config.libraries
         # load the runtime outside of a timing block (because it recurses into the compiler)
         if !uses_julia_runtime(job)
-            runtime = load_runtime(job)
+            runtime, runtime_relocs = load_runtime(job)
         end
 
         @tracepoint "Library linking" begin
@@ -301,7 +305,8 @@ const __llvm_initialized = Ref(false)
 
             # GPU run-time library
             if !uses_julia_runtime(job)
-                @tracepoint "runtime library" link!(ir, runtime; only_needed=true)
+                @tracepoint "runtime library" link_relocatable!(
+                    ir, relocations, runtime, runtime_relocs; only_needed=true)
             end
         end
     end
@@ -334,13 +339,16 @@ const __llvm_initialized = Ref(false)
 
             finish_linked_module!(job, ir)
 
-            # Materialize isbits and Bool boxes; bake addresses for other objects.
-            portable = relocate_gvs!(ir, gv_to_value)
-            portable || mark_session_dependent!(job)
+            # Resolve early so optimization sees concrete values.
+            resolve_early = resolve_relocations && relocation_lowering(job) === :bake
+            if resolve_early
+                prune_dead_relocations!(ir, relocations)
+                bake_relocations!(ir, relocations)
+            end
 
             if job.config.optimize
                 @tracepoint "optimization" begin
-                    optimize!(job, ir; job.config.opt_level)
+                    optimize!(job, ir, relocations; job.config.opt_level)
 
                     # deferred codegen has some special optimization requirements,
                     # which also need to happen _after_ regular optimization.
@@ -362,6 +370,12 @@ const __llvm_initialized = Ref(false)
                 end
             end
 
+            # Runtime linking during optimization can add relocations.
+            if resolve_early && !isempty(relocations.sites)
+                prune_dead_relocations!(ir, relocations)
+                bake_relocations!(ir, relocations)
+            end
+
             if job.config.cleanup
                 @tracepoint "clean-up" begin
                     @dispose pb=NewPMPassBuilder() begin
@@ -375,6 +389,9 @@ const __llvm_initialized = Ref(false)
                 end
             end
 
+            # Do not expose dead sites to loaders.
+            resolve_early || prune_dead_relocations!(ir, relocations)
+
             # optimization may have replaced functions, so look the entry point up again
             entry = functions(ir)[entry_fn]
 
@@ -405,26 +422,31 @@ const __llvm_initialized = Ref(false)
 
     if job.config.toplevel && job.config.validate
         @tracepoint "validation" begin
-            check_ir(job, ir)
+            check_ir(job, ir, relocations)
         end
     end
 
     if should_verify()
         @tracepoint "verification" verify(ir)
     end
 
-    return ir, (; entry, compiled, gv_to_value)
+    return ir, (; entry, compiled, relocations)
 end
 
+# Compatibility for back-ends that resolve relocations during `emit_llvm`.
+emit_asm(@nospecialize(job::CompilerJob), ir::LLVM.Module,
+         format::LLVM.API.LLVMCodeGenFileType) =
+    emit_asm(job, ir, Relocations(), format)
+
 @locked function emit_asm(@nospecialize(job::CompilerJob), ir::LLVM.Module,
-                          format::LLVM.API.LLVMCodeGenFileType)
+                          relocs::Relocations, format::LLVM.API.LLVMCodeGenFileType)
     # NOTE: strip after validation to get better errors
     if job.config.strip
         @tracepoint "Debug info removal" strip_debuginfo!(ir)
     end
 
     @tracepoint "LLVM back-end" begin
-        @tracepoint "preparation" prepare_execution!(job, ir)
+        @tracepoint "preparation" prepare_execution!(job, ir, relocs)
 
         code = @tracepoint "machine-code generation" mcgen(job, ir, format)
     end