[Support] Use block numbers for LoopInfo BBMap (#103400)

Replace the DenseMap from blocks to their innermost loop a vector
indexed by block numbers, when possible. Supporting number updates is
not trivial as we don't store a list of basic blocks, so this is not
implemented.

NB: I'm generally not happy with the way loops are stored. As I think
that there's room for improvement, I don't want to touch the
representation at this point.

Pull Request: https://github.com/llvm/llvm-project/pull/103400

llvm/include/llvm/Support/GenericLoopInfo.h

@@ -525,7 +525,14 @@ raw_ostream &operator<<(raw_ostream &OS, const LoopBase<BlockT, LoopT> &Loop) {
 
 template <class BlockT, class LoopT> class LoopInfoBase {
   // BBMap - Mapping of basic blocks to the inner most loop they occur in
-  DenseMap<const BlockT *, LoopT *> BBMap;
+  std::conditional_t<GraphHasNodeNumbers<const BlockT *>, SmallVector<LoopT *>,
+                     DenseMap<const BlockT *, LoopT *>>
+      BBMap;
+
+  using ParentT = decltype(std::declval<const BlockT *>()->getParent());
+  ParentT ParentPtr = nullptr;
+  unsigned BlockNumberEpoch;
+
   std::vector<LoopT *> TopLevelLoops;
   BumpPtrAllocator LoopAllocator;
 
@@ -543,11 +550,15 @@ public:
       : BBMap(std::move(Arg.BBMap)),
         TopLevelLoops(std::move(Arg.TopLevelLoops)),
         LoopAllocator(std::move(Arg.LoopAllocator)) {
+    ParentPtr = Arg.ParentPtr;
+    BlockNumberEpoch = Arg.BlockNumberEpoch;
     // We have to clear the arguments top level loops as we've taken ownership.
     Arg.TopLevelLoops.clear();
   }
   LoopInfoBase &operator=(LoopInfoBase &&RHS) {
     BBMap = std::move(RHS.BBMap);
+    ParentPtr = RHS.ParentPtr;
+    BlockNumberEpoch = RHS.BlockNumberEpoch;
 
     for (auto *L : TopLevelLoops)
       L->~LoopT();
@@ -601,9 +612,29 @@ public:
   /// reverse program order.
   SmallVector<LoopT *, 4> getLoopsInReverseSiblingPreorder() const;
 
+private:
+  /// Verify that used block numbers are still valid.
+  void verifyBlockNumberEpoch(ParentT BBParent) const {
+    if constexpr (GraphHasNodeNumbers<BlockT *>) {
+      assert(ParentPtr == BBParent &&
+             "loop info queried with block of other function");
+      assert(BlockNumberEpoch ==
+                 GraphTraits<ParentT>::getNumberEpoch(ParentPtr) &&
+             "loop info used with outdated block numbers");
+    }
+  }
+
+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 { return BBMap.lookup(BB); }
+  LoopT *getLoopFor(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);
+  }
 
   /// Same as getLoopFor.
   const LoopT *operator[](const BlockT *BB) const { return getLoopFor(BB); }
@@ -652,12 +683,23 @@ public:
   /// Change the top-level loop that contains BB to the specified loop.
   /// This should be used by transformations that restructure the loop hierarchy
   /// tree.
-  void changeLoopFor(BlockT *BB, LoopT *L) {
-    if (!L) {
-      BBMap.erase(BB);
-      return;
+  void changeLoopFor(const BlockT *BB, LoopT *L) {
+    if constexpr (GraphHasNodeNumbers<const BlockT *>) {
+      verifyBlockNumberEpoch(BB->getParent());
+      unsigned Number = GraphTraits<const BlockT *>::getNumber(BB);
+      if (Number >= BBMap.size()) {
+        unsigned Max = GraphTraits<decltype(BB->getParent())>::getMaxNumber(
+            BB->getParent());
+        BBMap.resize(Number >= Max ? Number + 1 : Max);
+      }
+      BBMap[Number] = L;
+    } else {
+      if (!L) {
+        BBMap.erase(BB);
+        return;
+      }
+      BBMap[BB] = L;
     }
-    BBMap[BB] = L;
   }
 
   /// Replace the specified loop in the top-level loops list with the indicated
@@ -680,12 +722,23 @@ public:
   /// including all of the Loop objects it is nested in and our mapping from
   /// BasicBlocks to loops.
   void removeBlock(BlockT *BB) {
-    auto I = BBMap.find(BB);
-    if (I != BBMap.end()) {
-      for (LoopT *L = I->second; L; L = L->getParentLoop())
-        L->removeBlockFromLoop(BB);
+    if constexpr (GraphHasNodeNumbers<BlockT *>) {
+      verifyBlockNumberEpoch(BB->getParent());
+      unsigned Number = GraphTraits<BlockT *>::getNumber(BB);
+      if (Number >= BBMap.size())
+        return;
 
-      BBMap.erase(I);
+      for (LoopT *L = BBMap[Number]; L; L = L->getParentLoop())
+        L->removeBlockFromLoop(BB);
+      BBMap[Number] = nullptr;
+    } else {
+      auto I = BBMap.find(BB);
+      if (I != BBMap.end()) {
+        for (LoopT *L = I->second; L; L = L->getParentLoop())
+          L->removeBlockFromLoop(BB);
+
+        BBMap.erase(I);
+      }
     }
   }
 

llvm/include/llvm/Support/GenericLoopInfoImpl.h

@@ -295,7 +295,7 @@ void LoopBase<BlockT, LoopT>::addBasicBlockToLoop(
   LoopT *L = static_cast<LoopT *>(this);
 
   // Add the loop mapping to the LoopInfo object...
-  LIB.BBMap[NewBB] = L;
+  LIB.changeLoopFor(NewBB, L);
 
   // Add the basic block to this loop and all parent loops...
   while (L) {
@@ -578,6 +578,12 @@ template <class BlockT, class LoopT>
 void LoopInfoBase<BlockT, LoopT>::analyze(const DomTreeBase<BlockT> &DomTree) {
   // Postorder traversal of the dominator tree.
   const DomTreeNodeBase<BlockT> *DomRoot = DomTree.getRootNode();
+  if constexpr (GraphHasNodeNumbers<const BlockT *>) {
+    ParentPtr = DomRoot->getBlock()->getParent();
+    BlockNumberEpoch = GraphTraits<ParentT>::getNumberEpoch(ParentPtr);
+    unsigned Max = GraphTraits<ParentT>::getMaxNumber(ParentPtr);
+    BBMap.resize(Max);
+  }
   for (auto DomNode : post_order(DomRoot)) {
 
     BlockT *Header = DomNode->getBlock();
@@ -756,14 +762,33 @@ void LoopInfoBase<BlockT, LoopT>::verify(
 
 // Verify that blocks are mapped to valid loops.
 #ifndef NDEBUG
-  for (auto &Entry : BBMap) {
-    const BlockT *BB = Entry.first;
-    LoopT *L = Entry.second;
-    assert(Loops.count(L) && "orphaned loop");
-    assert(L->contains(BB) && "orphaned block");
-    for (LoopT *ChildLoop : *L)
-      assert(!ChildLoop->contains(BB) &&
-             "BBMap should point to the innermost loop containing BB");
+  if constexpr (GraphHasNodeNumbers<const BlockT *>) {
+    for (auto It : enumerate(BBMap)) {
+      LoopT *L = It.value();
+      unsigned Number = It.index();
+      if (!L)
+        continue;
+      assert(Loops.count(L) && "orphaned loop");
+      // We have no way to map block numbers back to blocks, so find it.
+      auto BBIt = find_if(L->Blocks, [&Number](BlockT *BB) {
+        return GraphTraits<BlockT *>::getNumber(BB) == Number;
+      });
+      BlockT *BB = BBIt != L->Blocks.end() ? *BBIt : nullptr;
+      assert(BB && "orphaned block");
+      for (LoopT *ChildLoop : *L)
+        assert(!ChildLoop->contains(BB) &&
+               "BBMap should point to the innermost loop containing BB");
+    }
+  } else {
+    for (auto &Entry : BBMap) {
+      const BlockT *BB = Entry.first;
+      LoopT *L = Entry.second;
+      assert(Loops.count(L) && "orphaned loop");
+      assert(L->contains(BB) && "orphaned block");
+      for (LoopT *ChildLoop : *L)
+        assert(!ChildLoop->contains(BB) &&
+               "BBMap should point to the innermost loop containing BB");
+    }
   }
 
   // Recompute LoopInfo to verify loops structure.

llvm/lib/CodeGen/BranchFolding.cpp

@@ -199,10 +199,10 @@ void BranchFolder::RemoveDeadBlock(MachineBasicBlock *MBB) {
       MF->eraseAdditionalCallInfo(&MI);
 
   // Remove the block.
-  MF->erase(MBB);
-  EHScopeMembership.erase(MBB);
   if (MLI)
     MLI->removeBlock(MBB);
+  MF->erase(MBB);
+  EHScopeMembership.erase(MBB);
 }
 
 bool BranchFolder::OptimizeFunction(MachineFunction &MF,

llvm/lib/CodeGen/EarlyIfConversion.cpp

@@ -1267,9 +1267,9 @@ bool EarlyIfConverter::tryConvertIf(MachineBasicBlock *MBB) {
     IfConv.convertIf(RemoveBlocks);
     Changed = true;
     updateDomTree(DomTree, IfConv, RemoveBlocks);
+    updateLoops(Loops, RemoveBlocks);
     for (MachineBasicBlock *MBB : RemoveBlocks)
       MBB->eraseFromParent();
-    updateLoops(Loops, RemoveBlocks);
   }
   return Changed;
 }
@@ -1440,9 +1440,9 @@ bool EarlyIfPredicator::tryConvertIf(MachineBasicBlock *MBB) {
     IfConv.convertIf(RemoveBlocks, /*Predicate*/ true);
     Changed = true;
     updateDomTree(DomTree, IfConv, RemoveBlocks);
+    updateLoops(Loops, RemoveBlocks);
     for (MachineBasicBlock *MBB : RemoveBlocks)
       MBB->eraseFromParent();
-    updateLoops(Loops, RemoveBlocks);
   }
   return Changed;
 }

llvm/lib/Target/AArch64/AArch64ConditionalCompares.cpp

@@ -924,9 +924,9 @@ bool AArch64ConditionalCompares::tryConvert(MachineBasicBlock *MBB) {
     CmpConv.convert(RemovedBlocks);
     Changed = true;
     updateDomTree(RemovedBlocks);
+    updateLoops(RemovedBlocks);
     for (MachineBasicBlock *MBB : RemovedBlocks)
       MBB->eraseFromParent();
-    updateLoops(RemovedBlocks);
   }
   return Changed;
 }

llvm/lib/Transforms/IPO/LoopExtractor.cpp

@@ -238,8 +238,11 @@ bool LoopExtractor::extractLoop(Loop *L, LoopInfo &LI, DominatorTree &DT) {
   AssumptionCache *AC = LookupAssumptionCache(Func);
   CodeExtractorAnalysisCache CEAC(Func);
   CodeExtractor Extractor(L->getBlocks(), &DT, false, nullptr, nullptr, AC);
-  if (Extractor.extractCodeRegion(CEAC)) {
+  if (Extractor.isEligible()) {
+    // Remove loop while blocks are still in the current function
     LI.erase(L);
+    [[maybe_unused]] Function *ExtrF = Extractor.extractCodeRegion(CEAC);
+    assert(ExtrF && "CodeExtractor didn't extact eligible loop");
     --NumLoops;
     ++NumExtracted;
     return true;

llvm/lib/Transforms/Utils/LoopUtils.cpp

@@ -682,19 +682,16 @@ void llvm::deleteDeadLoop(Loop *L, DominatorTree *DT, ScalarEvolution *SE,
     MSSA->verifyMemorySSA();
 
   if (LI) {
+    SmallPtrSet<BasicBlock *, 8> Blocks(llvm::from_range, L->blocks());
+
     // Erase the instructions and the blocks without having to worry
     // about ordering because we already dropped the references.
-    // NOTE: This iteration is safe because erasing the block does not remove
-    // its entry from the loop's block list.  We do that in the next section.
-    for (BasicBlock *BB : L->blocks())
-      BB->eraseFromParent();
-
-    // Finally, the blocks from loopinfo.  This has to happen late because
-    // otherwise our loop iterators won't work.
-
-    SmallPtrSet<BasicBlock *, 8> blocks(llvm::from_range, L->blocks());
-    for (BasicBlock *BB : blocks)
+    // Remove blocks from loopinfo before erasing them, otherwise the loopinfo
+    // cannot find the loop using block numbers.
+    for (BasicBlock *BB : Blocks) {
       LI->removeBlock(BB);
+      BB->eraseFromParent();
+    }
 
     // The last step is to update LoopInfo now that we've eliminated this loop.
     // Note: LoopInfo::erase remove the given loop and relink its subloops with

polly/lib/CodeGen/IslNodeBuilder.cpp

@@ -234,7 +234,8 @@ static void findReferencesInStmt(ScopStmt *Stmt, SetVector<Value *> &Values,
   LoopInfo *LI = Stmt->getParent()->getLI();
 
   BasicBlock *BB = Stmt->getBasicBlock();
-  Loop *Scope = LI->getLoopFor(BB);
+  // TODO: Should BB ever be null?
+  Loop *Scope = BB ? LI->getLoopFor(BB) : nullptr;
   for (Instruction *Inst : Stmt->getInstructions())
     findReferencesInInst(Inst, Stmt, Scope, GlobalMap, Values, SCEVs);