[llvm] AArch64: Fix assert on fold of undef operand (PR #218600)

[Matt Arsenault via llvm-commits]

https://github.com/arsenm created https://github.com/llvm/llvm-project/pull/218600

Co-Authored-By: Claude <noreply at anthropic.com> (Claude Opus 4.8)

>From fb8f0e6d66ec4de8cc0e793cada32a994d53e923 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault at amd.com>
Date: Mon, 24 Aug 2026 22:55:37 +0200
Subject: [PATCH] AArch64: Fix assert on fold of undef operand

Co-Authored-By: Claude <noreply at anthropic.com> (Claude Opus 4.8)
---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp  |   7 +-
 .../AArch64/optimize-cond-branch-undef.mir    | 104 ++++++++++++++++++
 2 files changed, 108 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index 522b263df5de1..9867164af2e41 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -10102,9 +10102,9 @@ bool AArch64InstrInfo::optimizeCondBranch(MachineInstr &MI) const {
       return false;
     if (!MRI->hasOneNonDBGUse(CopyVReg))
       return false;
-    if (!MRI->hasOneDef(CopyVReg))
-      return false;
     DefMI = MRI->getVRegDef(CopyVReg);
+    if (!DefMI)
+      return false;
   }
 
   switch (DefMI->getOpcode()) {
@@ -10131,7 +10131,8 @@ bool AArch64InstrInfo::optimizeCondBranch(MachineInstr &MI) const {
     if (!NewReg.isVirtual())
       return false;
 
-    assert(!MRI->def_empty(NewReg) && "Register must be defined.");
+    if (!MRI->getVRegDef(NewReg))
+      return false;
 
     MachineBasicBlock &RefToMBB = *MBB;
     MachineBasicBlock *TBB = MI.getOperand(1).getMBB();
diff --git a/llvm/test/CodeGen/AArch64/optimize-cond-branch-undef.mir b/llvm/test/CodeGen/AArch64/optimize-cond-branch-undef.mir
index 7a4d9b98ca1a1..40cc6b826faf8 100644
--- a/llvm/test/CodeGen/AArch64/optimize-cond-branch-undef.mir
+++ b/llvm/test/CodeGen/AArch64/optimize-cond-branch-undef.mir
@@ -30,3 +30,107 @@ body: |
   bb.2:
     RET_ReallyLR
 ...
+
+---
+name: cbz_and_undef_source
+tracksRegLiveness: true
+body: |
+  ; CHECK-LABEL: name: cbz_and_undef_source
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[ANDWri:%[0-9]+]]:gpr32sp = ANDWri undef %1:gpr32, 0
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gpr32 = COPY [[ANDWri]]
+  ; CHECK-NEXT:   CBZW [[COPY]], %bb.2
+  ; CHECK-NEXT:   B %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   RET_ReallyLR
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   RET_ReallyLR
+  bb.0:
+    successors: %bb.1, %bb.2
+
+    %0:gpr32sp = ANDWri undef %1:gpr32, 0
+    %2:gpr32 = COPY %0
+    CBZW %2, %bb.2
+    B %bb.1
+
+  bb.1:
+    RET_ReallyLR
+
+  bb.2:
+    RET_ReallyLR
+...
+
+# The branch condition is a COPY of an undef register.
+---
+name: cbz_copy_of_undef
+tracksRegLiveness: true
+body: |
+  ; CHECK-LABEL: name: cbz_copy_of_undef
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gpr32 = COPY undef %1:gpr32
+  ; CHECK-NEXT:   CBZW [[COPY]], %bb.2
+  ; CHECK-NEXT:   B %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   RET_ReallyLR
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   RET_ReallyLR
+  bb.0:
+    successors: %bb.1, %bb.2
+
+    %0:gpr32 = COPY undef %1:gpr32
+    CBZW %0, %bb.2
+    B %bb.1
+
+  bb.1:
+    RET_ReallyLR
+
+  bb.2:
+    RET_ReallyLR
+...
+
+# A defined AND source is still folded into a TBZ through the copy
+# chain: the undef guard must not block the ordinary case.
+---
+name: tbz_and_defined_source
+tracksRegLiveness: true
+body: |
+  ; CHECK-LABEL: name: tbz_and_defined_source
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK-NEXT:   liveins: $w0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gpr32 = COPY $w0
+  ; CHECK-NEXT:   [[ANDWri:%[0-9]+]]:gpr32sp = ANDWri [[COPY]], 0
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:gpr32 = COPY [[ANDWri]]
+  ; CHECK-NEXT:   TBZW [[COPY]], 0, %bb.2
+  ; CHECK-NEXT:   B %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   RET_ReallyLR
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   RET_ReallyLR
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: $w0
+
+    %3:gpr32 = COPY $w0
+    %0:gpr32sp = ANDWri %3, 0
+    %2:gpr32 = COPY %0
+    CBZW %2, %bb.2
+    B %bb.1
+
+  bb.1:
+    RET_ReallyLR
+
+  bb.2:
+    RET_ReallyLR
+...