From 548f540a9e5f6a7c3491a06ae2a6272d2fbfc127 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 5 Nov 2026 11:29:57 +0100 Subject: [PATCH] Metal: give every loop a single exit block `lower_unreachable_control_flow!` turns every exception into a branch to the kernel's return, which makes each bounds check an exit of the loops around it. Apple's back-end, which has to structurize control flow, hangs on loops with many such exits, or silently skips code after them (Metal.jl#1127). Run FixIrreducible and UnifyLoopExits after the last SimplifyCFG, as LLVM does before structurizing for AMDGPU. Before LLVM 24 these passes only handle branches, so switches are lowered first, and LLVM 20-22 crash on loops without exits, so functions with blocks that can't return are left as they are. --- src/metal.jl | 41 +++++++++++++++++ test/metal.jl | 121 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 162 insertions(+) diff --git a/src/metal.jl b/src/metal.jl index 0f3b6f36..1cd16dc6 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -778,6 +778,9 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L # narrow the promoted operations back, but before merging the byte GEPs this introduces. legalize_odd_width_integers!(mod) + # after the last SimplifyCFG above, which would undo it + unify_loop_exits!(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. @@ -789,6 +792,44 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L return end +# The Metal back-end can hang, or skip code, on loops with many exits (Metal.jl#1127), and +# `lower_unreachable_control_flow!` makes an exit of every exception. Route each loop's exits +# through one block, as LLVM does before structurizing for AMDGPU. +function unify_loop_exits!(mod::LLVM.Module) + @dispose pb=PassBuilder() begin + add!(pb, FunctionPassManager()) do fpm + # these passes only handle branches before LLVM 24 (llvm/llvm-project#206567) + LLVM.version() < v"24" && add!(fpm, LowerSwitchPass()) + add!(fpm, FixIrreduciblePass()) + end + run!(pb, mod) + end + + LLVM.library!(mod.library) do + for f in mod.functions + isdeclaration(f) && continue + # LLVM 20-22 crash on a loop without exits (llvm/llvm-project#165343) + v"20" <= LLVM.version() < v"23" && has_blocks_without_return_path(f) && continue + @dispose pb=PassBuilder() begin + add!(pb, UnifyLoopExitsPass()) + run!(pb, f) + 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) + while !isempty(work) + for pred in pop!(work).predecessors + pred in returning || (push!(returning, pred); push!(work, pred)) + end + end + return length(returning) < length(f.blocks) +end + # The AGX back-end keeps only the low 32 bits of a thread (AS 0) pointer that is loaded from # another address space. That breaks the host addresses Julia keeps by reference, e.g. of a # `Symbol` stored in an array (Metal.jl#1017), so load such pointers as words; pointers made diff --git a/test/metal.jl b/test/metal.jl index 496ae34d..c1da4c2c 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -951,6 +951,127 @@ end end end +@testset "loop exits" begin + # the blocks outside the loop with header `name` that the loop branches to + function exit_blocks(f, name) + header = only(filter(bb -> bb.name == name, collect(f.blocks))) + function reach(next) + seen = Set([header]) + work = [header] + while !isempty(work) + for bb in next(pop!(work)) + bb in seen || (push!(seen, bb); push!(work, bb)) + end + end + return seen + end + body = intersect(reach(bb -> bb.successors), reach(bb -> bb.predecessors)) + return unique(succ for bb in body for succ in bb.successors if !(succ in body)) + end + + JuliaContext() do ctx + # every exception leaves the loop it is in, which the back-end can hang or + # miscompile on (Metal.jl#1127) + mod = parse(LLVM.Module, """ + define i32 @checked(i32 %n, i32 %m) { + entry: + br label %loop + loop: + %i = phi i32 [ 0, %entry ], [ %next, %latch ] + %oob = icmp ugt i32 %i, %m + br i1 %oob, label %throw, label %latch + latch: + %next = add i32 %i, 1 + %more = icmp ult i32 %next, %n + br i1 %more, label %loop, label %done + throw: + ret i32 -1 + done: + ret i32 %next + } + + define i32 @switched(i32 %n, i32 %k) { + entry: + br label %loop + loop: + %i = phi i32 [ 0, %entry ], [ %next, %latch ] + switch i32 %k, label %latch [ i32 1, label %one + i32 2, label %two ] + latch: + %next = add i32 %i, 1 + %more = icmp ult i32 %next, %n + br i1 %more, label %loop, label %done + one: + ret i32 1 + two: + ret i32 2 + done: + ret i32 %next + } + + define i32 @irreducible(i1 %c, i32 %n, i32 %m) { + entry: + br i1 %c, label %a, label %b + a: + %i = phi i32 [ 0, %entry ], [ %j, %b ] + %next = add i32 %i, 1 + %oob = icmp ugt i32 %next, %m + br i1 %oob, label %throw, label %b + b: + %j = phi i32 [ 0, %entry ], [ %next, %a ] + %more = icmp ult i32 %j, %n + br i1 %more, label %a, label %done + throw: + ret i32 -1 + done: + ret i32 %j + } + + define void @endless() { + entry: + br label %loop + loop: + br label %loop + } + + define void @endless_after_lowering() { + entry: + br label %loop + loop: + switch i32 0, label %done [ i32 0, label %loop ] + done: + ret void + } + """) + GPUCompiler.unify_loop_exits!(mod) + @test isnothing(verify(mod)) + @test length(exit_blocks(mod.functions["checked"], "loop")) == 1 + @test length(exit_blocks(mod.functions["switched"], "loop")) == 1 + @test length(exit_blocks(mod.functions["irreducible"], "a")) == 1 + dispose(mod) + end + + # the same for exceptions in a kernel, after the rest of the lowering + mod = @eval module $(gensym()) + function kernel(ptr, n, m) + acc = 0f0 + i = 1 + while i <= n + i > m && throw(BoundsError()) + acc += unsafe_load(ptr, i) + i + 1 > m && throw(BoundsError()) + acc += unsafe_load(ptr, i + 1) + i += 2 + end + unsafe_store!(ptr, acc) + return + end + end + air = sprint(io -> Metal.code_native(io, mod.kernel, (Core.LLVMPtr{Float32,1}, Int, Int); + kernel=true)) + @test occursin("loop.exit", air) +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