CodeGen: Fix machine sink critical edge crash on an undef register

[arsenm]

Found by AI while working on something else.

Co-authored-by: Claude (Opus 4.8) noreply@anthropic.com

[arsenm]

• CodeGen: Fix machine sink critical edge crash on an undef register #216870 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmorg-github-actions]

@llvm/pr-subscribers-backend-x86

Author: Matt Arsenault (arsenm)

Changes

Found by AI while working on something else.

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

---

Full diff: https://github.com/llvm/llvm-project/pull/216870.diff

2 Files Affected:

• (modified) llvm/lib/CodeGen/MachineSink.cpp (+1-2)
• (modified) llvm/test/CodeGen/X86/machinesink-coalesce-undef.mir (+59)

diff --git a/llvm/lib/CodeGen/MachineSink.cpp b/llvm/lib/CodeGen/MachineSink.cpp
index 2472005369c0a..161407cf2cd24 100644
--- a/llvm/lib/CodeGen/MachineSink.cpp
+++ b/llvm/lib/CodeGen/MachineSink.cpp
@@ -1102,8 +1102,7 @@ bool MachineSinking::isWorthBreakingCriticalEdge(
       // claim it's likely we can sink these together.
       // If definition resides elsewhere, we aren't
       // blocking it from being sunk so don't break the edge.
-      MachineInstr *DefMI = MRI->getVRegDef(Reg);
-      if (DefMI->getParent() == MI.getParent())
+      if (MRI->getDefBlock(Reg) == MI.getParent())
         return true;
     }
   }
diff --git a/llvm/test/CodeGen/X86/machinesink-coalesce-undef.mir b/llvm/test/CodeGen/X86/machinesink-coalesce-undef.mir
index 54e0c1d701321..37110754de27b 100644
--- a/llvm/test/CodeGen/X86/machinesink-coalesce-undef.mir
+++ b/llvm/test/CodeGen/X86/machinesink-coalesce-undef.mir
@@ -37,3 +37,62 @@ body: |
   bb.2:
     RET64
 ...
+
+# A copy sunk across a critical edge whose source is an undef register
+# with no def must not crash isWorthBreakingCriticalEdge.
+---
+name: crit_edge_undef_copy
+tracksRegLiveness: true
+body: |
+  ; CHECK-LABEL: name: crit_edge_undef_copy
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.3(0x40000000)
+  ; CHECK-NEXT:   liveins: $edi
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gr32 = COPY $edi
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:gr32 = COPY undef %2:gr32
+  ; CHECK-NEXT:   TEST32rr [[COPY]], [[COPY]], implicit-def $eflags
+  ; CHECK-NEXT:   JCC_1 %bb.1, 4, implicit $eflags
+  ; CHECK-NEXT:   JMP_1 %bb.3
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   successors: %bb.2(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   MOV32mr $rip, 1, $noreg, 0, $noreg, [[COPY1]] :: (store (s32))
+  ; CHECK-NEXT:   JMP_1 %bb.2
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.3(0x40000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   TEST32rr [[COPY]], [[COPY]], implicit-def $eflags
+  ; CHECK-NEXT:   JCC_1 %bb.1, 4, implicit $eflags
+  ; CHECK-NEXT:   JMP_1 %bb.3
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.3:
+  ; CHECK-NEXT:   RET64
+  bb.0:
+    successors: %bb.1(0x40000000), %bb.3(0x40000000)
+    liveins: $edi
+
+    %0:gr32 = COPY $edi
+    %1:gr32 = COPY undef %2:gr32
+    TEST32rr %0, %0, implicit-def $eflags
+    JCC_1 %bb.1, 4, implicit $eflags
+    JMP_1 %bb.3
+
+  bb.1:
+    successors: %bb.2(0x80000000)
+
+    MOV32mr $rip, 1, $noreg, 0, $noreg, %1 :: (store (s32))
+    JMP_1 %bb.2
+
+  bb.2:
+    successors: %bb.1(0x40000000), %bb.3(0x40000000)
+
+    TEST32rr %0, %0, implicit-def $eflags
+    JCC_1 %bb.1, 4, implicit $eflags
+    JMP_1 %bb.3
+
+  bb.3:
+    RET64
+...

[lukel97]

LGTM