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

[via llvm-commits]

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

>From 3e8b5ba2fa00704cdd3ff0424fe5a4075c5af76b 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 | 56 +++++++++++++++++++-----
 llvm/include/llvm/ADT/GenericCycleInfo.h |  1 +
 2 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/llvm/include/llvm/ADT/GenericCycleImpl.h b/llvm/include/llvm/ADT/GenericCycleImpl.h
index 2adec725cd3bf7d..8f5ed090ae5d675 100644
--- a/llvm/include/llvm/ADT/GenericCycleImpl.h
+++ b/llvm/include/llvm/ADT/GenericCycleImpl.h
@@ -368,17 +368,24 @@ 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;
+
+  // 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);
+
+  CycleT *ParentCycle = Cycle->getParentCycle();
+  while (ParentCycle) {
+    Cycle = ParentCycle;
+    Cycle->appendBlock(NewBlock);
+    ParentCycle = Cycle->getParentCycle();
   }
+
+  BlockMapTopLevel.try_emplace(NewBlock, Cycle);
   assert(validateTree());
 }
 
@@ -392,6 +399,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..613053f02b7aded 100644
--- a/llvm/include/llvm/ADT/GenericCycleInfo.h
+++ b/llvm/include/llvm/ADT/GenericCycleInfo.h
@@ -261,6 +261,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);