[llvm] CycleInfo: Fix splitCriticalEdge (PR #68584)

[via llvm-commits]

https://github.com/petar-avramovic updated https://github.com/llvm/llvm-project/pull/68584

>From b005774ecf9920928672913f5430c797e34bb3dc Mon Sep 17 00:00:00 2001
From: Petar Avramovic <Petar.Avramovic at amd.com>
Date: Mon, 9 Oct 2023 15:01:06 +0200
Subject: [PATCH] CycleInfo: Fix splitCriticalEdge

There are cases when cycle that contains both Pred and Succ is not found:
Pred is in a smaller cycle that is contained in a larger cycle that also
contains Succ, or Pred and Succ are in disjointed innermost cycles but
there is a larger cycle that contains both.
Add efficient helper function that finds innermost cycle that contains
both Pred and Succ if it exists.
---
 llvm/include/llvm/ADT/GenericCycleImpl.h | 63 +++++++++++++++++++-----
 llvm/include/llvm/ADT/GenericCycleInfo.h |  7 +++
 2 files changed, 59 insertions(+), 11 deletions(-)

diff --git a/llvm/include/llvm/ADT/GenericCycleImpl.h b/llvm/include/llvm/ADT/GenericCycleImpl.h
index 2adec725cd3bf7d..74faff98b903386 100644
--- a/llvm/include/llvm/ADT/GenericCycleImpl.h
+++ b/llvm/include/llvm/ADT/GenericCycleImpl.h
@@ -184,6 +184,24 @@ void GenericCycleInfo<ContextT>::moveTopLevelCycleToNewParent(CycleT *NewParent,
       It.second = NewParent;
 }
 
+template <typename ContextT>
+void GenericCycleInfo<ContextT>::addBlockToCycle(BlockT *Block, CycleT *Cycle) {
+  // FixMe: Appending NewBlock is fine as a set of blocks in a cycle. When
+  // printing, cycle NewBlock is at the end of list but it should be in the
+  // middle to represent actual traversal of a cycle.
+  Cycle->appendBlock(Block);
+  BlockMap.try_emplace(Block, Cycle);
+
+  CycleT *ParentCycle = Cycle->getParentCycle();
+  while (ParentCycle) {
+    Cycle = ParentCycle;
+    Cycle->appendBlock(Block);
+    ParentCycle = Cycle->getParentCycle();
+  }
+
+  BlockMapTopLevel.try_emplace(Block, Cycle);
+}
+
 /// \brief Main function of the cycle info computations.
 template <typename ContextT>
 void GenericCycleInfoCompute<ContextT>::run(BlockT *EntryBlock) {
@@ -368,17 +386,11 @@ void GenericCycleInfo<ContextT>::splitCriticalEdge(BlockT *Pred, BlockT *Succ,
                                                    BlockT *NewBlock) {
   // Edge Pred-Succ is replaced by edges Pred-NewBlock and NewBlock-Succ, all
   // cycles that had blocks Pred and Succ also get NewBlock.
-  CycleT *Cycle = this->getCycle(Pred);
-  if (Cycle && Cycle->contains(Succ)) {
-    while (Cycle) {
-      // FixMe: Appending NewBlock is fine as a set of blocks in a cycle. When
-      // printing cycle NewBlock is at the end of list but it should be in the
-      // middle to represent actual traversal of a cycle.
-      Cycle->appendBlock(NewBlock);
-      BlockMap.try_emplace(NewBlock, Cycle);
-      Cycle = Cycle->getParentCycle();
-    }
-  }
+  CycleT *Cycle = getSmallestCommonCycle(getCycle(Pred), getCycle(Succ));
+  if (!Cycle)
+    return;
+
+  addBlockToCycle(NewBlock, Cycle);
   assert(validateTree());
 }
 
@@ -392,6 +404,35 @@ auto GenericCycleInfo<ContextT>::getCycle(const BlockT *Block) const
   return BlockMap.lookup(Block);
 }
 
+/// \brief Find the innermost cycle containing both given cycles.
+///
+/// \returns the innermost cycle containing both \p A and \p B
+///          or nullptr if there is no such cycle.
+template <typename ContextT>
+auto GenericCycleInfo<ContextT>::getSmallestCommonCycle(CycleT *A,
+                                                        CycleT *B) const
+    -> CycleT * {
+  if (!A || !B)
+    return nullptr;
+
+  // If cycles A and B have different depth replace them with parent cycle
+  // until they have the same depth.
+  while (A->getDepth() > B->getDepth())
+    A = A->getParentCycle();
+  while (B->getDepth() > A->getDepth())
+    B = B->getParentCycle();
+
+  // Cycles A and B are at same depth but may be disjoint, replace them with
+  // parent cycles until we find cycle that contains both or we run out of
+  // parent cycles.
+  while (A != B) {
+    A = A->getParentCycle();
+    B = B->getParentCycle();
+  }
+
+  return A;
+}
+
 /// \brief get the depth for the cycle which containing a given block.
 ///
 /// \returns the depth for the innermost cycle containing \p Block or 0 if it is
diff --git a/llvm/include/llvm/ADT/GenericCycleInfo.h b/llvm/include/llvm/ADT/GenericCycleInfo.h
index 18be7eb4a6cca9d..83c4c2759d46891 100644
--- a/llvm/include/llvm/ADT/GenericCycleInfo.h
+++ b/llvm/include/llvm/ADT/GenericCycleInfo.h
@@ -248,6 +248,12 @@ template <typename ContextT> class GenericCycleInfo {
   /// the subtree.
   void moveTopLevelCycleToNewParent(CycleT *NewParent, CycleT *Child);
 
+  /// Assumes that \p Cycle is the innermost cycle containing \p Block.
+  /// \p Block will be appended to \p Cycle and all of its parent cycles.
+  /// \p Block will be added to BlockMap with \p Cycle and
+  /// BlockMapTopLevel with \p Cycle's top level parent cycle.
+  void addBlockToCycle(BlockT *Block, CycleT *Cycle);
+
 public:
   GenericCycleInfo() = default;
   GenericCycleInfo(GenericCycleInfo &&) = default;
@@ -261,6 +267,7 @@ template <typename ContextT> class GenericCycleInfo {
   const ContextT &getSSAContext() const { return Context; }
 
   CycleT *getCycle(const BlockT *Block) const;
+  CycleT *getSmallestCommonCycle(CycleT *A, CycleT *B) const;
   unsigned getCycleDepth(const BlockT *Block) const;
   CycleT *getTopLevelParentCycle(BlockT *Block);