From cd738f9ed39225103558fb954058aa847d46cfa0 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 7 Nov 2026 16:28:51 +0100 Subject: [PATCH] Metal: define undefined phi values The AGX back-end of M3 and newer GPUs miscompiles some phis that have an undef incoming value, even when that edge is never taken at run time. KernelAbstractions kernels routinely contain such phis: a variable assigned before a `@synchronize` and used after it is undefined on the path that skips the out-of-range threads. KernelAbstractions' `forloop` test then computed with the same local index in every thread, so all threads updated the same element, and similar kernels threw spurious bounds errors. Replace undef and poison phi operands with zero, after the loop-exit unification that introduces more of them. --- src/metal.jl | 24 ++++++++++++++++++++++++ test/metal.jl | 30 ++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+) 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