[llvm] [LoopInfo] Derive Loop::contains(BlockT*) from the block-to-loop map (PR #207613)

via llvm-commits llvm-commits at lists.llvm.org
Sun Jul 5 14:33:24 PDT 2026


llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-support

Author: Fangrui Song (MaskRay)

<details>
<summary>Changes</summary>

Each Loop stored its blocks twice: the Blocks vector and a DenseBlockSet used
only for O(1) contains(BlockT *). Block lists are inclusive, so building
LoopInfo inserted every block into that set once per enclosing loop -- an
O(blocks * depth) cost.

Drop DenseBlockSet and answer contains(BlockT *) from the innermost loop
in LoopInfoBase's block-number-indexed map. This removes the per-build
set maintenance and shrinks sizeof(Loop) from 160 to 72 bytes. Each Loop
keeps a back-pointer to its LoopInfo for the lookup; contains() stays
total for a block from another function (e.g. a global's use).

Note: A pass restructuring the loop nest must not rely on
`contains(BlockT *)` while the block lists and the map are transiently
out of sync, and should scan getBlocks() instead (as FixIrreducible and
LoopSimplifyCFG do). LoopInfoBase::verify() checks the block lists
rather than contains() for the same reason.

Aided by Claude Opus 4.8


---
Full diff: https://github.com/llvm/llvm-project/pull/207613.diff


5 Files Affected:

- (modified) llvm/include/llvm/Support/GenericLoopInfo.h (+46-27) 
- (modified) llvm/include/llvm/Support/GenericLoopInfoImpl.h (+10-11) 
- (modified) llvm/lib/Transforms/Scalar/LoopSimplifyCFG.cpp (+5-1) 
- (modified) llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp (-13) 
- (modified) llvm/lib/Transforms/Utils/FixIrreducible.cpp (+5-1) 


``````````diff
diff --git a/llvm/include/llvm/Support/GenericLoopInfo.h b/llvm/include/llvm/Support/GenericLoopInfo.h
index 9b3bcd09e289d..7b43eca8ba782 100644
--- a/llvm/include/llvm/Support/GenericLoopInfo.h
+++ b/llvm/include/llvm/Support/GenericLoopInfo.h
@@ -64,7 +64,10 @@ template <class BlockT, class LoopT> class LoopBase {
   // The list of blocks in this loop. First entry is the header node.
   std::vector<BlockT *> Blocks;
 
-  SmallPtrSet<const BlockT *, 8> DenseBlockSet;
+  // The LoopInfo that owns this loop. Used to answer contains(BlockT *) from
+  // the central block-number-indexed block-to-loop map rather than a per-loop
+  // block set.
+  LoopInfoBase<BlockT, LoopT> *LI = nullptr;
 
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   /// Indicator that this loop is no longer a valid loop.
@@ -130,10 +133,21 @@ template <class BlockT, class LoopT> class LoopBase {
     return contains(L->getParentLoop());
   }
 
-  /// Return true if the specified basic block is in this loop.
+  /// Return true if the specified basic block is in this loop, using LoopInfo's
+  /// block-to-loop map.
+  ///
+  /// This is only valid when that map agrees with the block lists. Avoid when
+  /// the loop nest is being restructured, when a block may appear in a loop's
+  /// block list before it is mapped to that loop. Code in such a transient
+  /// state must scan getBlocks() directly instead.
   bool contains(const BlockT *BB) const {
     assert(!isInvalid() && "Loop not in a valid state!");
-    return DenseBlockSet.count(BB);
+    // A block from another function is never contained, and its number would
+    // otherwise index this function's map.
+    if constexpr (GraphHasNodeNumbers<const BlockT *>)
+      if (BB->getParent() != LI->ParentPtr)
+        return false;
+    return contains(LI->lookupLoopFor(BB));
   }
 
   /// Return true if the specified instruction is in this loop.
@@ -195,19 +209,6 @@ template <class BlockT, class LoopT> class LoopBase {
     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!");
-    return DenseBlockSet;
-  }
-
   /// Return true if this loop is no longer valid.  The only valid use of this
   /// helper is "assert(L.isInvalid())" or equivalent, since IsInvalid is set to
   /// true by the destructor.  In other words, if this accessor returns true,
@@ -419,7 +420,6 @@ template <class BlockT, class LoopT> class LoopBase {
   void addBlockEntry(BlockT *BB) {
     assert(!isInvalid() && "Loop not in a valid state!");
     Blocks.push_back(BB);
-    DenseBlockSet.insert(BB);
   }
 
   /// interface to reverse Blocks[from, end of loop] in this loop
@@ -458,8 +458,6 @@ template <class BlockT, class LoopT> class LoopBase {
     auto I = find(Blocks, BB);
     assert(I != Blocks.end() && "N is not in this list!");
     Blocks.erase(I);
-
-    DenseBlockSet.erase(BB);
   }
 
   /// Verify loop structure
@@ -486,7 +484,6 @@ template <class BlockT, class LoopT> class LoopBase {
 
   explicit LoopBase(BlockT *BB) : ParentLoop(nullptr) {
     Blocks.push_back(BB);
-    DenseBlockSet.insert(BB);
   }
 
   // Since loop passes like SCEV are allowed to key analysis results off of
@@ -507,7 +504,6 @@ template <class BlockT, class LoopT> class LoopBase {
 #endif
     SubLoops.clear();
     Blocks.clear();
-    DenseBlockSet.clear();
     ParentLoop = nullptr;
   }
 };
@@ -552,6 +548,7 @@ template <class BlockT, class LoopT> class LoopInfoBase {
         LoopAllocator(std::move(Arg.LoopAllocator)) {
     ParentPtr = Arg.ParentPtr;
     BlockNumberEpoch = Arg.BlockNumberEpoch;
+    resetLoopInfoOwners();
     // We have to clear the arguments top level loops as we've taken ownership.
     Arg.TopLevelLoops.clear();
   }
@@ -565,6 +562,7 @@ template <class BlockT, class LoopT> class LoopInfoBase {
 
     TopLevelLoops = std::move(RHS.TopLevelLoops);
     LoopAllocator = std::move(RHS.LoopAllocator);
+    resetLoopInfoOwners();
     RHS.TopLevelLoops.clear();
     return *this;
   }
@@ -580,7 +578,9 @@ template <class BlockT, class LoopT> class LoopInfoBase {
 
   template <typename... ArgsTy> LoopT *AllocateLoop(ArgsTy &&...Args) {
     LoopT *Storage = LoopAllocator.Allocate<LoopT>();
-    return new (Storage) LoopT(std::forward<ArgsTy>(Args)...);
+    LoopT *L = new (Storage) LoopT(std::forward<ArgsTy>(Args)...);
+    L->LI = this;
+    return L;
   }
 
   /// iterator/begin/end - The interface to the top-level loops in the current
@@ -613,6 +613,18 @@ template <class BlockT, class LoopT> class LoopInfoBase {
   SmallVector<LoopT *, 4> getLoopsInReverseSiblingPreorder() const;
 
 private:
+  // Point every loop's owning-LoopInfo back-pointer at this object. Called
+  // after a move.
+  void resetLoopInfoOwners() {
+    SmallVector<LoopT *, 8> Worklist(TopLevelLoops.begin(),
+                                     TopLevelLoops.end());
+    while (!Worklist.empty()) {
+      LoopT *L = Worklist.pop_back_val();
+      L->LI = this;
+      Worklist.append(L->begin(), L->end());
+    }
+  }
+
   /// Verify that used block numbers are still valid.
   void verifyBlockNumberEpoch(ParentT BBParent) const {
     if constexpr (GraphHasNodeNumbers<BlockT *>) {
@@ -624,18 +636,25 @@ template <class BlockT, class LoopT> class LoopInfoBase {
     }
   }
 
-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.
-  LoopT *getLoopFor(const BlockT *BB) const {
+  // Look up BB's innermost loop in the block-to-loop map; BB must belong to
+  // this function.
+  LoopT *lookupLoopFor(const BlockT *BB) const {
     if constexpr (GraphHasNodeNumbers<const BlockT *>) {
-      verifyBlockNumberEpoch(BB->getParent());
       unsigned Number = GraphTraits<const BlockT *>::getNumber(BB);
       return Number < BBMap.size() ? BBMap[Number] : nullptr;
     } else
       return BBMap.lookup(BB);
   }
 
+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.
+  LoopT *getLoopFor(const BlockT *BB) const {
+    if constexpr (GraphHasNodeNumbers<const BlockT *>)
+      verifyBlockNumberEpoch(BB->getParent());
+    return lookupLoopFor(BB);
+  }
+
   /// Same as getLoopFor.
   const LoopT *operator[](const BlockT *BB) const { return getLoopFor(BB); }
 
diff --git a/llvm/include/llvm/Support/GenericLoopInfoImpl.h b/llvm/include/llvm/Support/GenericLoopInfoImpl.h
index 4666f4dac9cb6..0facd3341e0d2 100644
--- a/llvm/include/llvm/Support/GenericLoopInfoImpl.h
+++ b/llvm/include/llvm/Support/GenericLoopInfoImpl.h
@@ -17,7 +17,6 @@
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SetOperations.h"
 #include "llvm/Support/GenericLoopInfo.h"
 
 namespace llvm {
@@ -734,13 +733,6 @@ static void compareLoops(const LoopT *L, const LoopT *OtherL,
   std::vector<BlockT *> OtherBBs = OtherL->getBlocks();
   assert(compareVectors(BBs, OtherBBs) &&
          "Mismatched basic blocks in the loops!");
-
-  const SmallPtrSetImpl<const BlockT *> &BlocksSet = L->getBlocksSet();
-  const SmallPtrSetImpl<const BlockT *> &OtherBlocksSet =
-      OtherL->getBlocksSet();
-  assert(BlocksSet.size() == OtherBlocksSet.size() &&
-         llvm::set_is_subset(BlocksSet, OtherBlocksSet) &&
-         "Mismatched basic blocks in BlocksSets!");
 }
 #endif
 
@@ -755,6 +747,10 @@ void LoopInfoBase<BlockT, LoopT>::verify(
 
 // Verify that blocks are mapped to valid loops.
 #ifndef NDEBUG
+  // Every loop must point back at this LoopInfo (see resetLoopInfoOwners).
+  for (const LoopT *L : Loops)
+    assert(L->LI == this && "Loop has a stale owning-LoopInfo back-pointer");
+
   if constexpr (GraphHasNodeNumbers<const BlockT *>) {
     for (auto It : enumerate(BBMap)) {
       LoopT *L = It.value();
@@ -768,8 +764,11 @@ void LoopInfoBase<BlockT, LoopT>::verify(
       });
       BlockT *BB = BBIt != L->Blocks.end() ? *BBIt : nullptr;
       assert(BB && "orphaned block");
+      // Check the map against the (independent) block lists: L is its innermost
+      // loop (not in a deeper loop). Using contains() here would derive from
+      // BBMap itself and check nothing.
       for (LoopT *ChildLoop : *L)
-        assert(!ChildLoop->contains(BB) &&
+        assert(!llvm::is_contained(ChildLoop->getBlocks(), BB) &&
                "BBMap should point to the innermost loop containing BB");
     }
   } else {
@@ -777,9 +776,9 @@ void LoopInfoBase<BlockT, LoopT>::verify(
       const BlockT *BB = Entry.first;
       LoopT *L = Entry.second;
       assert(Loops.count(L) && "orphaned loop");
-      assert(L->contains(BB) && "orphaned block");
+      assert(llvm::is_contained(L->getBlocks(), BB) && "orphaned block");
       for (LoopT *ChildLoop : *L)
-        assert(!ChildLoop->contains(BB) &&
+        assert(!llvm::is_contained(ChildLoop->getBlocks(), BB) &&
                "BBMap should point to the innermost loop containing BB");
     }
   }
diff --git a/llvm/lib/Transforms/Scalar/LoopSimplifyCFG.cpp b/llvm/lib/Transforms/Scalar/LoopSimplifyCFG.cpp
index 328d842243f19..32620ba14e0e3 100644
--- a/llvm/lib/Transforms/Scalar/LoopSimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopSimplifyCFG.cpp
@@ -77,7 +77,11 @@ static void removeBlockFromLoops(BasicBlock *BB, Loop *FirstLoop,
                                  Loop *LastLoop = nullptr) {
   assert((!LastLoop || LastLoop->contains(FirstLoop->getHeader())) &&
          "First loop is supposed to be inside of last loop!");
-  assert(FirstLoop->contains(BB) && "Must be a loop block!");
+  // BB's innermost-loop mapping may already have been updated by the caller
+  // while it is still present in the block lists being cleaned up here, so
+  // check the block list directly rather than contains().
+  assert(llvm::is_contained(FirstLoop->getBlocks(), BB) &&
+         "Must be a loop block!");
   for (Loop *Current = FirstLoop; Current != LastLoop;
        Current = Current->getParentLoop())
     Current->removeBlockFromLoop(BB);
diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index abd5d30ab9c35..5e63338b3f63c 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -515,10 +515,6 @@ static void hoistLoopToNewParent(Loop &L, BasicBlock &Preheader,
                      return BB == &Preheader || L.contains(BB);
                    });
 
-    OldContainingL->getBlocksSet().erase(&Preheader);
-    for (BasicBlock *BB : L.blocks())
-      OldContainingL->getBlocksSet().erase(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
     // nodes for values used in the no-longer-nested loop.
@@ -1836,8 +1832,6 @@ static void deleteDeadBlocksFromLoop(Loop &L,
 
   // 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); });
   }
@@ -2036,9 +2030,6 @@ static bool rebuildLoopAfterUnswitch(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
     // 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) {
         return BB == PH || L.contains(BB);
       });
@@ -2067,8 +2058,6 @@ static bool rebuildLoopAfterUnswitch(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
     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
@@ -2083,8 +2072,6 @@ static bool rebuildLoopAfterUnswitch(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
 
   auto RemoveUnloopedBlocksFromLoop =
       [](Loop &L, SmallPtrSetImpl<BasicBlock *> &UnloopedBlocks) {
-        for (auto *BB : UnloopedBlocks)
-          L.getBlocksSet().erase(BB);
         llvm::erase_if(L.getBlocksVector(), [&](BasicBlock *BB) {
           return UnloopedBlocks.count(BB);
         });
diff --git a/llvm/lib/Transforms/Utils/FixIrreducible.cpp b/llvm/lib/Transforms/Utils/FixIrreducible.cpp
index b4f6262f7b309..df0fb72f3189a 100644
--- a/llvm/lib/Transforms/Utils/FixIrreducible.cpp
+++ b/llvm/lib/Transforms/Utils/FixIrreducible.cpp
@@ -183,8 +183,12 @@ static void reconnectChildLoops(LoopInfo &LI, Loop *ParentLoop, Loop *NewLoop,
                                     : LI.getTopLevelLoopsVector();
   // Any candidate is a child iff its header is owned by the new loop. Move all
   // the children to a new vector.
+  // The new loop's block list is already populated but its subloops are not yet
+  // attached, so query the block list directly rather than contains(), which
+  // reflects the not-yet-updated loop nesting.
   auto FirstChild = llvm::partition(CandidateLoops, [&](Loop *L) {
-    return NewLoop == L || !NewLoop->contains(L->getHeader());
+    return NewLoop == L ||
+           !llvm::is_contained(NewLoop->getBlocks(), L->getHeader());
   });
   SmallVector<Loop *, 8> ChildLoops(FirstChild, CandidateLoops.end());
   CandidateLoops.erase(FirstChild, CandidateLoops.end());

``````````

</details>


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


More information about the llvm-commits mailing list