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

[shchenz]

Summary

This solves the big compile time regression in https://github.com/llvm/llvm-project/issues/57664

The hot function is getTopLevelParentCycle.

Too many redundant computations to get a top level cycle for a block.

This patch improves compile time about X15 on PPC64 Linux.

[nikic]

llvm/include/llvm/ADT/GenericCycleImpl.h

185

This assumes the new parent is a top level cycle. Is this guaranteed?

[sameerds]

llvm/include/llvm/ADT/GenericCycleImpl.h

185

At the only call to this function, it seems to be true that NewParent is a newly discovered top-level cycle, and Child is an existing top-level cycle. This loop then updates the cache to replace Child with NewParent. But that makes the whole thing a bit fragile if this function is ever called from a different place.

Maybe we should make this a private function "moveTopLevelCycleToNewParent", with suitable assertions?

[shchenz]

address comments:
1: change moveToNewParent to only apply top level parameters

[shchenz]

llvm/include/llvm/ADT/GenericCycleImpl.h

185

Thanks. Updated!

[avogelsgesang]

Thanks for the quick fix!
I unfortunately don't understand this code here, and can't give you a review, though

[avogelsgesang]

does it make sense to add a regression test?

[shchenz]

In D134019#3804673, @avogelsgesang wrote:

does it make sense to add a regression test?

This is a compile time improvement, no compiler functionality change. AFAIK, we don't have regression test for the compile time changes.

llvm/include/llvm/ADT/GenericCycleImpl.h

public:
 
   static void updateDepth(CycleT *SubTree);
 
 private:
   void dfs(BlockT *EntryBlock);
 };
 
 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,
+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 {
     return Child == Ptr.get();
   });
   assert(Pos != CurrentContainer.end());
   NewParent->Children.push_back(std::move(*Pos));
   *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.
 template <typename ContextT>
 void GenericCycleInfoCompute<ContextT>::run(BlockT *EntryBlock) {
   LLVM_DEBUG(errs() << "Entry block: " << Info.Context.print(EntryBlock)
                     << "\n");
   dfs(EntryBlock);
do {
       if (auto *BlockParent = Info.getTopLevelParentCycle(Block)) {
         LLVM_DEBUG(errs() << "  block " << Info.Context.print(Block) << ": ");
 
         if (BlockParent != NewCycle.get()) {
           LLVM_DEBUG(errs()
                      << "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);
         } else {
           LLVM_DEBUG(errs()
                      << "known child cycle "
                      << Info.Context.print(BlockParent->getHeader()) << "\n");
         }
       } else {
         Info.BlockMap.try_emplace(Block, NewCycle.get());
         assert(!is_contained(NewCycle->Blocks, Block));
         NewCycle->Blocks.push_back(Block);
         ProcessPredecessors(Block);
+        Info.BlockMapTopLevel.try_emplace(Block, NewCycle.get());
       }
     } while (!Worklist.empty());
 
     Info.TopLevelCycles.push_back(std::move(NewCycle));
   }
 
   // Fix top-level cycle links and compute cycle depths.
   for (auto *TLC : Info.toplevel_cycles()) {
LLVM_DEBUG(
     }
   );
 }
 
 /// \brief Reset the object to its initial state.
 template <typename ContextT> void GenericCycleInfo<ContextT>::clear() {
   TopLevelCycles.clear();
   BlockMap.clear();
+  BlockMapTopLevel.clear();
 }
 
 /// \brief Compute the cycle info for a function.
 template <typename ContextT>
 void GenericCycleInfo<ContextT>::compute(FunctionT &F) {
   GenericCycleInfoCompute<ContextT> Compute(*this);
   Context.setFunction(F);
 

llvm/include/llvm/ADT/GenericCycleInfo.h

public:
   using CycleT = GenericCycle<ContextT>;
   using FunctionT = typename ContextT::FunctionT;
   template <typename> friend class GenericCycle;
   template <typename> friend class GenericCycleInfoCompute;
 
 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;
   GenericCycleInfo &operator=(GenericCycleInfo &&) = default;
 
   void clear();
   void compute(FunctionT &F);
 
   FunctionT *getFunction() const { return Context.getFunction(); }
   const ContextT &getSSAContext() const { return Context; }
 
   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.
   //@{
 #ifndef NDEBUG
   bool validateTree() const;
 #endif
   void print(raw_ostream &Out) const;
   void dump() const { print(dbgs()); }