Make GPUCompiler IR relocatable across Julia sessions

Author: maleadt

src/driver.jl

@@ -93,7 +93,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.host_references, format)
 
     if output == :asm || output == :obj
         return asm, (; asm_meta..., ir_meta..., ir)
@@ -192,6 +192,7 @@ const __llvm_initialized = Ref(false)
     # finalize the current module. this needs to happen before linking deferred modules,
     # since those modules have been finalized themselves, and we don't want to re-finalize.
     entry = finish_module!(job, ir, entry)
+    host_references = classify_gvs!(ir, gv_to_value)
 
     # deferred code generation
     has_deferred_jobs = job.config.toplevel && !job.config.only_entry &&
@@ -245,11 +246,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_with_host_references!(ir, host_references, dyn_ir,
+                                               dyn_meta.host_references)
                     changed = true
                     dyn_entry_fn
                 end
@@ -292,7 +291,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_references = load_runtime(job)
         end
 
         @tracepoint "Library linking" begin
@@ -301,7 +300,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_with_host_references!(
+                    ir, host_references, runtime, runtime_references; only_needed=true)
             end
         end
     end
@@ -334,13 +334,12 @@ 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)
+            host_references.embedded_pointer |= !materialize_bool_singletons!(ir)
+            host_references.embedded_pointer && mark_session_dependent!(job)
 
             if job.config.optimize
                 @tracepoint "optimization" begin
-                    optimize!(job, ir; job.config.opt_level)
+                    optimize!(job, ir, host_references; job.config.opt_level)
 
                     # deferred codegen has some special optimization requirements,
                     # which also need to happen _after_ regular optimization.
@@ -405,26 +404,27 @@ const __llvm_initialized = Ref(false)
 
     if job.config.toplevel && job.config.validate
         @tracepoint "validation" begin
-            check_ir(job, ir)
+            check_ir(job, ir, host_references)
         end
     end
 
     if should_verify()
         @tracepoint "verification" verify(ir)
     end
 
-    return ir, (; entry, compiled, gv_to_value)
+    return ir, (; entry, compiled, host_references)
 end
 
 @locked function emit_asm(@nospecialize(job::CompilerJob), ir::LLVM.Module,
-                          format::LLVM.API.LLVMCodeGenFileType)
+                          refs::HostReferences, 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, refs)
+        refs.embedded_pointer && mark_session_dependent!(job)
 
         code = @tracepoint "machine-code generation" mcgen(job, ir, format)
     end

src/interface.jl

@@ -429,6 +429,82 @@ the check-and-compile sequence (all back-ends already do, e.g. `mtlfunction_lock
 """
 function cached_results end
 
+"""
+    JuliaValueRef(value)
+
+A Julia value used as the serializable identity of a host reference. Resolve it in the active
+session with [`resolve_host_reference`](@ref); a loader that writes the resulting address into
+device storage must keep `value` rooted for as long as that storage remains reachable.
+"""
+struct JuliaValueRef
+    value::Any
+end
+
+"""
+    CGlobalRef(symbol)
+
+A libjulia C data global. Resolving this reference calls `jl_cglobal` for `symbol` and loads
+the word stored at that address.
+"""
+struct CGlobalRef
+    symbol::Symbol
+end
+
+"""
+    HostReference
+
+A serializable source for a host-derived word: either a [`JuliaValueRef`](@ref) or a
+[`CGlobalRef`](@ref).
+"""
+const HostReference = Union{JuliaValueRef,CGlobalRef}
+
+"""
+    HostReferences(slots, embedded_pointer)
+
+Host-reference metadata accompanying a module. `slots` maps each object-code symbol to the
+source of the word that must be written there. `embedded_pointer` conservatively records that
+an unavoidable session pointer has been embedded, so the artifact cannot be serialized into a
+package image even if `slots` is empty.
+"""
+mutable struct HostReferences
+    slots::Dict{String,HostReference}
+    embedded_pointer::Bool
+end
+
+HostReferences() = HostReferences(Dict{String,HostReference}(), false)
+copy_host_references(refs::HostReferences) =
+    HostReferences(copy(refs.slots), refs.embedded_pointer)
+
+same_host_reference(a::JuliaValueRef, b::JuliaValueRef) = a.value === b.value
+same_host_reference(a::CGlobalRef, b::CGlobalRef) = a.symbol === b.symbol
+same_host_reference(::HostReference, ::HostReference) = false
+
+"""
+    resolve_host_reference(ref) -> UInt
+
+Resolve `ref` to its current-session word. Loader-based backends use this after loading the
+object and before exposing a callable function.
+"""
+resolve_host_reference(ref::JuliaValueRef) =
+    UInt(ccall(:jl_value_ptr, Ptr{Cvoid}, (Any,), ref.value))
+function resolve_host_reference(ref::CGlobalRef)
+    address = ccall(:jl_cglobal, Any, (Any, Any), String(ref.symbol), UInt)
+    return unsafe_load(address)
+end
+
+"""
+    lower_host_references!(job, mod, refs)
+
+Lower compiler-managed Julia value globals and Julia runtime globals for `job`'s final
+backend representation. The default lowering resolves every live reference in the current
+Julia session. Backends with a loader relocation mechanism may call
+[`emit_host_reference_slots!`](@ref) instead, then resolve and patch those slots after loading.
+"""
+function lower_host_references!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
+                                refs::HostReferences)
+    resolve_host_reference_slots!(mod, refs)
+end
+
 @static if HAS_INTEGRATED_CACHE
 
 """
@@ -550,6 +626,8 @@ end
 end # HAS_INTEGRATED_CACHE
 
 @public GPUCompilerCacheToken, cache_owner, cached_results
+@public JuliaValueRef, CGlobalRef, HostReference, HostReferences
+@public lower_host_references!, emit_host_reference_slots!, resolve_host_reference
 
 # the method table to use
 #

src/irgen.jl

@@ -43,21 +43,36 @@ function irgen(@nospecialize(job::CompilerJob))
         end
     end
 
-    # sanitize global values (Julia doesn't when using the external codegen policy)
-    for val in [collect(globals(mod)); collect(functions(mod))]
-        isdeclaration(val) && continue
+    # Julia value globals can be declarations. Sanitize every definition as well as every
+    # mapped global, but leave unrelated external declarations alone because their names are
+    # part of their ABI.
+    for val in collect(globals(mod))
         old_name = LLVM.name(val)
+        mapped = haskey(gv_to_value, old_name)
+        isdeclaration(val) && !mapped && continue
         new_name = safe_name(old_name)
         if old_name != new_name
             LLVM.name!(val, new_name)
-            val = get(gv_to_value, old_name, nothing)
-            if val !== nothing
-                delete!(gv_to_value, old_name)
-                gv_to_value[new_name] = val
+            actual_name = LLVM.name(val)
+            if mapped
+                init = pop!(gv_to_value, old_name)
+                haskey(gv_to_value, actual_name) &&
+                    error("Duplicate Julia value global '$actual_name'")
+                gv_to_value[actual_name] = init
             end
         end
     end
 
+    # sanitize defined functions (Julia doesn't when using the external codegen policy)
+    for val in collect(functions(mod))
+        isdeclaration(val) && continue
+        old_name = LLVM.name(val)
+        new_name = safe_name(old_name)
+        if old_name != new_name
+            LLVM.name!(val, new_name)
+        end
+    end
+
     # rename and process the entry point
     if job.config.name !== nothing
         LLVM.name!(entry, safe_name(job.config.name))