Materialize isbits boxed constants as device globals

src/driver.jl

@@ -334,11 +334,11 @@ const __llvm_initialized = Ref(false)
 
             finish_linked_module!(job, ir)
 
-            relocate_gvs!(ir, gv_to_value)
-
-            # the IR (and any artifact derived from it) now embeds session-absolute
-            # addresses; keep results for this job out of a package image
-            isempty(gv_to_value) || mark_session_dependent!(job)
+            # resolve `julia.constgv` slots: isbits constants are materialized as
+            # device-resident globals, everything else gets this session's absolute
+            # address baked in; only the latter makes the result session-dependent
+            portable = relocate_gvs!(ir, gv_to_value)
+            portable || mark_session_dependent!(job)
 
             if job.config.optimize
                 @tracepoint "optimization" begin

src/interface.jl

@@ -508,7 +508,8 @@ end
 ## session-dependent results
 #
 # Some compilation results embed session-specific data: `relocate_gvs!` bakes absolute
-# pointers into the IR of toplevel jobs that reference `julia.constgv` globals, and any
+# pointers into the IR of toplevel jobs that reference `julia.constgv` globals (except
+# for slots it can materialize as session-portable device constants), and any
 # artifact a back-end derives from that IR (metallib, SPIR-V, ...) inherits them. Such
 # results must not survive into a package image, while remaining available for
 # within-session lookups during the precompilation process itself. Julia wipes its own

src/jlgen.jl

@@ -755,32 +755,100 @@ end
 """
     relocate_gvs!(mod::LLVM.Module, gv_to_value::Dict{String, Ptr{Cvoid}})
 
-Bake absolute pointer values into the initializers of `julia.constgv`-tagged
-global variables, matching them by name against `gv_to_value`. Only GVs that
-are declarations (as produced by `compile_method_instance`, which strips the
-initializers for session portability) or still have a null initializer are
-touched: preserving existing initializers covers the older-Julia path where
-Julia itself emits pointer values directly.
+Resolve `julia.constgv`-tagged global variable slots, matching them by name
+against `gv_to_value`. Slots whose object is a non-ghost, non-`Bool` isbits
+value are *materialized*: a device-resident constant replica of the box
+(header word + payload bytes) is emitted and the slot points at its payload,
+so device code can dereference it and LLVM can constant-fold it. All other
+slots get the object's absolute host address baked in, as before; device code
+only uses those as opaque identity tokens.
+
+Only GVs that are declarations (as produced by `compile_method_instance`,
+which strips the initializers for session portability) or still have a null
+initializer are touched: preserving existing initializers covers the
+older-Julia path where Julia itself emits pointer values directly.
 
 GVs present in `mod` but missing from `gv_to_value` remain declarations, which
 back-ends will reject loudly (undefined symbol) rather than silently folding
 to null.
+
+Returns `true` if the module is session-portable afterwards: no absolute host
+address was written (neither a baked slot nor a materialized header carrying
+a non-smalltag type pointer).
 """
 function relocate_gvs!(mod::LLVM.Module, gv_to_value::Dict{String, Ptr{Cvoid}})
+    portable = true
     mod_gvs = globals(mod)
     for (name, init) in gv_to_value
         haskey(mod_gvs, name) || continue
         gv = mod_gvs[name]
         cur = initializer(gv)
-        if cur === nothing || LLVM.isnull(cur)
+        if !(cur === nothing || LLVM.isnull(cur))
+            # pre-baked by Julia itself (pre-1.13): also session-absolute
+            portable = false
+            continue
+        end
+
+        val = nothing
+        if init != C_NULL
+            obj = Base.unsafe_pointer_to_objref(init)
+            # ghosts/singletons and Bools are compared by pointer; keep their
+            # identity by baking the host address (they are never dereferenced)
+            if isbitstype(typeof(obj)) && sizeof(obj) > 0 && !(obj isa Bool)
+                val, hdr = materialize_box!(mod, gv, obj, init)
+                # non-smalltag headers carry a host DataType pointer
+                portable &= hdr < UInt(64 << 4)   # jl_max_tags << 4
+            end
+        end
+        if val === nothing
             val = const_inttoptr(ConstantInt(Int64(init)), global_value_type(gv))
-            initializer!(gv, val)
-            # re-internalize what compile_method_instance demoted to an external
-            # declaration; with the address baked in, the optimizer can now fold it
-            linkage!(gv, LLVM.API.LLVMPrivateLinkage)
+            portable = false
         end
+        initializer!(gv, val)
+        # re-internalize what compile_method_instance demoted to an external
+        # declaration; with the value in place, the optimizer can now fold it
+        linkage!(gv, LLVM.API.LLVMPrivateLinkage)
     end
-    return
+    return portable
+end
+
+# emit a device-resident constant replica of the box holding `obj`; returns
+# the constant to store in the slot, and the (gcbits-masked) header word
+function materialize_box!(mod::LLVM.Module, gv::GlobalVariable, @nospecialize(obj),
+                          init::Ptr{Cvoid})
+    W = sizeof(Int)
+    hdr, bytes = GC.@preserve obj begin
+        # the header word transparently yields the smalltag immediate for
+        # smalltag types and the host type pointer otherwise; drop the gcbits
+        hdr = unsafe_load(Ptr{UInt}(init - W)) & ~UInt(15)
+        bytes = [unsafe_load(Ptr{UInt8}(init), i) for i in 1:sizeof(obj)]
+        hdr, bytes
+    end
+
+    T_word = LLVM.IntType(8W)
+    T_byte = LLVM.Int8Type()
+    fields = LLVM.Constant[ConstantInt(T_word, hdr), ConstantDataArray(T_byte, bytes)]
+    payload_idx = 1
+    if Base.datatype_alignment(typeof(obj)) > W
+        # pad so the payload lands at a 16-byte offset (JL_HEAP_ALIGNMENT max)
+        pushfirst!(fields, ConstantDataArray(T_byte, zeros(UInt8, 16 - W)))
+        payload_idx = 2
+    end
+    boxinit = ConstantStruct(fields)
+    boxty = value_type(boxinit)
+
+    box = GlobalVariable(mod, boxty, safe_name(LLVM.name(gv)) * "_box")
+    initializer!(box, boxinit)
+    constant!(box, true)
+    linkage!(box, LLVM.API.LLVMPrivateLinkage)
+    alignment!(box, 16)
+    unnamed_addr!(box, true)
+
+    idx(i) = ConstantInt(LLVM.Int32Type(), i)
+    payload = const_gep(boxty, box, LLVM.Constant[idx(0), idx(payload_idx)])
+    slotty = global_value_type(gv)
+    val = value_type(payload) == slotty ? payload : const_addrspacecast(payload, slotty)
+    return val, hdr
 end
 
 # partially revert JuliaLang/julia#49391 — see #527

src/rtlib.jl

@@ -102,12 +102,13 @@ function emit_function!(mod, config::CompilerConfig, source::MethodInstance, met
     # runtime functions may reference Julia objects through `julia.constgv` globals (e.g.
     # `box_bool` returning the `jl_true`/`jl_false` singletons). Kernels get theirs
     # relocated when the fully-linked toplevel module is finalized, but the runtime's
-    # mappings would be dropped along with the rest of its per-function metadata: bake
-    # the session-absolute addresses into the cached bitcode instead, and keep such
-    # functions out of package images (their pointers don't survive into other sessions).
+    # mappings would be dropped along with the rest of its per-function metadata: resolve
+    # the slots into the cached bitcode instead, and keep functions whose resolution baked
+    # session-absolute addresses (rather than materialized session-portable constants)
+    # out of package images.
     if !isempty(meta.gv_to_value)
-        relocate_gvs!(new_mod, meta.gv_to_value)
-        mark_session_dependent!(rt_job)
+        portable = relocate_gvs!(new_mod, meta.gv_to_value)
+        portable || mark_session_dependent!(rt_job)
     end
 
     # rename to the final `gpu_*` name on the per-function module, so the cached bitcode

test/native.jl

@@ -299,6 +299,100 @@ end
         end
     end
 
+    @testset "boxed constant materialization" begin
+        # since JuliaLang/julia#55045, isbits union constants stay fully boxed; the
+        # box must be replicated on device instead of referencing a host address
+        mod = @eval module $(gensym())
+            union_smalltag(cond::Bool, a::Int32) = cond ? a : Int64(0)
+            union_float(cond::Bool, a::Float32) = cond ? a : 1.0
+            union_bool(cond::Bool, a::Int32) = cond ? a : true
+            union_ghost(cond::Bool, a::Int32) = cond ? a : nothing
+            function kernel(p::Ptr{Int64}, cond::Bool, a::Int32)
+                x = cond ? a : Int64(0)
+                unsafe_store!(p, Int64(x))
+                return
+            end
+            function egal_kernel(p::Ptr{Bool}, cond::Bool, a::Int32)
+                x = cond ? a : Int64(0)
+                unsafe_store!(p, x === Int64(0))
+                return
+            end
+        end
+
+        # smalltag constants materialize fully session-portably
+        ir = sprint(io->Native.code_llvm(io, mod.union_smalltag, Tuple{Bool, Int32};
+                                         dump_module=true, validate=true))
+        @static if VERSION >= v"1.14.0-DEV.1348"
+            @test occursin("_box", ir)
+            @test !occursin("inttoptr", ir)
+        end
+
+        # non-smalltag constants carry a host type pointer in the box header,
+        # but the payload is still device-resident
+        ir = sprint(io->Native.code_llvm(io, mod.union_float, Tuple{Bool, Float32};
+                                         dump_module=true, validate=true))
+        @static if VERSION >= v"1.14.0-DEV.1348"
+            @test occursin("_box", ir)
+        end
+
+        # identity objects (Bool singletons, ghosts) are never replicated
+        ir = sprint(io->Native.code_llvm(io, mod.union_bool, Tuple{Bool, Int32};
+                                         dump_module=true, validate=true))
+        @test !occursin("_box", ir)
+        ir = sprint(io->Native.code_llvm(io, mod.union_ghost, Tuple{Bool, Int32};
+                                         dump_module=true, validate=true))
+        @test !occursin("_box", ir)
+
+        # kernel compilation, including bits-egal on the materialized leaf
+        Native.code_execution(mod.kernel, (Ptr{Int64}, Bool, Int32))
+        Native.code_execution(mod.egal_kernel, (Ptr{Bool}, Bool, Int32))
+
+        # relocate_gvs! reports whether the module stayed session-portable
+        JuliaContext() do ctx
+            # Unlike Int128, vector-shaped tuples are 16-byte aligned on all
+            # supported architectures and Julia versions.
+            aligned = (VecElement(Int64(1)), VecElement(Int64(2)))
+            @test Base.datatype_alignment(typeof(aligned)) > sizeof(Int)
+            objs = Any[Int64(42), 1.25, :sym, aligned]
+            # pointers to the heap boxes rooted in `objs` (passing an element
+            # through a specialized function would re-box, possibly on the stack)
+            ptrs = [ccall(:jl_value_ptr, Ptr{Cvoid}, (Any,), x) for x in objs]
+            function slot_module(ptr::Ptr{Cvoid})
+                llvm_mod = LLVM.Module("test")
+                name = "jl_global#0"
+                LLVM.GlobalVariable(llvm_mod, LLVM.PointerType(LLVM.Int8Type()), name)
+                llvm_mod, Dict(name => ptr)
+            end
+            GC.@preserve objs begin
+                # smalltag isbits: materialized, portable
+                m, map = slot_module(ptrs[1])
+                @test GPUCompiler.relocate_gvs!(m, map)
+                @test haskey(globals(m), "jl_global_0_box")
+                dispose(m)
+
+                # Float64: materialized, but the header carries a type pointer
+                m, map = slot_module(ptrs[2])
+                @test !GPUCompiler.relocate_gvs!(m, map)
+                @test haskey(globals(m), "jl_global_0_box")
+                dispose(m)
+
+                # Symbol: baked address
+                m, map = slot_module(ptrs[3])
+                @test !GPUCompiler.relocate_gvs!(m, map)
+                @test !haskey(globals(m), "jl_global_0_box")
+                @test occursin("inttoptr", string(m))
+                dispose(m)
+
+                # 16-byte-aligned payloads get padded past the header word
+                m, map = slot_module(ptrs[4])
+                GPUCompiler.relocate_gvs!(m, map)
+                box = globals(m)["jl_global_0_box"]
+                @test length(elements(LLVM.global_value_type(box))) == 3
+                dispose(m)
+            end
+        end
+    end
+
     @testset "allowed mutable types" begin
         # when types have no fields, we should always allow them
         mod = @eval module $(gensym())