[llvm] Branch folding: Avoid infinite loop with indirect target blocks (PR #96956)

[via llvm-commits]

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

>From a37b892f8555cee33d16ae2d53f16e1bed41421a Mon Sep 17 00:00:00 2001
From: v01dxyz <v01dxyz at v01d.xyz>
Date: Thu, 27 Jun 2024 21:33:57 +0200
Subject: [PATCH 1/2] Branch folding: Avoid infinite loop with indirect target
 blocks

---
 llvm/lib/CodeGen/BranchFolding.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index 1dc278586f1178..bf351efaf56fee 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -1767,8 +1767,8 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
       // possible and not remove the "!FallThrough()->isEHPad" condition below.
       MachineBasicBlock *PrevTBB = nullptr, *PrevFBB = nullptr;
       SmallVector<MachineOperand, 4> PrevCond;
-      if (FallThrough != MF.end() &&
-          !FallThrough->isEHPad() &&
+      if (FallThrough != MF.end() && !FallThrough->isEHPad() &&
+          !FallThrough->isInlineAsmBrIndirectTarget() &&
           !TII->analyzeBranch(PrevBB, PrevTBB, PrevFBB, PrevCond, true) &&
           PrevBB.isSuccessor(&*FallThrough)) {
         MBB->moveAfter(&MF.back());

>From ad95db2fa830f375bf203371b52b254183637c9a Mon Sep 17 00:00:00 2001
From: v01dxyz <v01dxyz at v01d.xyz>
Date: Fri, 28 Jun 2024 17:58:09 +0200
Subject: [PATCH 2/2] Branch folding/OptimizeBlock: Add support for CALLBR to
 avoid infinite loops

Do not move the MBB to the end if fallthrough is an indirect
target block.

Improvement based on the commit ff6a1ed.

CALLBR RFC:
https://lists.llvm.org/pipermail/llvm-dev/2018-October/127239.html
---
 llvm/lib/CodeGen/BranchFolding.cpp            | 20 ++++++++++++-------
 ...no-inf-loop-with-callbr-indirect-blocks.ll | 18 +++++++++++++++++
 2 files changed, 31 insertions(+), 7 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-blocks.ll

diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index bf351efaf56fee..3a03a86cb44ada 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -1757,14 +1757,20 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
       // 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
+      // Also, some constructs introduced the possibility of an arbitrary number
+      // of successors to a given block.
+      //
+      // - Windows EH: with EH blocks
+      // - CALLBR instruction: with Indirect Target blocks
+      //
+      // The analyzeBranch call does not consider exception handling or inline
+      // asm with control flow instructions. So we can get in a state where a
+      // block containing a call is followed by multiple 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.
+      // below were performed for such blocks. Therefore, we should assume
+      // that it is possible and not remove the "!FallThrough()->isEHPad &&
+      // !FallThrough->isInlineAsmIndirectTarget()" condition below (even if
+      // in the case of EHPad that appears not to be happening anymore).
       MachineBasicBlock *PrevTBB = nullptr, *PrevFBB = nullptr;
       SmallVector<MachineOperand, 4> PrevCond;
       if (FallThrough != MF.end() && !FallThrough->isEHPad() &&
diff --git a/llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-blocks.ll b/llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-blocks.ll
new file mode 100644
index 00000000000000..5a9557c3e343f8
--- /dev/null
+++ b/llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-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
+}