[llvm] [polly] [LoopInfo] Add an opaque LoopRef handle and stop exposing loop containers (PR #210653)

via llvm-commits llvm-commits at lists.llvm.org
Sun Jul 19 23:42:06 PDT 2026


llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Fangrui Song (MaskRay)

<details>
<summary>Changes</summary>

Introduce llvm::LoopRef, the LoopInfo counterpart of CycleRef
(cc98c1eb84c6), wrapping the loop pointer rather than a preorder index
since passes hold loop identities across loop-forest mutation.

getBlocksVector, mutable getBlocksSet, getSubLoopsVector, and
getTopLevelLoopsVector hand out mutable references to the internal
containers, so every caller re-implements loop-nest surgery by hand.
Delete them in favor of two operations, making representation changes
(removing the block set, maintaining a loop numbering) local to
LoopInfoBase.

Also move hasNoExitBlocks, getExitEdges, and getUniqueLatchExitBlock
from LoopBase members to LoopRef queries.

Two behavior changes: SimpleLoopUnswitch's deleteDeadBlocksFromLoop now
runs the per-child deletion callbacks while the forest is still fully
consistent and detaches dead children before destroying them
(previously the callbacks could observe a destroyed sibling still
linked into the forest); FixIrreducible's reconnectChildLoops takes
children in their original relative order rather than std::partition's
arbitrary order.

Aided by Claude Fable 5


---

Patch is 27.74 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/210653.diff


13 Files Affected:

- (modified) bolt/lib/Core/BinaryFunction.cpp (+1-1) 
- (modified) llvm/include/llvm/Support/GenericLoopInfo.h (+83-30) 
- (modified) llvm/include/llvm/Support/GenericLoopInfoImpl.h (+19-18) 
- (modified) llvm/lib/Target/AMDGPU/AMDGPUNextUseAnalysis.cpp (+1-1) 
- (modified) llvm/lib/Transforms/Scalar/LoopDeletion.cpp (+1-1) 
- (modified) llvm/lib/Transforms/Scalar/LoopPredication.cpp (+1-1) 
- (modified) llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp (+38-65) 
- (modified) llvm/lib/Transforms/Utils/FixIrreducible.cpp (+8-15) 
- (modified) llvm/lib/Transforms/Utils/LoopPeel.cpp (+1-1) 
- (modified) llvm/lib/Transforms/Utils/LoopUtils.cpp (+1-1) 
- (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+2-1) 
- (modified) llvm/lib/Transforms/Vectorize/VPlan.cpp (+2-1) 
- (modified) polly/lib/Analysis/ScopDetection.cpp (+1-1) 


``````````diff
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index f580edf2626d2..e0a76c1462654 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -4501,7 +4501,7 @@ void BinaryFunction::calculateLoopInfo() {
 
     // Compute exit count.
     SmallVector<BinaryLoop::Edge, 1> ExitEdges;
-    L->getExitEdges(ExitEdges);
+    BLI->getExitEdges(BLI->ref(L), ExitEdges);
     for (BinaryLoop::Edge &Exit : ExitEdges) {
       const BinaryBasicBlock *Exiting = Exit.first;
       const BinaryBasicBlock *ExitTarget = Exit.second;
diff --git a/llvm/include/llvm/Support/GenericLoopInfo.h b/llvm/include/llvm/Support/GenericLoopInfo.h
index 77f9cd731f71f..e83d1fdccfbe0 100644
--- a/llvm/include/llvm/Support/GenericLoopInfo.h
+++ b/llvm/include/llvm/Support/GenericLoopInfo.h
@@ -51,6 +51,24 @@ namespace llvm {
 
 template <class N, class M> class LoopInfoBase;
 template <class N, class M> class LoopBase;
+template <class N, class M> class PopulateLoopsDFS;
+
+/// Opaque handle to a loop within a LoopInfoBase. A handle stays valid for as
+/// long as the loop itself. The long-term goal is to move LoopBase operations
+/// to LoopInfoBase.
+class LoopRef {
+  void *Ptr = nullptr;
+
+  explicit LoopRef(void *Ptr) : Ptr(Ptr) {}
+  template <class N, class M> friend class LoopInfoBase;
+
+public:
+  LoopRef() = default;
+  bool isValid() const { return Ptr != nullptr; }
+  explicit operator bool() const { return isValid(); }
+  bool operator==(LoopRef O) const { return Ptr == O.Ptr; }
+  bool operator!=(LoopRef O) const { return Ptr != O.Ptr; }
+};
 
 //===----------------------------------------------------------------------===//
 /// Instances of this class are used to represent loops that are detected in the
@@ -146,10 +164,6 @@ template <class BlockT, class LoopT> class LoopBase {
     assert(!isInvalid() && "Loop not in a valid state!");
     return SubLoops;
   }
-  std::vector<LoopT *> &getSubLoopsVector() {
-    assert(!isInvalid() && "Loop not in a valid state!");
-    return SubLoops;
-  }
   using iterator = typename std::vector<LoopT *>::const_iterator;
   using reverse_iterator =
       typename std::vector<LoopT *>::const_reverse_iterator;
@@ -189,19 +203,6 @@ template <class BlockT, class LoopT> class LoopBase {
     return Blocks.size();
   }
 
-  /// Return a direct, mutable handle to the blocks vector so that we can
-  /// mutate it efficiently with techniques like `std::remove`.
-  std::vector<BlockT *> &getBlocksVector() {
-    assert(!isInvalid() && "Loop not in a valid state!");
-    return Blocks;
-  }
-  /// Return a direct, mutable handle to the blocks set so that we can
-  /// mutate it efficiently.
-  SmallPtrSetImpl<const BlockT *> &getBlocksSet() {
-    assert(!isInvalid() && "Loop not in a valid state!");
-    return DenseBlockSet;
-  }
-
   /// Return a direct, immutable handle to the blocks set.
   const SmallPtrSetImpl<const BlockT *> &getBlocksSet() const {
     assert(!isInvalid() && "Loop not in a valid state!");
@@ -294,19 +295,9 @@ template <class BlockT, class LoopT> class LoopBase {
   /// Otherwise return null.
   BlockT *getUniqueExitBlock() const;
 
-  /// Return the unique exit block for the latch, or null if there are multiple
-  /// different exit blocks or the latch is not exiting.
-  BlockT *getUniqueLatchExitBlock() const;
-
-  /// Return true if this loop does not have any exit blocks.
-  bool hasNoExitBlocks() const;
-
   /// Edge type.
   using Edge = std::pair<BlockT *, BlockT *>;
 
-  /// Return all pairs of (_inside_block_,_outside_block_).
-  void getExitEdges(SmallVectorImpl<Edge> &ExitEdges) const;
-
   /// If there is a preheader for this loop, return it. A loop has a preheader
   /// if there is only one edge to the header of the loop from outside of the
   /// loop. If this is the case, the block branching to the header of the loop
@@ -434,6 +425,16 @@ template <class BlockT, class LoopT> class LoopBase {
     Blocks.reserve(size);
   }
 
+  /// interface to do reserve() for SubLoops
+  void reserveSubLoops(unsigned Size) {
+    assert(!isInvalid() && "Loop not in a valid state!");
+    SubLoops.reserve(Size);
+  }
+
+  /// Capacity of the block list; input to an enclosing loop's reserveBlocks()
+  /// during construction, when this loop's list is not yet fully populated.
+  unsigned getBlocksCapacity() const { return Blocks.capacity(); }
+
   /// This method is used to move BB (which must be part of this loop) to be the
   /// loop header of the loop (the block that dominates all others).
   void moveToHeader(BlockT *BB) {
@@ -480,6 +481,7 @@ template <class BlockT, class LoopT> class LoopBase {
 
 protected:
   friend class LoopInfoBase<BlockT, LoopT>;
+  friend class PopulateLoopsDFS<BlockT, LoopT>;
 
   /// This creates an empty loop.
   LoopBase() : ParentLoop(nullptr) {}
@@ -625,6 +627,12 @@ template <class BlockT, class LoopT> class LoopInfoBase {
            "loop info used with outdated block numbers");
   }
 
+  /// Resolve a handle to its loop.
+  LoopT *deref(LoopRef L) const {
+    assert(L.isValid() && "dereferencing an invalid loop handle");
+    return static_cast<LoopT *>(L.Ptr);
+  }
+
 public:
   /// Return the inner most loop that BB lives in. If a basic block is in no
   /// loop (for example the entry node), null is returned.
@@ -644,6 +652,54 @@ template <class BlockT, class LoopT> class LoopInfoBase {
     return L ? L->getLoopDepth() : 0;
   }
 
+  /// The handle for a loop of this LoopInfo; an invalid handle for null.
+  LoopRef ref(const LoopT *L) const { return LoopRef(const_cast<LoopT *>(L)); }
+
+  using Edge = std::pair<BlockT *, BlockT *>;
+
+  /// Return true if \p L does not have any exit blocks.
+  bool hasNoExitBlocks(LoopRef L) const;
+
+  /// Return all pairs of (_inside_block_,_outside_block_).
+  void getExitEdges(LoopRef L, SmallVectorImpl<Edge> &ExitEdges) const;
+
+  /// Return the unique exit block for the latch of \p L, or null if there are
+  /// multiple different exit blocks or the latch is not exiting.
+  BlockT *getUniqueLatchExitBlock(LoopRef L) const;
+
+  /// Remove every block satisfying \p Pred from \p L's block list, preserving
+  /// the order of the remaining blocks. Only \p L itself is updated, not its
+  /// ancestors or descendants, and not the block-to-loop mapping.
+  template <typename PredicateT>
+  void removeBlocksIf(LoopRef L, PredicateT Pred) {
+    LoopT *Lp = deref(L);
+    llvm::erase_if(Lp->Blocks, [&](BlockT *BB) {
+      if (!Pred(BB))
+        return false;
+      Lp->DenseBlockSet.erase(BB);
+      return true;
+    });
+  }
+
+  /// Detach and return the children of \p Parent (the top-level loops if
+  /// \p Parent is invalid) that satisfy \p Pred, clearing their parent
+  /// pointers. Both the remaining and the returned children keep their
+  /// relative order.
+  template <typename PredicateT>
+  SmallVector<LoopT *, 4> takeChildrenIf(LoopRef Parent, PredicateT Pred) {
+    std::vector<LoopT *> &List =
+        Parent ? deref(Parent)->SubLoops : TopLevelLoops;
+    SmallVector<LoopT *, 4> Taken;
+    llvm::erase_if(List, [&](LoopT *Child) {
+      if (!Pred(Child))
+        return false;
+      Child->ParentLoop = nullptr;
+      Taken.push_back(Child);
+      return true;
+    });
+    return Taken;
+  }
+
   /// \brief Find the innermost loop containing both given loops.
   ///
   /// \returns the innermost loop containing both \p A and \p B
@@ -664,9 +720,6 @@ template <class BlockT, class LoopT> class LoopInfoBase {
   /// Return the top-level loops.
   const std::vector<LoopT *> &getTopLevelLoops() const { return TopLevelLoops; }
 
-  /// Return the top-level loops.
-  std::vector<LoopT *> &getTopLevelLoopsVector() { return TopLevelLoops; }
-
   /// This removes the specified top-level loop from this loop info object.
   /// The loop is not deleted, as it will presumably be inserted into
   /// another loop.
diff --git a/llvm/include/llvm/Support/GenericLoopInfoImpl.h b/llvm/include/llvm/Support/GenericLoopInfoImpl.h
index e0a06ad75842c..434dc6feaa545 100644
--- a/llvm/include/llvm/Support/GenericLoopInfoImpl.h
+++ b/llvm/include/llvm/Support/GenericLoopInfoImpl.h
@@ -92,8 +92,8 @@ std::pair<BlockT *, bool> getExitBlockHelper(const LoopBase<BlockT, LoopT> *L,
 }
 
 template <class BlockT, class LoopT>
-bool LoopBase<BlockT, LoopT>::hasNoExitBlocks() const {
-  auto RC = getExitBlockHelper(this, false);
+bool LoopInfoBase<BlockT, LoopT>::hasNoExitBlocks(LoopRef L) const {
+  auto RC = getExitBlockHelper(deref(L), false);
   if (RC.second)
     // found multiple exit blocks
     return false;
@@ -160,24 +160,26 @@ BlockT *LoopBase<BlockT, LoopT>::getUniqueExitBlock() const {
 }
 
 template <class BlockT, class LoopT>
-BlockT *LoopBase<BlockT, LoopT>::getUniqueLatchExitBlock() const {
-  BlockT *Latch = getLoopLatch();
+BlockT *LoopInfoBase<BlockT, LoopT>::getUniqueLatchExitBlock(LoopRef L) const {
+  const LoopT *Lp = deref(L);
+  BlockT *Latch = Lp->getLoopLatch();
   assert(Latch && "Latch block must exists");
-  auto IsExitBlock = [this](BlockT *BB, bool AllowRepeats) -> BlockT * {
+  auto IsExitBlock = [Lp](BlockT *BB, bool AllowRepeats) -> BlockT * {
     assert(!AllowRepeats && "Unexpected parameter value.");
-    return !contains(BB) ? BB : nullptr;
+    return !Lp->contains(BB) ? BB : nullptr;
   };
-  return find_singleton<BlockT>(children<BlockT *>(Latch), IsExitBlock);
+  return find_singleton<BlockT>(llvm::children<BlockT *>(Latch), IsExitBlock);
 }
 
 /// getExitEdges - Return all pairs of (_inside_block_,_outside_block_).
 template <class BlockT, class LoopT>
-void LoopBase<BlockT, LoopT>::getExitEdges(
-    SmallVectorImpl<Edge> &ExitEdges) const {
-  assert(!isInvalid() && "Loop not in a valid state!");
-  for (const auto BB : blocks())
-    for (auto *Succ : children<BlockT *>(BB))
-      if (!contains(Succ))
+void LoopInfoBase<BlockT, LoopT>::getExitEdges(
+    LoopRef L, SmallVectorImpl<Edge> &ExitEdges) const {
+  const LoopT *Lp = deref(L);
+  assert(!Lp->isInvalid() && "Loop not in a valid state!");
+  for (const auto BB : Lp->blocks())
+    for (auto *Succ : llvm::children<BlockT *>(BB))
+      if (!Lp->contains(Succ))
         // Not in current loop? It must be an exit block.
         ExitEdges.emplace_back(BB, Succ);
 }
@@ -495,7 +497,7 @@ static void discoverAndMapSubloop(LoopT *L, ArrayRef<BlockT *> Backedges,
       // Discover a subloop of this loop.
       Subloop->setParentLoop(L);
       ++NumSubloops;
-      NumBlocks += Subloop->getBlocksVector().capacity();
+      NumBlocks += Subloop->getBlocksCapacity();
       PredBB = Subloop->getHeader();
       // Continue traversal along predecessors that are not loop-back edges from
       // within this subloop tree itself. Note that a predecessor may directly
@@ -507,7 +509,7 @@ static void discoverAndMapSubloop(LoopT *L, ArrayRef<BlockT *> Backedges,
       }
     }
   }
-  L->getSubLoopsVector().reserve(NumSubloops);
+  L->reserveSubLoops(NumSubloops);
   L->reserveBlocks(NumBlocks);
 }
 
@@ -544,15 +546,14 @@ void PopulateLoopsDFS<BlockT, LoopT>::insertIntoLoop(BlockT *Block) {
     // We reach this point once per subloop after processing all the blocks in
     // the subloop.
     if (!Subloop->isOutermost())
-      Subloop->getParentLoop()->getSubLoopsVector().push_back(Subloop);
+      Subloop->getParentLoop()->SubLoops.push_back(Subloop);
     else
       LI->addTopLevelLoop(Subloop);
 
     // For convenience, Blocks and Subloops are inserted in postorder. Reverse
     // the lists, except for the loop header, which is always at the beginning.
     Subloop->reverseBlock(1);
-    std::reverse(Subloop->getSubLoopsVector().begin(),
-                 Subloop->getSubLoopsVector().end());
+    std::reverse(Subloop->SubLoops.begin(), Subloop->SubLoops.end());
 
     Subloop = Subloop->getParentLoop();
   }
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUNextUseAnalysis.cpp b/llvm/lib/Target/AMDGPU/AMDGPUNextUseAnalysis.cpp
index cad1513f280fb..16b928594b7eb 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUNextUseAnalysis.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUNextUseAnalysis.cpp
@@ -1122,7 +1122,7 @@ class llvm::AMDGPUNextUseAnalysisImpl {
   MBBDistPair calcShortestDistanceToExit(const MachineBasicBlock *CurMBB,
                                          const MachineLoop *CurLoop) const {
     SmallVector<std::pair<MachineBasicBlock *, MachineBasicBlock *>> ExitEdges;
-    CurLoop->getExitEdges(ExitEdges);
+    MLI->getExitEdges(MLI->ref(CurLoop), ExitEdges);
     MBBDistPair LD;
 
     for (auto [Exit, Dest] : ExitEdges) {
diff --git a/llvm/lib/Transforms/Scalar/LoopDeletion.cpp b/llvm/lib/Transforms/Scalar/LoopDeletion.cpp
index 0664eed072a9a..4ce52179992dd 100644
--- a/llvm/lib/Transforms/Scalar/LoopDeletion.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopDeletion.cpp
@@ -488,7 +488,7 @@ static LoopDeletionResult deleteLoopIfDead(Loop *L, DominatorTree &DT,
   // the situation of needing to be able to solve statically which exit block
   // will be branched to, or trying to preserve the branching logic in a loop
   // invariant manner.
-  if (!ExitBlock && !L->hasNoExitBlocks()) {
+  if (!ExitBlock && !LI.hasNoExitBlocks(LI.ref(L))) {
     LLVM_DEBUG(dbgs() << "Deletion requires at most one exit block.\n");
     return LoopDeletionResult::Unmodified;
   }
diff --git a/llvm/lib/Transforms/Scalar/LoopPredication.cpp b/llvm/lib/Transforms/Scalar/LoopPredication.cpp
index de5365271e233..2487a5786c8e1 100644
--- a/llvm/lib/Transforms/Scalar/LoopPredication.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopPredication.cpp
@@ -877,7 +877,7 @@ bool LoopPredication::isLoopProfitableToPredicate() {
     return true;
 
   SmallVector<std::pair<BasicBlock *, BasicBlock *>, 8> ExitEdges;
-  L->getExitEdges(ExitEdges);
+  LI->getExitEdges(LI->ref(L), ExitEdges);
   // If there is only one exiting edge in the loop, it is always profitable to
   // predicate the loop.
   if (ExitEdges.size() == 1)
diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index 89c0288140237..60a2151e26c94 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -509,14 +509,9 @@ static void hoistLoopToNewParent(Loop &L, BasicBlock &Preheader,
   // no-longer-containing loops to reflect the nesting change.
   for (Loop *OldContainingL = OldParentL; OldContainingL != NewParentL;
        OldContainingL = OldContainingL->getParentLoop()) {
-    llvm::erase_if(OldContainingL->getBlocksVector(),
-                   [&](const BasicBlock *BB) {
-                     return BB == &Preheader || L.contains(BB);
-                   });
-
-    OldContainingL->getBlocksSet().erase(&Preheader);
-    for (BasicBlock *BB : L.blocks())
-      OldContainingL->getBlocksSet().erase(BB);
+    LI.removeBlocksIf(LI.ref(OldContainingL), [&](const BasicBlock *BB) {
+      return BB == &Preheader || L.contains(BB);
+    });
 
     // Because we just hoisted a loop out of this one, we have essentially
     // created new exit paths from it. That means we need to form LCSSA PHI
@@ -593,7 +588,7 @@ static bool unswitchTrivialBranch(Loop &L, CondBrInst &BI, DominatorTree &DT,
 
   std::optional<int> LatchIdx = std::nullopt;
   auto *LoopLatch = L.getLoopLatch();
-  auto *ULExit = L.getUniqueLatchExitBlock();
+  auto *ULExit = LI.getUniqueLatchExitBlock(LI.ref(&L));
   if (SE && FullUnswitch && ULExit) {
     if (BI.getSuccessor(0) == LoopLatch && L.contains(BI.getSuccessor(1)))
       LatchIdx = 0;
@@ -1845,18 +1840,17 @@ static void deleteDeadBlocksFromLoop(Loop &L,
                  [&](BasicBlock *BB) { return DeadBlockSet.count(BB); });
 
   // Walk from this loop up through its parents removing all of the dead blocks.
-  for (Loop *ParentL = &L; ParentL; ParentL = ParentL->getParentLoop()) {
-    for (auto *BB : DeadBlockSet)
-      ParentL->getBlocksSet().erase(BB);
-    llvm::erase_if(ParentL->getBlocksVector(),
-                   [&](BasicBlock *BB) { return DeadBlockSet.count(BB); });
-  }
-
-  // Now delete the dead child loops. This raw delete will clear them
-  // recursively.
-  llvm::erase_if(L.getSubLoopsVector(), [&](Loop *ChildL) {
+  for (Loop *ParentL = &L; ParentL; ParentL = ParentL->getParentLoop())
+    LI.removeBlocksIf(LI.ref(ParentL),
+                      [&](BasicBlock *BB) { return DeadBlockSet.count(BB); });
+
+  // Now delete the dead child loops. Run the per-child deletion callbacks
+  // first, while the loop forest is still fully consistent (markLoopAsDeleted
+  // checks the child's position in it), then detach the dead children and
+  // destroy them.
+  for (Loop *ChildL : L) {
     if (!DeadBlockSet.count(ChildL->getHeader()))
-      return false;
+      continue;
 
     assert(llvm::all_of(ChildL->blocks(),
                         [&](BasicBlock *ChildBB) {
@@ -1867,9 +1861,11 @@ static void deleteDeadBlocksFromLoop(Loop &L,
     LoopUpdater.markLoopAsDeleted(*ChildL, ChildL->getName());
     if (SE)
       SE->forgetBlockAndLoopDispositions();
+  }
+  for (Loop *ChildL : LI.takeChildrenIf(LI.ref(&L), [&](Loop *ChildL) {
+         return DeadBlockSet.count(ChildL->getHeader());
+       }))
     LI.destroy(ChildL);
-    return true;
-  });
 
   // Remove the loop mappings for the dead blocks and drop all the references
   // from these blocks to others to handle cyclic references as we start
@@ -2044,15 +2040,10 @@ static bool rebuildLoopAfterUnswitch(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
   // *up* the nest.
   if (!LoopBlockSet.empty() && L.getParentLoop() != ParentL) {
     // Remove this loop's (original) blocks from all of the intervening loops.
-    for (Loop *IL = L.getParentLoop(); IL != ParentL;
-         IL = IL->getParentLoop()) {
-      IL->getBlocksSet().erase(PH);
-      for (auto *BB : L.blocks())
-        IL->getBlocksSet().erase(BB);
-      llvm::erase_if(IL->getBlocksVector(), [&](BasicBlock *BB) {
+    for (Loop *IL = L.getParentLoop(); IL != ParentL; IL = IL->getParentLoop())
+      LI.removeBlocksIf(LI.ref(IL), [&](BasicBlock *BB) {
         return BB == PH || L.contains(BB);
       });
-    }
 
     LI.changeLoopFor(PH, ParentL);
     L.getParentLoop()->removeChildLoop(&L);
@@ -2062,25 +2053,18 @@ static bool rebuildLoopAfterUnswitch(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
       LI.addTopLevelLoop(&L);
   }
 
-  // Now we update all the blocks which are no longer within the loop.
-  auto &Blocks = L.getBlocksVector();
-  auto BlocksSplitI =
-      LoopBlockSet.empty()
-          ? Blocks.begin()
-          : std::stable_partition(
-                Blocks.begin(), Blocks.end(),
-                [&](BasicBlock *BB) { return LoopBlockSet.count(BB); });
-
-  // Before we erase the list of unlooped blocks, build a set of them.
-  SmallPtrSet<BasicBlock *, 16> UnloopedBlocks(BlocksSplitI, Blocks.end());
+  // Now we update all the blocks which are no longer within the loop, building
+  // the set of them as they are removed.
+  SmallPtrSet<BasicBlock *, 16> UnloopedBlocks;
+  LI.removeBlocksIf(LI.ref(&L), [&](BasicBlock *BB) {
+    if (LoopBlockSet.count(BB))
+      return false;
+    UnloopedBlocks.insert(BB);
+    return true;
+  });
   if (LoopBlockSet.empty())
     UnloopedBlocks.insert(PH);
 
-  // Now erase these blocks from the loop.
-  for (auto *BB : make_range(BlocksSplitI, Blocks.end()))
-    L.getBlocksSet().erase(BB);
-  Blocks.erase(BlocksSplitI, Blocks.end());
-
   // Sort the exits in ascending loop depth, we'll work backwards across these
   // to process them inside out.
   llvm::stable_sort(ExitsInLoops, [&](BasicBlock *LHS, BasicBlock *RHS) {
@@ -2092,10 +2076,8 @@ static bool rebuildLoopAfterUnswitch(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
   Loop *PrevExitL = L.getParentLoop(); // The deepest possible exit loop.
 
   auto RemoveUnloopedBlocksFromLoop =
-      [](Loop &L, SmallPtrSetImpl<BasicBlock *> &UnloopedBlocks) {
-        for (auto *BB : UnloopedBlocks)
-          L.getBlocksSet().erase(BB);
-        llvm::erase_if(L.getBlocksVector(), [&](BasicBlock *BB) {
+      [&LI](Loop &L, SmallPtrSetImpl<BasicBlock *> &UnloopedBlocks) {
+        LI.removeBlocksIf(LI.ref(&L), [&](BasicBlock *BB) {
           return UnloopedBlocks.count(BB);
         });
       };
@@ -2171,19 +2153,11 @@ static bool rebuildLoopAfterUnswitch(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
         LI.changeLoopFor(BB, nullptr);
...
[truncated]

``````````

</details>


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


More information about the llvm-commits mailing list