[SimplifyCFG] Start redesigning `FoldTwoEntryPHINode()`.

The current `FoldTwoEntryPHINode()` is not quite designed correctly.
It starts from the merge point, and then tries to detect
the 'divergence' point.

Because of that, it is limited to the simple two-predecessor case,
where the PHI completely goes away. but that is rather pessimistic,
and it doesn't make much sense from the costmodel side of things.

For example if there is some other unrelated predecessor of
the merge point,  we could split the merge point so that
the then/else blocks first branch to an empty block
and then to the merge point, and then we'd be able to speculate
the then/else code.

But if we'd instead simply start at the divergence point,
and look for the merge point, then we'll just natively support this case.

There's also the fact that `SpeculativelyExecuteBB()` already does
just that, but only if there is a single block to speculate,
and with a much more restrictive cost model.
But that also means we have code duplication.

Now, sadly, while this is as much NFCI as possible,
there is just no way to cleanly migrate to
the proper implementation. The results *are* going to be different
somewhat because of various phase ordering effects and SimplifyCFG
block iteration strategy.

1 files changed, 80 insertions, 9 deletions

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 335ac03..b16fbd1 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -2960,6 +2960,81 @@ static bool FoldTwoEntryPHINode(PHINode *PN, const TargetTransformInfo &TTI,
   return true;
 }
 
+static bool SpeculativelyExecuteThenElseCode(BranchInst *BI,
+                                             const TargetTransformInfo &TTI,
+                                             DomTreeUpdater *DTU,
+                                             const DataLayout &DL) {
+  assert(BI->isConditional() && !isa<ConstantInt>(BI->getCondition()) &&
+         BI->getSuccessor(0) != BI->getSuccessor(1) &&
+         "Only for truly conditional branches.");
+  BasicBlock *BB = BI->getParent();
+
+  // Which ones of our successors end up with an unconditional branch?
+  SmallVector<BasicBlock *, 2> UncondSuccessors;
+  SmallVector<BasicBlock *, 2> OtherSuccessors;
+  for (BasicBlock *Succ : successors(BI)) {
+    auto *SuccBI = dyn_cast<BranchInst>(Succ->getTerminator());
+    if (SuccBI && SuccBI->isUnconditional())
+      UncondSuccessors.emplace_back(Succ);
+    else
+      OtherSuccessors.emplace_back(Succ);
+  }
+  assert(UncondSuccessors.size() + OtherSuccessors.size() == 2 &&
+         "Can not have more than two successors!");
+
+  // If none do, then we can't do anything.
+  if (UncondSuccessors.empty())
+    return false;
+
+  // We want to hoist code from the unconditional block[s] and eliminate them,
+  // but if they have their address taken, then we essentially can't do this.
+  for (BasicBlock *UncondSucc : UncondSuccessors)
+    if (UncondSucc->hasAddressTaken())
+      return false;
+
+  // All unconditional successors must have a single (and the same) predecessor.
+  // FIXME: lift this restriction.
+  for (BasicBlock *UncondSucc : UncondSuccessors)
+    if (!UncondSucc->getSinglePredecessor())
+      return false;
+
+  // Now, what is the merge point?
+  BasicBlock *MergeBB = nullptr;
+  // If there was only a single unconditional successor,
+  // then the other successor *must* be the merge point.
+  if (UncondSuccessors.size() == 1)
+    MergeBB = OtherSuccessors.front();
+
+  // All unconditional successors must have the same successor themselves.
+  for (BasicBlock *UncondSucc : UncondSuccessors) {
+    auto *SuccBI = cast<BranchInst>(UncondSucc->getTerminator());
+    assert(SuccBI->isUnconditional() && "Should be an unconditional branch.");
+    BasicBlock *SuccOfSucc = SuccBI->getSuccessor(0);
+    if (!MergeBB) // First unconditional successor, record it's successor.
+      MergeBB = SuccOfSucc;
+    else if (SuccOfSucc != MergeBB) // Do all succs have the same successor?
+      return false;
+  }
+
+  assert(MergeBB && "Should have found the merge point.");
+  assert(all_of(UncondSuccessors,
+                [MergeBB](BasicBlock *UncondSucc) {
+                  return is_contained(predecessors(MergeBB), UncondSucc);
+                }) &&
+         "All unconditional successors must be predecessors of merge block.");
+  assert((UncondSuccessors.size() != 1 ||
+          is_contained(predecessors(MergeBB), BB)) &&
+         "If there is only a single unconditional successor, then the dispatch "
+         "block must also be merge block's predecessor.");
+
+  if (auto *PN = dyn_cast<PHINode>(MergeBB->begin()))
+    // FIXME: lift this restriction.
+    if (PN->getNumIncomingValues() == 2)
+      return FoldTwoEntryPHINode(PN, TTI, DTU, DL);
+
+  return false;
+}
+
 static Value *createLogicalOp(IRBuilderBase &Builder,
                               Instruction::BinaryOps Opc, Value *LHS,
                               Value *RHS, const Twine &Name = "") {
@@ -6521,6 +6596,11 @@ bool SimplifyCFGOpt::simplifyCondBranch(BranchInst *BI, IRBuilder<> &Builder) {
         return requestResimplify();
   }
 
+  if (Options.FoldTwoEntryPHINode) {
+    if (SpeculativelyExecuteThenElseCode(BI, TTI, DTU, DL))
+      return true;
+  }
+
   // If this is a branch on a phi node in the current block, thread control
   // through this block if any PHI node entries are constants.
   if (PHINode *PN = dyn_cast<PHINode>(BI->getCondition()))
@@ -6736,15 +6816,6 @@ bool SimplifyCFGOpt::simplifyOnce(BasicBlock *BB) {
 
   IRBuilder<> Builder(BB);
 
-  if (Options.FoldTwoEntryPHINode) {
-    // If there is a trivial two-entry PHI node in this basic block, and we can
-    // eliminate it, do so now.
-    if (auto *PN = dyn_cast<PHINode>(BB->begin()))
-      if (PN->getNumIncomingValues() == 2)
-        if (FoldTwoEntryPHINode(PN, TTI, DTU, DL))
-          return true;
-  }
-
   Instruction *Terminator = BB->getTerminator();
   Builder.SetInsertPoint(Terminator);
   switch (Terminator->getOpcode()) {