From 2faa2af3e53eaad8a5f0d674b71790f56688fe77 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 18:13:10 +0200 Subject: [PATCH 1/3] Handle loads from merged relocation slot addresses. LLVM can merge the loads from several relocation slots into a single load from a phi or select of the slots' addresses, e.g. when a boxed value is compared against the singletons it may be. The :table lowering only redirected direct loads of a slot, so such a slot kept a use and compilation failed with "still has uses after redirection". Rebuild those phis and selects over the slots' table offsets instead, and load the word at the resulting offset. Merging a slot with any other address, or using a merged address for anything but loading its word, still errors. Libjulia globals can be merged the same way, e.g. jl_nothing with the slots of other singletons. Only their direct loads were relocated, leaving the phi referencing the global itself. As such a phi is only ever loaded from, give it the address of the global's slot, which holds the same word, and report merged loads that cannot be redirected as unresolved. Fixes #959. --- src/relocation.jl | 207 ++++++++++++++++++++++++++++++++++++++++++---- test/metal.jl | 38 +++++++++ test/native.jl | 170 +++++++++++++++++++++++++++++++++++++ 3 files changed, 399 insertions(+), 16 deletions(-) diff --git a/src/relocation.jl b/src/relocation.jl index f7ba7ef9..637484e7 100644 --- a/src/relocation.jl +++ b/src/relocation.jl @@ -515,18 +515,93 @@ function rewrite_word_loads!(produce_word, @nospecialize(value), what::String; elseif isa(val, LLVM.LoadInst) offset === nothing && error("Unsupported $what load through constant expression $(operands(val)[1])") - T = value_type(val) - (T isa LLVM.PointerType || - (T isa LLVM.IntegerType && width(T) == 8sizeof(UInt))) || - error("Unsupported $what load of LLVM type $T") - @dispose builder=IRBuilder() begin - position!(builder, val) - replacement = produce_word(builder, offset) - T isa LLVM.PointerType && - (replacement = inttoptr!(builder, replacement, T)) - replace_uses!(val, replacement) + replace_word_load!(builder -> produce_word(builder, offset), val, what) + changed = true + end + end + return changed +end + +is_word_type(T::LLVMType) = + T isa LLVM.PointerType || (T isa LLVM.IntegerType && width(T) == 8sizeof(UInt)) + +# Replace a word-sized `load` with the word `produce_word(builder)` emits in its place. +function replace_word_load!(produce_word, load::LLVM.LoadInst, what::String) + T = value_type(load) + is_word_type(T) || error("Unsupported $what load of LLVM type $T") + @dispose builder=IRBuilder() begin + position!(builder, load) + replacement = produce_word(builder) + T isa LLVM.PointerType && (replacement = inttoptr!(builder, replacement, T)) + replace_uses!(load, replacement) + end + erase!(load) + return +end + +# The addresses an instruction merely forwards: those a `phi` or `select` picks from, or the +# operand of a pointer cast. `nothing` for any other instruction. +function forwarded_addresses(inst::LLVM.Instruction) + if inst isa LLVM.PHIInst + return LLVM.Value[value for (value, _) in incoming(inst)] + elseif inst isa LLVM.SelectInst + return LLVM.Value[operands(inst)[2], operands(inst)[3]] + elseif inst isa LLVM.BitCastInst || inst isa LLVM.AddrSpaceCastInst + return LLVM.Value[operands(inst)[1]] + end + return nothing +end + +# LLVM merges loads from different addresses into one load from a `phi` or `select` of those +# addresses, e.g. when sinking the loads of a boxed value's possible singletons out of their +# branches. Starting from such an instruction, collect into `merged` the instructions the +# address flows through, and into `loads` the loads it ends up in. Returns the first user that +# does anything else with it, or `nothing`. +function merged_address_loads!(merged::Vector{LLVM.Instruction}, loads::Vector{LLVM.LoadInst}, + inst::LLVM.Instruction) + push!(merged, inst) + worklist = LLVM.Instruction[inst] + while !isempty(worklist) + for use in uses(pop!(worklist)) + val = user(use) + if val isa LLVM.LoadInst + push!(loads, val) + elseif val isa LLVM.Instruction && forwarded_addresses(val) !== nothing + val in merged && continue + push!(merged, val) + push!(worklist, val) + else + return val + end + end + end + return nothing +end + +# Where `value`'s address is merged with others before a word is loaded from it, substitute +# the address of a slot that holds the same word, `slot_address(offset)` for the word at byte +# `offset`. Other uses of `value` are left alone. +function redirect_merged_addresses!(slot_address, @nospecialize(value); offset::Int=0, + dl::DataLayout=datalayout(LLVM.parent(value)::LLVM.Module)) + changed = false + for use in collect(uses(value)) + val = user(use) + if isa(val, LLVM.ConstantExpr) + delta = constexpr_byte_offset(val, dl) + delta === nothing && continue + changed |= redirect_merged_addresses!(slot_address, val; offset=offset + delta, dl) + elseif isa(val, LLVM.Instruction) && forwarded_addresses(val) !== nothing + # slots live in the default address space + T = value_type(value) + addrspace(T) == 0 || continue + loads = LLVM.LoadInst[] + merged_address_loads!(LLVM.Instruction[], loads, val) === nothing || continue + all(load -> is_word_type(value_type(load)), loads) || continue + slot = slot_address(offset) + ops = operands(val) + for i in 1:length(ops) + ops[i] == value && (ops[i] = const_pointercast(slot, T)) end - erase!(val) changed = true end end @@ -567,17 +642,25 @@ function collect_cglobal_relocations!(@nospecialize(job::CompilerJob), mod::LLVM changed |= rewrite_word_loads!(f, "cglobal '$fn'") do builder, offset load!(builder, relocation_word_type(), cglobal_slot(offset)) end + # e.g. `jl_nothing` merged with the relocation slots of other singletons + changed |= redirect_merged_addresses!(cglobal_slot, f) end return changed end function has_unresolved_cglobal_loads(mod::LLVM.Module, relocs::Relocations) - function has_load(value) + # also through merged addresses that `redirect_merged_addresses!` had to leave alone + function has_load(value, seen=Set{LLVM.Value}()) for use in uses(value) val = user(use) val isa LLVM.LoadInst && return true - val isa LLVM.ConstantExpr && has_load(val) && return true + if val isa LLVM.ConstantExpr || + (val isa LLVM.Instruction && forwarded_addresses(val) !== nothing) + val in seen && continue + push!(seen, val) + has_load(val, seen) && return true + end end return false end @@ -831,14 +914,20 @@ function emit_table_relocations!(@nospecialize(job::CompilerJob), mod::LLVM.Modu end end end - function table_word(builder::IRBuilder, index::Int) + table_offset(index::Int) = ConstantInt(LLVM.Int32Type(), index - 1) + function table_word(builder::IRBuilder, offset::LLVM.Value) f = LLVM.parent(position(builder)) - ptr = inbounds_gep!(builder, T_word, table_base(f), - [ConstantInt(LLVM.Int32Type(), index - 1)]) + ptr = inbounds_gep!(builder, T_word, table_base(f), [offset]) load!(builder, T_word, ptr) end + table_word(builder::IRBuilder, index::Int) = table_word(builder, table_offset(index)) mod_gvs = globals(mod) + rewrite_merged_slot_loads!(table_word, mod, + Pair{LLVM.Value,LLVM.Value}[mod_gvs[rec.name] => table_offset(index) + for (index, rec) in enumerate(relocs.records) + if rec.kind === SlotSite && haskey(mod_gvs, rec.name)]) + for (index, rec) in enumerate(relocs.records) haskey(mod_gvs, rec.name) || error("Missing relocation global '$(rec.name)'") gv = mod_gvs[rec.name] @@ -861,6 +950,92 @@ function emit_table_relocations!(@nospecialize(job::CompilerJob), mod::LLVM.Modu return end +# Loads from a `phi` or `select` of slot addresses (see `merged_address_loads!`) cannot be +# redirected one slot at a time, so merge the slots' table offsets the same way instead and +# load the word at the merged offset. `slots` maps each slot to its offset, in table order. +# The offsets are constants and each merge is rebuilt next to the one it replaces, so the new +# values dominate wherever the old ones did. +function rewrite_merged_slot_loads!(table_word, mod::LLVM.Module, + slots::Vector{Pair{LLVM.Value,LLVM.Value}}) + dl = datalayout(mod) + offsets = Dict{LLVM.Value,LLVM.Value}(slots) + function slot_offset(value) + while value isa LLVM.ConstantExpr && constexpr_byte_offset(value, dl) == 0 + value = first(operands(value)) + end + return get(offsets, value, nothing) + end + + merged = LLVM.Instruction[] + loads = LLVM.LoadInst[] + function find_merges(value) + for use in uses(value) + val = user(use) + if val isa LLVM.ConstantExpr + constexpr_byte_offset(val, dl) == 0 && find_merges(val) + elseif val isa LLVM.Instruction && forwarded_addresses(val) !== nothing && + !(val in merged) + other = merged_address_loads!(merged, loads, val) + other === nothing || + error("Unsupported use of merged relocation slot addresses: $other") + end + end + end + foreach(find_merges ∘ first, slots) + isempty(merged) && return + + # A slot can only be merged with other slots: the table holds no word for other addresses. + for inst in merged, value in forwarded_addresses(inst) + value in merged || slot_offset(value) !== nothing || + error("Relocation slot address merged with unsupported address $value in $inst") + end + + T_offset = LLVM.Int32Type() + merged_offsets = Dict{LLVM.Value,LLVM.Value}() + @dispose builder=IRBuilder() begin + # `phi`s first, as they may be merged with themselves through a loop + for inst in merged + inst isa LLVM.PHIInst || continue + position!(builder, inst) + merged_offsets[inst] = phi!(builder, T_offset) + end + function merged_offset(value) + haskey(merged_offsets, value) && return merged_offsets[value] + offset = slot_offset(value) + offset === nothing || return offset + if value isa LLVM.SelectInst + cond, a, b = operands(value) + a, b = merged_offset(a), merged_offset(b) + position!(builder, value) + offset = select!(builder, cond, a, b) + else + offset = merged_offset(first(forwarded_addresses(value))) + end + merged_offsets[value] = offset + end + for inst in merged + if inst isa LLVM.PHIInst + append!(incoming(merged_offsets[inst]), + Tuple{LLVM.Value,LLVM.BasicBlock}[(merged_offset(value), block) + for (value, block) in incoming(inst)]) + else + merged_offset(inst) + end + end + end + + for load in loads + offset = merged_offsets[first(operands(load))] + replace_word_load!(builder -> table_word(builder, offset), load, + "merged relocation slot") + end + for inst in merged + replace_uses!(inst, PoisonValue(value_type(inst))) + end + foreach(erase!, merged) + return +end + # Copy a relocatable box into a per-function stack slot and fill its header from the # relocation table. Sound because a box address carries no identity of its own: `isbits` egal # compares by content, so a per-invocation copy is indistinguishable from a shared one. diff --git a/test/metal.jl b/test/metal.jl index 9b69b434..e1cca4ec 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -435,6 +435,44 @@ 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 + # has to turn into a `phi` of table offsets (#959). + 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) : + i == 4 ? Val(4) : Val(5) + function kernel(ptr, i) + v = pick(unsafe_load(i)) + unsafe_store!(ptr, v === Val(1) ? 1f0 : v === Val(2) ? 2f0 : 3f0) + return + end + # `nothing` is merged in through `jl_nothing`, which needs a slot of its own + function maybe_kernel(ptr, i) + x = unsafe_load(i) + v = x > 5 ? Base.inferencebarrier(nothing) : pick(x) + unsafe_store!(ptr, v === Val(1) ? 1f0 : v === nothing ? 2f0 : 3f0) + return + end + end + tt = (Core.LLVMPtr{Float32,1}, Core.LLVMPtr{Int,1}) + + @test @filecheck begin + @check "phi i32" + @check "load i64" + Metal.code_native_table(mod.kernel, tt; kernel=true) + end + for f in (mod.kernel, mod.maybe_kernel) + air = sprint(io -> Metal.code_native_table(io, f, tt; kernel=true)) + @test occursin("reloc_table", air) + @test !occursin("jl_global", air) + @test !occursin("jl_nothing", air) + end + end +end + @testset "codegen counter normalization" begin # Julia's per-session codegen counter has to be scrubbed from everything that reaches the # bitcode, or the metallib is not reproducible: symbol names, the block labels inlining diff --git a/test/native.jl b/test/native.jl index 52eea89e..759d79a2 100644 --- a/test/native.jl +++ b/test/native.jl @@ -1276,6 +1276,129 @@ end end end +@testset "tabulated relocation of merged slots" begin + # LLVM merges loads from different slots into one load from a `phi` or `select` of their + # addresses (#959). The table has no address to take their place, so the lowering has to + # merge their table offsets instead. + if GPUCompiler.supports_relocatable_ir() && LLVM.version() >= v"17" + mod = @eval module $(gensym()) + f() = nothing + end + job, _ = Native.create_job(mod.f, Tuple{}; relocations=:table, jlruntime=false) + slots = ("merged_a", "merged_b", "merged_c") + relocations() = GPUCompiler.Relocations( + [GPUCompiler.Relocation(GPUCompiler.SlotSite, name, 0, + GPUCompiler.JuliaValueRef(Symbol(name))) + for name in slots]) + word(name) = + GPUCompiler.resolve_relocation_target(GPUCompiler.JuliaValueRef(Symbol(name))) + JuliaContext() do ctx + m = parse(LLVM.Module, """ + @merged_a = external global ptr + @merged_b = external global ptr + @merged_c = external global ptr + @jl_nothing = external global ptr + + define i64 @pick(i64 %i) { + top: + switch i64 %i, label %other [ i64 1, label %one + i64 2, label %two ] + one: + br label %join + two: + br label %join + other: + %large = icmp sgt i64 %i, 3 + %other.addr = select i1 %large, ptr @jl_nothing, ptr @merged_c + br label %join + join: + %addr = phi ptr [ @merged_a, %one ], [ @merged_b, %two ], + [ %other.addr, %other ] + %word = load i64, ptr %addr + ret i64 %word + } + + define i64 @flip(i64 %n) { + top: + br label %loop + loop: + %k = phi i64 [ 0, %top ], [ %k.next, %loop ] + %addr = phi ptr [ @merged_a, %top ], [ %addr.next, %loop ] + %flip = icmp eq i64 %k, 2 + %addr.next = select i1 %flip, ptr @merged_b, ptr %addr + %k.next = add i64 %k, 1 + %done = icmp sge i64 %k.next, %n + br i1 %done, label %exit, label %loop + exit: + %word = load ptr, ptr %addr.next + %int = ptrtoint ptr %word to i64 + ret i64 %int + } + + define i64 @cast(i1 %cond) { + %addr = select i1 %cond, ptr @merged_a, ptr @merged_b + %cast = addrspacecast ptr %addr to ptr addrspace(1) + %word = load i64, ptr addrspace(1) %cast + ret i64 %word + }""") + relocs = relocations() + obj, _ = GPUCompiler.emit_asm(job, m, relocs, LLVM.API.LLVMObjectFile) + # `jl_nothing` became a slot too, and all of them were replaced by the table + @test length(relocs) == length(slots) + 1 + @test any(rec -> rec.target == GPUCompiler.CGlobalRef(:jl_nothing), + relocs.records) + for rec in relocs.records + @test !haskey(globals(m), rec.name) + end + @test occursin("phi i32", string(m)) + + fptr, lljit, table = Native.load(Vector{UInt8}(codeunits(obj)), "pick", relocs; + table=true) + try + GC.@preserve table begin + nothing_word = GPUCompiler.resolve_relocation_target( + GPUCompiler.CGlobalRef(:jl_nothing)) + @test [ccall(fptr, UInt, (Int,), i) for i in 1:4] == + [word("merged_a"), word("merged_b"), word("merged_c"), nothing_word] + flip = pointer(lookup(lljit, "flip")) + @test [ccall(flip, UInt, (Int,), n) for n in 1:4] == + word.(["merged_a", "merged_a", "merged_b", "merged_b"]) + cast = pointer(lookup(lljit, "cast")) + @test [ccall(cast, UInt, (Bool,), c) for c in (true, false)] == + word.(["merged_a", "merged_b"]) + end + finally + dispose(lljit) + end + + # the table holds no word for any other address... + m = parse(LLVM.Module, """ + @merged_a = external global ptr + + define i64 @mixed(i1 %cond, ptr %other) { + %addr = select i1 %cond, ptr @merged_a, ptr %other + %word = load i64, ptr %addr + ret i64 %word + }""") + @test_throws "merged with unsupported address" GPUCompiler.emit_asm( + job, m, relocations(), LLVM.API.LLVMObjectFile) + + # ...and has no address to give out + m = parse(LLVM.Module, """ + @merged_a = external global ptr + @merged_b = external global ptr + + define i1 @compare(i1 %cond) { + %addr = select i1 %cond, ptr @merged_a, ptr @merged_b + %same = icmp eq ptr %addr, @merged_a + ret i1 %same + }""") + @test_throws "Unsupported use of merged relocation slot addresses" GPUCompiler.emit_asm( + job, m, relocations(), LLVM.API.LLVMObjectFile) + 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 @@ -1470,6 +1593,53 @@ end GPUCompiler.emit_patchable_relocations!(mod, relocs) @test occursin("externally_initialized global i64 0", string(mod)) + # A global whose address is merged with others before being loaded from (e.g. + # `jl_nothing` with the relocation slots of other singletons) takes its slot's address. + merged_ir = """ + @jl_float32_type = external global $word_ptr + @jl_float64_type = external global $word_ptr + + define $word_ptr @entry(i1 %cond) { + %addr = select i1 %cond, $word_ptr_ptr @jl_float32_type, + $word_ptr_ptr @jl_float64_type + %value = load $word_ptr, $word_ptr_ptr %addr + ret $word_ptr %value + }""" + mod = parse(LLVM.Module, merged_ir) + relocs = GPUCompiler.Relocations() + @test GPUCompiler.collect_cglobal_relocations!(job, mod, relocs) + @test [rec.target for rec in relocs.records] == + [GPUCompiler.CGlobalRef(:jl_float32_type), GPUCompiler.CGlobalRef(:jl_float64_type)] + addr = first(instructions(first(blocks(functions(mod)["entry"])))) + @test !occursin(r"@jl_float(32|64)_type\b", string(addr)) + for rec in relocs.records + @test occursin("@$(rec.name)", string(addr)) + end + mod = parse(LLVM.Module, merged_ir) + GPUCompiler.prepare_execution!(job, mod) + ir = string(mod) + for T in (Float32, Float64) + expected = GPUCompiler.resolve_relocation_target( + GPUCompiler.CGlobalRef(Symbol("jl_$(lowercase(string(T)))_type"))) + @test expected == UInt(pointer_from_objref(T)) + @test occursin("inttoptr (i64 $expected to $word_ptr)", ir) + end + # a slot only holds a word, so anything else loaded is left, and reported, as is + mod = parse(LLVM.Module, """ + @jl_float32_type = external global i32 + @jl_float64_type = external global i32 + + define i32 @entry(i1 %cond) { + %addr = select i1 %cond, $(ptr("i32")) @jl_float32_type, + $(ptr("i32")) @jl_float64_type + %value = load i32, $(ptr("i32")) %addr + ret i32 %value + }""") + relocs = GPUCompiler.Relocations() + @test !GPUCompiler.collect_cglobal_relocations!(job, mod, relocs) + @test isempty(relocs) + @test GPUCompiler.has_unresolved_cglobal_loads(mod, relocs) + # Fold aggregate GEPs according to the module data layout. mod = parse(LLVM.Module, """ target datalayout = "e-p:64:64-i64:64" From f61f8f8c1760128c5e921eb252750cd8aa9c96b6 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 26 Sep 2026 10:04:45 +0200 Subject: [PATCH 2/3] Preserve loads when relocating cglobal addresses. Redirect both direct and merged cglobal addresses to their word slots. This keeps pointer loads intact instead of rebuilding them as integer loads followed by inttoptr, which can miscompile mixed singleton comparisons on Metal. --- src/relocation.jl | 76 ++++++++++++++++++++++++++++++----------------- test/native.jl | 16 +++++++++- 2 files changed, 64 insertions(+), 28 deletions(-) diff --git a/src/relocation.jl b/src/relocation.jl index 637484e7..b0e4678d 100644 --- a/src/relocation.jl +++ b/src/relocation.jl @@ -578,32 +578,58 @@ function merged_address_loads!(merged::Vector{LLVM.Instruction}, loads::Vector{L return nothing end -# Where `value`'s address is merged with others before a word is loaded from it, substitute -# the address of a slot that holds the same word, `slot_address(offset)` for the word at byte -# `offset`. Other uses of `value` are left alone. -function redirect_merged_addresses!(slot_address, @nospecialize(value); offset::Int=0, - dl::DataLayout=datalayout(LLVM.parent(value)::LLVM.Module)) +# Check forwarded addresses and relax load alignment to what a word slot guarantees. +# Cycles can arise from loop PHIs. +function check_word_loads!(value, seen=Set{LLVM.Value}()) + value in seen && return true + push!(seen, value) + for use in uses(value) + val = user(use) + if val isa LLVM.LoadInst + is_word_type(value_type(val)) || return false + alignment(val) > sizeof(UInt) && alignment!(val, sizeof(UInt)) + elseif val isa LLVM.Instruction && forwarded_addresses(val) !== nothing + check_word_loads!(val, seen) || return false + else + return false + end + end + return true +end + +# Substitute a slot holding the same word for each loaded cglobal address. Replacing the +# address preserves the loads and any PHIs/selects LLVM has introduced between them. +function redirect_word_addresses!(slot_address, @nospecialize(value), what::String; + offset::Union{Int,Nothing}=0, + dl::DataLayout=datalayout(LLVM.parent(value)::LLVM.Module)) changed = false for use in collect(uses(value)) val = user(use) - if isa(val, LLVM.ConstantExpr) + if val isa LLVM.ConstantExpr delta = constexpr_byte_offset(val, dl) - delta === nothing && continue - changed |= redirect_merged_addresses!(slot_address, val; offset=offset + delta, dl) - elseif isa(val, LLVM.Instruction) && forwarded_addresses(val) !== nothing - # slots live in the default address space - T = value_type(value) - addrspace(T) == 0 || continue - loads = LLVM.LoadInst[] - merged_address_loads!(LLVM.Instruction[], loads, val) === nothing || continue - all(load -> is_word_type(value_type(load)), loads) || continue - slot = slot_address(offset) - ops = operands(val) - for i in 1:length(ops) - ops[i] == value && (ops[i] = const_pointercast(slot, T)) - end - changed = true + inner = (offset === nothing || delta === nothing) ? nothing : offset + delta + changed |= redirect_word_addresses!(slot_address, val, what; offset=inner, dl) + continue + elseif val isa LLVM.LoadInst + offset === nothing && + error("Unsupported $what load through constant expression $(operands(val)[1])") + is_word_type(value_type(val)) || + error("Unsupported $what load of LLVM type $(value_type(val))") + alignment(val) > sizeof(UInt) && alignment!(val, sizeof(UInt)) + elseif val isa LLVM.Instruction && forwarded_addresses(val) !== nothing + offset === nothing && continue + # Slots live in the default address space. + addrspace(value_type(value)) == 0 || continue + check_word_loads!(val) || continue + else + continue end + slot = const_pointercast(slot_address(offset), value_type(value)) + ops = operands(val) + for i in 1:length(ops) + ops[i] == value && (ops[i] = slot) + end + changed = true end return changed end @@ -639,18 +665,14 @@ function collect_cglobal_relocations!(@nospecialize(job::CompilerJob), mod::LLVM end end - changed |= rewrite_word_loads!(f, "cglobal '$fn'") do builder, offset - load!(builder, relocation_word_type(), cglobal_slot(offset)) - end - # e.g. `jl_nothing` merged with the relocation slots of other singletons - changed |= redirect_merged_addresses!(cglobal_slot, f) + changed |= redirect_word_addresses!(cglobal_slot, f, "cglobal '$fn'") end return changed end function has_unresolved_cglobal_loads(mod::LLVM.Module, relocs::Relocations) - # also through merged addresses that `redirect_merged_addresses!` had to leave alone + # also through merged addresses that `redirect_word_addresses!` had to leave alone function has_load(value, seen=Set{LLVM.Value}()) for use in uses(value) val = user(use) diff --git a/test/native.jl b/test/native.jl index 759d79a2..77dd2e51 100644 --- a/test/native.jl +++ b/test/native.jl @@ -1599,10 +1599,15 @@ end @jl_float32_type = external global $word_ptr @jl_float64_type = external global $word_ptr + define $word_ptr @direct() { + %value = load $word_ptr, $word_ptr_ptr @jl_float32_type, align 16 + ret $word_ptr %value + } + define $word_ptr @entry(i1 %cond) { %addr = select i1 %cond, $word_ptr_ptr @jl_float32_type, $word_ptr_ptr @jl_float64_type - %value = load $word_ptr, $word_ptr_ptr %addr + %value = load $word_ptr, $word_ptr_ptr %addr, align 16 ret $word_ptr %value }""" mod = parse(LLVM.Module, merged_ir) @@ -1615,6 +1620,15 @@ 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. + for f in ("direct", "entry") + load = only(inst for bb in blocks(functions(mod)[f]) for inst in instructions(bb) + if inst isa LLVM.LoadInst) + @test value_type(load) isa LLVM.PointerType + @test alignment(load) == sizeof(UInt) + end + LLVM.verify(mod) mod = parse(LLVM.Module, merged_ir) GPUCompiler.prepare_execution!(job, mod) ir = string(mod) From 9d808bc2a927204810a4a503313fa3c8c694c303 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 26 Sep 2026 10:04:56 +0200 Subject: [PATCH 3/3] Lower relocation slots by substituting table addresses. Use replace_global_with_local! and LLVM address-space inference instead of rebuilding pointer PHIs and selects as an offset graph. Keep the word-load checks, cap alignment at the table word size, and discard session-specific load names. Expand constant users before choosing insertion points so later substitutions cannot introduce uses before their definitions. --- src/relocation.jl | 225 ++++++++++++++-------------------------------- test/metal.jl | 23 +++-- test/native.jl | 43 +++++++-- 3 files changed, 119 insertions(+), 172 deletions(-) diff --git a/src/relocation.jl b/src/relocation.jl index b0e4678d..2371fa83 100644 --- a/src/relocation.jl +++ b/src/relocation.jl @@ -500,45 +500,9 @@ function constexpr_byte_offset(ce::LLVM.ConstantExpr, dl::DataLayout) return nothing end -# Rewrite word-sized loads derived through constant casts or GEPs from `value`. The producer -# receives the byte offset; reject paths whose offset is not static. -function rewrite_word_loads!(produce_word, @nospecialize(value), what::String; - offset::Union{Int,Nothing}=0, - dl::DataLayout=datalayout(LLVM.parent(value)::LLVM.Module)) - changed = false - for use in collect(uses(value)) - val = user(use) - if isa(val, LLVM.ConstantExpr) - delta = constexpr_byte_offset(val, dl) - inner = (offset === nothing || delta === nothing) ? nothing : offset + delta - changed |= rewrite_word_loads!(produce_word, val, what; offset=inner, dl) - elseif isa(val, LLVM.LoadInst) - offset === nothing && - error("Unsupported $what load through constant expression $(operands(val)[1])") - replace_word_load!(builder -> produce_word(builder, offset), val, what) - changed = true - end - end - return changed -end - is_word_type(T::LLVMType) = T isa LLVM.PointerType || (T isa LLVM.IntegerType && width(T) == 8sizeof(UInt)) -# Replace a word-sized `load` with the word `produce_word(builder)` emits in its place. -function replace_word_load!(produce_word, load::LLVM.LoadInst, what::String) - T = value_type(load) - is_word_type(T) || error("Unsupported $what load of LLVM type $T") - @dispose builder=IRBuilder() begin - position!(builder, load) - replacement = produce_word(builder) - T isa LLVM.PointerType && (replacement = inttoptr!(builder, replacement, T)) - replace_uses!(load, replacement) - end - erase!(load) - return -end - # The addresses an instruction merely forwards: those a `phi` or `select` picks from, or the # operand of a pointer cast. `nothing` for any other instruction. function forwarded_addresses(inst::LLVM.Instruction) @@ -552,32 +516,6 @@ function forwarded_addresses(inst::LLVM.Instruction) return nothing end -# LLVM merges loads from different addresses into one load from a `phi` or `select` of those -# addresses, e.g. when sinking the loads of a boxed value's possible singletons out of their -# branches. Starting from such an instruction, collect into `merged` the instructions the -# address flows through, and into `loads` the loads it ends up in. Returns the first user that -# does anything else with it, or `nothing`. -function merged_address_loads!(merged::Vector{LLVM.Instruction}, loads::Vector{LLVM.LoadInst}, - inst::LLVM.Instruction) - push!(merged, inst) - worklist = LLVM.Instruction[inst] - while !isempty(worklist) - for use in uses(pop!(worklist)) - val = user(use) - if val isa LLVM.LoadInst - push!(loads, val) - elseif val isa LLVM.Instruction && forwarded_addresses(val) !== nothing - val in merged && continue - push!(merged, val) - push!(worklist, val) - else - return val - end - end - end - return nothing -end - # Check forwarded addresses and relax load alignment to what a word slot guarantees. # Cycles can arise from loop PHIs. function check_word_loads!(value, seen=Set{LLVM.Value}()) @@ -927,28 +865,32 @@ function emit_table_relocations!(@nospecialize(job::CompilerJob), mod::LLVM.Modu # One base pointer per function, materialized at the top of its entry block (the state # it derives from is a function argument, so it dominates every use). - bases = Dict{LLVM.Function, LLVM.Value}() + bases = Dict{LLVM.Function, Tuple{LLVM.Value, LLVM.Instruction}}() function table_base(f::LLVM.Function) get!(bases, f) do + entry = first(instructions(first(blocks(f)))) @dispose builder=IRBuilder() begin - position!(builder, first(instructions(first(blocks(f))))) - relocation_table_pointer(job, builder, f) + position!(builder, entry) + relocation_table_pointer(job, builder, f), entry end end end - table_offset(index::Int) = ConstantInt(LLVM.Int32Type(), index - 1) - function table_word(builder::IRBuilder, offset::LLVM.Value) + function table_address(builder::IRBuilder, base::LLVM.Value, index::Int) + inbounds_gep!(builder, T_word, base, [ConstantInt(LLVM.Int32Type(), index - 1)]) + end + function table_word(builder::IRBuilder, index::Int) f = LLVM.parent(position(builder)) - ptr = inbounds_gep!(builder, T_word, table_base(f), [offset]) - load!(builder, T_word, ptr) + base, _ = table_base(f) + load!(builder, T_word, table_address(builder, base, index)) end - table_word(builder::IRBuilder, index::Int) = table_word(builder, table_offset(index)) mod_gvs = globals(mod) - rewrite_merged_slot_loads!(table_word, mod, - Pair{LLVM.Value,LLVM.Value}[mod_gvs[rec.name] => table_offset(index) - for (index, rec) in enumerate(relocs.records) - if rec.kind === SlotSite && haskey(mod_gvs, rec.name)]) + 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) + # 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) for (index, rec) in enumerate(relocs.records) haskey(mod_gvs, rec.name) || error("Missing relocation global '$(rec.name)'") @@ -956,105 +898,72 @@ function emit_table_relocations!(@nospecialize(job::CompilerJob), mod::LLVM.Modu check_relocation(mod, rec, gv) if rec.kind === SlotSite - rewrite_word_loads!(gv, "relocation slot '$(rec.name)'") do builder, offset - offset == 0 || - error("Relocation slot '$(rec.name)' is loaded at offset $offset") - table_word(builder, index) + addresses = Dict{LLVM.Function, LLVM.Value}() + function slot_address(f::LLVM.Function) + get!(addresses, f) do + base, entry = table_base(f) + @dispose builder=IRBuilder() begin + # After the base, but before any original instruction or PHI edge use. + position!(builder, entry) + ptr = table_address(builder, base, index) + pointercast!(builder, ptr, value_type(gv)) + end + end end - prune_constexpr_uses!(gv) - isempty(uses(gv)) || - error("Relocation slot '$(rec.name)' still has uses after redirection") - erase!(gv) + replace_global_with_local!(gv, slot_address) else demote_relocatable_box!(mod, gv, rec, table_word, index) end end - return -end -# Loads from a `phi` or `select` of slot addresses (see `merged_address_loads!`) cannot be -# redirected one slot at a time, so merge the slots' table offsets the same way instead and -# load the word at the merged offset. `slots` maps each slot to its offset, in table order. -# The offsets are constants and each merge is rebuilt next to the one it replaces, so the new -# values dominate wherever the old ones did. -function rewrite_merged_slot_loads!(table_word, mod::LLVM.Module, - slots::Vector{Pair{LLVM.Value,LLVM.Value}}) - dl = datalayout(mod) - offsets = Dict{LLVM.Value,LLVM.Value}(slots) - function slot_offset(value) - while value isa LLVM.ConstantExpr && constexpr_byte_offset(value, dl) == 0 - value = first(operands(value)) + # Table addresses may be in a different address space from the original slots. + # Let LLVM propagate that space through the existing PHIs, selects and casts. + @dispose pb=NewPMPassBuilder() begin + tti = llvm_targetinfo(job.config.target) + tti === nothing || LLVM.target_transform_info!(pb, tti) + add!(pb, NewPMFunctionPassManager()) do fpm + add!(fpm, InferAddressSpacesPass()) end - return get(offsets, value, nothing) + run!(pb, mod, llvm_machine(job.config.target)) end + return +end - merged = LLVM.Instruction[] - loads = LLVM.LoadInst[] - function find_merges(value) - for use in uses(value) +# 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. +function check_relocation_slot_uses!(mod::LLVM.Module, slots::Vector{LLVM.GlobalVariable}) + dl = datalayout(mod) + seen = Set{LLVM.Value}(slots) + worklist = LLVM.Value[slots...] + while !isempty(worklist) + for use in uses(pop!(worklist)) val = user(use) - if val isa LLVM.ConstantExpr - constexpr_byte_offset(val, dl) == 0 && find_merges(val) - elseif val isa LLVM.Instruction && forwarded_addresses(val) !== nothing && - !(val in merged) - other = merged_address_loads!(merged, loads, val) - other === nothing || - error("Unsupported use of merged relocation slot addresses: $other") + if val isa LLVM.LoadInst + is_word_type(value_type(val)) || + error("Unsupported relocation slot load of LLVM type $(value_type(val))") + # Julia names these loads after globals with session-specific counters. + LLVM.name!(val, "") + # The packed table guarantees word alignment, even if the old global had more. + alignment(val) > sizeof(UInt) && alignment!(val, sizeof(UInt)) + continue + elseif val isa LLVM.ConstantExpr + constexpr_byte_offset(val, dl) == 0 || + error("Unsupported relocation slot address $val") + elseif !(val isa LLVM.Instruction && forwarded_addresses(val) !== nothing) + error("Unsupported use of relocation slot address: $val") end + val in seen && continue + push!(seen, val) + push!(worklist, val) end end - foreach(find_merges ∘ first, slots) - isempty(merged) && return - - # A slot can only be merged with other slots: the table holds no word for other addresses. - for inst in merged, value in forwarded_addresses(inst) - value in merged || slot_offset(value) !== nothing || - error("Relocation slot address merged with unsupported address $value in $inst") - end - - T_offset = LLVM.Int32Type() - merged_offsets = Dict{LLVM.Value,LLVM.Value}() - @dispose builder=IRBuilder() begin - # `phi`s first, as they may be merged with themselves through a loop - for inst in merged - inst isa LLVM.PHIInst || continue - position!(builder, inst) - merged_offsets[inst] = phi!(builder, T_offset) - end - function merged_offset(value) - haskey(merged_offsets, value) && return merged_offsets[value] - offset = slot_offset(value) - offset === nothing || return offset - if value isa LLVM.SelectInst - cond, a, b = operands(value) - a, b = merged_offset(a), merged_offset(b) - position!(builder, value) - offset = select!(builder, cond, a, b) - else - offset = merged_offset(first(forwarded_addresses(value))) - end - merged_offsets[value] = offset + for val in seen + val isa LLVM.Instruction || continue + for address in forwarded_addresses(val) + address in seen || + error("Relocation slot address merged with unsupported address $address in $val") end - for inst in merged - if inst isa LLVM.PHIInst - append!(incoming(merged_offsets[inst]), - Tuple{LLVM.Value,LLVM.BasicBlock}[(merged_offset(value), block) - for (value, block) in incoming(inst)]) - else - merged_offset(inst) - end - end - end - - for load in loads - offset = merged_offsets[first(operands(load))] - replace_word_load!(builder -> table_word(builder, offset), load, - "merged relocation slot") - end - for inst in merged - replace_uses!(inst, PoisonValue(value_type(inst))) end - foreach(erase!, merged) return end diff --git a/test/metal.jl b/test/metal.jl index e1cca4ec..2a35f40c 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -427,7 +427,7 @@ 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, (i64 addrspace\(1\)\*|ptr addrspace\(1\))", air) + @test occursin(r"load (i64|ptr), (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) @@ -439,7 +439,7 @@ end # 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 - # has to turn into a `phi` of table offsets (#959). + # redirects to the device-space table (#959). 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) : @@ -456,19 +456,28 @@ end unsafe_store!(ptr, v === Val(1) ? 1f0 : v === nothing ? 2f0 : 3f0) return end + function loop_kernel(ptr, i) + n = unsafe_load(i) + v = pick(n) + for k in 1:n + v = k == 3 ? pick(k) : v + end + unsafe_store!(ptr, v === Val(1) ? 1f0 : v === Val(2) ? 2f0 : 3f0) + return + end end tt = (Core.LLVMPtr{Float32,1}, Core.LLVMPtr{Int,1}) @test @filecheck begin - @check "phi i32" - @check "load i64" + @check "phi ptr addrspace(1)" + @check "load ptr, ptr addrspace(1)" Metal.code_native_table(mod.kernel, tt; kernel=true) end - for f in (mod.kernel, mod.maybe_kernel) + 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("jl_global", air) - @test !occursin("jl_nothing", air) + @test !occursin(r"(?m)^@.*jl_global", air) + @test !occursin(r"(?m)^@.*jl_nothing", air) end end end diff --git a/test/native.jl b/test/native.jl index 77dd2e51..79a7b68c 100644 --- a/test/native.jl +++ b/test/native.jl @@ -1278,14 +1278,14 @@ end @testset "tabulated relocation of merged slots" begin # LLVM merges loads from different slots into one load from a `phi` or `select` of their - # addresses (#959). The table has no address to take their place, so the lowering has to - # merge their table offsets instead. + # addresses (#959). The lowering substitutes table entry addresses while preserving + # the existing PHIs and selects. if GPUCompiler.supports_relocatable_ir() && LLVM.version() >= v"17" mod = @eval module $(gensym()) f() = nothing end job, _ = Native.create_job(mod.f, Tuple{}; relocations=:table, jlruntime=false) - slots = ("merged_a", "merged_b", "merged_c") + slots = ("merged_a", "merged_b", "merged_c", "merged_d") relocations() = GPUCompiler.Relocations( [GPUCompiler.Relocation(GPUCompiler.SlotSite, name, 0, GPUCompiler.JuliaValueRef(Symbol(name))) @@ -1297,6 +1297,7 @@ end @merged_a = external global ptr @merged_b = external global ptr @merged_c = external global ptr + @merged_d = external addrspace(1) global ptr @jl_nothing = external global ptr define i64 @pick(i64 %i) { @@ -1314,6 +1315,13 @@ end join: %addr = phi ptr [ @merged_a, %one ], [ @merged_b, %two ], [ %other.addr, %other ] + %word = load i64, ptr %addr, align 16 + ret i64 %word + } + + define i64 @constant_cast(i1 %cond) { + %addr = select i1 %cond, ptr @merged_a, + ptr addrspacecast (ptr addrspace(1) @merged_d to ptr) %word = load i64, ptr %addr ret i64 %word } @@ -1350,7 +1358,10 @@ end for rec in relocs.records @test !haskey(globals(m), rec.name) end - @test occursin("phi i32", string(m)) + LLVM.verify(m) + @test all(alignment(inst) <= sizeof(UInt) + for bb in blocks(functions(m)["pick"]) for inst in instructions(bb) + if inst isa LLVM.LoadInst) fptr, lljit, table = Native.load(Vector{UInt8}(codeunits(obj)), "pick", relocs; table=true) @@ -1360,6 +1371,9 @@ end GPUCompiler.CGlobalRef(:jl_nothing)) @test [ccall(fptr, UInt, (Int,), i) for i in 1:4] == [word("merged_a"), word("merged_b"), word("merged_c"), nothing_word] + constant_cast = pointer(lookup(lljit, "constant_cast")) + @test [ccall(constant_cast, UInt, (Bool,), c) for c in (true, false)] == + word.(["merged_a", "merged_d"]) flip = pointer(lookup(lljit, "flip")) @test [ccall(flip, UInt, (Int,), n) for n in 1:4] == word.(["merged_a", "merged_a", "merged_b", "merged_b"]) @@ -1371,7 +1385,7 @@ end dispose(lljit) end - # the table holds no word for any other address... + # Mixed addresses need not share an address space on the target. m = parse(LLVM.Module, """ @merged_a = external global ptr @@ -1383,7 +1397,22 @@ end @test_throws "merged with unsupported address" GPUCompiler.emit_asm( job, m, relocations(), LLVM.API.LLVMObjectFile) - # ...and has no address to give out + # The table contains whole, read-only words. + for (body, message) in ( + ("%word = load i32, ptr %addr; ret i32 %word", "Unsupported relocation slot load"), + ("store i64 0, ptr %addr; ret i32 0", "Unsupported use of relocation slot address")) + m = parse(LLVM.Module, """ + @merged_a = external global ptr + @merged_b = external global ptr + define i32 @unsupported(i1 %cond) { + %addr = select i1 %cond, ptr @merged_a, ptr @merged_b + $(replace(body, "; " => "\n")) + }""") + @test_throws message GPUCompiler.emit_asm( + job, m, relocations(), LLVM.API.LLVMObjectFile) + end + + # Slot addresses are not exposed as general storage. m = parse(LLVM.Module, """ @merged_a = external global ptr @merged_b = external global ptr @@ -1393,7 +1422,7 @@ end %same = icmp eq ptr %addr, @merged_a ret i1 %same }""") - @test_throws "Unsupported use of merged relocation slot addresses" GPUCompiler.emit_asm( + @test_throws "Unsupported use of relocation slot address" GPUCompiler.emit_asm( job, m, relocations(), LLVM.API.LLVMObjectFile) end end