diff --git a/src/metal.jl b/src/metal.jl index 014e610c..44c2d6a3 100644 --- a/src/metal.jl +++ b/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 diff --git a/src/relocation.jl b/src/relocation.jl index ed4bd8f6..8cc8f13c 100644 --- a/src/relocation.jl +++ b/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 diff --git a/src/utils.jl b/src/utils.jl index 663864d3..3fc15d79 100644 --- a/src/utils.jl +++ b/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 diff --git a/test/metal.jl b/test/metal.jl index 479e5e88..e20d24c7 100644 --- a/test/metal.jl +++ b/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")