[llvm] [IR] Use block numbers in PredIteratorCache (PR #101885)

[Alexis Engelke via llvm-commits]

https://github.com/aengelke created https://github.com/llvm/llvm-project/pull/101885

Avoid using a DenseMap of BasicBlocks when a vector is sufficient. This only has a few users, so gains are low, but so is the effort.

There's no longer a separate map which only holds sizes. All users that used size() also use get() afterwards, so it's always beneficial to compute both.

[c-t-t](http://llvm-compile-time-tracker.com/compare.php?from=e7f9d8e5c3e49e729c69aaa9be3322f7902370b8&to=ad71aea8ebad4ea5d0ae34c537e0cc899064db6b&stat=instructions:u) -0.08% in stage2-O3

>From ad71aea8ebad4ea5d0ae34c537e0cc899064db6b Mon Sep 17 00:00:00 2001
From: Alexis Engelke <engelke at in.tum.de>
Date: Sun, 4 Aug 2024 11:38:42 +0000
Subject: [PATCH] [IR] Improve PredIteratorCache to use block numbers

Avoid using DenseMap of BasicBlocks when a vector is sufficient. This
only has a few users, so gains are low, but so is the effort.

There's no longer a separate map which only holds sizes. All users that
used size() also use get() afterwards, so it's always beneficial to
compute both.
---
 llvm/include/llvm/IR/PredIteratorCache.h | 60 ++++++++++--------------
 1 file changed, 25 insertions(+), 35 deletions(-)

diff --git a/llvm/include/llvm/IR/PredIteratorCache.h b/llvm/include/llvm/IR/PredIteratorCache.h
index fc8cf20e9f759..c9603c870820b 100644
--- a/llvm/include/llvm/IR/PredIteratorCache.h
+++ b/llvm/include/llvm/IR/PredIteratorCache.h
@@ -25,52 +25,42 @@ namespace llvm {
 /// predecessor iterator queries.  This is useful for code that repeatedly
 /// wants the predecessor list for the same blocks.
 class PredIteratorCache {
-  /// BlockToPredsMap - Pointer to null-terminated list.
-  mutable DenseMap<BasicBlock *, BasicBlock **> BlockToPredsMap;
-  mutable DenseMap<BasicBlock *, unsigned> BlockToPredCountMap;
+  /// Storage, indexed by block number.
+  SmallVector<ArrayRef<BasicBlock *>> Storage;
+  /// Block number epoch to guard against renumberings.
+  unsigned BlockNumberEpoch;
 
   /// Memory - This is the space that holds cached preds.
   BumpPtrAllocator Memory;
 
-private:
-  /// GetPreds - Get a cached list for the null-terminated predecessor list of
-  /// the specified block.  This can be used in a loop like this:
-  ///   for (BasicBlock **PI = PredCache->GetPreds(BB); *PI; ++PI)
-  ///      use(*PI);
-  /// instead of:
-  /// for (pred_iterator PI = pred_begin(BB), E = pred_end(BB); PI != E; ++PI)
-  BasicBlock **GetPreds(BasicBlock *BB) {
-    BasicBlock **&Entry = BlockToPredsMap[BB];
-    if (Entry)
-      return Entry;
-
-    SmallVector<BasicBlock *, 32> PredCache(predecessors(BB));
-    PredCache.push_back(nullptr); // null terminator.
+public:
+  size_t size(BasicBlock *BB) { return get(BB).size(); }
+  ArrayRef<BasicBlock *> get(BasicBlock *BB) {
+#ifndef NDEBUG
+    // In debug builds, verify that no renumbering has occured.
+    if (Storage.empty())
+      BlockNumberEpoch = BB->getParent()->getBlockNumberEpoch();
+    else
+      assert(BlockNumberEpoch == BB->getParent()->getBlockNumberEpoch() &&
+             "Blocks renumbered during lifetime of PredIteratorCache");
+#endif
 
-    BlockToPredCountMap[BB] = PredCache.size() - 1;
+    if (LLVM_LIKELY(BB->getNumber() < Storage.size()))
+      if (auto Res = Storage[BB->getNumber()]; Res.data())
+        return Res;
 
-    Entry = Memory.Allocate<BasicBlock *>(PredCache.size());
-    std::copy(PredCache.begin(), PredCache.end(), Entry);
-    return Entry;
-  }
+    if (BB->getNumber() >= Storage.size())
+      Storage.resize(BB->getParent()->getMaxBlockNumber());
 
-  unsigned GetNumPreds(BasicBlock *BB) const {
-    auto Result = BlockToPredCountMap.find(BB);
-    if (Result != BlockToPredCountMap.end())
-      return Result->second;
-    return BlockToPredCountMap[BB] = pred_size(BB);
-  }
-
-public:
-  size_t size(BasicBlock *BB) const { return GetNumPreds(BB); }
-  ArrayRef<BasicBlock *> get(BasicBlock *BB) {
-    return ArrayRef(GetPreds(BB), GetNumPreds(BB));
+    SmallVector<BasicBlock *, 32> PredCache(predecessors(BB));
+    BasicBlock **Data = Memory.Allocate<BasicBlock *>(PredCache.size());
+    std::copy(PredCache.begin(), PredCache.end(), Data);
+    return Storage[BB->getNumber()] = ArrayRef(Data, PredCache.size());
   }
 
   /// clear - Remove all information.
   void clear() {
-    BlockToPredsMap.clear();
-    BlockToPredCountMap.clear();
+    Storage.clear();
     Memory.Reset();
   }
 };