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

[via llvm-commits]

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

>From 9d152c95d3b9610088cfde5ff3c684603a9b3a7b 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 | 48 +++++++++++++++++++-----
 llvm/include/llvm/ADT/GenericCycleInfo.h |  1 +
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/llvm/include/llvm/ADT/GenericCycleImpl.h b/llvm/include/llvm/ADT/GenericCycleImpl.h
index 2adec725cd3bf7d..e072447600c058e 100644
--- a/llvm/include/llvm/ADT/GenericCycleImpl.h
+++ b/llvm/include/llvm/ADT/GenericCycleImpl.h
@@ -368,16 +368,14 @@ 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(Pred, 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();
   }
   assert(validateTree());
 }
@@ -392,6 +390,36 @@ auto GenericCycleInfo<ContextT>::getCycle(const BlockT *Block) const
   return BlockMap.lookup(Block);
 }
 
+/// \brief Find the innermost cycle containing both given blocks.
+///
+/// \returns the innermost cycle containing both \p A and \p B or nullptr if
+///          they are not both contained in any cycle.
+template <typename ContextT>
+auto GenericCycleInfo<ContextT>::getSmallestCommonCycle(const BlockT *A,
+                                                        const BlockT *B) const
+    -> CycleT * {
+  CycleT *CycleA = getCycle(A);
+  CycleT *CycleB = getCycle(B);
+  if (!CycleA || !CycleB)
+    return nullptr;
+
+  // If CycleA and CycleB have different depth replace them with parents until
+  // they have the same depth.
+  while (CycleA->getDepth() > CycleB->getDepth())
+    CycleA = CycleA->getParentCycle();
+  while (CycleB->getDepth() > CycleA->getDepth())
+    CycleB = CycleB->getParentCycle();
+
+  // CycleA and CycleB are at same depth but may be disjoint, replace them with
+  // parents until we find cycle that contains both or we run out of parents.
+  while (CycleA != CycleB) {
+    CycleA = CycleA->getParentCycle();
+    CycleB = CycleB->getParentCycle();
+  }
+
+  return CycleA;
+}
+
 /// \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..7ddb2f349f062d1 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(const BlockT *A, const BlockT *B) const;
   unsigned getCycleDepth(const BlockT *Block) const;
   CycleT *getTopLevelParentCycle(BlockT *Block);