Metal: fix byte-GEP coalescing crash on cross-block chains (#863)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: maleadt

src/metal.jl

@@ -407,18 +407,25 @@ function merge_byte_gep_chains!(mod::LLVM.Module)
         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[]
+    # the first `gep i8, (gep i8, p, A), B` in `f`, or `nothing`
+    function next_chain(f)
         for bb in blocks(f), inst in instructions(bb)
-            is_byte_gep(inst) && is_byte_gep(operands(inst)[1]) && push!(worklist, inst)
+            is_byte_gep(inst) && is_byte_gep(operands(inst)[1]) && return 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]
+        return nothing
+    end
+
+    for f in functions(mod)
+        isdeclaration(f) && continue
+        # Rescan after each merge instead of keeping a worklist: a merge can erase a *different*
+        # chained GEP (a now-dead inner GEP whose only use was the one we just folded), and a
+        # block layout that doesn't follow dominance can place that inner GEP after the outer one,
+        # so stale worklist entries would be use-after-free. Rescanning only ever inspects live
+        # instructions; each merge replaces a chained GEP with a shallower one (and shortens its
+        # users), so the total chain depth strictly decreases and this terminates.
+        while (gep = next_chain(f)) !== nothing
+            src  = operands(gep)[1]
+            base = operands(src)[1]
             inbounds = LLVM.API.LLVMIsInBounds(gep) != 0 && LLVM.API.LLVMIsInBounds(src) != 0
             @dispose builder=IRBuilder() begin
                 position!(builder, gep)

test/metal.jl

@@ -159,6 +159,38 @@ end
     @test count("getelementptr", air) == 1
 end
 
+@testset "byte GEP coalescing across blocks" begin
+    # Regression: when the outer GEP's block is laid out before its (dominating) inner base's
+    # block, the inner GEP can be erased before it is reached. A stale worklist would then
+    # dereference a freed instruction (`BoundsError` on a 0-element operand set), so the pass
+    # must rescan live instructions instead. It must also still fully coalesce the 3-level chain.
+    p = opaque_ptrs ? "ptr addrspace(1)" : "i8 addrspace(1)*"
+    ir = """
+    define void @k($p %p, i64 %i, i64 %a, i64 %b) {
+    entry:
+      br label %bbB
+    bbA:
+      %g3 = getelementptr i8, $p %g2, i64 %b
+      store i8 0, $p %g3, align 1
+      br label %end
+    bbB:
+      %g1 = getelementptr i8, $p %p, i64 %i
+      %g2 = getelementptr i8, $p %g1, i64 %a
+      br label %bbA
+    end:
+      ret void
+    }
+    """
+    @dispose ctx=Context() begin
+        mod = parse(LLVM.Module, ir)
+        GPUCompiler.merge_byte_gep_chains!(mod)
+        f = functions(mod)["k"]
+        ngep = sum(bb -> count(i -> i isa LLVM.GetElementPtrInst, collect(instructions(bb))),
+                   collect(blocks(f)))
+        @test ngep == 1
+    end
+end
+
 @testset "GC runtime input arguments" begin
     mod = @eval module $(gensym())
         mutable struct PleaseAllocate