From 6233346895abfb57782511cddc263d439fdd537b Mon Sep 17 00:00:00 2001
From: Chengjun <chengjunp@Nvidia.com>
Date: Mon, 28 Oct 2024 17:22:48 -0700
Subject: [PATCH] [GenericCycle] Add a Cache for getExitBlocks in GenericCycle
 (#112290)

In `UniformityAnalysis`, we need to get the exit blocks of cycles in the
`DivergencePropagator` and currently, we have to do a search for the
exit blocks every time. In this change, we add a cache of the results in
the `GenericCycle` so that it can save the compile time. By testing, for
some large cases, this can save about 60% compile time in the
`UniformityAnalysis`.
---
 llvm/include/llvm/ADT/GenericCycleImpl.h |  9 +++++++++
 llvm/include/llvm/ADT/GenericCycleInfo.h | 21 +++++++++++++++++++--
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/ADT/GenericCycleImpl.h b/llvm/include/llvm/ADT/GenericCycleImpl.h
index 3d2c5f42883558..41ba8bf8fde14b 100644
--- a/llvm/include/llvm/ADT/GenericCycleImpl.h
+++ b/llvm/include/llvm/ADT/GenericCycleImpl.h
@@ -47,6 +47,11 @@ bool GenericCycle<ContextT>::contains(const GenericCycle *C) const {
 template <typename ContextT>
 void GenericCycle<ContextT>::getExitBlocks(
     SmallVectorImpl<BlockT *> &TmpStorage) const {
+  if (!ExitBlocksCache.empty()) {
+    TmpStorage = ExitBlocksCache;
+    return;
+  }
+
   TmpStorage.clear();
 
   size_t NumExitBlocks = 0;
@@ -65,6 +70,7 @@ void GenericCycle<ContextT>::getExitBlocks(
 
     TmpStorage.resize(NumExitBlocks);
   }
+  ExitBlocksCache.append(TmpStorage.begin(), TmpStorage.end());
 }
 
 template <typename ContextT>
@@ -298,6 +304,8 @@ void GenericCycleInfo<ContextT>::moveTopLevelCycleToNewParent(CycleT *NewParent,
   for (auto &It : BlockMapTopLevel)
     if (It.second == Child)
       It.second = NewParent;
+  NewParent->clearCache();
+  Child->clearCache();
 }
 
 template <typename ContextT>
@@ -316,6 +324,7 @@ void GenericCycleInfo<ContextT>::addBlockToCycle(BlockT *Block, CycleT *Cycle) {
   }
 
   BlockMapTopLevel.try_emplace(Block, Cycle);
+  Cycle->clearCache();
 }
 
 /// \brief Main function of the cycle info computations.
diff --git a/llvm/include/llvm/ADT/GenericCycleInfo.h b/llvm/include/llvm/ADT/GenericCycleInfo.h
index 8c2fa0490e638a..b8b6e3e9967a4a 100644
--- a/llvm/include/llvm/ADT/GenericCycleInfo.h
+++ b/llvm/include/llvm/ADT/GenericCycleInfo.h
@@ -74,16 +74,27 @@ template <typename ContextT> class GenericCycle {
   ///       always have the same depth.
   unsigned Depth = 0;
 
+  /// Cache for the results of GetExitBlocks
+  mutable SmallVector<BlockT *, 4> ExitBlocksCache;
+
   void clear() {
     Entries.clear();
     Children.clear();
     Blocks.clear();
     Depth = 0;
     ParentCycle = nullptr;
+    clearCache();
+  }
+
+  void appendEntry(BlockT *Block) {
+    Entries.push_back(Block);
+    clearCache();
   }
 
-  void appendEntry(BlockT *Block) { Entries.push_back(Block); }
-  void appendBlock(BlockT *Block) { Blocks.insert(Block); }
+  void appendBlock(BlockT *Block) {
+    Blocks.insert(Block);
+    clearCache();
+  }
 
   GenericCycle(const GenericCycle &) = delete;
   GenericCycle &operator=(const GenericCycle &) = delete;
@@ -102,6 +113,11 @@ template <typename ContextT> class GenericCycle {
     return Entries;
   }
 
+  /// Clear the cache of the cycle.
+  /// This should be run in all non-const function in GenericCycle
+  /// and GenericCycleInfo.
+  void clearCache() const { ExitBlocksCache.clear(); }
+
   /// \brief Return whether \p Block is an entry block of the cycle.
   bool isEntry(const BlockT *Block) const {
     return is_contained(Entries, Block);
@@ -112,6 +128,7 @@ template <typename ContextT> class GenericCycle {
     assert(contains(Block));
     Entries.clear();
     Entries.push_back(Block);
+    clearCache();
   }
 
   /// \brief Return whether \p Block is contained in the cycle.