diff --git a/src/relocation.jl b/src/relocation.jl index f7ba7ef9..2371fa83 100644 --- a/src/relocation.jl +++ b/src/relocation.jl @@ -500,35 +500,74 @@ 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)) +is_word_type(T::LLVMType) = + T isa LLVM.PointerType || (T isa LLVM.IntegerType && width(T) == 8sizeof(UInt)) + +# 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 + +# 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) inner = (offset === nothing || delta === nothing) ? nothing : offset + delta - changed |= rewrite_word_loads!(produce_word, val, what; offset=inner, dl) - elseif isa(val, LLVM.LoadInst) + 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])") - 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) - end - erase!(val) - changed = true + 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 @@ -564,20 +603,24 @@ 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 + changed |= redirect_word_addresses!(cglobal_slot, f, "cglobal '$fn'") end return changed end function has_unresolved_cglobal_loads(mod::LLVM.Module, relocs::Relocations) - function has_load(value) + # 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) 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 @@ -822,42 +865,105 @@ 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 + 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), - [ConstantInt(LLVM.Int32Type(), index - 1)]) - load!(builder, T_word, ptr) + base, _ = table_base(f) + load!(builder, T_word, table_address(builder, base, index)) end mod_gvs = globals(mod) + 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)'") gv = mod_gvs[rec.name] 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 + + # 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 + run!(pb, mod, llvm_machine(job.config.target)) + end + return +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. +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.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 + 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 + end return end diff --git a/test/metal.jl b/test/metal.jl index 9b69b434..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) @@ -435,6 +435,53 @@ 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). + 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 + 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 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, 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) + 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..79a7b68c 100644 --- a/test/native.jl +++ b/test/native.jl @@ -1276,6 +1276,158 @@ 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 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", "merged_d") + 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 + @merged_d = external addrspace(1) 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, 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 + } + + 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 + 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) + 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] + 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"]) + 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 + + # Mixed addresses need not share an address space on the target. + 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) + + # 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 + + 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 relocation slot address" 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 +1622,67 @@ 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 @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, align 16 + 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 + # 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) + 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"