[llvm] [BranchFolding] Add an optional target hook to skip branch folding when it's unsafe (PR #147029)

[via llvm-commits]

https://github.com/csstormq created https://github.com/llvm/llvm-project/pull/147029

Assume the indirect branch implementation is as follows:

```
mov r0.lo BB0.label_lo
mov r0.hi BB0.label_hi
add r2 r1 r0
jr r2
```

Here, the base address and offset are stored in `r0`  and `r1` respectively, with `r2 = r1 + r0` as the jump target. Analyzing the target of the indirect branch `jr`  is often challenging, particularly when register `r2` is spilled.

When basic block `BB0` is empty, `BranchFolding` may optimize it, causing a dangling label. To address this, the target can override  `isMBBSafeToBranchFolding`  and return  `false`.
 
This issue was found in an out-of-tree target, hence the lack of testcase in this commit.

>From a034b82a528346dea1f1a61dde950d1e483ab073 Mon Sep 17 00:00:00 2001
From: csstormq <swust_xiaoqiangxu at 163.com>
Date: Fri, 4 Jul 2025 17:42:58 +0800
Subject: [PATCH] [BranchFolding] Add an optional target hook to skip branch
 folding when it's unsafe

---
 llvm/include/llvm/CodeGen/TargetInstrInfo.h |  6 ++++++
 llvm/lib/CodeGen/BranchFolding.cpp          | 16 ++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/llvm/include/llvm/CodeGen/TargetInstrInfo.h b/llvm/include/llvm/CodeGen/TargetInstrInfo.h
index b5b83c7ff1164..0d6976e91d245 100644
--- a/llvm/include/llvm/CodeGen/TargetInstrInfo.h
+++ b/llvm/include/llvm/CodeGen/TargetInstrInfo.h
@@ -2335,6 +2335,12 @@ class LLVM_ABI TargetInstrInfo : public MCInstrInfo {
     llvm_unreachable("unknown number of operands necessary");
   }
 
+  /// Return false if \p MBB whose any predecessor is not analyzable is not safe
+  /// to do \p BranchFolding.
+  virtual bool isMBBSafeToBranchFolding(const MachineBasicBlock &MBB) const {
+    return true;
+  }
+
 private:
   mutable std::unique_ptr<MIRFormatter> Formatter;
   unsigned CallFrameSetupOpcode, CallFrameDestroyOpcode;
diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index 3b3e7a418feb5..8a46b42c47542 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -1346,6 +1346,19 @@ static void salvageDebugInfoFromEmptyBlock(const TargetInstrInfo *TII,
       copyDebugInfoToPredecessor(TII, MBB, *PredBB);
 }
 
+static bool areAllPredsAnalyzable(const llvm::TargetInstrInfo *TII,
+                                  MachineBasicBlock *MBB) {
+  for (auto itr = MBB->pred_begin(); itr != MBB->pred_end(); ++itr) {
+    MachineBasicBlock *CurTBB = nullptr, *CurFBB = nullptr;
+    SmallVector<MachineOperand, 4> CurCond;
+    bool CantAnalyzable = TII->analyzeBranch(**itr, CurTBB, CurFBB, CurCond,
+                                             /*AllowModify*/ false);
+    if (CantAnalyzable)
+      return false;
+  }
+  return true;
+}
+
 bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
   bool MadeChange = false;
   MachineFunction &MF = *MBB->getParent();
@@ -1389,6 +1402,9 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
       // TODO: Is it ever worth rewriting predecessors which don't already
       // jump to a landing pad, and so can safely jump to the fallthrough?
     } else if (MBB->isSuccessor(&*FallThrough)) {
+      if (!TII->isMBBSafeToBranchFolding(*MBB) &&
+          !areAllPredsAnalyzable(TII, MBB))
+        return MadeChange;
       // Rewrite all predecessors of the old block to go to the fallthrough
       // instead.
       while (!MBB->pred_empty()) {