[Transforms] Refactor CreateControlFlowHub (#103013)

CreateControlFlowHub is a method that redirects control flow edges from a set of
incoming blocks to a set of outgoing blocks through a new set of "guard" blocks.
This is now refactored into a separate file with one enhancement: The input to
the method is now a set of branches rather than two sets of blocks.

The original implementation reroutes every edge from incoming blocks to outgoing
blocks. But it is possible that for some incoming block InBB, some successor S
might be in the set of outgoing blocks, but that particular edge should not be
rerouted. The new implementation makes this possible by allowing the user to
specify the targets of each branch that need to be rerouted.

This is needed when improving the implementation of FixIrreducible #101386.
Current use in FixIrreducible does not demonstrate this finer control over the
edges being rerouted. But in UnifyLoopExits, when only one successor of an
exiting block is an exit block, this refinement now reroutes only the relevant
control-flow through the edge; the non-exit successor is not rerouted. This
results in fewer branches and PHI nodes in the hub.

1 files changed, 0 insertions, 314 deletions

diff --git a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
index c78965a..4144c79 100644
--- a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
+++ b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
@@ -1891,320 +1891,6 @@ BranchInst *llvm::GetIfCondition(BasicBlock *BB, BasicBlock *&IfTrue,
   return BI;
 }
 
-// After creating a control flow hub, the operands of PHINodes in an outgoing
-// block Out no longer match the predecessors of that block. Predecessors of Out
-// that are incoming blocks to the hub are now replaced by just one edge from
-// the hub. To match this new control flow, the corresponding values from each
-// PHINode must now be moved a new PHINode in the first guard block of the hub.
-//
-// This operation cannot be performed with SSAUpdater, because it involves one
-// new use: If the block Out is in the list of Incoming blocks, then the newly
-// created PHI in the Hub will use itself along that edge from Out to Hub.
-static void reconnectPhis(BasicBlock *Out, BasicBlock *GuardBlock,
-                          const SetVector<BasicBlock *> &Incoming,
-                          BasicBlock *FirstGuardBlock) {
-  auto I = Out->begin();
-  while (I != Out->end() && isa<PHINode>(I)) {
-    auto Phi = cast<PHINode>(I);
-    auto NewPhi =
-        PHINode::Create(Phi->getType(), Incoming.size(),
-                        Phi->getName() + ".moved", FirstGuardBlock->begin());
-    bool AllUndef = true;
-    for (auto *In : Incoming) {
-      Value *V = UndefValue::get(Phi->getType());
-      if (In == Out) {
-        V = NewPhi;
-      } else if (Phi->getBasicBlockIndex(In) != -1) {
-        V = Phi->removeIncomingValue(In, false);
-        AllUndef &= isa<UndefValue>(V);
-      }
-      NewPhi->addIncoming(V, In);
-    }
-    assert(NewPhi->getNumIncomingValues() == Incoming.size());
-    Value *NewV = NewPhi;
-    if (AllUndef) {
-      NewPhi->eraseFromParent();
-      NewV = UndefValue::get(Phi->getType());
-    }
-    if (Phi->getNumOperands() == 0) {
-      Phi->replaceAllUsesWith(NewV);
-      I = Phi->eraseFromParent();
-      continue;
-    }
-    Phi->addIncoming(NewV, GuardBlock);
-    ++I;
-  }
-}
-
-using BBPredicates = DenseMap<BasicBlock *, Instruction *>;
-using BBSetVector = SetVector<BasicBlock *>;
-
-// Redirects the terminator of the incoming block to the first guard
-// block in the hub. The condition of the original terminator (if it
-// was conditional) and its original successors are returned as a
-// tuple <condition, succ0, succ1>. The function additionally filters
-// out successors that are not in the set of outgoing blocks.
-//
-// - condition is non-null iff the branch is conditional.
-// - Succ1 is non-null iff the sole/taken target is an outgoing block.
-// - Succ2 is non-null iff condition is non-null and the fallthrough
-//         target is an outgoing block.
-static std::tuple<Value *, BasicBlock *, BasicBlock *>
-redirectToHub(BasicBlock *BB, BasicBlock *FirstGuardBlock,
-              const BBSetVector &Outgoing) {
-  assert(isa<BranchInst>(BB->getTerminator()) &&
-         "Only support branch terminator.");
-  auto Branch = cast<BranchInst>(BB->getTerminator());
-  auto Condition = Branch->isConditional() ? Branch->getCondition() : nullptr;
-
-  BasicBlock *Succ0 = Branch->getSuccessor(0);
-  BasicBlock *Succ1 = nullptr;
-  Succ0 = Outgoing.count(Succ0) ? Succ0 : nullptr;
-
-  if (Branch->isUnconditional()) {
-    Branch->setSuccessor(0, FirstGuardBlock);
-    assert(Succ0);
-  } else {
-    Succ1 = Branch->getSuccessor(1);
-    Succ1 = Outgoing.count(Succ1) ? Succ1 : nullptr;
-    assert(Succ0 || Succ1);
-    if (Succ0 && !Succ1) {
-      Branch->setSuccessor(0, FirstGuardBlock);
-    } else if (Succ1 && !Succ0) {
-      Branch->setSuccessor(1, FirstGuardBlock);
-    } else {
-      Branch->eraseFromParent();
-      BranchInst::Create(FirstGuardBlock, BB);
-    }
-  }
-
-  assert(Succ0 || Succ1);
-  return std::make_tuple(Condition, Succ0, Succ1);
-}
-// Setup the branch instructions for guard blocks.
-//
-// Each guard block terminates in a conditional branch that transfers
-// control to the corresponding outgoing block or the next guard
-// block. The last guard block has two outgoing blocks as successors
-// since the condition for the final outgoing block is trivially
-// true. So we create one less block (including the first guard block)
-// than the number of outgoing blocks.
-static void setupBranchForGuard(SmallVectorImpl<BasicBlock *> &GuardBlocks,
-                                const BBSetVector &Outgoing,
-                                BBPredicates &GuardPredicates) {
-  // To help keep the loop simple, temporarily append the last
-  // outgoing block to the list of guard blocks.
-  GuardBlocks.push_back(Outgoing.back());
-
-  for (int i = 0, e = GuardBlocks.size() - 1; i != e; ++i) {
-    auto Out = Outgoing[i];
-    assert(GuardPredicates.count(Out));
-    BranchInst::Create(Out, GuardBlocks[i + 1], GuardPredicates[Out],
-                       GuardBlocks[i]);
-  }
-
-  // Remove the last block from the guard list.
-  GuardBlocks.pop_back();
-}
-
-/// We are using one integer to represent the block we are branching to. Then at
-/// each guard block, the predicate was calcuated using a simple `icmp eq`.
-static void calcPredicateUsingInteger(
-    const BBSetVector &Incoming, const BBSetVector &Outgoing,
-    SmallVectorImpl<BasicBlock *> &GuardBlocks, BBPredicates &GuardPredicates) {
-  auto &Context = Incoming.front()->getContext();
-  auto FirstGuardBlock = GuardBlocks.front();
-
-  auto Phi = PHINode::Create(Type::getInt32Ty(Context), Incoming.size(),
-                             "merged.bb.idx", FirstGuardBlock);
-
-  for (auto In : Incoming) {
-    Value *Condition;
-    BasicBlock *Succ0;
-    BasicBlock *Succ1;
-    std::tie(Condition, Succ0, Succ1) =
-        redirectToHub(In, FirstGuardBlock, Outgoing);
-    Value *IncomingId = nullptr;
-    if (Succ0 && Succ1) {
-      // target_bb_index = Condition ? index_of_succ0 : index_of_succ1.
-      auto Succ0Iter = find(Outgoing, Succ0);
-      auto Succ1Iter = find(Outgoing, Succ1);
-      Value *Id0 = ConstantInt::get(Type::getInt32Ty(Context),
-                                    std::distance(Outgoing.begin(), Succ0Iter));
-      Value *Id1 = ConstantInt::get(Type::getInt32Ty(Context),
-                                    std::distance(Outgoing.begin(), Succ1Iter));
-      IncomingId = SelectInst::Create(Condition, Id0, Id1, "target.bb.idx",
-                                      In->getTerminator()->getIterator());
-    } else {
-      // Get the index of the non-null successor.
-      auto SuccIter = Succ0 ? find(Outgoing, Succ0) : find(Outgoing, Succ1);
-      IncomingId = ConstantInt::get(Type::getInt32Ty(Context),
-                                    std::distance(Outgoing.begin(), SuccIter));
-    }
-    Phi->addIncoming(IncomingId, In);
-  }
-
-  for (int i = 0, e = Outgoing.size() - 1; i != e; ++i) {
-    auto Out = Outgoing[i];
-    auto Cmp = ICmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_EQ, Phi,
-                                ConstantInt::get(Type::getInt32Ty(Context), i),
-                                Out->getName() + ".predicate", GuardBlocks[i]);
-    GuardPredicates[Out] = Cmp;
-  }
-}
-
-/// We record the predicate of each outgoing block using a phi of boolean.
-static void calcPredicateUsingBooleans(
-    const BBSetVector &Incoming, const BBSetVector &Outgoing,
-    SmallVectorImpl<BasicBlock *> &GuardBlocks, BBPredicates &GuardPredicates,
-    SmallVectorImpl<WeakVH> &DeletionCandidates) {
-  auto &Context = Incoming.front()->getContext();
-  auto BoolTrue = ConstantInt::getTrue(Context);
-  auto BoolFalse = ConstantInt::getFalse(Context);
-  auto FirstGuardBlock = GuardBlocks.front();
-
-  // The predicate for the last outgoing is trivially true, and so we
-  // process only the first N-1 successors.
-  for (int i = 0, e = Outgoing.size() - 1; i != e; ++i) {
-    auto Out = Outgoing[i];
-    LLVM_DEBUG(dbgs() << "Creating guard for " << Out->getName() << "\n");
-
-    auto Phi =
-        PHINode::Create(Type::getInt1Ty(Context), Incoming.size(),
-                        StringRef("Guard.") + Out->getName(), FirstGuardBlock);
-    GuardPredicates[Out] = Phi;
-  }
-
-  for (auto *In : Incoming) {
-    Value *Condition;
-    BasicBlock *Succ0;
-    BasicBlock *Succ1;
-    std::tie(Condition, Succ0, Succ1) =
-        redirectToHub(In, FirstGuardBlock, Outgoing);
-
-    // Optimization: Consider an incoming block A with both successors
-    // Succ0 and Succ1 in the set of outgoing blocks. The predicates
-    // for Succ0 and Succ1 complement each other. If Succ0 is visited
-    // first in the loop below, control will branch to Succ0 using the
-    // corresponding predicate. But if that branch is not taken, then
-    // control must reach Succ1, which means that the incoming value of
-    // the predicate from `In` is true for Succ1.
-    bool OneSuccessorDone = false;
-    for (int i = 0, e = Outgoing.size() - 1; i != e; ++i) {
-      auto Out = Outgoing[i];
-      PHINode *Phi = cast<PHINode>(GuardPredicates[Out]);
-      if (Out != Succ0 && Out != Succ1) {
-        Phi->addIncoming(BoolFalse, In);
-      } else if (!Succ0 || !Succ1 || OneSuccessorDone) {
-        // Optimization: When only one successor is an outgoing block,
-        // the incoming predicate from `In` is always true.
-        Phi->addIncoming(BoolTrue, In);
-      } else {
-        assert(Succ0 && Succ1);
-        if (Out == Succ0) {
-          Phi->addIncoming(Condition, In);
-        } else {
-          auto Inverted = invertCondition(Condition);
-          DeletionCandidates.push_back(Condition);
-          Phi->addIncoming(Inverted, In);
-        }
-        OneSuccessorDone = true;
-      }
-    }
-  }
-}
-
-// Capture the existing control flow as guard predicates, and redirect
-// control flow from \p Incoming block through the \p GuardBlocks to the
-// \p Outgoing blocks.
-//
-// There is one guard predicate for each outgoing block OutBB. The
-// predicate represents whether the hub should transfer control flow
-// to OutBB. These predicates are NOT ORTHOGONAL. The Hub evaluates
-// them in the same order as the Outgoing set-vector, and control
-// branches to the first outgoing block whose predicate evaluates to true.
-static void
-convertToGuardPredicates(SmallVectorImpl<BasicBlock *> &GuardBlocks,
-                         SmallVectorImpl<WeakVH> &DeletionCandidates,
-                         const BBSetVector &Incoming,
-                         const BBSetVector &Outgoing, const StringRef Prefix,
-                         std::optional<unsigned> MaxControlFlowBooleans) {
-  BBPredicates GuardPredicates;
-  auto F = Incoming.front()->getParent();
-
-  for (int i = 0, e = Outgoing.size() - 1; i != e; ++i)
-    GuardBlocks.push_back(
-        BasicBlock::Create(F->getContext(), Prefix + ".guard", F));
-
-  // When we are using an integer to record which target block to jump to, we
-  // are creating less live values, actually we are using one single integer to
-  // store the index of the target block. When we are using booleans to store
-  // the branching information, we need (N-1) boolean values, where N is the
-  // number of outgoing block.
-  if (!MaxControlFlowBooleans || Outgoing.size() <= *MaxControlFlowBooleans)
-    calcPredicateUsingBooleans(Incoming, Outgoing, GuardBlocks, GuardPredicates,
-                               DeletionCandidates);
-  else
-    calcPredicateUsingInteger(Incoming, Outgoing, GuardBlocks, GuardPredicates);
-
-  setupBranchForGuard(GuardBlocks, Outgoing, GuardPredicates);
-}
-
-BasicBlock *llvm::CreateControlFlowHub(
-    DomTreeUpdater *DTU, SmallVectorImpl<BasicBlock *> &GuardBlocks,
-    const BBSetVector &Incoming, const BBSetVector &Outgoing,
-    const StringRef Prefix, std::optional<unsigned> MaxControlFlowBooleans) {
-  if (Outgoing.size() < 2)
-    return Outgoing.front();
-
-  SmallVector<DominatorTree::UpdateType, 16> Updates;
-  if (DTU) {
-    for (auto *In : Incoming) {
-      for (auto Succ : successors(In))
-        if (Outgoing.count(Succ))
-          Updates.push_back({DominatorTree::Delete, In, Succ});
-    }
-  }
-
-  SmallVector<WeakVH, 8> DeletionCandidates;
-  convertToGuardPredicates(GuardBlocks, DeletionCandidates, Incoming, Outgoing,
-                           Prefix, MaxControlFlowBooleans);
-  auto FirstGuardBlock = GuardBlocks.front();
-  
-  // Update the PHINodes in each outgoing block to match the new control flow.
-  for (int i = 0, e = GuardBlocks.size(); i != e; ++i)
-    reconnectPhis(Outgoing[i], GuardBlocks[i], Incoming, FirstGuardBlock);
-
-  reconnectPhis(Outgoing.back(), GuardBlocks.back(), Incoming, FirstGuardBlock);
-
-  if (DTU) {
-    int NumGuards = GuardBlocks.size();
-    assert((int)Outgoing.size() == NumGuards + 1);
-
-    for (auto In : Incoming)
-      Updates.push_back({DominatorTree::Insert, In, FirstGuardBlock});
-
-    for (int i = 0; i != NumGuards - 1; ++i) {
-      Updates.push_back({DominatorTree::Insert, GuardBlocks[i], Outgoing[i]});
-      Updates.push_back(
-          {DominatorTree::Insert, GuardBlocks[i], GuardBlocks[i + 1]});
-    }
-    Updates.push_back({DominatorTree::Insert, GuardBlocks[NumGuards - 1],
-                       Outgoing[NumGuards - 1]});
-    Updates.push_back({DominatorTree::Insert, GuardBlocks[NumGuards - 1],
-                       Outgoing[NumGuards]});
-    DTU->applyUpdates(Updates);
-  }
-
-  for (auto I : DeletionCandidates) {
-    if (I->use_empty())
-      if (auto Inst = dyn_cast_or_null<Instruction>(I))
-        Inst->eraseFromParent();
-  }
-
-  return FirstGuardBlock;
-}
-
 void llvm::InvertBranch(BranchInst *PBI, IRBuilderBase &Builder) {
   Value *NewCond = PBI->getCondition();
   // If this is a "cmp" instruction, only used for branching (and nowhere