diff --git a/src/metal.jl b/src/metal.jl index dc6f32ed..0bf75392 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -390,6 +390,51 @@ function split_aggregate_loads!(mod::LLVM.Module) return changed end +# Flatten chained single-index byte `getelementptr`s into one: `gep i8, (gep i8, p, A), B` +# -> `gep i8, p, (A + B)`. The AGX back-end miscompiles a 1-byte load/store made through a +# chained GEP (a byte GEP whose base is another byte GEP) when the grid has exactly two +# threadgroups -- the first threadgroup's access is silently dropped. LLVM deliberately keeps +# such chains split: `InstCombine`'s `visitGEPOfGEP` only merges when the combined index folds +# to an existing value (it bails on variable-plus-constant to avoid materializing an extra +# `add`), and the Metal driver never coalesces them either. On a normal back-end the split form +# is free (the constant GEP folds into the addressing mode), so this is purely an AGX defect; we +# work around it by force-merging here, which only costs a cheap `add` and makes the back-end +# emit correct code. Runs on the optimized, opaque-pointer IR, before AIR lowering. +function merge_byte_gep_chains!(mod::LLVM.Module) + changed = false + i8 = LLVM.Int8Type() + is_byte_gep(v) = + v isa LLVM.GetElementPtrInst && length(operands(v)) == 2 && + LLVM.LLVMType(LLVM.API.LLVMGetGEPSourceElementType(v)) == i8 + + for f in functions(mod) + isdeclaration(f) && continue + worklist = LLVM.Instruction[] + for bb in blocks(f), inst in instructions(bb) + is_byte_gep(inst) && is_byte_gep(operands(inst)[1]) && push!(worklist, inst) + end + # process defs-before-uses (instruction order) so longer chains collapse fully: + # an inner gep is rewritten before the outer one that consumes it. + for gep in worklist + src = operands(gep)[1] + (is_byte_gep(gep) && is_byte_gep(src)) || continue # may already be merged + base = operands(src)[1] + inbounds = LLVM.API.LLVMIsInBounds(gep) != 0 && LLVM.API.LLVMIsInBounds(src) != 0 + @dispose builder=IRBuilder() begin + position!(builder, gep) + sum = add!(builder, operands(src)[2], operands(gep)[2]) + newgep = inbounds ? inbounds_gep!(builder, i8, base, [sum]) : + gep!(builder, i8, base, [sum]) + replace_uses!(gep, newgep) + end + erase!(gep) + isempty(uses(src)) && erase!(src) + changed = true + end + end + return changed +end + function finish_ir!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::LLVM.Module, entry::LLVM.Function) entry_fn = LLVM.name(entry) @@ -494,6 +539,11 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L end end + # flatten chained byte GEPs that the AGX back-end miscompiles for 1-byte accesses on a + # 2-threadgroup grid (see `merge_byte_gep_chains!`). run last, after the intrinsic-lowering + # cleanup above, so the merged form is what reaches the AIR downgrader / back-end. + merge_byte_gep_chains!(mod) + return end diff --git a/test/metal.jl b/test/metal.jl index 1173514f..fb18b6ee 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -139,6 +139,26 @@ end end end +@testset "byte GEP coalescing" begin + # The AGX back-end miscompiles a 1-byte load/store addressed through a *chained* byte + # `getelementptr` (`gep i8, (gep i8, p, i), 1`) on a grid of exactly two threadgroups, + # silently dropping the first threadgroup's access. LLVM deliberately leaves such chains + # split (`InstCombine`'s `visitGEPOfGEP` only merges when the combined index folds, to avoid + # an extra `add`), so `merge_byte_gep_chains!` coalesces them into a single GEP during AIR + # lowering. Without that pass the kernel below keeps two GEPs. + mod = @eval module $(gensym()) + function kernel(ptr::Core.LLVMPtr{Int8,1}, i::Int64) + p = ptr + i + Base.unsafe_store!(p + Int64(1), Int8(0)) + return + end + end + + air = sprint(io -> Metal.code_native(io, mod.kernel, + Tuple{Core.LLVMPtr{Int8,1}, Int64}; kernel=true)) + @test count("getelementptr", air) == 1 +end + @testset "GC runtime input arguments" begin mod = @eval module $(gensym()) mutable struct PleaseAllocate