Materialize isbits boxed constants as device globals

src/driver.jl

@@ -334,11 +334,9 @@ 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)
+            # Materialize isbits and Bool boxes; bake addresses for other objects.
+            portable = relocate_gvs!(ir, gv_to_value)
+            portable || mark_session_dependent!(job)
 
             if job.config.optimize
                 @tracepoint "optimization" begin

src/interface.jl

@@ -507,7 +507,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

@@ -759,32 +759,127 @@ 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.
-
-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.
+Resolve globals that refer to Julia objects. `jl_true`/`jl_false`, which are
+absent from `gv_to_value`, become canonical module-local boxes. Ordinary
+non-ghost isbits objects are also materialized; other objects keep their host
+address as an opaque identity token.
+
+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.
+
+Apart from the dedicated Bool globals, 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 = materialize_bool_singletons!(mod)
     mod_gvs = globals(mod)
     for (name, init) in gv_to_value
+        # Bools are resolved by name above.
+        name in ("jl_true", "jl_false") && continue
+
         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)
+            # Zero-sized objects remain identity tokens.
+            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
+
+# Bool JuliaVariables are absent from `gv_to_value`; define one device box per name.
+function materialize_bool_singletons!(mod::LLVM.Module)
+    portable = true
+    mod_gvs = globals(mod)
+    for (name, obj) in ("jl_true" => true, "jl_false" => false)
+        haskey(mod_gvs, name) || continue
+        gv = mod_gvs[name]
+        cur = initializer(gv)
+        if !(cur === nothing || LLVM.isnull(cur))
+            # Existing definitions may contain session-specific addresses.
+            portable = false
+            continue
+        end
+
+        init = ccall(:jl_value_ptr, Ptr{Cvoid}, (Any,), obj)
+        val, hdr = materialize_box!(mod, gv, obj, init)
+        initializer!(gv, val)
+        constant!(gv, true)
+        linkage!(gv, LLVM.API.LLVMPrivateLinkage)
+
+        # Stay conservative if Bool stops using a smalltag.
+        portable &= hdr < UInt(64 << 4)   # jl_max_tags << 4
+    end
+    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})
+    @assert isbitstype(typeof(obj)) && sizeof(obj) > 0
+
+    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/metal.jl

@@ -775,6 +775,8 @@ function add_global_address_spaces!(@nospecialize(job::CompilerJob), mod::LLVM.M
     # delete old globals
     for (old, new) in global_map
         prune_constexpr_uses!(old)
+        # Rewrite constant-expression uses left after cloning.
+        replace_uses!(old, new)
         @assert isempty(uses(old))
         replace_metadata_uses!(old, new)
         erase!(old)

src/rtlib.jl

@@ -98,15 +98,11 @@ function emit_function!(mod, config::CompilerConfig, source::MethodInstance, met
     # recent Julia versions include prototypes for all runtime functions, even if unused
     run!(StripDeadPrototypesPass(), new_mod, llvm_machine(config.target))
 
-    # 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).
+    # Resolve constgv mappings before their metadata is discarded. Dedicated Bool
+    # globals are resolved after linking into the toplevel module.
     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/metal.jl

@@ -368,6 +368,31 @@ end
     end
 end
 
+@testset "boxed Bool singleton relocation" begin
+    @static if VERSION >= v"1.14.0-DEV.1348"
+        mod = @eval module $(gensym())
+            @noinline produce(cond::Bool, a::Int32) = cond ? a : true
+            function consume(cond::Bool, a::Int32)
+                x = produce(cond, a)
+                x isa Bool && x && return Int32(1)
+                return Int32(0)
+            end
+            function kernel(ptr, cond::Bool, a::Int32)
+                unsafe_store!(ptr, consume(cond, a))
+                return
+            end
+        end
+
+        ir = sprint() do io
+            Metal.code_llvm(io, mod.kernel,
+                            Tuple{Core.LLVMPtr{Int32,1}, Bool, Int32};
+                            dump_module=true, kernel=true)
+        end
+        @test occursin("@jl_true_box = private unnamed_addr addrspace(2) constant", ir)
+        @test !occursin("@jl_true = external", ir)
+    end
+end
+
 # Tuples with a dynamic index are lowered to an addrspace(2) constant plus a
 # GEP+load. Without InferAddressSpaces propagating AS 2 through the cast to
 # the generic AS introduced during `add_global_address_spaces!`, the load

test/native.jl

@@ -301,6 +301,137 @@ 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_ghost(cond::Bool, a::Int32) = cond ? a : nothing
+            @noinline produce_true(cond::Bool, a::Int32) = cond ? a : true
+            @noinline produce_false(cond::Bool, a::Int32) = cond ? a : false
+            function consume_true(cond::Bool, a::Int32)
+                x = produce_true(cond, a)
+                x isa Bool && x && return Int32(1)
+                return Int32(0)
+            end
+            function consume_false(cond::Bool, a::Int32)
+                x = produce_false(cond, a)
+                x isa Bool && !x && return Int32(1)
+                return Int32(0)
+            end
+            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
+
+        # Boxed Bool leaves use canonical device boxes.
+        for (f, name) in ((mod.consume_true, "jl_true"),
+                          (mod.consume_false, "jl_false"))
+            ir = sprint(io->Native.code_llvm(io, f, Tuple{Bool, Int32};
+                                             dump_module=true, validate=true))
+            @static if VERSION >= v"1.14.0-DEV.1348"
+                @test occursin("@$(name)_box = private unnamed_addr constant", ir)
+                @test !occursin("@$name = external", ir)
+                @test !occursin("inttoptr", ir)
+            end
+        end
+
+        # zero-sized identity objects remain opaque host tokens
+        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
+
+            # Bool JuliaVariables are absent from `gv_to_value`.
+            m = LLVM.Module("bool singletons")
+            for name in ("jl_true", "jl_false")
+                gv = LLVM.GlobalVariable(m, LLVM.PointerType(LLVM.Int8Type()), name)
+                constant!(gv, true)
+            end
+            @test GPUCompiler.relocate_gvs!(m, Dict{String, Ptr{Cvoid}}())
+            bool_ir = string(m)
+            for name in ("jl_true", "jl_false")
+                @test haskey(globals(m), "$(name)_box")
+                @test occursin("@$name = private constant", bool_ir)
+            end
+            @test !occursin("external", bool_ir)
+            @test !occursin("inttoptr", bool_ir)
+            dispose(m)
+
+            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())