[MachineCycle][NFC] add a cache for block and its top level cycle

This solves llvm#57664

Reviewed By: sameerds

Differential Revision: https://reviews.llvm.org/D134019

Author: Chen Zheng

llvm/include/llvm/ADT/GenericCycleImpl.h

@@ -144,21 +144,28 @@ template <typename ContextT> class GenericCycleInfoCompute {
 };
 
 template <typename ContextT>
-auto GenericCycleInfo<ContextT>::getTopLevelParentCycle(
-    const BlockT *Block) const -> CycleT * {
+auto GenericCycleInfo<ContextT>::getTopLevelParentCycle(BlockT *Block)
+    -> CycleT * {
+  auto Cycle = BlockMapTopLevel.find(Block);
+  if (Cycle != BlockMapTopLevel.end())
+    return Cycle->second;
+
   auto MapIt = BlockMap.find(Block);
   if (MapIt == BlockMap.end())
     return nullptr;
 
   auto *C = MapIt->second;
   while (C->ParentCycle)
     C = C->ParentCycle;
+  BlockMapTopLevel.try_emplace(Block, C);
   return C;
 }
 
 template <typename ContextT>
-void GenericCycleInfo<ContextT>::moveToNewParent(CycleT *NewParent,
-                                                 CycleT *Child) {
+void GenericCycleInfo<ContextT>::moveTopLevelCycleToNewParent(CycleT *NewParent,
+                                                              CycleT *Child) {
+  assert((!Child->ParentCycle && !NewParent->ParentCycle) &&
+         "NewParent and Child must be both top level cycle!\n");
   auto &CurrentContainer =
       Child->ParentCycle ? Child->ParentCycle->Children : TopLevelCycles;
   auto Pos = llvm::find_if(CurrentContainer, [=](const auto &Ptr) -> bool {
@@ -169,6 +176,13 @@ void GenericCycleInfo<ContextT>::moveToNewParent(CycleT *NewParent,
   *Pos = std::move(CurrentContainer.back());
   CurrentContainer.pop_back();
   Child->ParentCycle = NewParent;
+
+  NewParent->Blocks.insert(NewParent->Blocks.end(), Child->block_begin(),
+                           Child->block_end());
+
+  for (auto &It : BlockMapTopLevel)
+    if (It.second == Child)
+      It.second = NewParent;
 }
 
 /// \brief Main function of the cycle info computations.
@@ -240,10 +254,7 @@ void GenericCycleInfoCompute<ContextT>::run(BlockT *EntryBlock) {
                      << "discovered child cycle "
                      << Info.Context.print(BlockParent->getHeader()) << "\n");
           // Make BlockParent the child of NewCycle.
-          Info.moveToNewParent(NewCycle.get(), BlockParent);
-          NewCycle->Blocks.insert(NewCycle->Blocks.end(),
-                                  BlockParent->block_begin(),
-                                  BlockParent->block_end());
+          Info.moveTopLevelCycleToNewParent(NewCycle.get(), BlockParent);
 
           for (auto *ChildEntry : BlockParent->entries())
             ProcessPredecessors(ChildEntry);
@@ -257,6 +268,7 @@ void GenericCycleInfoCompute<ContextT>::run(BlockT *EntryBlock) {
         assert(!is_contained(NewCycle->Blocks, Block));
         NewCycle->Blocks.push_back(Block);
         ProcessPredecessors(Block);
+        Info.BlockMapTopLevel.try_emplace(Block, NewCycle.get());
       }
     } while (!Worklist.empty());
 
@@ -336,6 +348,7 @@ void GenericCycleInfoCompute<ContextT>::dfs(BlockT *EntryBlock) {
 template <typename ContextT> void GenericCycleInfo<ContextT>::clear() {
   TopLevelCycles.clear();
   BlockMap.clear();
+  BlockMapTopLevel.clear();
 }
 
 /// \brief Compute the cycle info for a function.

llvm/include/llvm/ADT/GenericCycleInfo.h

@@ -232,15 +232,24 @@ template <typename ContextT> class GenericCycleInfo {
 private:
   ContextT Context;
 
-  /// Map basic blocks to their inner-most containing loop.
+  /// Map basic blocks to their inner-most containing cycle.
   DenseMap<BlockT *, CycleT *> BlockMap;
 
+  /// Map basic blocks to their top level containing cycle.
+  DenseMap<BlockT *, CycleT *> BlockMapTopLevel;
+
   /// Top-level cycles discovered by any DFS.
   ///
   /// Note: The implementation treats the nullptr as the parent of
   /// every top-level cycle. See \ref contains for an example.
   std::vector<std::unique_ptr<CycleT>> TopLevelCycles;
 
+  /// Move \p Child to \p NewParent by manipulating Children vectors.
+  ///
+  /// Note: This is an incomplete operation that does not update the depth of
+  /// the subtree.
+  void moveTopLevelCycleToNewParent(CycleT *NewParent, CycleT *Child);
+
 public:
   GenericCycleInfo() = default;
   GenericCycleInfo(GenericCycleInfo &&) = default;
@@ -254,13 +263,7 @@ template <typename ContextT> class GenericCycleInfo {
 
   CycleT *getCycle(const BlockT *Block) const;
   unsigned getCycleDepth(const BlockT *Block) const;
-  CycleT *getTopLevelParentCycle(const BlockT *Block) const;
-
-  /// Move \p Child to \p NewParent by manipulating Children vectors.
-  ///
-  /// Note: This is an incomplete operation that does not update the
-  /// list of blocks in the new parent or the depth of the subtree.
-  void moveToNewParent(CycleT *NewParent, CycleT *Child);
+  CycleT *getTopLevelParentCycle(BlockT *Block);
 
   /// Methods for debug and self-test.
   //@{