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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-x86

Author: None (v01dXYZ)

<details>
<summary>Changes</summary>

Fixes #<!-- -->96893.

Tasks remaining:

- [x] Add in the comment about EHPad it's also necessary for indirect targets.
- [x] Add a sanity test with a timeout (I could not set a timeout duration for a specific test only so I add the test file).

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


2 Files Affected:

- (modified) llvm/lib/CodeGen/BranchFolding.cpp (+15-9) 
- (added) llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-blocks.ll (+18) 


``````````diff
diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index 1dc278586f1178..3a03a86cb44ada 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -1757,18 +1757,24 @@ 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() &&
+      if (FallThrough != MF.end() && !FallThrough->isEHPad() &&
+          !FallThrough->isInlineAsmBrIndirectTarget() &&
           !TII->analyzeBranch(PrevBB, PrevTBB, PrevFBB, PrevCond, true) &&
           PrevBB.isSuccessor(&*FallThrough)) {
         MBB->moveAfter(&MF.back());
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
+}

``````````

</details>


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