[llvm] [SelectionDAGBuilder] Look for apropriate INLINEASM_BR instruction to verify (PR #152591)

[via llvm-commits]

https://github.com/XChy created https://github.com/llvm/llvm-project/pull/152591

Partially fix #149023.
The original code `MRI.def_begin(Reg)->getParent()` may return the incorrect MI, as the physical register `Reg` may have multiple definitions. 
This patch selects the correct MI to verify by comparing the MBB of each definition.

New testcase hangs with -O1/2/3 enabled. The BranchFolding may be to blame.

>From 02989d77071cabb556962df8339ab687e64fdd5c Mon Sep 17 00:00:00 2001
From: XChy <xxs_chy at outlook.com>
Date: Fri, 8 Aug 2025 04:15:08 +0800
Subject: [PATCH] [SelectionDAGBuilder] Look for apropriate INLINEASM_BR
 instruction to verify

---
 .../SelectionDAG/SelectionDAGBuilder.cpp      | 21 ++++++--
 llvm/test/CodeGen/X86/callbr-asm-loop.ll      | 50 +++++++++++++++++++
 2 files changed, 68 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/callbr-asm-loop.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index d0815e9f51822..c345c6207dfee 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -12731,23 +12731,37 @@ void SelectionDAGBuilder::visitVectorSplice(const CallInst &I) {
 // increase in virtreg number from there. If a callbr has no outputs, then it
 // should not have a corresponding callbr landingpad; in fact, the callbr
 // landingpad would not even be able to refer to such a callbr.
-static Register FollowCopyChain(MachineRegisterInfo &MRI, Register Reg) {
+static Register FollowCopyChain(MachineRegisterInfo &MRI, Register Reg,
+                                MachineBasicBlock *CallBrMBB) {
   MachineInstr *MI = MRI.def_begin(Reg)->getParent();
   // There is definitely at least one copy.
   assert(MI->getOpcode() == TargetOpcode::COPY &&
          "start of copy chain MUST be COPY");
   Reg = MI->getOperand(1).getReg();
+
+  // If the copied register in the first copy must be virtual.
+  assert(Reg.isVirtual() && "expected COPY of virtual register");
   MI = MRI.def_begin(Reg)->getParent();
+
   // There may be an optional second copy.
   if (MI->getOpcode() == TargetOpcode::COPY) {
     assert(Reg.isVirtual() && "expected COPY of virtual register");
     Reg = MI->getOperand(1).getReg();
     assert(Reg.isPhysical() && "expected COPY of physical register");
-    MI = MRI.def_begin(Reg)->getParent();
+
+    // Look for the machine callbr instruction.
+    for (auto Def : MRI.def_operands(Reg))
+      if (Def.getParent()->getOpcode() == TargetOpcode::INLINEASM_BR &&
+          Def.getParent()->getParent() == CallBrMBB) {
+        MI = Def.getParent();
+        break;
+      }
   }
+
   // The start of the chain must be an INLINEASM_BR.
   assert(MI->getOpcode() == TargetOpcode::INLINEASM_BR &&
          "end of copy chain MUST be INLINEASM_BR");
+
   return Reg;
 }
 
@@ -12766,6 +12780,7 @@ void SelectionDAGBuilder::visitCallBrLandingPad(const CallInst &I) {
   MachineRegisterInfo &MRI = DAG.getMachineFunction().getRegInfo();
 
   Register InitialDef = FuncInfo.ValueMap[CBR];
+  MachineBasicBlock *CallBrMBB = FuncInfo.getMBB(CBR->getParent());
   SDValue Chain = DAG.getRoot();
 
   // Re-parse the asm constraints string.
@@ -12789,7 +12804,7 @@ void SelectionDAGBuilder::visitCallBrLandingPad(const CallInst &I) {
       // getRegistersForValue may produce 1 to many registers based on whether
       // the OpInfo.ConstraintVT is legal on the target or not.
       for (Register &Reg : OpInfo.AssignedRegs.Regs) {
-        Register OriginalDef = FollowCopyChain(MRI, InitialDef++);
+        Register OriginalDef = FollowCopyChain(MRI, InitialDef++, CallBrMBB);
         if (OriginalDef.isPhysical())
           FuncInfo.MBB->addLiveIn(OriginalDef);
         // Update the assigned registers to use the original defs.
diff --git a/llvm/test/CodeGen/X86/callbr-asm-loop.ll b/llvm/test/CodeGen/X86/callbr-asm-loop.ll
new file mode 100644
index 0000000000000..f2a63db9107f2
--- /dev/null
+++ b/llvm/test/CodeGen/X86/callbr-asm-loop.ll
@@ -0,0 +1,50 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+
+; RUN: llc -O0 -mtriple=i686-- -verify-machineinstrs < %s | FileCheck %s
+
+; Test that causes multiple defs of %eax.
+; FIXME: The testcase hangs with -O1/2/3 enabled.
+define i32 @loop1() {
+; CHECK-LABEL: loop1:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    pushl %esi
+; CHECK-NEXT:    .cfi_def_cfa_offset 8
+; CHECK-NEXT:    .cfi_offset %esi, -8
+; CHECK-NEXT:    jmp .LBB0_1
+; CHECK-NEXT:  .LBB0_1: # %tailrecurse
+; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    xorl %eax, %eax
+; CHECK-NEXT:    movl $1, %edx
+; CHECK-NEXT:    #APP
+; CHECK-NEXT:    #NO_APP
+; CHECK-NEXT:    movl %eax, %ecx
+; CHECK-NEXT:    movl %edx, %esi
+; CHECK-NEXT:    jmp .LBB0_3
+; CHECK-NEXT:  .LBB0_2: # Inline asm indirect target
+; CHECK-NEXT:    # %tailrecurse.tailrecurse.backedge_crit_edge
+; CHECK-NEXT:    # in Loop: Header=BB0_1 Depth=1
+; CHECK-NEXT:    # Label of block must be emitted
+; CHECK-NEXT:  .LBB0_3: # %tailrecurse.backedge
+; CHECK-NEXT:    # in Loop: Header=BB0_1 Depth=1
+; CHECK-NEXT:    jmp .LBB0_1
+; CHECK-NEXT:  .LBB0_4: # Inline asm indirect target
+; CHECK-NEXT:    # %lab2.split
+; CHECK-NEXT:    # Label of block must be emitted
+; CHECK-NEXT:    movl %edx, %eax
+; CHECK-NEXT:    popl %esi
+; CHECK-NEXT:    .cfi_def_cfa_offset 4
+; CHECK-NEXT:    retl
+entry:
+  br label %tailrecurse
+
+tailrecurse:
+  %0 = callbr { i32, i32 } asm "", "={ax},={dx},0,1,!i,!i"(i32 0, i32 1) #1
+          to label %tailrecurse.backedge [label %tailrecurse.backedge, label %lab2.split]
+
+tailrecurse.backedge:
+  br label %tailrecurse
+
+lab2.split:
+  %asmresult5 = extractvalue { i32, i32 } %0, 1
+  ret i32 %asmresult5
+}