[llvm] [CodeMoverUtils] Enhance CodeMoverUtils to sink an entire BB (PR #87857)

[via llvm-commits]

https://github.com/CongzheUalberta updated https://github.com/llvm/llvm-project/pull/87857

>From 3924918eb7ef1068f07cf8bdb0efff529d041046 Mon Sep 17 00:00:00 2001
From: Congzhe Cao <c00548033 at china.huawei.com>
Date: Fri, 5 Apr 2024 22:39:09 -0400
Subject: [PATCH]  [CodeMoverUtils] Enhance CodeMoverUtils to sink entire BB

When moving an entire basic block BB after InsertPoint, currently
we check for all instructions whether their users are dominated by
InsertPoint, however, this can be improved such that even a user
is not dominated InsertPoint, as long as it appears as a subsequent
instruction in the same BB, it is safe to move.

This patch is similar to commit 751be2a064f119af74c7b9b1e52bc904d8aa114d
that enhanced hoisting an entire BB, and this patch enhances sinking an
entire BB.
---
 llvm/lib/Transforms/Utils/CodeMoverUtils.cpp    | 17 +++++++++++++++--
 .../Transforms/Utils/CodeMoverUtilsTest.cpp     |  5 +++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/CodeMoverUtils.cpp b/llvm/lib/Transforms/Utils/CodeMoverUtils.cpp
index 6a2dae5bab68ee..1db63c9a197b3a 100644
--- a/llvm/lib/Transforms/Utils/CodeMoverUtils.cpp
+++ b/llvm/lib/Transforms/Utils/CodeMoverUtils.cpp
@@ -336,9 +336,22 @@ bool llvm::isSafeToMoveBefore(Instruction &I, Instruction &InsertPoint,
 
   if (isReachedBefore(&I, &InsertPoint, &DT, PDT))
     for (const Use &U : I.uses())
-      if (auto *UserInst = dyn_cast<Instruction>(U.getUser()))
-        if (UserInst != &InsertPoint && !DT.dominates(&InsertPoint, U))
+      if (auto *UserInst = dyn_cast<Instruction>(U.getUser())) {
+        // If InsertPoint is in a BB that comes after I, then we cannot move if
+        // I is used in the terminator of the current BB.
+        if (I.getParent() != InsertPoint.getParent() &&
+              UserInst == I.getParent()->getTerminator())
           return false;
+        if (UserInst != &InsertPoint && !DT.dominates(&InsertPoint, U)) {
+          // If UserInst is an instruction that appears later in the same BB as
+          // I, then it is okay to move since I will still be available when
+          // UserInst is executed.
+          if (CheckForEntireBlock && I.getParent() == UserInst->getParent() &&
+              DT.dominates(&I, UserInst))
+            continue;
+          return false;
+        }
+      }
   if (isReachedBefore(&InsertPoint, &I, &DT, PDT))
     for (const Value *Op : I.operands())
       if (auto *OpInst = dyn_cast<Instruction>(Op)) {
diff --git a/llvm/unittests/Transforms/Utils/CodeMoverUtilsTest.cpp b/llvm/unittests/Transforms/Utils/CodeMoverUtilsTest.cpp
index 8554d1a33cadee..cbddc6d92772ab 100644
--- a/llvm/unittests/Transforms/Utils/CodeMoverUtilsTest.cpp
+++ b/llvm/unittests/Transforms/Utils/CodeMoverUtilsTest.cpp
@@ -673,6 +673,11 @@ TEST(CodeMoverUtils, IsSafeToMoveTest4) {
         // Can move as %add2 and %sub2 are control flow equivalent,
         // although %add2 does not strictly dominate %sub2.
         EXPECT_TRUE(isSafeToMoveBefore(*SubInst2, *AddInst2, DT, &PDT, &DI));
+
+        BasicBlock *BB0 = getBasicBlockByName(F, "if.then.first");
+        BasicBlock *BB1 = getBasicBlockByName(F, "if.end.first");
+        EXPECT_TRUE(
+            isSafeToMoveBefore(*BB0, *BB1->getTerminator(), DT, &PDT, &DI));
       });
 }