Metal: define undefined phi values

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)

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