implement feedback

Author: ro-i

llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h

@@ -608,17 +608,10 @@ LLVM_ABI bool SplitIndirectBrCriticalEdges(Function &F,
 // successors
 LLVM_ABI void InvertBranch(BranchInst *PBI, IRBuilderBase &Builder);
 
-template <typename... TermInst>
-LLVM_ABI bool hasOnlyGivenTerminators(const Function &F);
-
-// Check whether the function only has blocks with simple terminators:
+// Check whether the function only has simple terminator:
 // br/brcond/unreachable/ret
 LLVM_ABI bool hasOnlySimpleTerminator(const Function &F);
 
-// Check whether the function only has blocks with simple terminators
-// (br/brcond/unreachable/ret) or callbr.
-LLVM_ABI bool hasOnlySimpleTerminatorOrCallBr(const Function &F);
-
 LLVM_ABI Printable printBBPtr(const BasicBlock *BB);
 
 } // end namespace llvm

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

@@ -1766,25 +1766,16 @@ void llvm::InvertBranch(BranchInst *PBI, IRBuilderBase &Builder) {
   PBI->swapSuccessors();
 }
 
-template <typename... TermInst>
-bool llvm::hasOnlyGivenTerminators(const Function &F) {
+bool llvm::hasOnlySimpleTerminator(const Function &F) {
   for (auto &BB : F) {
     auto *Term = BB.getTerminator();
-    if (!(isa<TermInst>(Term) || ...))
+    if (!(isa<ReturnInst>(Term) || isa<UnreachableInst>(Term) ||
+          isa<BranchInst>(Term)))
       return false;
   }
   return true;
 }
 
-bool llvm::hasOnlySimpleTerminator(const Function &F) {
-  return hasOnlyGivenTerminators<ReturnInst, UnreachableInst, BranchInst>(F);
-}
-
-bool llvm::hasOnlySimpleTerminatorOrCallBr(const Function &F) {
-  return hasOnlyGivenTerminators<ReturnInst, UnreachableInst, BranchInst,
-                                 CallBrInst>(F);
-}
-
 Printable llvm::printBBPtr(const BasicBlock *BB) {
   return Printable([BB](raw_ostream &OS) {
     if (BB)

llvm/lib/Transforms/Utils/ControlFlowUtils.cpp

@@ -346,6 +346,53 @@ std::pair<BasicBlock *, bool> ControlFlowHub::finalize(
   return {FirstGuardBlock, true};
 }
 
+// Check if the given cycle is disjoint with the cycle of the given basic block.
+// If the basic block does not belong to any cycle, this is regarded as
+// disjoint, too.
+static bool disjointWithBlock(CycleInfo *CI, Cycle *C, BasicBlock *BB) {
+  Cycle *CC = CI->getCycle(BB);
+  if (!CC)
+    return true;
+  Cycle *CommonC = CI->getSmallestCommonCycle(C, CC);
+  return CommonC != C && CommonC != CC;
+}
+
+// Check if the given loop is disjoint with the loop of the given basic block.
+// If the basic block does not belong to any loop, this is regarded as
+// disjoint, too.
+static bool disjointWithBlock(LoopInfo *LI, Loop *L, BasicBlock *BB) {
+  Loop *LL = LI->getLoopFor(BB);
+  return LL && !L->contains(LL) && !LL->contains(L);
+}
+
+template <typename TI, typename T>
+static void updateCycleLoopInfo(TI *LCI, bool AttachToCallBr,
+                                BasicBlock *CallBrBlock,
+                                BasicBlock *CallBrTarget, BasicBlock *Succ) {
+  static_assert(std::is_same_v<TI, CycleInfo> || std::is_same_v<TI, LoopInfo>,
+                "type must be CycleInfo or LoopInfo");
+  if (!LCI)
+    return;
+
+  T *LC;
+  if constexpr (std::is_same_v<TI, CycleInfo>)
+    LC = LCI->getCycle(AttachToCallBr ? CallBrBlock : Succ);
+  else
+    LC = LCI->getLoopFor(AttachToCallBr ? CallBrBlock : Succ);
+  if (!LC)
+    return;
+
+  // Check if the cycles/loops are disjoint. In that case, we do not add the
+  // intermediate target to any cycle/loop.
+  if (AttachToCallBr && disjointWithBlock(LCI, LC, Succ))
+    return;
+
+  if constexpr (std::is_same_v<TI, CycleInfo>)
+    LCI->addBlockToCycle(CallBrTarget, LC);
+  else
+    LC->addBasicBlockToLoop(CallBrTarget, *LCI);
+}
+
 BasicBlock *ControlFlowHub::createCallBrTarget(
     CallBrInst *CallBr, BasicBlock *Succ, unsigned SuccIdx, bool AttachToCallBr,
     CycleInfo *CI, DomTreeUpdater *DTU, LoopInfo *LI) {
@@ -359,36 +406,10 @@ BasicBlock *ControlFlowHub::createCallBrTarget(
   CallBr->setSuccessor(SuccIdx, CallBrTarget);
   // Jump from the new target block to the original successor.
   BranchInst::Create(Succ, CallBrTarget);
-  if (LI) {
-    if (Loop *L = LI->getLoopFor(AttachToCallBr ? CallBrBlock : Succ)) {
-      bool AddToLoop = true;
-      if (AttachToCallBr) {
-        // Check if the loops are disjoint. In that case, we do not add the
-        // intermediate target to any loop.
-        if (auto *LL = LI->getLoopFor(Succ);
-            LL && !L->contains(LL) && !LL->contains(L))
-          AddToLoop = false;
-      }
-      if (AddToLoop)
-        L->addBasicBlockToLoop(CallBrTarget, *LI);
-    }
-  }
-  if (CI) {
-    if (auto *C = CI->getCycle(AttachToCallBr ? CallBrBlock : Succ); C) {
-      bool AddToCycle = true;
-      if (AttachToCallBr) {
-        // Check if the cycles are disjoint. In that case, we do not add the
-        // intermediate target to any cycle.
-        if (auto *CC = CI->getCycle(Succ); CC) {
-          auto *CommonC = CI->getSmallestCommonCycle(C, CC);
-          if (CommonC != C && CommonC != CC)
-            AddToCycle = false;
-        }
-      }
-      if (AddToCycle)
-        CI->addBlockToCycle(CallBrTarget, C);
-    }
-  }
+  updateCycleLoopInfo<LoopInfo, Loop>(LI, AttachToCallBr, CallBrBlock,
+                                      CallBrTarget, Succ);
+  updateCycleLoopInfo<CycleInfo, Cycle>(CI, AttachToCallBr, CallBrBlock,
+                                        CallBrTarget, Succ);
   if (DTU) {
     DTU->applyUpdates({{DominatorTree::Insert, CallBrBlock, CallBrTarget}});
     if (DTU->getDomTree().dominates(CallBrBlock, Succ))

llvm/lib/Transforms/Utils/FixIrreducible.cpp

@@ -308,7 +308,7 @@ static bool fixIrreducible(Cycle &C, CycleInfo &CI, DominatorTree &DT,
         if (Succ != Header)
           continue;
         BasicBlock *NewSucc = ControlFlowHub::createCallBrTarget(
-            CallBr, Succ, I, false, &CI, &DTU, LI);
+            CallBr, Succ, I, /* AttachToCallBr = */ false, &CI, &DTU, LI);
         CHub.addBranch(NewSucc, Succ);
         LLVM_DEBUG(dbgs() << "Added internal branch: " << printBBPtr(NewSucc)
                           << " -> " << printBBPtr(Succ) << "\n");
@@ -345,7 +345,7 @@ static bool fixIrreducible(Cycle &C, CycleInfo &CI, DominatorTree &DT,
         if (!C.contains(Succ))
           continue;
         BasicBlock *NewSucc = ControlFlowHub::createCallBrTarget(
-            CallBr, Succ, I, true, &CI, &DTU, LI);
+            CallBr, Succ, I, /* AttachToCallBr = */ true, &CI, &DTU, LI);
         CHub.addBranch(NewSucc, Succ);
         LLVM_DEBUG(dbgs() << "Added external branch: " << printBBPtr(NewSucc)
                           << " -> " << printBBPtr(Succ) << "\n");
@@ -400,8 +400,6 @@ static bool FixIrreducibleImpl(Function &F, CycleInfo &CI, DominatorTree &DT,
   LLVM_DEBUG(dbgs() << "===== Fix irreducible control-flow in function: "
                     << F.getName() << "\n");
 
-  assert(hasOnlySimpleTerminatorOrCallBr(F) && "Unsupported block terminator.");
-
   bool Changed = false;
   for (Cycle *TopCycle : CI.toplevel_cycles()) {
     for (Cycle *C : depth_first(TopCycle)) {

llvm/lib/Transforms/Utils/UnifyLoopExits.cpp

@@ -179,7 +179,7 @@ static bool unifyLoopExits(DominatorTree &DT, LoopInfo &LI, Loop *L) {
         if (L->contains(Succ))
           continue;
         BasicBlock *NewSucc = ControlFlowHub::createCallBrTarget(
-            CallBr, Succ, J, false, nullptr, &DTU, &LI);
+            CallBr, Succ, J, /* AttachToCallBr = */ false, nullptr, &DTU, &LI);
         // ExitingBlocks is later used to restore SSA, so we need to make sure
         // that the blocks used for phi nodes in the guard blocks match the
         // predecessors of the guard blocks, which, in the case of callbr, are
@@ -219,7 +219,7 @@ static bool unifyLoopExits(DominatorTree &DT, LoopInfo &LI, Loop *L) {
   // members of the parent loop. Same goes for the callbr target blocks if they
   // have not already been added to the respective parent loop by adding them to
   // the original callbr target's loop.
-  if (auto ParentLoop = L->getParentLoop()) {
+  if (auto *ParentLoop = L->getParentLoop()) {
     for (auto *G : GuardBlocks) {
       ParentLoop->addBasicBlockToLoop(G, LI);
     }
@@ -254,8 +254,6 @@ bool UnifyLoopExitsLegacyPass::runOnFunction(Function &F) {
   auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
   auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
 
-  assert(hasOnlySimpleTerminatorOrCallBr(F) && "Unsupported block terminator.");
-
   return runImpl(LI, DT);
 }