diff --git a/src/metal.jl b/src/metal.jl index 1cd16dc6..6b5b6855 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -781,6 +781,9 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L # after the last SimplifyCFG above, which would undo it unify_loop_exits!(mod) + # after the exit unification, which introduces such values + define_undef_phi_values!(mod) + # 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. @@ -819,6 +822,27 @@ function unify_loop_exits!(mod::LLVM.Module) return end +# The AGX back-end of Apple9 and later GPUs (M3 and newer) can miscompile a `phi` with an +# undefined incoming value, even when the edge carrying it is never taken: a kernel then +# computes with the same value in every thread instead of the defined one. KernelAbstractions +# kernels are prone to this, as a variable assigned in one `@synchronize`-separated section of +# a kernel and used in the next is undefined on the path that skips the out-of-range threads. +# Define those values; any value is a valid refinement, so pick zero. +function define_undef_phi_values!(mod::LLVM.Module) + for f in mod.functions, bb in f.blocks + for inst in bb.instructions + inst isa LLVM.PHIInst || break + ops = LLVM.operands(inst) + for i in 1:length(ops) + if ops[i] isa Union{UndefValue, PoisonValue} + ops[i] = null(inst.value_type) + end + end + end + end + return +end + function has_blocks_without_return_path(f::LLVM.Function) returning = Set(bb for bb in f.blocks if bb.terminator isa LLVM.RetInst) work = collect(returning) diff --git a/test/metal.jl b/test/metal.jl index c1da4c2c..5aef587b 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -1072,6 +1072,36 @@ end @test occursin("loop.exit", air) end +@testset "undefined phi values" begin + # the back-end of M3 and newer GPUs can miscompile these, even on a path that is never + # taken (a KernelAbstractions kernel with `@synchronize` returned wrong results) + Context(; opaque_pointers=true) do ctx + mod = parse(LLVM.Module, """ + define i64 @f(i1 %c, i64 %x, ptr %p) { + entry: + br i1 %c, label %def, label %join + def: + br label %join + join: + %v = phi i64 [ %x, %def ], [ undef, %entry ] + %w = phi ptr [ %p, %def ], [ poison, %entry ] + %u = phi { i32, float } [ zeroinitializer, %def ], [ undef, %entry ] + store i64 %v, ptr %w + ret i64 %v + } + """) + GPUCompiler.define_undef_phi_values!(mod) + @test isnothing(verify(mod)) + phis = filter(inst -> inst isa LLVM.PHIInst, + collect(LLVM.instructions(mod.functions["f"].blocks[3]))) + @test length(phis) == 3 + for phi in phis + @test !any(op -> op isa Union{UndefValue, PoisonValue}, LLVM.operands(phi)) + end + dispose(mod) + end +end + @testset "reproducible output" begin # Compiling the same kernel twice in one session has to yield byte-identical AIR: the two # compilations see different LLVM wrapper addresses and codegen counters, so any pass that