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

[Matt Arsenault via llvm-branch-commits]

================
@@ -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),
----------------
arsenm wrote:

Avoid repeated getInt32Ty calls 

https://github.com/llvm/llvm-project/pull/103013