[BranchFolding] Add an optional target hook to skip branch folding when it's unsafe

Author: csstormq

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;

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()) {