diff --git a/src/relocation.jl b/src/relocation.jl index c3e9b6f4..10a7685f 100644 --- a/src/relocation.jl +++ b/src/relocation.jl @@ -31,17 +31,23 @@ struct JuliaValueRef end """ - CGlobalRef(symbol, library=nothing) + CGlobalRef(symbol, library=nothing; offset=0) A named C data global. With `library === nothing`, resolution uses `jl_cglobal`'s process-wide lookup. Otherwise it looks up `symbol` in `library`. Resolution returns the -word stored in the global. +word stored at byte `offset`. """ struct CGlobalRef symbol::Symbol library::Union{Nothing,String} + offset::Int + + function CGlobalRef(symbol::Symbol, library::Union{Nothing,String}=nothing; + offset::Integer=0) + offset >= 0 || throw(ArgumentError("cglobal offset must be nonnegative")) + new(symbol, library, Int(offset)) + end end -CGlobalRef(symbol::Symbol) = CGlobalRef(symbol, nothing) """ RelocationTarget @@ -53,7 +59,7 @@ const RelocationTarget = Union{JuliaValueRef,CGlobalRef} same_relocation_target(a::JuliaValueRef, b::JuliaValueRef) = a.value === b.value same_relocation_target(a::CGlobalRef, b::CGlobalRef) = - a.symbol === b.symbol && a.library == b.library + a.symbol === b.symbol && a.library == b.library && a.offset == b.offset same_relocation_target(::RelocationTarget, ::RelocationTarget) = false # Permanently root a value in the current process and return the canonical rooted @@ -86,11 +92,11 @@ function resolve_relocation_target(target::CGlobalRef) if target.library === nothing # `jl_cglobal` accepts the symbol directly and does the process-wide `jl_dlfind`. address = ccall(:jl_cglobal, Any, (Any, Any), target.symbol, UInt) - return unsafe_load(address) + return unsafe_load(address + target.offset) end handle = Libdl.dlopen(target.library) address = Libdl.dlsym(handle, target.symbol) - return unsafe_load(Ptr{UInt}(address)) + return unsafe_load(Ptr{UInt}(address) + target.offset) end @@ -466,25 +472,56 @@ function materialize_box!(mod::LLVM.Module, relocs::Relocations, namespace::Stri return val end -# Rewrite every load of `value` into the word `produce_word(builder)` emits at the load's -# position, restoring a pointer with `inttoptr` where the original load produced one. -# Constant expressions (typed-pointer bitcasts) are recursed through, so both `i64` and -# pointer-typed words are handled. Shared by every producer and lowering that replaces a -# word-sized global with a run-time value. -function rewrite_word_loads!(produce_word, @nospecialize(value), what::String) +# Return the byte offset added by a constant cast or GEP, or `nothing` if it is not static. +function constexpr_byte_offset(ce::LLVM.ConstantExpr, dl::DataLayout) + op = opcode(ce) + if op == LLVM.API.LLVMBitCast || op == LLVM.API.LLVMAddrSpaceCast + return 0 + elseif op == LLVM.API.LLVMGetElementPtr + ops = operands(ce) + indices = ops[2:end] + all(idx -> idx isa LLVM.ConstantInt, indices) || return nothing + T = LLVMType(LLVM.API.LLVMGetGEPSourceElementType(ce)) + offset = convert(Int, indices[1]) * Int(abi_size(dl, T)) + for idx in indices[2:end] + i = convert(Int, idx) + if T isa LLVM.StructType + offset += Int(offsetof(dl, T, i)) + T = elements(T)[i+1] + elseif T isa LLVM.ArrayType || T isa LLVM.VectorType + T = eltype(T) + offset += i * Int(abi_size(dl, T)) + else + return nothing + end + end + return offset + end + 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) - changed |= rewrite_word_loads!(produce_word, val, what) + 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])") 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) + replacement = produce_word(builder, offset) T isa LLVM.PointerType && (replacement = inttoptr!(builder, replacement, T)) replace_uses!(val, replacement) @@ -496,8 +533,7 @@ function rewrite_word_loads!(produce_word, @nospecialize(value), what::String) return changed end -# Some Julia code loads words from libjulia C globals, for example type tags. Record those -# loads as dedicated zero-offset relocations immediately before object emission. +# Record loaded words from libjulia globals as one relocation slot per symbol and offset. function is_cglobal_candidate(value, relocs::Relocations) name = LLVM.name(value) value isa LLVM.GlobalVariable && @@ -515,20 +551,21 @@ function collect_cglobal_relocations!(@nospecialize(job::CompilerJob), mod::LLVM for f in [collect(functions(mod)); collect(globals(mod))] is_cglobal_candidate(f, relocs) || continue fn = LLVM.name(f) - slot = nothing - function cglobal_slot() - if slot === nothing - name = namespaced_name(namespace, "gpu_" * fn) + slots = Dict{Int,GlobalVariable}() + function cglobal_slot(offset::Int) + get!(slots, offset) do + # Including zero distinguishes `symbol` at N from `symbol_N` at zero. + name = namespaced_name(namespace, "gpu_$(fn)_$(offset)") slot = GlobalVariable(mod, relocation_word_type(), name) LLVM.name(slot) == name || error("cglobal slot name '$name' is already in use") - add_relocation!(relocs, SlotSite, name, 0, CGlobalRef(Symbol(fn))) + add_relocation!(relocs, SlotSite, name, 0, CGlobalRef(Symbol(fn); offset)) + slot end - slot end - changed |= rewrite_word_loads!(f, "cglobal '$fn'") do builder - load!(builder, relocation_word_type(), cglobal_slot()) + changed |= rewrite_word_loads!(f, "cglobal '$fn'") do builder, offset + load!(builder, relocation_word_type(), cglobal_slot(offset)) end end @@ -808,7 +845,9 @@ 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 + 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) end prune_constexpr_uses!(gv) diff --git a/test/metal.jl b/test/metal.jl index 18dbe2b2..daa15ec7 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -1489,6 +1489,34 @@ end end end +@testset "small typeof table relocation" begin + # Julia codegen references small-tagged datatypes like `Bool` through a constant offset + # into `jl_small_typeof`. The relocation must resolve to the actual type pointer, not to + # the table's (null) first word. `pointer_from_objref` lowers inline, keeping the load + # live in the kernel without any allocation or exception machinery. + function kernel(a::Core.LLVMPtr{UInt,1}) + unsafe_store!(a, UInt(pointer_from_objref(Bool))) + return + end + source = methodinstance(typeof(kernel), Tuple{Core.LLVMPtr{UInt,1}}, + Base.get_world_counter()) + target = MetalCompilerTarget(; macos=v"12.2", metal=v"3.0", air=v"3.0") + config = CompilerConfig(target, Metal.CompilerParams(); kernel=true) + job = CompilerJob(source, config) + + # Julia 1.10 still emits small-tagged types as literal pointers. + if VERSION >= v"1.11-" + # precondition: the unoptimized IR loads the tag through a nonzero table offset + ir = sprint(io->GPUCompiler.code_llvm(io, job; dump_module=true, optimize=false)) + @test occursin(r"getelementptr[^(]*\(i8, [^@]*@jl_small_typeof, i64 [1-9]\d*\)", ir) + + air = sprint(io->GPUCompiler.code_native(io, job; dump_module=true)) + @test occursin(Regex("store i64 $(UInt(pointer_from_objref(Bool))),"), air) + @test !occursin("jl_small_typeof", air) + @test !occursin(r"store i64 0,", air) + end +end + # byval lowering must strip the (non-IPO-safe) Julia const-region metadata off loads derived # from the materialized argument; check the helper walks gep/addrspacecast chains and removes it. @testset "const-region metadata stripping for materialized args" begin diff --git a/test/native.jl b/test/native.jl index af7a766e..6bd7a413 100644 --- a/test/native.jl +++ b/test/native.jl @@ -1244,6 +1244,78 @@ end @test occursin("@$(rec.name) = external global i64", string(mod)) GPUCompiler.emit_patchable_relocations!(mod, relocs) @test occursin("externally_initialized global i64 0", string(mod)) + + # Fold aggregate GEPs according to the module data layout. + mod = parse(LLVM.Module, """ + target datalayout = "e-p:64:64-i64:64" + @jl_layout = external global { i8, i64, [3 x i32] } + + define i32 @entry() { + %value = load i32, $(ptr("i32")) getelementptr ( + { i8, i64, [3 x i32] }, $(ptr("{ i8, i64, [3 x i32] }")) @jl_layout, + i64 0, i32 2, i64 1) + ret i32 %value + }""") + load = first(instructions(first(blocks(functions(mod)["entry"])))) + @test GPUCompiler.constexpr_byte_offset(operands(load)[1], datalayout(mod)) == 20 + @test_throws ArgumentError GPUCompiler.CGlobalRef(:jl_layout; offset=-1) + + # Julia codegen references small-tagged DataTypes through `jl_small_typeof` offsets. + small_typeof = cglobal(:jl_small_typeof, Ptr{Cvoid}) + bool_index = findfirst(i -> unsafe_load(small_typeof, i) == pointer_from_objref(Bool), + 1:(64 << 4) รท sizeof(Ptr{Cvoid})) + bool_index === nothing && error("Bool is absent from jl_small_typeof") + bool_offset = (bool_index - 1) * sizeof(Ptr{Cvoid}) + @test bool_offset > 0 + table_entry(offset) = GPUCompiler.supports_typed_pointers(ctx) ? + "bitcast (i8* getelementptr (i8, i8* @jl_small_typeof, i64 $offset) to $word_ptr_ptr)" : + "getelementptr (i8, ptr @jl_small_typeof, i64 $offset)" + table_ir = """ + @jl_small_typeof = external global i8 + + define $word_ptr @entry() { + %bool = load $word_ptr, $word_ptr_ptr $(table_entry(bool_offset)) + %first = load $word_ptr, $word_ptr_ptr $(table_entry(0)) + %junk = ptrtoint $word_ptr %first to i64 + ret $word_ptr %bool + }""" + + mod = parse(LLVM.Module, table_ir) + relocs = GPUCompiler.Relocations() + @test GPUCompiler.collect_cglobal_relocations!(job, mod, relocs) + @test !GPUCompiler.has_unresolved_cglobal_loads(mod, relocs) + @test length(relocs.records) == 2 + @test allunique(rec.name for rec in relocs.records) + @test all(rec -> endswith(rec.name, "_$(rec.target.offset)"), relocs.records) + @test Set(rec.target for rec in relocs.records) == + Set([GPUCompiler.CGlobalRef(:jl_small_typeof), + GPUCompiler.CGlobalRef(:jl_small_typeof; offset=bool_offset)]) + @test all(rec -> rec.kind === GPUCompiler.SlotSite && rec.offset == 0, relocs.records) + @test GPUCompiler.resolve_relocation_target( + GPUCompiler.CGlobalRef(:jl_small_typeof; offset=bool_offset)) == + UInt(pointer_from_objref(Bool)) + + mod = parse(LLVM.Module, table_ir) + GPUCompiler.prepare_execution!(job, mod) + @test occursin("inttoptr (i64 $(UInt(pointer_from_objref(Bool))) to $word_ptr)", + string(mod)) + + # Never silently relocate a dynamically-indexed load to the base word. + mod = parse(LLVM.Module, """ + @jl_small_typeof = external global i8 + @jl_float32_type = external global i8 + + define $word_ptr @entry() { + %value = load $word_ptr, $word_ptr_ptr $(GPUCompiler.supports_typed_pointers(ctx) ? + "bitcast (i8* getelementptr (i8, i8* @jl_small_typeof, i64 ptrtoint (i8* @jl_float32_type to i64)) to $word_ptr_ptr)" : + "getelementptr (i8, ptr @jl_small_typeof, i64 ptrtoint (ptr @jl_float32_type to i64))") + ret $word_ptr %value + }""") + relocs = GPUCompiler.Relocations() + @test_throws_message(ErrorException, + GPUCompiler.collect_cglobal_relocations!(job, mod, relocs)) do msg + occursin("Unsupported cglobal 'jl_small_typeof' load", msg) + end end end