Load relocation table entries as words

Project.toml

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

src/metal.jl

@@ -10,11 +10,13 @@ const LLVMDowngrader_jll =
 # Metal has no target machine, so provide our own TTI
 struct MetalTTI <: LLVM.AbstractTargetTransformInfo end
 
-# teache LLVM about Metal's address-space hierarchy:
+# teach LLVM about Metal's address-space hierarchy (named as in Metal.jl's `AS`):
 #   0: Generic    1: Device       2: Constant
 #   3: ThreadGroup 4: Thread      5: ThreadGroup_ImgBlock  6: Ray
-# AS 0 is the flat/generic space; only casts involving it are legal, and the
-# specific spaces are mutually disjoint.
+# The AIR we emit has no generic space, though: AS 0 is thread memory. We only present it
+# as the flat space so that `InferAddressSpacesPass` moves pointers derived from a specific
+# space into that space; a pointer still in AS 0 afterwards is a thread pointer. Only casts
+# involving AS 0 are legal, and the specific spaces are mutually disjoint.
 LLVM.flat_address_space(::MetalTTI) = UInt(0)
 LLVM.is_noop_addr_space_cast(::MetalTTI, from::Unsigned, to::Unsigned) =
     from == 0 || to == 0

src/relocation.jl

@@ -818,9 +818,10 @@ Rewrite every record into an indexed load from a back-end-provided table of word
 copies into `relocs.table` so that [`resolved_relocation_table`](@ref) delivers the words in
 that same order regardless of what happens to the records afterwards.
 
-Slots become `load(gep(base, index))` and are erased. Interior boxes cannot be patched
-after load — the platforms needing this have no writable program-scope storage — so each is
-demoted to a per-function stack copy whose header word comes from the table.
+Slots become word loads `load(gep(base, index))`, converted back with `inttoptr` where a
+slot was loaded as a pointer, and are erased. Interior boxes cannot be patched after load —
+the platforms needing this have no writable program-scope storage — so each is demoted to a
+per-function stack copy whose header word comes from the table.
 [`relocation_table_pointer`](@ref) supplies the base pointer; since it can only do so where
 the state is available, callees still holding a relocation use are inlined first.
 """
@@ -864,7 +865,8 @@ function emit_table_relocations!(@nospecialize(job::CompilerJob), mod::LLVM.Modu
     mod_gvs = mod.globals
     slots = LLVM.GlobalVariable[mod_gvs[rec.name] for rec in relocs.records
                        if rec.kind === SlotSite && haskey(mod_gvs, rec.name)]
-    check_relocation_slot_uses!(mod, slots)
+    loads = check_relocation_slot_uses!(mod, slots)
+    load_relocation_words!(loads, T_word)
     # Expand all constant users before choosing entry insertion points. Expanding a later
     # slot could otherwise insert a use before the entry instruction saved for an earlier one.
     convert_users_to_instructions!(slots)
@@ -910,8 +912,10 @@ end
 
 # Slots denote read-only words, not general storage. In particular, don't merge a table
 # address with an unrelated pointer: a back-end may not have a common address space for them.
+# Returns the loads of the words, however their address was forwarded.
 function check_relocation_slot_uses!(mod::LLVM.Module, slots::Vector{LLVM.GlobalVariable})
     dl = mod.datalayout
+    loads = LLVM.LoadInst[]
     seen = Set{LLVM.Value}(slots)
     worklist = LLVM.Value[slots...]
     while !isempty(worklist)
@@ -920,6 +924,7 @@ function check_relocation_slot_uses!(mod::LLVM.Module, slots::Vector{LLVM.Global
             if val isa LLVM.LoadInst
                 is_word_type(val.value_type) ||
                     error("Unsupported relocation slot load of LLVM type $(val.value_type)")
+                push!(loads, val)
                 # Julia names these loads after globals with session-specific counters.
                 val.name = ""
                 # The packed table guarantees word alignment, even if the old global had more.
@@ -943,6 +948,36 @@ function check_relocation_slot_uses!(mod::LLVM.Module, slots::Vector{LLVM.Global
                 error("Relocation slot address merged with unsupported address $address in $val")
         end
     end
+    return loads
+end
+
+# Metadata that describes a loaded pointer, and so does not apply to a loaded word.
+const PointerLoadMetadataKinds =
+    (MD_nonnull, MD_dereferenceable, MD_dereferenceable_or_null, MD_align)
+
+# Load the slots' values as the words the table holds. A pointer-typed load would read a
+# host address as a pointer into the target's default address space, which on Metal is
+# thread memory, and shader validation does not preserve such a value's bits. Converting
+# the word back with `inttoptr` keeps the IR valid for existing users, while comparisons
+# fold to integer ones.
+function load_relocation_words!(loads::Vector{LLVM.LoadInst}, T_word::LLVMType)
+    @dispose builder=IRBuilder() begin
+        for load in loads
+            load.value_type isa LLVM.PointerType || continue
+            position!(builder, LLVM.before(load))
+            word = load!(builder, T_word, load.pointer_operand;
+                         align=load.alignment, volatile=load.volatile)
+            if isatomic(load)
+                word.ordering = load.ordering
+                word.syncscope = load.syncscope
+            end
+            for (kind, md) in load.metadata
+                kind in PointerLoadMetadataKinds || (word.metadata[kind] = md)
+            end
+            replace_uses!(load, inttoptr!(builder, word, load.value_type))
+            erase!(load)
+        end
+    end
     return
 end
 

test/metal.jl

@@ -427,19 +427,44 @@ end
         # the table base is loaded out of the kernel-state argument, and the words out of the
         # table -- a bake would instead leave a private constant holding the resolved address
         @test occursin("reloc_table", air)
-        @test occursin(r"load (i64|ptr), (i64 addrspace\(1\)\*|ptr addrspace\(1\))", air)
+        @test occursin(r"load i64, (i64 addrspace\(1\)\*|ptr addrspace\(1\))", air)
         # nothing is left of the site globals the records named
         for rec in relocs.records
             @test !occursin("@$(rec.name) ", air)
         end
     end
 end
 
+@testset "relocation words" begin
+    # The table holds host words, such as the address of a Symbol, that the device only
+    # compares. Loaded as a pointer, a word would be a pointer into thread memory, whose bits
+    # Metal's shader validation does not preserve; so it is loaded as an integer, just like
+    # the Symbol argument it is compared against.
+    if GPUCompiler.supports_relocatable_ir() && LLVM.version() >= v"17"
+        mod = @eval module $(gensym())
+            function kernel(ptr, sym::Symbol)
+                unsafe_store!(ptr, sym === :foo ? 1f0 : sym === :bar ? 2f0 : 3f0)
+                return
+            end
+        end
+        tt = (Core.LLVMPtr{Float32,1}, Symbol)
+
+        air = sprint(io -> Metal.code_native_table(io, mod.kernel, tt; kernel=true))
+        @test occursin("reloc_table", air)
+        @test !occursin("load ptr, ptr addrspace(1)", air)
+        # LLVM 20 folds the comparison of the converted words into one of the words
+        if LLVM.version() >= v"20"
+            @test occursin("icmp eq i64", air)
+            @test !occursin("icmp eq ptr", air)
+        end
+    end
+end
+
 @testset "merged relocation slots" begin
     # With five possible values, inference widens `pick`'s result to `Val`, so `v` is boxed
     # and `===` compares its address with those of the singletons. LLVM merges the loads of
     # those addresses into one load from a `phi` of their slots, which the table lowering
-    # redirects to the device-space table (#959).
+    # redirects to the device-space table (#959), still reading the words as integers.
     if GPUCompiler.supports_relocatable_ir() && LLVM.version() >= v"17"
         mod = @eval module $(gensym())
             pick(i) = i == 1 ? Val(1) : i == 2 ? Val(2) : i == 3 ? Val(3) :
@@ -470,14 +495,15 @@ end
 
         @test @filecheck begin
             @check "phi ptr addrspace(1)"
-            @check "load ptr, ptr addrspace(1)"
+            @check "load i64, ptr addrspace(1)"
             Metal.code_native_table(mod.kernel, tt; kernel=true)
         end
         for f in (mod.kernel, mod.maybe_kernel, mod.loop_kernel)
             air = sprint(io -> Metal.code_native_table(io, f, tt; kernel=true))
             @test occursin("reloc_table", air)
             @test !occursin(r"(?m)^@.*jl_global", air)
             @test !occursin(r"(?m)^@.*jl_nothing", air)
+            @test !occursin("load ptr, ptr addrspace(1)", air)
         end
     end
 end

test/native.jl

@@ -1565,6 +1565,70 @@ end
     end
 end
 
+@testset "pointer-valued relocation words" begin
+    # Slots are generally loaded as pointers, which code may also return or dereference.
+    # Whether a strategy patches, bakes or tabulates the word, that pointer must be the same.
+    if GPUCompiler.supports_relocatable_ir() && LLVM.version() >= v"17"
+        target = GPUCompiler.JuliaValueRef(:pointer_probe)
+        expected = GPUCompiler.resolve_relocation_target(target)
+        for strategy in (:bake, :patch, :table)
+            job, _ = Native.create_job(identity, (Nothing,); relocations=strategy,
+                                       jlruntime=false)
+            JuliaContext() do ctx
+                m = parse(LLVM.Module, """
+                    @slot = external global ptr
+                    @jl_float32_type = external global ptr
+
+                    define ptr @value() {
+                        %value = load ptr, ptr @slot, align 8, !nonnull !0, !invariant.load !0
+                        ret ptr %value
+                    }
+
+                    define i64 @header() {
+                        %value = load ptr, ptr @slot, align 8, !nonnull !0, !invariant.load !0
+                        %header = getelementptr i64, ptr %value, i64 -1
+                        %word = load i64, ptr %header
+                        ret i64 %word
+                    }
+
+                    define ptr @type() {
+                        %type = load ptr, ptr @jl_float32_type, align 8, !nonnull !0
+                        ret ptr %type
+                    }
+
+                    !0 = !{}""")
+                relocs = GPUCompiler.Relocations(
+                    [GPUCompiler.Relocation(GPUCompiler.SlotSite, "slot", 0, target)])
+                obj, _ = GPUCompiler.emit_asm(job, m, relocs, LLVM.CodeGenFileType.Object)
+                @test verification_error(m) === nothing
+                if strategy === :table
+                    # the words are loaded as such, keeping what is known about the access,
+                    # but not what was known about the pointer
+                    loads = [inst for f in ("value", "type")
+                                  for inst in m.functions[f].entry.instructions
+                                  if inst isa LLVM.LoadInst]
+                    @test all(load -> load.value_type == LLVM.Int64Type(), loads)
+                    @test !any(load -> haskey(load.metadata, LLVM.MD_nonnull), loads)
+                    @test any(load -> haskey(load.metadata, LLVM.MD_invariant_load), loads)
+                end
+
+                fptr, lljit, table = Native.load(Vector{UInt8}(codeunits(obj)), "value",
+                                                 relocs; table=strategy === :table)
+                @dispose lljit=lljit begin
+                    GC.@preserve table begin
+                        @test UInt(ccall(fptr, Ptr{Cvoid}, ())) == expected
+                        header = pointer(lookup(lljit, "header"))
+                        @test ccall(header, UInt, ()) ==
+                              unsafe_load(Ptr{UInt}(expected - sizeof(UInt)))
+                        type = pointer(lookup(lljit, "type"))
+                        @test ccall(type, Ptr{Cvoid}, ()) == pointer_from_objref(Float32)
+                    end
+                end
+            end
+        end
+    end
+end
+
 @testset "unlowered relocation table" begin
     # Emitting a `:table` module through the 3-argument `emit_asm` hands the lowering an
     # empty manifest, leaving the real one unlowered and the module's slots stranded. The
@@ -1786,8 +1850,8 @@ end
         for rec in relocs.records
             @test occursin("@$(rec.name)", string(addr))
         end
-        # Direct and merged references retain the same load type. Mixing pointer loads
-        # with rebuilt integer loads/inttoptr miscompiles the `nothing` case on Metal.
+        # Direct and merged references retain their load type, so that later optimization
+        # sees consistent pointer expressions; only the `:table` lowering loads words.
         for f in ("direct", "entry")
             load = only(inst for bb in mod.functions[f].blocks for inst in bb.instructions
                         if inst isa LLVM.LoadInst)