[llvm] [BranchFolding] Avoid moving blocks to fall through to an indirect target (PR #152916)

[via llvm-commits]

https://github.com/XChy updated https://github.com/llvm/llvm-project/pull/152916

>From 2913dc66e0684edd3596d06939b13cbb7c04bcee 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 1/2] [BranchFolding] Avoid moving blocks to fall through to an
 indirect target

---
 llvm/lib/CodeGen/BranchFolding.cpp       | 12 ++++++++++--
 llvm/test/CodeGen/X86/callbr-asm-loop.ll | 16 +++++-----------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index dcfd9aad70fc5..2f6e3b1f76b11 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -1787,10 +1787,18 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
       // 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.
+      //
+      // Similarly, the analyzeBranch call does not consider callbr, which also
+      // introduces the possibility of infinite rotation, as there may be
+      // multiple successors of PrevBB. Thus we checks such case by
+      // !FallThrough->isInlineAsmBrIndirectTarget().
+      // NOTE:  Checking if PrevBB contains callbr is more precise, but
+      // much more expensive.
       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/callbr-asm-loop.ll b/llvm/test/CodeGen/X86/callbr-asm-loop.ll
index 999b04c4f4838..fab167b449de9 100644
--- a/llvm/test/CodeGen/X86/callbr-asm-loop.ll
+++ b/llvm/test/CodeGen/X86/callbr-asm-loop.ll
@@ -1,35 +1,29 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
 
-; RUN: llc -O0 -mtriple=i686-- < %s | FileCheck %s
+; RUN: llc -O1 -mtriple=i686-- < %s | FileCheck %s
 
 ; Test that causes multiple defs of %eax.
 ; FIXME: The testcase hangs with -O1/2/3 enabled.
 define i32 @loop1() nounwind {
 ; CHECK-LABEL: loop1:
 ; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    pushl %esi
-; CHECK-NEXT:    jmp .LBB0_1
+; CHECK-NEXT:    .p2align 4
 ; 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:    jmp .LBB0_1
 ; CHECK-NEXT:  .LBB0_2: # Inline asm indirect target
-; CHECK-NEXT:    # %tailrecurse.tailrecurse.backedge_crit_edge
+; CHECK-NEXT:    # %tailrecurse.tailrecurse_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:  .LBB0_3: # 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:    retl
 entry:
   br label %tailrecurse

>From 818ab09d4184ce481d3d1aa3943187ea253885af Mon Sep 17 00:00:00 2001
From: XChy <xxs_chy at outlook.com>
Date: Thu, 14 Aug 2025 01:34:57 +0800
Subject: [PATCH 2/2] fix comment

---
 llvm/lib/CodeGen/BranchFolding.cpp       | 8 ++++----
 llvm/test/CodeGen/X86/callbr-asm-loop.ll | 1 -
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index 2f6e3b1f76b11..7292bc2be0df2 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -1790,10 +1790,10 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
       //
       // Similarly, the analyzeBranch call does not consider callbr, which also
       // introduces the possibility of infinite rotation, as there may be
-      // multiple successors of PrevBB. Thus we checks such case by
-      // !FallThrough->isInlineAsmBrIndirectTarget().
-      // NOTE:  Checking if PrevBB contains callbr is more precise, but
-      // much more expensive.
+      // multiple successors of PrevBB. Thus we check such case by
+      // FallThrough->isInlineAsmBrIndirectTarget().
+      // NOTE: Checking if PrevBB contains callbr is more precise, but much
+      // more expensive.
       MachineBasicBlock *PrevTBB = nullptr, *PrevFBB = nullptr;
       SmallVector<MachineOperand, 4> PrevCond;
 
diff --git a/llvm/test/CodeGen/X86/callbr-asm-loop.ll b/llvm/test/CodeGen/X86/callbr-asm-loop.ll
index fab167b449de9..0b6898815f8cc 100644
--- a/llvm/test/CodeGen/X86/callbr-asm-loop.ll
+++ b/llvm/test/CodeGen/X86/callbr-asm-loop.ll
@@ -3,7 +3,6 @@
 ; RUN: llc -O1 -mtriple=i686-- < %s | FileCheck %s
 
 ; Test that causes multiple defs of %eax.
-; FIXME: The testcase hangs with -O1/2/3 enabled.
 define i32 @loop1() nounwind {
 ; CHECK-LABEL: loop1:
 ; CHECK:       # %bb.0: # %entry