[llvm-branch-commits] [llvm] [Transforms] Refactor CreateControlFlowHub (PR #103013)

Sameer Sahasrabuddhe via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Mon Aug 12 23:53:43 PDT 2024


https://github.com/ssahasra created https://github.com/llvm/llvm-project/pull/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 uses in FixIrreducible and UnifyLoopExits do not demonstrate this finer control over the edges being rerouted.

>From 9b8d7f65680155f04bc754aebd4d820bad743581 Mon Sep 17 00:00:00 2001
From: Sameer Sahasrabuddhe <sameer.sahasrabuddhe at amd.com>
Date: Tue, 13 Aug 2024 12:07:00 +0530
Subject: [PATCH] [Transforms] Refactor CreateControlFlowHub

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
uses in FixIrreducible and UnifyLoopExits do not demonstrate this finer control
over the edges being rerouted.
---
 .../llvm/Transforms/Utils/BasicBlockUtils.h   |  75 ----
 .../llvm/Transforms/Utils/ControlFlowUtils.h  | 121 +++++++
 llvm/lib/Transforms/Utils/BasicBlockUtils.cpp | 314 ----------------
 llvm/lib/Transforms/Utils/CMakeLists.txt      |   1 +
 .../lib/Transforms/Utils/ControlFlowUtils.cpp | 341 ++++++++++++++++++
 llvm/lib/Transforms/Utils/FixIrreducible.cpp  |  35 +-
 llvm/lib/Transforms/Utils/UnifyLoopExits.cpp  |  69 ++--
 .../CodeGen/AMDGPU/local-atomicrmw-fadd.ll    |  32 +-
 llvm/test/Transforms/FixIrreducible/basic.ll  |   4 +-
 .../Transforms/FixIrreducible/bug45623.ll     |   3 +-
 llvm/test/Transforms/FixIrreducible/nested.ll |   3 +-
 llvm/test/Transforms/FixIrreducible/switch.ll |   3 +-
 .../Transforms/FixIrreducible/unreachable.ll  |   4 +-
 13 files changed, 545 insertions(+), 460 deletions(-)
 create mode 100644 llvm/include/llvm/Transforms/Utils/ControlFlowUtils.h
 create mode 100644 llvm/lib/Transforms/Utils/ControlFlowUtils.cpp

diff --git a/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h b/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
index c99df6bf94d025..b447942ffbd676 100644
--- a/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
+++ b/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
@@ -602,81 +602,6 @@ bool SplitIndirectBrCriticalEdges(Function &F, bool IgnoreBlocksWithoutPHI,
                                   BranchProbabilityInfo *BPI = nullptr,
                                   BlockFrequencyInfo *BFI = nullptr);
 
-/// Given a set of incoming and outgoing blocks, create a "hub" such that every
-/// edge from an incoming block InBB to an outgoing block OutBB is now split
-/// into two edges, one from InBB to the hub and another from the hub to
-/// OutBB. The hub consists of a series of guard blocks, one for each outgoing
-/// block. Each guard block conditionally branches to the corresponding outgoing
-/// block, or the next guard block in the chain. These guard blocks are returned
-/// in the argument vector.
-///
-/// Since the control flow edges from InBB to OutBB have now been replaced, the
-/// function also updates any PHINodes in OutBB. For each such PHINode, the
-/// operands corresponding to incoming blocks are moved to a new PHINode in the
-/// hub, and the hub is made an operand of the original PHINode.
-///
-/// Input CFG:
-/// ----------
-///
-///                    Def
-///                     |
-///                     v
-///           In1      In2
-///            |        |
-///            |        |
-///            v        v
-///  Foo ---> Out1     Out2
-///                     |
-///                     v
-///                    Use
-///
-///
-/// Create hub: Incoming = {In1, In2}, Outgoing = {Out1, Out2}
-/// ----------------------------------------------------------
-///
-///             Def
-///              |
-///              v
-///  In1        In2          Foo
-///   |    Hub   |            |
-///   |    + - - | - - +      |
-///   |    '     v     '      V
-///   +------> Guard1 -----> Out1
-///        '     |     '
-///        '     v     '
-///        '   Guard2 -----> Out2
-///        '           '      |
-///        + - - - - - +      |
-///                           v
-///                          Use
-///
-/// Limitations:
-/// -----------
-/// 1. This assumes that all terminators in the CFG are direct branches (the
-///    "br" instruction). The presence of any other control flow such as
-///    indirectbr, switch or callbr will cause an assert.
-///
-/// 2. The updates to the PHINodes are not sufficient to restore SSA
-///    form. Consider a definition Def, its use Use, incoming block In2 and
-///    outgoing block Out2, such that:
-///    a. In2 is reachable from D or contains D.
-///    b. U is reachable from Out2 or is contained in Out2.
-///    c. U is not a PHINode if U is contained in Out2.
-///
-///    Clearly, Def dominates Out2 since the program is valid SSA. But when the
-///    hub is introduced, there is a new path through the hub along which Use is
-///    reachable from entry without passing through Def, and SSA is no longer
-///    valid. To fix this, we need to look at all the blocks post-dominated by
-///    the hub on the one hand, and dominated by Out2 on the other. This is left
-///    for the caller to accomplish, since each specific use of this function
-///    may have additional information which simplifies this fixup. For example,
-///    see restoreSSA() in the UnifyLoopExits pass.
-BasicBlock *CreateControlFlowHub(
-    DomTreeUpdater *DTU, SmallVectorImpl<BasicBlock *> &GuardBlocks,
-    const SetVector<BasicBlock *> &Predecessors,
-    const SetVector<BasicBlock *> &Successors, const StringRef Prefix,
-    std::optional<unsigned> MaxControlFlowBooleans = std::nullopt);
-
 // Utility function for inverting branch condition and for swapping its
 // successors
 void InvertBranch(BranchInst *PBI, IRBuilderBase &Builder);
diff --git a/llvm/include/llvm/Transforms/Utils/ControlFlowUtils.h b/llvm/include/llvm/Transforms/Utils/ControlFlowUtils.h
new file mode 100644
index 00000000000000..f5d237bc96627f
--- /dev/null
+++ b/llvm/include/llvm/Transforms/Utils/ControlFlowUtils.h
@@ -0,0 +1,121 @@
+//===- Transforms/Utils/ControlFlowUtils.h --------------------*- C++ -*---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Utilities to manipulate the CFG and restore SSA for the new control flow.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_UTILS_CONTROLFLOWUTILS_H
+#define LLVM_TRANSFORMS_UTILS_CONTROLFLOWUTILS_H
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+
+namespace llvm {
+
+class BasicBlock;
+class DomTreeUpdater;
+
+/// Given a set of branch descriptors [BB, Succ0, Succ1], create a "hub" such
+/// that the control flow from each BB to a successor is now split into two
+/// edges, one from BB to the hub and another from the hub to the successor. The
+/// hub consists of a series of guard blocks, one for each outgoing block. Each
+/// guard block conditionally branches to the corresponding outgoing block, or
+/// the next guard block in the chain. These guard blocks are returned in the
+/// argument vector.
+///
+/// This also updates any PHINodes in the successor. For each such PHINode, the
+/// operands corresponding to incoming blocks are moved to a new PHINode in the
+/// hub, and the hub is made an operand of the original PHINode.
+///
+/// Note that for some block BB with a conditional branch, it is not necessary
+/// that both successors are rerouted. The client specifies this by setting
+/// either Succ0 or Succ1 to nullptr, in which case, the corresponding successor
+/// is not rerouted.
+///
+/// Input CFG:
+/// ----------
+///
+///                    Def
+///                     |
+///                     v
+///           In1      In2
+///            |        |
+///            |        |
+///            v        v
+///  Foo ---> Out1     Out2
+///                     |
+///                     v
+///                    Use
+///
+///
+/// Create hub: Incoming = {In1, In2}, Outgoing = {Out1, Out2}
+/// ----------------------------------------------------------
+///
+///             Def
+///              |
+///              v
+///  In1        In2          Foo
+///   |    Hub   |            |
+///   |    + - - | - - +      |
+///   |    '     v     '      V
+///   +------> Guard1 -----> Out1
+///        '     |     '
+///        '     v     '
+///        '   Guard2 -----> Out2
+///        '           '      |
+///        + - - - - - +      |
+///                           v
+///                          Use
+///
+/// Limitations:
+/// -----------
+/// 1. This assumes that all terminators in the CFG are direct branches (the
+///    "br" instruction). The presence of any other control flow such as
+///    indirectbr, switch or callbr will cause an assert.
+///
+/// 2. The updates to the PHINodes are not sufficient to restore SSA
+///    form. Consider a definition Def, its use Use, incoming block In2 and
+///    outgoing block Out2, such that:
+///    a. In2 is reachable from D or contains D.
+///    b. U is reachable from Out2 or is contained in Out2.
+///    c. U is not a PHINode if U is contained in Out2.
+///
+///    Clearly, Def dominates Out2 since the program is valid SSA. But when the
+///    hub is introduced, there is a new path through the hub along which Use is
+///    reachable from entry without passing through Def, and SSA is no longer
+///    valid. To fix this, we need to look at all the blocks post-dominated by
+///    the hub on the one hand, and dominated by Out2 on the other. This is left
+///    for the caller to accomplish, since each specific use of this function
+///    may have additional information which simplifies this fixup. For example,
+///    see restoreSSA() in the UnifyLoopExits pass.
+struct ControlFlowHub {
+  struct BranchDescriptor {
+    BasicBlock *BB;
+    BasicBlock *Succ0;
+    BasicBlock *Succ1;
+
+    BranchDescriptor(BasicBlock *BB, BasicBlock *Succ0, BasicBlock *Succ1)
+        : BB(BB), Succ0(Succ0), Succ1(Succ1) {}
+  };
+
+  void addBranch(BasicBlock *BB, BasicBlock *Succ0, BasicBlock *Succ1) {
+    Branches.emplace_back(BB, Succ0, Succ1);
+  }
+
+  BasicBlock *
+  finalize(DomTreeUpdater *DTU, SmallVectorImpl<BasicBlock *> &GuardBlocks,
+           const StringRef Prefix,
+           std::optional<unsigned> MaxControlFlowBooleans = std::nullopt);
+
+  SmallVector<BranchDescriptor> Branches;
+};
+
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_UTILS_CONTROLFLOWUTILS_H
diff --git a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
index c78965aaeba3ca..4144c7993b7e42 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
diff --git a/llvm/lib/Transforms/Utils/CMakeLists.txt b/llvm/lib/Transforms/Utils/CMakeLists.txt
index 51e8821773c3af..b5a7eedbac62b1 100644
--- a/llvm/lib/Transforms/Utils/CMakeLists.txt
+++ b/llvm/lib/Transforms/Utils/CMakeLists.txt
@@ -16,6 +16,7 @@ add_llvm_component_library(LLVMTransformUtils
   CodeExtractor.cpp
   CodeLayout.cpp
   CodeMoverUtils.cpp
+  ControlFlowUtils.cpp
   CtorUtils.cpp
   CountVisits.cpp
   Debugify.cpp
diff --git a/llvm/lib/Transforms/Utils/ControlFlowUtils.cpp b/llvm/lib/Transforms/Utils/ControlFlowUtils.cpp
new file mode 100644
index 00000000000000..89359ac6f762d6
--- /dev/null
+++ b/llvm/lib/Transforms/Utils/ControlFlowUtils.cpp
@@ -0,0 +1,341 @@
+//===- ControlFlowUtils.cpp - Control Flow Utilities -----------------------==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Utilities to manipulate the CFG and restore SSA for the new control flow.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/ControlFlowUtils.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/Analysis/DomTreeUpdater.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/ValueHandle.h"
+#include "llvm/Transforms/Utils/Local.h"
+
+#define DEBUG_TYPE "control-flow-hub"
+
+using namespace llvm;
+
+using BBPredicates = DenseMap<BasicBlock *, Instruction *>;
+using EdgeDescriptor = ControlFlowHub::BranchDescriptor;
+
+// Redirects the terminator of the incoming block to the first guard block in
+// the hub. Returns the branch condition from `BB` if it exits.
+// - If only one of Succ0 or Succ1 is not null, the corresponding branch
+//   successor is redirected to the FirstGuardBlock.
+// - Else both are not null, and branch is replaced with an unconditional
+//   branch to the FirstGuardBlock.
+static Value *redirectToHub(BasicBlock *BB, BasicBlock *Succ0,
+                            BasicBlock *Succ1, BasicBlock *FirstGuardBlock) {
+  assert(isa<BranchInst>(BB->getTerminator()) &&
+         "Only support branch terminator.");
+  auto *Branch = cast<BranchInst>(BB->getTerminator());
+  auto *Condition = Branch->isConditional() ? Branch->getCondition() : nullptr;
+
+  assert(Succ0 || Succ1);
+
+  if (Branch->isUnconditional()) {
+    assert(Succ0 == Branch->getSuccessor(0));
+    assert(!Succ1);
+    Branch->setSuccessor(0, FirstGuardBlock);
+  } else {
+    assert(!Succ1 || Succ1 == Branch->getSuccessor(1));
+    if (Succ0 && !Succ1) {
+      Branch->setSuccessor(0, FirstGuardBlock);
+    } else if (Succ1 && !Succ0) {
+      Branch->setSuccessor(1, FirstGuardBlock);
+    } else {
+      Branch->eraseFromParent();
+      BranchInst::Create(FirstGuardBlock, BB);
+    }
+  }
+
+  return Condition;
+}
+
+// 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.
+static void setupBranchForGuard(ArrayRef<BasicBlock *> GuardBlocks,
+                                ArrayRef<BasicBlock *> Outgoing,
+                                BBPredicates &GuardPredicates) {
+  assert(Outgoing.size() > 1);
+  assert(GuardBlocks.size() == Outgoing.size() - 1);
+  int I = 0;
+  for (int E = GuardBlocks.size() - 1; I != E; ++I) {
+    BasicBlock *Out = Outgoing[I];
+    BranchInst::Create(Out, GuardBlocks[I + 1], GuardPredicates[Out],
+                       GuardBlocks[I]);
+  }
+  BasicBlock *Out = Outgoing[I];
+  BranchInst::Create(Out, Outgoing[I + 1], GuardPredicates[Out],
+                     GuardBlocks[I]);
+}
+
+// Assign an index to each outgoing block. At the corresponding guard
+// block, compute the branch condition by comparing this index.
+static void calcPredicateUsingInteger(ArrayRef<EdgeDescriptor> Branches,
+                                      ArrayRef<BasicBlock *> Outgoing,
+                                      ArrayRef<BasicBlock *> GuardBlocks,
+                                      BBPredicates &GuardPredicates) {
+  LLVMContext &Context = GuardBlocks.front()->getContext();
+  BasicBlock *FirstGuardBlock = GuardBlocks.front();
+
+  auto *Phi = PHINode::Create(Type::getInt32Ty(Context), Branches.size(),
+                              "merged.bb.idx", FirstGuardBlock);
+
+  for (auto [BB, Succ0, Succ1] : Branches) {
+    Value *Condition = redirectToHub(BB, Succ0, Succ1, FirstGuardBlock);
+    Value *IncomingId = nullptr;
+    if (Succ0 && 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",
+                                      BB->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, BB);
+  }
+
+  for (int I = 0, E = Outgoing.size() - 1; I != E; ++I) {
+    BasicBlock *Out = Outgoing[I];
+    LLVM_DEBUG(dbgs() << "Creating integer guard for " << Out->getName()
+                      << "\n");
+    auto *Cmp = ICmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_EQ, Phi,
+                                 ConstantInt::get(Type::getInt32Ty(Context), I),
+                                 Out->getName() + ".predicate", GuardBlocks[I]);
+    GuardPredicates[Out] = Cmp;
+  }
+}
+
+// Determine the branch condition to be used at each guard block from the
+// original boolean values.
+static void calcPredicateUsingBooleans(
+    ArrayRef<EdgeDescriptor> Branches, ArrayRef<BasicBlock *> Outgoing,
+    SmallVectorImpl<BasicBlock *> &GuardBlocks, BBPredicates &GuardPredicates,
+    SmallVectorImpl<WeakVH> &DeletionCandidates) {
+  LLVMContext &Context = GuardBlocks.front()->getContext();
+  auto *BoolTrue = ConstantInt::getTrue(Context);
+  auto *BoolFalse = ConstantInt::getFalse(Context);
+  BasicBlock *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) {
+    BasicBlock *Out = Outgoing[I];
+    LLVM_DEBUG(dbgs() << "Creating boolean guard for " << Out->getName()
+                      << "\n");
+
+    auto *Phi =
+        PHINode::Create(Type::getInt1Ty(Context), Branches.size(),
+                        StringRef("Guard.") + Out->getName(), FirstGuardBlock);
+    GuardPredicates[Out] = Phi;
+  }
+
+  for (auto [BB, Succ0, Succ1] : Branches) {
+    Value *Condition = redirectToHub(BB, Succ0, Succ1, FirstGuardBlock);
+
+    // 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 `BB` is true for Succ1.
+    bool OneSuccessorDone = false;
+    for (int I = 0, E = Outgoing.size() - 1; I != E; ++I) {
+      BasicBlock *Out = Outgoing[I];
+      PHINode *Phi = cast<PHINode>(GuardPredicates[Out]);
+      if (Out != Succ0 && Out != Succ1) {
+        Phi->addIncoming(BoolFalse, BB);
+      } else if (!Succ0 || !Succ1 || OneSuccessorDone) {
+        // Optimization: When only one successor is an outgoing block,
+        // the incoming predicate from `BB` is always true.
+        Phi->addIncoming(BoolTrue, BB);
+      } else {
+        assert(Succ0 && Succ1);
+        if (Out == Succ0) {
+          Phi->addIncoming(Condition, BB);
+        } else {
+          Value *Inverted = invertCondition(Condition);
+          DeletionCandidates.push_back(Condition);
+          Phi->addIncoming(Inverted, BB);
+        }
+        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.
+//
+// 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 convertToGuardPredicates(
+    ArrayRef<EdgeDescriptor> Branches, ArrayRef<BasicBlock *> Outgoing,
+    SmallVectorImpl<BasicBlock *> &GuardBlocks,
+    SmallVectorImpl<WeakVH> &DeletionCandidates, const StringRef Prefix,
+    std::optional<unsigned> MaxControlFlowBooleans) {
+  BBPredicates GuardPredicates;
+  Function *F = Outgoing.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(Branches, Outgoing, GuardBlocks, GuardPredicates,
+                               DeletionCandidates);
+  else
+    calcPredicateUsingInteger(Branches, Outgoing, GuardBlocks, GuardPredicates);
+
+  setupBranchForGuard(GuardBlocks, Outgoing, GuardPredicates);
+}
+
+// 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,
+                          ArrayRef<EdgeDescriptor> 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 [BB, Succ0, Succ1] : Incoming) {
+      Value *V = UndefValue::get(Phi->getType());
+      if (BB == Out) {
+        V = NewPhi;
+      } else if (Phi->getBasicBlockIndex(BB) != -1) {
+        V = Phi->removeIncomingValue(BB, false);
+        AllUndef &= isa<UndefValue>(V);
+      }
+      NewPhi->addIncoming(V, BB);
+    }
+    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;
+  }
+}
+
+BasicBlock *ControlFlowHub::finalize(
+    DomTreeUpdater *DTU, SmallVectorImpl<BasicBlock *> &GuardBlocks,
+    const StringRef Prefix, std::optional<unsigned> MaxControlFlowBooleans) {
+#ifndef NDEBUG
+  SmallSet<BasicBlock *, 8> Incoming;
+#endif
+  SetVector<BasicBlock *> Outgoing;
+
+  for (auto [BB, Succ0, Succ1] : Branches) {
+#ifndef NDEBUG
+    assert(Incoming.insert(BB).second && "Duplicate entry for incoming block.");
+#endif
+    if (Succ0)
+      Outgoing.insert(Succ0);
+    if (Succ1)
+      Outgoing.insert(Succ1);
+  }
+
+  if (Outgoing.size() < 2)
+    return Outgoing.front();
+
+  SmallVector<DominatorTree::UpdateType, 16> Updates;
+  if (DTU) {
+    for (auto [BB, Succ0, Succ1] : Branches) {
+      if (Succ0)
+        Updates.push_back({DominatorTree::Delete, BB, Succ0});
+      if (Succ1)
+        Updates.push_back({DominatorTree::Delete, BB, Succ1});
+    }
+  }
+
+  SmallVector<WeakVH, 8> DeletionCandidates;
+  convertToGuardPredicates(Branches, Outgoing.getArrayRef(), GuardBlocks,
+                           DeletionCandidates, Prefix, MaxControlFlowBooleans);
+  BasicBlock *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], Branches, FirstGuardBlock);
+  // Process the Nth (last) outgoing block with the (N-1)th (last) guard block.
+  reconnectPhis(Outgoing.back(), GuardBlocks.back(), Branches, FirstGuardBlock);
+
+  if (DTU) {
+    int NumGuards = GuardBlocks.size();
+
+    for (auto [BB, Succ0, Succ1] : Branches)
+      Updates.push_back({DominatorTree::Insert, BB, 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]});
+    }
+    // The second successor of the last guard block is an outgoing block instead
+    // of having a "next" guard block.
+    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;
+}
diff --git a/llvm/lib/Transforms/Utils/FixIrreducible.cpp b/llvm/lib/Transforms/Utils/FixIrreducible.cpp
index 11e24d0585be48..30075af2ffc654 100644
--- a/llvm/lib/Transforms/Utils/FixIrreducible.cpp
+++ b/llvm/lib/Transforms/Utils/FixIrreducible.cpp
@@ -51,15 +51,14 @@
 // including inside any newly created loops. This ensures that any SCC hidden
 // inside a maximal SCC is also transformed.
 //
-// The actual transformation is handled by function CreateControlFlowHub, which
-// takes a set of incoming blocks (the predecessors) and outgoing blocks (the
-// headers). The function also moves every PHINode in an outgoing block to the
-// hub. Since the hub dominates all the outgoing blocks, each such PHINode
-// continues to dominate its uses. Since every header in an SCC has at least two
-// predecessors, every value used in the header (or later) but defined in a
-// predecessor (or earlier) is represented by a PHINode in a header. Hence the
-// above handling of PHINodes is sufficient and no further processing is
-// required to restore SSA.
+// The actual transformation is handled by the ControlFlowHub, which redirects
+// specified control flow edges through a set of guard blocks. This also moves
+// every PHINode in an outgoing block to the hub. Since the hub dominates all
+// the outgoing blocks, each such PHINode continues to dominate its uses. Since
+// every header in an SCC has at least two predecessors, every value used in the
+// header (or later) but defined in a predecessor (or earlier) is represented by
+// a PHINode in a header. Hence the above handling of PHINodes is sufficient and
+// no further processing is required to restore SSA.
 //
 // Limitation: The pass cannot handle switch statements and indirect
 //             branches. Both must be lowered to plain branches first.
@@ -74,6 +73,7 @@
 #include "llvm/Pass.h"
 #include "llvm/Transforms/Utils.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/ControlFlowUtils.h"
 
 #define DEBUG_TYPE "fix-irreducible"
 
@@ -189,9 +189,24 @@ static void createNaturalLoopInternal(LoopInfo &LI, DominatorTree &DT,
   // Redirect all the backedges through a "hub" consisting of a series
   // of guard blocks that manage the flow of control from the
   // predecessors to the headers.
+  ControlFlowHub CHub;
+  for (BasicBlock *P : Predecessors) {
+    auto *Branch = cast<BranchInst>(P->getTerminator());
+    BasicBlock *Succ0 = Branch->getSuccessor(0);
+    Succ0 = Headers.count(Succ0) ? Succ0 : nullptr;
+    BasicBlock *Succ1 =
+        Branch->isUnconditional() ? nullptr : Branch->getSuccessor(1);
+    Succ1 = Succ1 && Headers.count(Succ1) ? Succ1 : nullptr;
+    CHub.addBranch(P, Succ0, Succ1);
+
+    LLVM_DEBUG(dbgs() << "Added internal branch: " << P->getName() << " -> "
+                      << (Succ0 ? Succ0->getName() : "") << " "
+                      << (Succ1 ? Succ1->getName() : "") << "\n");
+  }
+
   SmallVector<BasicBlock *, 8> GuardBlocks;
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
-  CreateControlFlowHub(&DTU, GuardBlocks, Predecessors, Headers, "irr");
+  CHub.finalize(&DTU, GuardBlocks, "irr");
 #if defined(EXPENSIVE_CHECKS)
   assert(DT.verify(DominatorTree::VerificationLevel::Full));
 #else
diff --git a/llvm/lib/Transforms/Utils/UnifyLoopExits.cpp b/llvm/lib/Transforms/Utils/UnifyLoopExits.cpp
index 1d51f61351fe27..e4f9600ac6c984 100644
--- a/llvm/lib/Transforms/Utils/UnifyLoopExits.cpp
+++ b/llvm/lib/Transforms/Utils/UnifyLoopExits.cpp
@@ -26,6 +26,7 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/Utils.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/ControlFlowUtils.h"
 
 #define DEBUG_TYPE "unify-loop-exits"
 
@@ -34,7 +35,7 @@ using namespace llvm;
 static cl::opt<unsigned> MaxBooleansInControlFlowHub(
     "max-booleans-in-control-flow-hub", cl::init(32), cl::Hidden,
     cl::desc("Set the maximum number of outgoing blocks for using a boolean "
-             "value to record the exiting block in CreateControlFlowHub."));
+             "value to record the exiting block in the ControlFlowHub."));
 
 namespace {
 struct UnifyLoopExitsLegacyPass : public FunctionPass {
@@ -86,7 +87,7 @@ INITIALIZE_PASS_END(UnifyLoopExitsLegacyPass, "unify-loop-exits",
 // for creating the new PHI is well-known, and also the set of incoming blocks
 // to the new PHI.
 static void restoreSSA(const DominatorTree &DT, const Loop *L,
-                       const SetVector<BasicBlock *> &Incoming,
+                       SmallVectorImpl<BasicBlock *> &Incoming,
                        BasicBlock *LoopExitBlock) {
   using InstVector = SmallVector<Instruction *, 8>;
   using IIMap = MapVector<Instruction *, InstVector>;
@@ -140,53 +141,43 @@ static void restoreSSA(const DominatorTree &DT, const Loop *L,
   }
 }
 
+static bool isExitBlock(Loop *L, BasicBlock *Succ, LoopInfo &LI) {
+  Loop *SL = LI.getLoopFor(Succ);
+  if (SL == L || L->contains(SL))
+    return false;
+  return true;
+}
+
 static bool unifyLoopExits(DominatorTree &DT, LoopInfo &LI, Loop *L) {
   // To unify the loop exits, we need a list of the exiting blocks as
   // well as exit blocks. The functions for locating these lists both
   // traverse the entire loop body. It is more efficient to first
   // locate the exiting blocks and then examine their successors to
   // locate the exit blocks.
-  SetVector<BasicBlock *> ExitingBlocks;
-  SetVector<BasicBlock *> Exits;
-
-  // We need SetVectors, but the Loop API takes a vector, so we use a temporary.
-  SmallVector<BasicBlock *, 8> Temp;
-  L->getExitingBlocks(Temp);
-  for (auto *BB : Temp) {
-    ExitingBlocks.insert(BB);
-    for (auto *S : successors(BB)) {
-      auto SL = LI.getLoopFor(S);
-      // A successor is not an exit if it is directly or indirectly in the
-      // current loop.
-      if (SL == L || L->contains(SL))
-        continue;
-      Exits.insert(S);
-    }
-  }
-
-  LLVM_DEBUG(
-      dbgs() << "Found exit blocks:";
-      for (auto Exit : Exits) {
-        dbgs() << " " << Exit->getName();
-      }
-      dbgs() << "\n";
-
-      dbgs() << "Found exiting blocks:";
-      for (auto EB : ExitingBlocks) {
-        dbgs() << " " << EB->getName();
-      }
-      dbgs() << "\n";);
-
-  if (Exits.size() <= 1) {
-    LLVM_DEBUG(dbgs() << "loop does not have multiple exits; nothing to do\n");
-    return false;
+  SmallVector<BasicBlock *, 8> ExitingBlocks;
+  L->getExitingBlocks(ExitingBlocks);
+
+  // Redirect exiting edges through a control flow hub.
+  ControlFlowHub CHub;
+  for (auto *BB : ExitingBlocks) {
+    auto *Branch = cast<BranchInst>(BB->getTerminator());
+    BasicBlock *Succ0 = Branch->getSuccessor(0);
+    Succ0 = isExitBlock(L, Succ0, LI) ? Succ0 : nullptr;
+
+    BasicBlock *Succ1 =
+        Branch->isUnconditional() ? nullptr : Branch->getSuccessor(1);
+    Succ1 = Succ1 && isExitBlock(L, Succ1, LI) ? Succ1 : nullptr;
+    CHub.addBranch(BB, Succ0, Succ1);
+
+    LLVM_DEBUG(dbgs() << "Added internal branch: " << BB->getName() << " -> "
+                      << (Succ0 ? Succ0->getName() : "") << " "
+                      << (Succ1 ? Succ1->getName() : "") << "\n");
   }
 
   SmallVector<BasicBlock *, 8> GuardBlocks;
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
-  auto LoopExitBlock =
-      CreateControlFlowHub(&DTU, GuardBlocks, ExitingBlocks, Exits, "loop.exit",
-                           MaxBooleansInControlFlowHub.getValue());
+  BasicBlock *LoopExitBlock = CHub.finalize(
+      &DTU, GuardBlocks, "loop.exit", MaxBooleansInControlFlowHub.getValue());
 
   restoreSSA(DT, L, ExitingBlocks, LoopExitBlock);
 
diff --git a/llvm/test/CodeGen/AMDGPU/local-atomicrmw-fadd.ll b/llvm/test/CodeGen/AMDGPU/local-atomicrmw-fadd.ll
index 65614a17fc0114..0894d3251423d6 100644
--- a/llvm/test/CodeGen/AMDGPU/local-atomicrmw-fadd.ll
+++ b/llvm/test/CodeGen/AMDGPU/local-atomicrmw-fadd.ll
@@ -7644,9 +7644,9 @@ define amdgpu_kernel void @local_ds_fadd(ptr addrspace(1) %out, ptr addrspace(3)
 ; GFX7-NEXT:    s_or_b64 s[8:9], s[0:1], s[8:9]
 ; GFX7-NEXT:    s_andn2_b64 exec, exec, s[8:9]
 ; GFX7-NEXT:    s_cbranch_execnz .LBB28_2
-; GFX7-NEXT:  ; %bb.3: ; %Flow22
+; GFX7-NEXT:  ; %bb.3: ; %Flow23
 ; GFX7-NEXT:    s_or_b64 exec, exec, s[8:9]
-; GFX7-NEXT:  .LBB28_4: ; %Flow23
+; GFX7-NEXT:  .LBB28_4: ; %Flow24
 ; GFX7-NEXT:    s_or_b64 exec, exec, s[6:7]
 ; GFX7-NEXT:    s_mov_b64 s[8:9], exec
 ; GFX7-NEXT:    v_readfirstlane_b32 s10, v1
@@ -7674,7 +7674,7 @@ define amdgpu_kernel void @local_ds_fadd(ptr addrspace(1) %out, ptr addrspace(3)
 ; GFX7-NEXT:    v_mov_b32_e32 v3, v4
 ; GFX7-NEXT:    s_andn2_b64 exec, exec, s[8:9]
 ; GFX7-NEXT:    s_cbranch_execnz .LBB28_6
-; GFX7-NEXT:  .LBB28_7: ; %Flow21
+; GFX7-NEXT:  .LBB28_7: ; %Flow22
 ; GFX7-NEXT:    s_or_b64 exec, exec, s[6:7]
 ; GFX7-NEXT:    v_cvt_f32_ubyte0_e32 v0, v0
 ; GFX7-NEXT:    v_mul_f32_e32 v0, 0x42280000, v0
@@ -7723,7 +7723,7 @@ define amdgpu_kernel void @local_ds_fadd(ptr addrspace(1) %out, ptr addrspace(3)
 ; GFX7-NEXT:    s_cbranch_execnz .LBB28_11
 ; GFX7-NEXT:  ; %bb.12: ; %Flow
 ; GFX7-NEXT:    s_or_b64 exec, exec, s[4:5]
-; GFX7-NEXT:  .LBB28_13: ; %Flow19
+; GFX7-NEXT:  .LBB28_13: ; %Flow20
 ; GFX7-NEXT:    s_or_b64 exec, exec, s[6:7]
 ; GFX7-NEXT:    s_load_dwordx2 s[0:1], s[2:3], 0x0
 ; GFX7-NEXT:    v_readfirstlane_b32 s4, v2
@@ -7768,9 +7768,9 @@ define amdgpu_kernel void @local_ds_fadd(ptr addrspace(1) %out, ptr addrspace(3)
 ; GFX6-NEXT:    s_or_b64 s[8:9], s[0:1], s[8:9]
 ; GFX6-NEXT:    s_andn2_b64 exec, exec, s[8:9]
 ; GFX6-NEXT:    s_cbranch_execnz .LBB28_2
-; GFX6-NEXT:  ; %bb.3: ; %Flow20
+; GFX6-NEXT:  ; %bb.3: ; %Flow21
 ; GFX6-NEXT:    s_or_b64 exec, exec, s[8:9]
-; GFX6-NEXT:  .LBB28_4: ; %Flow21
+; GFX6-NEXT:  .LBB28_4: ; %Flow22
 ; GFX6-NEXT:    s_or_b64 exec, exec, s[6:7]
 ; GFX6-NEXT:    s_mov_b64 s[8:9], exec
 ; GFX6-NEXT:    v_readfirstlane_b32 s10, v1
@@ -7798,7 +7798,7 @@ define amdgpu_kernel void @local_ds_fadd(ptr addrspace(1) %out, ptr addrspace(3)
 ; GFX6-NEXT:    v_mov_b32_e32 v3, v4
 ; GFX6-NEXT:    s_andn2_b64 exec, exec, s[8:9]
 ; GFX6-NEXT:    s_cbranch_execnz .LBB28_6
-; GFX6-NEXT:  .LBB28_7: ; %Flow19
+; GFX6-NEXT:  .LBB28_7: ; %Flow20
 ; GFX6-NEXT:    s_or_b64 exec, exec, s[6:7]
 ; GFX6-NEXT:    v_cvt_f32_ubyte0_e32 v0, v0
 ; GFX6-NEXT:    v_mul_f32_e32 v0, 0x42280000, v0
@@ -7847,7 +7847,7 @@ define amdgpu_kernel void @local_ds_fadd(ptr addrspace(1) %out, ptr addrspace(3)
 ; GFX6-NEXT:    s_cbranch_execnz .LBB28_11
 ; GFX6-NEXT:  ; %bb.12: ; %Flow
 ; GFX6-NEXT:    s_or_b64 exec, exec, s[4:5]
-; GFX6-NEXT:  .LBB28_13: ; %Flow17
+; GFX6-NEXT:  .LBB28_13: ; %Flow18
 ; GFX6-NEXT:    s_or_b64 exec, exec, s[6:7]
 ; GFX6-NEXT:    s_load_dwordx2 s[0:1], s[2:3], 0x0
 ; GFX6-NEXT:    v_readfirstlane_b32 s4, v2
@@ -8484,9 +8484,9 @@ define amdgpu_kernel void @local_ds_fadd_one_as(ptr addrspace(1) %out, ptr addrs
 ; GFX7-NEXT:    s_or_b64 s[8:9], s[0:1], s[8:9]
 ; GFX7-NEXT:    s_andn2_b64 exec, exec, s[8:9]
 ; GFX7-NEXT:    s_cbranch_execnz .LBB29_2
-; GFX7-NEXT:  ; %bb.3: ; %Flow22
+; GFX7-NEXT:  ; %bb.3: ; %Flow23
 ; GFX7-NEXT:    s_or_b64 exec, exec, s[8:9]
-; GFX7-NEXT:  .LBB29_4: ; %Flow23
+; GFX7-NEXT:  .LBB29_4: ; %Flow24
 ; GFX7-NEXT:    s_or_b64 exec, exec, s[6:7]
 ; GFX7-NEXT:    s_mov_b64 s[8:9], exec
 ; GFX7-NEXT:    v_readfirstlane_b32 s10, v1
@@ -8514,7 +8514,7 @@ define amdgpu_kernel void @local_ds_fadd_one_as(ptr addrspace(1) %out, ptr addrs
 ; GFX7-NEXT:    v_mov_b32_e32 v3, v4
 ; GFX7-NEXT:    s_andn2_b64 exec, exec, s[8:9]
 ; GFX7-NEXT:    s_cbranch_execnz .LBB29_6
-; GFX7-NEXT:  .LBB29_7: ; %Flow21
+; GFX7-NEXT:  .LBB29_7: ; %Flow22
 ; GFX7-NEXT:    s_or_b64 exec, exec, s[6:7]
 ; GFX7-NEXT:    v_cvt_f32_ubyte0_e32 v0, v0
 ; GFX7-NEXT:    v_mul_f32_e32 v0, 0x42280000, v0
@@ -8563,7 +8563,7 @@ define amdgpu_kernel void @local_ds_fadd_one_as(ptr addrspace(1) %out, ptr addrs
 ; GFX7-NEXT:    s_cbranch_execnz .LBB29_11
 ; GFX7-NEXT:  ; %bb.12: ; %Flow
 ; GFX7-NEXT:    s_or_b64 exec, exec, s[4:5]
-; GFX7-NEXT:  .LBB29_13: ; %Flow19
+; GFX7-NEXT:  .LBB29_13: ; %Flow20
 ; GFX7-NEXT:    s_or_b64 exec, exec, s[6:7]
 ; GFX7-NEXT:    s_load_dwordx2 s[0:1], s[2:3], 0x0
 ; GFX7-NEXT:    v_readfirstlane_b32 s4, v2
@@ -8608,9 +8608,9 @@ define amdgpu_kernel void @local_ds_fadd_one_as(ptr addrspace(1) %out, ptr addrs
 ; GFX6-NEXT:    s_or_b64 s[8:9], s[0:1], s[8:9]
 ; GFX6-NEXT:    s_andn2_b64 exec, exec, s[8:9]
 ; GFX6-NEXT:    s_cbranch_execnz .LBB29_2
-; GFX6-NEXT:  ; %bb.3: ; %Flow20
+; GFX6-NEXT:  ; %bb.3: ; %Flow21
 ; GFX6-NEXT:    s_or_b64 exec, exec, s[8:9]
-; GFX6-NEXT:  .LBB29_4: ; %Flow21
+; GFX6-NEXT:  .LBB29_4: ; %Flow22
 ; GFX6-NEXT:    s_or_b64 exec, exec, s[6:7]
 ; GFX6-NEXT:    s_mov_b64 s[8:9], exec
 ; GFX6-NEXT:    v_readfirstlane_b32 s10, v1
@@ -8638,7 +8638,7 @@ define amdgpu_kernel void @local_ds_fadd_one_as(ptr addrspace(1) %out, ptr addrs
 ; GFX6-NEXT:    v_mov_b32_e32 v3, v4
 ; GFX6-NEXT:    s_andn2_b64 exec, exec, s[8:9]
 ; GFX6-NEXT:    s_cbranch_execnz .LBB29_6
-; GFX6-NEXT:  .LBB29_7: ; %Flow19
+; GFX6-NEXT:  .LBB29_7: ; %Flow20
 ; GFX6-NEXT:    s_or_b64 exec, exec, s[6:7]
 ; GFX6-NEXT:    v_cvt_f32_ubyte0_e32 v0, v0
 ; GFX6-NEXT:    v_mul_f32_e32 v0, 0x42280000, v0
@@ -8687,7 +8687,7 @@ define amdgpu_kernel void @local_ds_fadd_one_as(ptr addrspace(1) %out, ptr addrs
 ; GFX6-NEXT:    s_cbranch_execnz .LBB29_11
 ; GFX6-NEXT:  ; %bb.12: ; %Flow
 ; GFX6-NEXT:    s_or_b64 exec, exec, s[4:5]
-; GFX6-NEXT:  .LBB29_13: ; %Flow17
+; GFX6-NEXT:  .LBB29_13: ; %Flow18
 ; GFX6-NEXT:    s_or_b64 exec, exec, s[6:7]
 ; GFX6-NEXT:    s_load_dwordx2 s[0:1], s[2:3], 0x0
 ; GFX6-NEXT:    v_readfirstlane_b32 s4, v2
diff --git a/llvm/test/Transforms/FixIrreducible/basic.ll b/llvm/test/Transforms/FixIrreducible/basic.ll
index 7ba1360160db20..c682da1727e641 100644
--- a/llvm/test/Transforms/FixIrreducible/basic.ll
+++ b/llvm/test/Transforms/FixIrreducible/basic.ll
@@ -1,6 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt < %s -fix-irreducible -S | FileCheck %s -check-prefix=CHECK
-; RUN: opt < %s -passes=fix-irreducible -S | FileCheck %s -check-prefix=CHECK
+; RUN: opt < %s -fix-irreducible --verify-loop-info -S | FileCheck %s
+; RUN: opt < %s -passes="fix-irreducible,verify<loops>" -S | FileCheck %s -check-prefix=CHECK
 
 define i32 @basic(i1 %PredEntry, i1 %PredLeft, i1 %PredRight, i32 %X, i32 %Y) {
 ; CHECK-LABEL: @basic(
diff --git a/llvm/test/Transforms/FixIrreducible/bug45623.ll b/llvm/test/Transforms/FixIrreducible/bug45623.ll
index 89d31dd4fea6b8..ac188cb936c034 100644
--- a/llvm/test/Transforms/FixIrreducible/bug45623.ll
+++ b/llvm/test/Transforms/FixIrreducible/bug45623.ll
@@ -1,5 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt < %s -fix-irreducible -S | FileCheck %s
+; RUN: opt < %s -fix-irreducible --verify-loop-info -S | FileCheck %s
+; RUN: opt < %s -passes="fix-irreducible,verify<loops>" -S | FileCheck %s
 
 define dso_local void @tre_tnfa_run_backtrack() {
 ; CHECK-LABEL: @tre_tnfa_run_backtrack(
diff --git a/llvm/test/Transforms/FixIrreducible/nested.ll b/llvm/test/Transforms/FixIrreducible/nested.ll
index 85f7d39920b9d6..82c49e67a94842 100644
--- a/llvm/test/Transforms/FixIrreducible/nested.ll
+++ b/llvm/test/Transforms/FixIrreducible/nested.ll
@@ -1,5 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt < %s -fix-irreducible -S | FileCheck %s -check-prefix=CHECK
+; RUN: opt < %s -fix-irreducible --verify-loop-info -S | FileCheck %s
+; RUN: opt < %s -passes="fix-irreducible,verify<loops>" -S | FileCheck %s
 
 define void @nested_irr_top_level(i1 %Pred0, i1 %Pred1, i1 %Pred2, i1 %Pred3, i1 %Pred4, i1 %Pred5) {
 ; CHECK-LABEL: @nested_irr_top_level(
diff --git a/llvm/test/Transforms/FixIrreducible/switch.ll b/llvm/test/Transforms/FixIrreducible/switch.ll
index 27b9f9bf3b53be..8b30c9aaf8770b 100644
--- a/llvm/test/Transforms/FixIrreducible/switch.ll
+++ b/llvm/test/Transforms/FixIrreducible/switch.ll
@@ -1,5 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt < %s -passes='lower-switch,fix-irreducible' -S | FileCheck %s
+; RUN: opt < %s -lowerswitch -fix-irreducible --verify-loop-info -S | FileCheck %s
+; RUN: opt < %s -passes="lower-switch,fix-irreducible,verify<loops>" -S | FileCheck %s
 
 define void @loop_1(i32 %Value, i1 %PredEntry, i1 %PredD) {
 ; CHECK-LABEL: @loop_1(
diff --git a/llvm/test/Transforms/FixIrreducible/unreachable.ll b/llvm/test/Transforms/FixIrreducible/unreachable.ll
index 71cd81e01953ea..b7dcea03f2e22b 100644
--- a/llvm/test/Transforms/FixIrreducible/unreachable.ll
+++ b/llvm/test/Transforms/FixIrreducible/unreachable.ll
@@ -1,4 +1,6 @@
-; RUN: opt %s -fix-irreducible -S -o - | FileCheck %s
+; NOTE: Do not autogenerate
+; RUN: opt < %s -fix-irreducible --verify-loop-info -S | FileCheck %s
+; RUN: opt < %s -passes="fix-irreducible,verify<loops>" -S | FileCheck %s
 
 ; CHECK-LABEL: @unreachable(
 ; CHECK: entry:



More information about the llvm-branch-commits mailing list