[llvm] [MachineSink] Only add sink candidate if toBB is a successor of fromBB (PR #98540)

[via llvm-commits]

https://github.com/yozhu updated https://github.com/llvm/llvm-project/pull/98540

>From c230c574fc3c7f3ebd837cf04a8a3aba2abf67db Mon Sep 17 00:00:00 2001
From: YongKang Zhu <yongzhu at fb.com>
Date: Thu, 11 Jul 2024 13:17:43 -0700
Subject: [PATCH 1/2] [MachineSink] Only add sink candidate if toBB is a
 successor of fromBB

---
 llvm/lib/CodeGen/MachineSink.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/MachineSink.cpp b/llvm/lib/CodeGen/MachineSink.cpp
index bbc5ab13a0cd3..4d8de6765273f 100644
--- a/llvm/lib/CodeGen/MachineSink.cpp
+++ b/llvm/lib/CodeGen/MachineSink.cpp
@@ -913,7 +913,7 @@ bool MachineSinking::isWorthBreakingCriticalEdge(
     auto Res = CEMergeCandidates.try_emplace(Key, From);
     // We wanted to sink the same register into the same block, consider it to
     // be profitable.
-    if (!Res.second) {
+    if (!Res.second && Res.first->second->isSuccessor(To)) {
       // Return the source block that was previously held off.
       DeferredFromBlock = Res.first->second;
       return true;

>From 321c05138b4bdc4ba9faa0089233479b7b807a78 Mon Sep 17 00:00:00 2001
From: YongKang Zhu <yongzhu at fb.com>
Date: Fri, 12 Jul 2024 07:41:35 -0700
Subject: [PATCH 2/2] Move check to a more genetic place and add test

---
 llvm/lib/CodeGen/MachineSink.cpp              |  4 +--
 .../CodeGen/X86/MachineSink-Issue98477.ll     | 36 +++++++++++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/MachineSink-Issue98477.ll

diff --git a/llvm/lib/CodeGen/MachineSink.cpp b/llvm/lib/CodeGen/MachineSink.cpp
index 4d8de6765273f..57b71d1a26859 100644
--- a/llvm/lib/CodeGen/MachineSink.cpp
+++ b/llvm/lib/CodeGen/MachineSink.cpp
@@ -913,7 +913,7 @@ bool MachineSinking::isWorthBreakingCriticalEdge(
     auto Res = CEMergeCandidates.try_emplace(Key, From);
     // We wanted to sink the same register into the same block, consider it to
     // be profitable.
-    if (!Res.second && Res.first->second->isSuccessor(To)) {
+    if (!Res.second) {
       // Return the source block that was previously held off.
       DeferredFromBlock = Res.first->second;
       return true;
@@ -959,7 +959,7 @@ bool MachineSinking::isLegalToBreakCriticalEdge(MachineInstr &MI,
                                                 MachineBasicBlock *ToBB,
                                                 bool BreakPHIEdge) {
   // Avoid breaking back edge. From == To means backedge for single BB cycle.
-  if (!SplitEdges || FromBB == ToBB)
+  if (!SplitEdges || FromBB == ToBB || !FromBB->isSuccessor(ToBB))
     return false;
 
   MachineCycle *FromCycle = CI->getCycle(FromBB);
diff --git a/llvm/test/CodeGen/X86/MachineSink-Issue98477.ll b/llvm/test/CodeGen/X86/MachineSink-Issue98477.ll
new file mode 100644
index 0000000000000..abc41341043c8
--- /dev/null
+++ b/llvm/test/CodeGen/X86/MachineSink-Issue98477.ll
@@ -0,0 +1,36 @@
+; RUN: llc < %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define i32 @main(i1 %tobool.not, i32 %0) {
+entry:
+  br i1 %tobool.not, label %if.end13, label %j.preheader
+
+  j.preheader:       ; preds = %if.end13, %entry
+  %h.0.ph = phi i32 [ 1, %entry ], [ 0, %if.end13 ]
+  br label %j
+
+  j:                 ; preds = %if.then4, %j.preheader
+  %1 = phi i32 [ %div2, %if.then4 ], [ 0, %j.preheader ]
+  %rem1 = srem i32 1, %0
+  %cmp = icmp slt i32 %1, 0
+  %or.cond = select i1 false, i1 true, i1 %cmp
+  br i1 %or.cond, label %if.then4, label %if.end9
+
+  if.then4:          ; preds = %j
+  %div2 = sdiv i32 1, 0
+  %rem5 = srem i32 1, %h.0.ph
+  br i1 %tobool.not, label %if.end9, label %j
+
+  if.end9:           ; preds = %if.then4, %j
+  %2 = phi i32 [ 0, %j ], [ %rem5, %if.then4 ]
+  %tobool10.not = icmp eq i32 %2, 0
+  br i1 %tobool10.not, label %if.end13, label %while.body.lr.ph
+
+  while.body.lr.ph:  ; preds = %if.end9
+  ret i32 %rem1
+
+  if.end13:          ; preds = %if.end9, %entry
+  br label %j.preheader
+}