[llvm] Branch folding: Do not use both TII::analyzeBranch and MBB::isSuccessor (PR #104240)

[via llvm-commits]

https://github.com/v01dXYZ updated https://github.com/llvm/llvm-project/pull/104240

>From 0280f1405d26439a61b871b6407a9267a779b7ab Mon Sep 17 00:00:00 2001
From: v01dxyz <v01dxyz at v01d.xyz>
Date: Thu, 27 Jun 2024 21:33:57 +0200
Subject: [PATCH] Branch folding: Do not use both TII::analyzeBranch and
 MBB::isSuccessor

Avoid infinite loop with EHPad blocks or indirect target blocks.

After a ``!TII->analyzeBranch(PrevBB, PrevTBB, PrevFBB, PrevCond,
true)``, the code previously used ``PrevBB.isSuccessor(&*MBB)`` to
check if MBB is PrevTBB or PrevFBB.

But ``TII->analyzeBranch(...)`` only supports unconditional/conditional
branches while INVOKE or CALLBR could also add blocks as successors (EHPad
blocks for INVOKE, Indirect Target Blocks for CALLBR).
---
 llvm/lib/CodeGen/BranchFolding.cpp            | 32 +++++++------------
 ...loop-with-callbr-indirect-target-blocks.ll | 18 +++++++++++
 2 files changed, 29 insertions(+), 21 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-target-blocks.ll

diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index 1dc278586f1178..1d52a461d626bf 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -1343,13 +1343,6 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
     SameEHScope = MBBEHScope->second == FallThroughEHScope->second;
   }
 
-  // Analyze the branch in the current block. As a side-effect, this may cause
-  // the block to become empty.
-  MachineBasicBlock *CurTBB = nullptr, *CurFBB = nullptr;
-  SmallVector<MachineOperand, 4> CurCond;
-  bool CurUnAnalyzable =
-      TII->analyzeBranch(*MBB, CurTBB, CurFBB, CurCond, true);
-
   // If this block is empty, make everyone use its fall-through, not the block
   // explicitly.  Landing pads should not do this since the landing-pad table
   // points to this block.  Blocks with their addresses taken shouldn't be
@@ -1422,8 +1415,8 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
     // This has to check PrevBB->succ_size() because EH edges are ignored by
     // analyzeBranch.
     if (PriorCond.empty() && !PriorTBB && MBB->pred_size() == 1 &&
-        PrevBB.succ_size() == 1 && PrevBB.isSuccessor(MBB) &&
-        !MBB->hasAddressTaken() && !MBB->isEHPad()) {
+        PrevBB.succ_size() == 1 && ((PriorTBB == MBB) || (PriorFBB == MBB)) &&
+        !MBB->hasAddressTaken()) {
       LLVM_DEBUG(dbgs() << "\nMerging into block: " << PrevBB
                         << "From MBB: " << *MBB);
       // Remove redundant DBG_VALUEs first.
@@ -1568,6 +1561,13 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
     }
   }
 
+  // Analyze the branch in the current block. As a side-effect, this may cause
+  // the block to become empty.
+  MachineBasicBlock *CurTBB = nullptr, *CurFBB = nullptr;
+  SmallVector<MachineOperand, 4> CurCond;
+  bool CurUnAnalyzable =
+      TII->analyzeBranch(*MBB, CurTBB, CurFBB, CurCond, true);
+
   if (!CurUnAnalyzable) {
     // If this is a two-way branch, and the FBB branches to this block, reverse
     // the condition so the single-basic-block loop is faster.  Instead of:
@@ -1612,7 +1612,7 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
       if (MBB->empty()) {
         bool PredHasNoFallThrough = !PrevBB.canFallThrough();
         if (PredHasNoFallThrough || !PriorUnAnalyzable ||
-            !PrevBB.isSuccessor(MBB)) {
+            !((CurTBB == MBB) || (CurFBB == MBB))) {
           // If the prior block falls through into us, turn it into an
           // explicit branch to us to make updates simpler.
           if (!PredHasNoFallThrough && PrevBB.isSuccessor(MBB) &&
@@ -1756,21 +1756,11 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
       // removed, move this block to the end of the function. There is no real
       // advantage in "falling through" to an EH block, so we don't want to
       // perform this transformation for that case.
-      //
-      // Also, Windows EH introduced the possibility of an arbitrary number of
-      // successors to a given block.  The analyzeBranch call does not consider
-      // exception handling and so we can get in a state where a block
-      // containing a call is followed by multiple EH blocks that would be
-      // rotated infinitely at the end of the function if the transformation
-      // below were performed for EH "FallThrough" blocks.  Therefore, even if
-      // that appears not to be happening anymore, we should assume that it is
-      // possible and not remove the "!FallThrough()->isEHPad" condition below.
       MachineBasicBlock *PrevTBB = nullptr, *PrevFBB = nullptr;
       SmallVector<MachineOperand, 4> PrevCond;
       if (FallThrough != MF.end() &&
-          !FallThrough->isEHPad() &&
           !TII->analyzeBranch(PrevBB, PrevTBB, PrevFBB, PrevCond, true) &&
-          PrevBB.isSuccessor(&*FallThrough)) {
+          ((PrevTBB == &*FallThrough) || (PrevFBB == &*FallThrough))) {
         MBB->moveAfter(&MF.back());
         MadeChange = true;
         return MadeChange;
diff --git a/llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-target-blocks.ll b/llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-target-blocks.ll
new file mode 100644
index 00000000000000..5a9557c3e343f8
--- /dev/null
+++ b/llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-target-blocks.ll
@@ -0,0 +1,18 @@
+; RUN: llc -mtriple x86_64 %s -stop-after=branch-folder
+
+; Check branch-folder do not keep swapping %indirect.target.block.1 and
+; %indirect.target.block.2
+define void @no_inf_loop() {
+entry:
+  br label %bb1
+
+bb1:
+  callbr void asm "", "!i,!i"() to label %bb1
+    [label %indirect.target.block.1, label %indirect.target.block.2]
+
+indirect.target.block.1:
+  br label %bb1
+
+indirect.target.block.2:
+  br label %bb1
+}