[llvm] [SimplifyCFG] Add predecessor threshold for TryToSimplifyUncondBranchFromEmptyBlock optimization. (PR #110715)

[Amara Emerson via llvm-commits]

https://github.com/aemerson created https://github.com/llvm/llvm-project/pull/110715

In some pathological cases this optimization can spend an unreasonable amount of time.
Specifically, there could be a very large number of predecessor blocks of the unconditional
branch target. This change adds a threshold bailout to mitigate this.

rdar://137063034


>From 0637e4359e24c8119eedf60b2079d1b5ae7fcef4 Mon Sep 17 00:00:00 2001
From: Amara Emerson <amara at apple.com>
Date: Mon, 30 Sep 2024 14:51:05 -0700
Subject: [PATCH] [SimplifyCFG] Add predecessor threshold for
 TryToSimplifyUncondBranchFromEmptyBlock optimization.

In some pathological cases this optimization can spend an unreasonable amount of time.
Specifically, there could be a very large number of predecessor blocks of the unconditional
branch target. This change adds a threshold bailout to mitigate this.

rdar://137063034
---
 llvm/lib/Transforms/Utils/Local.cpp           | 10 +++++
 .../branch-common-pred-threshold.ll           | 44 +++++++++++++++++++
 2 files changed, 54 insertions(+)
 create mode 100644 llvm/test/Transforms/SimplifyCFG/branch-common-pred-threshold.ll

diff --git a/llvm/lib/Transforms/Utils/Local.cpp b/llvm/lib/Transforms/Utils/Local.cpp
index 7659fc69196151..3e00a63e5f2cf9 100644
--- a/llvm/lib/Transforms/Utils/Local.cpp
+++ b/llvm/lib/Transforms/Utils/Local.cpp
@@ -122,6 +122,11 @@ static cl::opt<unsigned> MaxPhiEntriesIncreaseAfterRemovingEmptyBlock(
 // bitreverse idioms.
 static const unsigned BitPartRecursionMaxDepth = 48;
 
+static cl::opt<unsigned> MaxNumPredsThreshold(
+    "simplifycfg-uncondbr-max-num-preds-threshold", cl::init(128), cl::Hidden,
+    cl::desc("The maximum number of predecessors a block can have to simplify "
+             "uncond-brs from empty blocks."));
+
 //===----------------------------------------------------------------------===//
 //  Local constant propagation.
 //
@@ -1165,6 +1170,11 @@ bool llvm::TryToSimplifyUncondBranchFromEmptyBlock(BasicBlock *BB,
   if (BB == Succ)
     return false;
 
+  // Some pathological IR can have extremely high numbers of predecessor blocks.
+  // Bail out early if so.
+  if (Succ->hasNPredecessorsOrMore(MaxNumPredsThreshold))
+    return false;
+
   SmallPtrSet<BasicBlock *, 16> BBPreds(pred_begin(BB), pred_end(BB));
   SmallPtrSet<BasicBlock *, 16> SuccPreds(pred_begin(Succ), pred_end(Succ));
 
diff --git a/llvm/test/Transforms/SimplifyCFG/branch-common-pred-threshold.ll b/llvm/test/Transforms/SimplifyCFG/branch-common-pred-threshold.ll
new file mode 100644
index 00000000000000..3f1c93762e34eb
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/branch-common-pred-threshold.ll
@@ -0,0 +1,44 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=simplifycfg -simplifycfg-uncondbr-max-num-preds-threshold=3 -S | FileCheck %s
+
+; This test checks that we don't optimize the edge from Pred -> BB to Pred -> Succ when
+; the number of predecessors of Succ is greater than the threshold.
+
+define i8 @succ_has_3_preds(i8 noundef %arg, i1 %c1, i1 %c2) {
+; CHECK-LABEL: @succ_has_3_preds(
+; CHECK-NEXT:  Pred:
+; CHECK-NEXT:    call void @dummy()
+; CHECK-NEXT:    br i1 [[C1:%.*]], label [[COMMONPRED:%.*]], label [[EXTRA_BB:%.*]]
+; CHECK:       extra_bb:
+; CHECK-NEXT:    br label [[SUCC:%.*]]
+; CHECK:       CommonPred:
+; CHECK-NEXT:    call void @dummy()
+; CHECK-NEXT:    br i1 [[C2:%.*]], label [[SUCC]], label [[BB:%.*]]
+; CHECK:       BB:
+; CHECK-NEXT:    [[PHI1:%.*]] = phi i8 [ 1, [[COMMONPRED]] ]
+; CHECK-NEXT:    br label [[SUCC]]
+; CHECK:       Succ:
+; CHECK-NEXT:    [[PHI2:%.*]] = phi i8 [ [[PHI1]], [[BB]] ], [ 4, [[COMMONPRED]] ], [ 0, [[EXTRA_BB]] ]
+; CHECK-NEXT:    ret i8 [[PHI2]]
+;
+Pred:
+call void @dummy()
+  br i1 %c1, label %CommonPred, label %extra_bb
+
+extra_bb:
+  br i1 %c1, label %CommonPred, label %Succ
+
+CommonPred:
+call void @dummy()
+  br i1 %c2, label %Succ, label %BB
+
+BB:
+  %phi1 = phi i8 [1, %CommonPred]
+  br label %Succ
+
+Succ:
+  %phi2 = phi i8 [ %phi1, %BB ], [ 4, %CommonPred ], [0, %extra_bb]
+  ret i8 %phi2
+}
+
+declare void @dummy()