Metal: load thread pointers from other address spaces as words

src/metal.jl

@@ -740,6 +740,27 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L
     # cleanup above, so the merged form is what reaches the AIR downgrader / back-end.
     merge_byte_gep_chains!(mod)
 
+    # run last too, as InstCombine would turn the word loads back into pointer loads
+    load_thread_pointers_as_words!(mod)
+
+    return
+end
+
+# The AGX back-end keeps only the low 32 bits of a thread (AS 0) pointer that is loaded from
+# another address space. That breaks the host addresses Julia keeps by reference, e.g. of a
+# `Symbol` stored in an array (Metal.jl#1017), so load such pointers as words; pointers made
+# with `inttoptr` keep all their bits.
+function load_thread_pointers_as_words!(mod::LLVM.Module)
+    T_word = LLVM.IntType(64)
+    @dispose builder=IRBuilder() begin
+        for f in mod.functions, bb in f.blocks, inst in collect(bb.instructions)
+            inst isa LLVM.LoadInst || continue
+            T = inst.value_type
+            T isa LLVM.PointerType && T.addrspace == 0 || continue
+            inst.pointer_operand.value_type.addrspace == 0 && continue
+            load_pointer_as_word!(builder, inst, T_word)
+        end
+    end
     return
 end
 
@@ -1518,8 +1539,10 @@ function add_parameter_address_spaces!(@nospecialize(job::CompilerJob), mod::LLV
     # we cannot simply remap the function arguments, because that will not propagate the
     # address space changes across, e.g, bitcasts (the dest would still be in AS 0).
     # using a type remapper on the other hand changes too much, including unrelated insts.
-    # so instead, we load the arguments in stack slots and dereference them so that we can
-    # keep on using the original IR that assumed pointers without address spaces
+    # so instead, arguments that are only read from are cast back to AS 0 for the address-
+    # space inference in `finish_ir!` to undo, while others are copied to a stack slot. that
+    # copy is a last resort: when it survives optimization, e.g. because of dynamic indexing,
+    # the AGX back-end can miscompile the kernel (Metal.jl#1018).
     new_args = LLVM.Value[]
     @dispose builder=IRBuilder() begin
         entry = BasicBlock(new_f, "conversion")
@@ -1537,12 +1560,16 @@ function add_parameter_address_spaces!(@nospecialize(job::CompilerJob), mod::LLV
                 end
                 push!(new_args, inttoptr!(builder, load!(builder, T_word, slot), param))
             elseif byref[i]
-                # load the argument in a stack slot
-                llvm_typ = convert(LLVMType, args[i].typ)
-                val = load!(builder, llvm_typ, new_f.parameters[i])
-                ptr = alloca!(builder, llvm_typ)
-                store!(builder, val, ptr)
-                push!(new_args, ptr)
+                if only_loaded_from(f.parameters[i])
+                    push!(new_args, addrspacecast!(builder, new_f.parameters[i], param))
+                else
+                    # load the argument in a stack slot
+                    llvm_typ = convert(LLVMType, args[i].typ)
+                    val = load!(builder, llvm_typ, new_f.parameters[i])
+                    ptr = alloca!(builder, llvm_typ)
+                    store!(builder, val, ptr)
+                    push!(new_args, ptr)
+                end
             else
                 push!(new_args, new_f.parameters[i])
             end
@@ -1590,6 +1617,20 @@ function add_parameter_address_spaces!(@nospecialize(job::CompilerJob), mod::LLV
     return new_f
 end
 
+# whether `ptr` is only used as the address of plain loads, directly or through GEPs and
+# bitcasts, all of which address-space inference can move to another address space
+function only_loaded_from(ptr::LLVM.Value)
+    for use in ptr.uses
+        u = LLVM.user(use)
+        if u isa LLVM.GetElementPtrInst || u isa LLVM.BitCastInst
+            only_loaded_from(u) || return false
+        elseif !(u isa LLVM.LoadInst) || u.volatile || isatomic(u)
+            return false
+        end
+    end
+    return true
+end
+
 # update address spaces of constant global objects
 #
 # global constant objects need to reside in address space 2, so we clone each function

src/relocation.jl

@@ -958,10 +958,6 @@ function check_relocation_slot_uses!(mod::LLVM.Module, slots::Vector{LLVM.Global
     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
@@ -971,18 +967,7 @@ 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)
+            load_pointer_as_word!(builder, load, T_word)
         end
     end
     return

src/utils.jl

@@ -495,3 +495,30 @@ else
     # declarations the relocation machinery collects. Only 1.11+ emits those declarations.
     VERSION >= v"1.11-" && HAS_LLVM_GVS_GLOBALS
 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)
+
+# Replace a pointer-typed `load` by a load of the `T_word` integer holding the pointer's bits,
+# converted back with `inttoptr`. Alignment, volatility, atomic ordering and the metadata that
+# applies to integers carry over.
+function load_pointer_as_word!(builder::IRBuilder, load::LLVM.LoadInst, T_word::LLVMType)
+    position!(builder, LLVM.before(load))
+    ptr = load.pointer_operand
+    if supports_typed_pointers(context())
+        ptr = bitcast!(builder, ptr, LLVM.PointerType(T_word, ptr.value_type.addrspace))
+    end
+    word = load!(builder, T_word, ptr; 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)
+    return word
+end

test/metal.jl

@@ -22,6 +22,39 @@
     end
 end
 
+@testset "byref aggregates read in place" begin
+    # A byref argument that is only read from is accessed in device memory. A copy on the
+    # stack would survive here because of the dynamic index, which the AGX back-end can
+    # miscompile (Metal.jl#1018).
+    mod = @eval module $(gensym())
+        function kernel(out, t, i)
+            unsafe_store!(out, @inbounds t[i])
+            return
+        end
+    end
+    tt = (Core.LLVMPtr{Float32,1}, NTuple{8,Float32}, Int)
+    air = sprint(io -> Metal.code_native(io, mod.kernel, tt; kernel=true))
+    @test !occursin("alloca", air)
+    @test !occursin("addrspacecast", air)
+end
+
+@testset "byref aggregates passed along" begin
+    # When the argument also flows into a call that is passed stack memory elsewhere, its
+    # address cannot be inferred to be in device memory, so it is still copied to the stack.
+    mod = @eval module $(gensym())
+        @noinline get(t, i) = @inbounds t[i]
+        function kernel(out, t, i)
+            local_t = ntuple(j -> Float32(i * j), Val(8))
+            unsafe_store!(out, get(t, i) + get(local_t, i))
+            return
+        end
+    end
+    tt = (Core.LLVMPtr{Float32,1}, NTuple{8,Float32}, Int)
+    air = sprint(io -> Metal.code_native(io, mod.kernel, tt; kernel=true))
+    @test !occursin("addrspacecast", air)
+    @test count("alloca [8 x float]", air) == 2  # the argument's copy and `local_t`
+end
+
 @testset "byref primitives" begin
     mod = @eval module $(gensym())
         kernel(x) = return
@@ -713,6 +746,25 @@ end
     @test occursin("!\"air.read\"", sym_arg)
 end
 
+@testset "thread pointers loaded from device memory" begin
+    # The AGX back-end truncates a thread (AS 0) pointer that is loaded from another address
+    # space to 32 bits, breaking e.g. the identity of a `Symbol` stored in an array
+    # (Metal.jl#1017). Such pointers are loaded as words instead.
+    mod = @eval module $(gensym())
+        function kernel(out, ptrs)
+            unsafe_store!(out, unsafe_load(ptrs, 1) === unsafe_load(ptrs, 2) ? 1 : 2)
+            return
+        end
+    end
+    tt = (Core.LLVMPtr{Int,1}, Core.LLVMPtr{Ptr{Cvoid},1})
+
+    air = sprint(io -> Metal.code_native(io, mod.kernel, tt; kernel=true))
+    @test !occursin(r"load (ptr|i8\*), ", air)
+    if VERSION >= v"1.12.0-DEV.225"     # older Julia passes `Ptr` as an integer
+        @test count("inttoptr i64", air) == 2
+    end
+end
+
 @testset "boxed constant classification" begin
     JuliaContext() do ctx
         mod = LLVM.Module("boxed-constant-classification")