implement feedback

Author: ro-i

llvm/include/llvm/ADT/GenericCycleImpl.h

@@ -561,6 +561,17 @@ auto GenericCycleInfo<ContextT>::getSmallestCommonCycle(CycleT *A,
   return A;
 }
 
+/// \brief Find the innermost cycle containing both given blocks.
+///
+/// \returns the innermost cycle containing both \p A and \p B
+///          or nullptr if there is no such cycle.
+template <typename ContextT>
+auto GenericCycleInfo<ContextT>::getSmallestCommonCycle(BlockT *A,
+                                                        BlockT *B) const
+    -> CycleT * {
+  return getSmallestCommonCycle(getCycle(A), getCycle(B));
+}
+
 /// \brief get the depth for the cycle which containing a given block.
 ///
 /// \returns the depth for the innermost cycle containing \p Block or 0 if it is

llvm/include/llvm/ADT/GenericCycleInfo.h

@@ -298,6 +298,7 @@ template <typename ContextT> class GenericCycleInfo {
 
   CycleT *getCycle(const BlockT *Block) const;
   CycleT *getSmallestCommonCycle(CycleT *A, CycleT *B) const;
+  CycleT *getSmallestCommonCycle(BlockT *A, BlockT *B) const;
   unsigned getCycleDepth(const BlockT *Block) const;
   CycleT *getTopLevelParentCycle(BlockT *Block);
 

llvm/include/llvm/Support/GenericLoopInfo.h

@@ -615,6 +615,17 @@ template <class BlockT, class LoopT> class LoopInfoBase {
     return L ? L->getLoopDepth() : 0;
   }
 
+  /// \brief Find the innermost loop containing both given loops.
+  ///
+  /// \returns the innermost loop containing both \p A and \p B
+  ///          or nullptr if there is no such loop.
+  LoopT *getSmallestCommonLoop(LoopT *A, LoopT *B) const;
+  /// \brief Find the innermost loop containing both given blocks.
+  ///
+  /// \returns the innermost loop containing both \p A and \p B
+  ///          or nullptr if there is no such loop.
+  LoopT *getSmallestCommonLoop(BlockT *A, BlockT *B) const;
+
   // True if the block is a loop header node
   bool isLoopHeader(const BlockT *BB) const {
     const LoopT *L = getLoopFor(BB);

llvm/include/llvm/Support/GenericLoopInfoImpl.h

@@ -645,6 +645,36 @@ LoopInfoBase<BlockT, LoopT>::getLoopsInReverseSiblingPreorder() const {
   return PreOrderLoops;
 }
 
+template <class BlockT, class LoopT>
+LoopT *LoopInfoBase<BlockT, LoopT>::getSmallestCommonLoop(LoopT *A,
+                                                          LoopT *B) const {
+  if (!A || !B)
+    return nullptr;
+
+  // If lops A and B have different depth replace them with parent loop
+  // until they have the same depth.
+  while (A->getLoopDepth() > B->getLoopDepth())
+    A = A->getParentLoop();
+  while (B->getLoopDepth() > A->getLoopDepth())
+    B = B->getParentLoop();
+
+  // Loops A and B are at same depth but may be disjoint, replace them with
+  // parent loops until we find loop that contains both or we run out of
+  // parent loops.
+  while (A != B) {
+    A = A->getParentLoop();
+    B = B->getParentLoop();
+  }
+
+  return A;
+}
+
+template <class BlockT, class LoopT>
+LoopT *LoopInfoBase<BlockT, LoopT>::getSmallestCommonLoop(BlockT *A,
+                                                          BlockT *B) const {
+  return getSmallestCommonLoop(getLoopFor(A), getLoopFor(B));
+}
+
 // Debugging
 template <class BlockT, class LoopT>
 void LoopInfoBase<BlockT, LoopT>::print(raw_ostream &OS) const {

llvm/include/llvm/Transforms/Utils/ControlFlowUtils.h

@@ -135,23 +135,21 @@ struct ControlFlowHub {
   /// \param Succ           The original successor basic block to be reached.
   /// \param SuccIdx        The index of the successor in the callbr
   ///                       instruction.
-  /// \param AttachToCallBr If true, the new block is associated with the
-  ///                       callbr's parent for loop/cycle info.
-  ///                       If false, the new block is associated with the
-  ///                       callbr's successor for loop/cycle info.
   /// \param CI             Optional CycleInfo for updating cycle membership.
   /// \param DTU            Optional DomTreeUpdater for updating the dominator
   ///                       tree.
   /// \param LI             Optional LoopInfo for updating loop membership.
+  /// \param UpdatedLI      Optional output flag indicating if LoopInfo has been
+  ///                       updated.
   ///
   /// \returns The newly created intermediate target block.
   ///
   /// \note This function updates PHI nodes, dominator tree, loop info, and
   /// cycle info as needed.
   static BasicBlock *
   createCallBrTarget(CallBrInst *CallBr, BasicBlock *Succ, unsigned SuccIdx,
-                     bool AttachToCallBr = true, CycleInfo *CI = nullptr,
-                     DomTreeUpdater *DTU = nullptr, LoopInfo *LI = nullptr);
+                     CycleInfo *CI = nullptr, DomTreeUpdater *DTU = nullptr,
+                     LoopInfo *LI = nullptr, bool *UpdatedLI = nullptr);
 };
 
 } // end namespace llvm

llvm/lib/Transforms/Utils/ControlFlowUtils.cpp

@@ -346,56 +346,35 @@ 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,
+static bool updateCycleLoopInfo(TI *LCI, 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;
+    return false;
 
   T *LC;
   if constexpr (std::is_same_v<TI, CycleInfo>)
-    LC = LCI->getCycle(AttachToCallBr ? CallBrBlock : Succ);
+    LC = LCI->getSmallestCommonCycle(CallBrBlock, Succ);
   else
-    LC = LCI->getLoopFor(AttachToCallBr ? CallBrBlock : Succ);
+    LC = LCI->getSmallestCommonLoop(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;
+    return false;
 
   if constexpr (std::is_same_v<TI, CycleInfo>)
     LCI->addBlockToCycle(CallBrTarget, LC);
   else
     LC->addBasicBlockToLoop(CallBrTarget, *LCI);
+
+  return true;
 }
 
-BasicBlock *ControlFlowHub::createCallBrTarget(
-    CallBrInst *CallBr, BasicBlock *Succ, unsigned SuccIdx, bool AttachToCallBr,
-    CycleInfo *CI, DomTreeUpdater *DTU, LoopInfo *LI) {
+BasicBlock *ControlFlowHub::createCallBrTarget(CallBrInst *CallBr,
+                                               BasicBlock *Succ,
+                                               unsigned SuccIdx, CycleInfo *CI,
+                                               DomTreeUpdater *DTU,
+                                               LoopInfo *LI, bool *UpdatedLI) {
   BasicBlock *CallBrBlock = CallBr->getParent();
   BasicBlock *CallBrTarget =
       BasicBlock::Create(CallBrBlock->getContext(),
@@ -406,10 +385,11 @@ BasicBlock *ControlFlowHub::createCallBrTarget(
   CallBr->setSuccessor(SuccIdx, CallBrTarget);
   // Jump from the new target block to the original successor.
   BranchInst::Create(Succ, CallBrTarget);
-  updateCycleLoopInfo<LoopInfo, Loop>(LI, AttachToCallBr, CallBrBlock,
-                                      CallBrTarget, Succ);
-  updateCycleLoopInfo<CycleInfo, Cycle>(CI, AttachToCallBr, CallBrBlock,
-                                        CallBrTarget, Succ);
+  bool Updated =
+      updateCycleLoopInfo<LoopInfo, Loop>(LI, CallBrBlock, CallBrTarget, Succ);
+  if (UpdatedLI)
+    *UpdatedLI = Updated;
+  updateCycleLoopInfo<CycleInfo, Cycle>(CI, CallBrBlock, CallBrTarget, Succ);
   if (DTU) {
     DTU->applyUpdates({{DominatorTree::Insert, CallBrBlock, CallBrTarget}});
     if (DTU->getDomTree().dominates(CallBrBlock, Succ))

llvm/lib/Transforms/Utils/FixIrreducible.cpp

@@ -308,8 +308,8 @@ static bool fixIrreducible(Cycle &C, CycleInfo &CI, DominatorTree &DT,
         BasicBlock *Succ = CallBr->getSuccessor(I);
         if (Succ != Header)
           continue;
-        BasicBlock *NewSucc = ControlFlowHub::createCallBrTarget(
-            CallBr, Succ, I, /* AttachToCallBr = */ false, &CI, &DTU, LI);
+        BasicBlock *NewSucc =
+            ControlFlowHub::createCallBrTarget(CallBr, Succ, I, &CI, &DTU, LI);
         CHub.addBranch(NewSucc, Succ);
         LLVM_DEBUG(dbgs() << "Added internal branch: "
                           << printBasicBlock(NewSucc) << " -> "
@@ -347,8 +347,8 @@ static bool fixIrreducible(Cycle &C, CycleInfo &CI, DominatorTree &DT,
         BasicBlock *Succ = CallBr->getSuccessor(I);
         if (!C.contains(Succ))
           continue;
-        BasicBlock *NewSucc = ControlFlowHub::createCallBrTarget(
-            CallBr, Succ, I, /* AttachToCallBr = */ true, &CI, &DTU, LI);
+        BasicBlock *NewSucc =
+            ControlFlowHub::createCallBrTarget(CallBr, Succ, I, &CI, &DTU, LI);
         CHub.addBranch(NewSucc, Succ);
         LLVM_DEBUG(dbgs() << "Added external branch: "
                           << printBasicBlock(NewSucc) << " -> "

llvm/lib/Transforms/Utils/UnifyLoopExits.cpp

@@ -155,7 +155,7 @@ static bool unifyLoopExits(DominatorTree &DT, LoopInfo &LI, Loop *L) {
   L->getExitingBlocks(ExitingBlocks);
 
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
-  SmallVector<BasicBlock *, 8> CallBrTargetBlocks;
+  SmallVector<BasicBlock *, 8> CallBrTargetBlocksToFix;
   // Redirect exiting edges through a control flow hub.
   ControlFlowHub CHub;
 
@@ -179,8 +179,13 @@ static bool unifyLoopExits(DominatorTree &DT, LoopInfo &LI, Loop *L) {
         BasicBlock *Succ = CallBr->getSuccessor(J);
         if (L->contains(Succ))
           continue;
+        bool UpdatedLI = false;
         BasicBlock *NewSucc = ControlFlowHub::createCallBrTarget(
-            CallBr, Succ, J, /* AttachToCallBr = */ false, nullptr, &DTU, &LI);
+            CallBr, Succ, J, nullptr, &DTU, &LI, &UpdatedLI);
+        // Even if CallBr and Succ do not have a common parent loop, we need to
+        // add the new target block to the parent loop of the current loop.
+        if (!UpdatedLI)
+          CallBrTargetBlocksToFix.push_back(NewSucc);
         // 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
@@ -191,9 +196,6 @@ static bool unifyLoopExits(DominatorTree &DT, LoopInfo &LI, Loop *L) {
         LLVM_DEBUG(dbgs() << "Added exiting branch: "
                           << printBasicBlock(NewSucc) << " -> "
                           << printBasicBlock(Succ) << "\n");
-        // Also add the new target block to the list of exiting blocks that
-        // should later be added to the parent loops.
-        CallBrTargetBlocks.push_back(NewSucc);
       }
     } else {
       llvm_unreachable("unsupported block terminator");
@@ -218,16 +220,19 @@ static bool unifyLoopExits(DominatorTree &DT, LoopInfo &LI, Loop *L) {
   L->verifyLoop();
 
   // The guard blocks were created outside the loop, so they need to become
-  // 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.
+  // members of the parent loop.
+  // Same goes for the callbr target blocks.  Although we try to add them to the
+  // smallest common parent loop of the callbr block and the corresponding
+  // original target block, there might not have been such a loop, in which case
+  // the newly created callbr target blocks are not part of any loop. For nested
+  // loops, this might result in them leading to a loop with multiple entry
+  // points.
   if (auto *ParentLoop = L->getParentLoop()) {
     for (auto *G : GuardBlocks) {
       ParentLoop->addBasicBlockToLoop(G, LI);
     }
-    for (auto *C : CallBrTargetBlocks) {
-      if (!ParentLoop->contains(C))
-        ParentLoop->addBasicBlockToLoop(C, LI);
+    for (auto *C : CallBrTargetBlocksToFix) {
+      ParentLoop->addBasicBlockToLoop(C, LI);
     }
     ParentLoop->verifyLoop();
   }

llvm/test/Transforms/UnifyLoopExits/nested.ll

@@ -168,3 +168,55 @@ exit:
   %exit.phi = phi i32 [%A4.phi, %A5], [%Z, %C]
   ret void
 }
+
+; Here, the newly created target loop that connects b to r1 needs to be part of
+; the parent loop (the outer loop b participates in). Otherwise, it will be
+; regarded as an additional loop entry point to this outer loop.
+define void @nested_callbr_multiple_exits() {
+; CHECK-LABEL: @nested_callbr_multiple_exits(
+; CHECK-NEXT:    br label [[A:%.*]]
+; CHECK:       a:
+; CHECK-NEXT:    callbr void asm "", ""()
+; CHECK-NEXT:            to label [[B:%.*]] []
+; CHECK:       b:
+; CHECK-NEXT:    callbr void asm "", "!i"()
+; CHECK-NEXT:            to label [[C:%.*]] [label %b.target.b.target.r1]
+; CHECK:       c:
+; CHECK-NEXT:    callbr void asm "", "!i"()
+; CHECK-NEXT:            to label [[C_TARGET_E:%.*]] [label %b]
+; CHECK:       e:
+; CHECK-NEXT:    callbr void asm "", "!i"()
+; CHECK-NEXT:            to label [[A]] [label %e.target.r2]
+; CHECK:       r1:
+; CHECK-NEXT:    ret void
+; CHECK:       r2:
+; CHECK-NEXT:    ret void
+; CHECK:       b.target.r1:
+; CHECK-NEXT:    br label [[LOOP_EXIT_GUARD:%.*]]
+; CHECK:       e.target.r2:
+; CHECK-NEXT:    br label [[LOOP_EXIT_GUARD]]
+; CHECK:       loop.exit.guard:
+; CHECK-NEXT:    [[GUARD_R1:%.*]] = phi i1 [ true, [[B_TARGET_R1:%.*]] ], [ false, [[E_TARGET_R2:%.*]] ]
+; CHECK-NEXT:    br i1 [[GUARD_R1]], label [[R1:%.*]], label [[R2:%.*]]
+; CHECK:       b.target.b.target.r1:
+; CHECK-NEXT:    br label [[LOOP_EXIT_GUARD1:%.*]]
+; CHECK:       c.target.e:
+; CHECK-NEXT:    br label [[LOOP_EXIT_GUARD1]]
+; CHECK:       loop.exit.guard1:
+; CHECK-NEXT:    [[GUARD_B_TARGET_R1:%.*]] = phi i1 [ true, [[B_TARGET_B_TARGET_R1:%.*]] ], [ false, [[C_TARGET_E]] ]
+; CHECK-NEXT:    br i1 [[GUARD_B_TARGET_R1]], label [[B_TARGET_R1]], label [[E:%.*]]
+;
+  br label %a
+a:
+  callbr void asm "", ""() to label %b []
+b:
+  callbr void asm "", "!i"() to label %c [label %r1]
+c:
+  callbr void asm "", "!i"() to label %e [label %b]
+e:
+  callbr void asm "", "!i"() to label %a [label %r2]
+r1:
+  ret void
+r2:
+  ret void
+}