[llvm] [MemCpyOpt] move SrcAlloca to the entry if transformation is performed (PR #67226)

[Kohei Asano via llvm-commits]

https://github.com/khei4 updated https://github.com/llvm/llvm-project/pull/67226

>From 4f5132a430f1ee402fbc0607d6ab410d99e09551 Mon Sep 17 00:00:00 2001
From: khei4 <kk.asano.luxy at gmail.com>
Date: Sun, 24 Sep 2023 01:20:25 +0900
Subject: [PATCH] [MemCpyOpt] move SrcAlloca to the entry if transformation is
 performed

---
 .../lib/Transforms/Scalar/MemCpyOptimizer.cpp | 22 +++++++++++--------
 llvm/test/Transforms/MemCpyOpt/stack-move.ll  |  7 ++----
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp b/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
index d87f2fb59814edf..44cfc80134ca18e 100644
--- a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
+++ b/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
@@ -1463,6 +1463,7 @@ bool MemCpyOptPass::performStackMoveOptzn(Instruction *Load, Instruction *Store,
 
   SmallVector<Instruction *, 4> LifetimeMarkers;
   SmallSet<Instruction *, 4> NoAliasInstrs;
+  bool SrcNotDom = false;
 
   // Recursively track the user and check whether modified alias exist.
   auto IsDereferenceableOrNull = [](Value *V, const DataLayout &DL) -> bool {
@@ -1483,14 +1484,11 @@ bool MemCpyOptPass::performStackMoveOptzn(Instruction *Load, Instruction *Store,
       Worklist.pop_back();
       for (const Use &U : I->uses()) {
         auto *UI = cast<Instruction>(U.getUser());
-        // TODO: We can perform the transformation if we move src alloca to
-        // before the dominator of all uses. If any use that isn't dominated by
-        // SrcAlloca exists, non-dominating uses will be produced.
-        if (!DT->dominates(SrcAlloca, UI)) {
-          LLVM_DEBUG(dbgs() << "Stack Move: SrcAlloca doesn't dominate all "
-                               "uses for the location, bailing\n");
-          return false;
-        }
+        // If any use that isn't dominated by SrcAlloca exists, we move src
+        // alloca to the entry before the transformation.
+        if (!DT->dominates(SrcAlloca, UI))
+          SrcNotDom = true;
+
         if (Visited.size() >= MaxUsesToExplore) {
           LLVM_DEBUG(
               dbgs()
@@ -1600,7 +1598,13 @@ bool MemCpyOptPass::performStackMoveOptzn(Instruction *Load, Instruction *Store,
   if (!CaptureTrackingWithModRef(SrcAlloca, SrcModRefCallback))
     return false;
 
-  // We can do the transformation. First, align the allocas appropriately.
+  // We can do the transformation. First, move the SrcAlloca to the entry point
+  // if it's not dominator for all uses. After that SrcAlloca becomes the
+  // dominator, because it's in the entry BB.
+  if (SrcNotDom)
+    SrcAlloca->moveBefore(*SrcAlloca->getParent(),
+                          SrcAlloca->getParent()->getFirstInsertionPt());
+  // Align the allocas appropriately.
   SrcAlloca->setAlignment(
       std::max(SrcAlloca->getAlign(), DestAlloca->getAlign()));
 
diff --git a/llvm/test/Transforms/MemCpyOpt/stack-move.ll b/llvm/test/Transforms/MemCpyOpt/stack-move.ll
index f6be486dce8ac1e..bae0a152116004d 100644
--- a/llvm/test/Transforms/MemCpyOpt/stack-move.ll
+++ b/llvm/test/Transforms/MemCpyOpt/stack-move.ll
@@ -986,12 +986,9 @@ bb2:
 ; when any use that isn't dominated by SrcAlloca exists.
 define i32 @use_not_dominated_by_src_alloca() {
 ; CHECK-LABEL: define i32 @use_not_dominated_by_src_alloca() {
-; CHECK-NEXT:    [[DEST:%.*]] = alloca i1, align 1
-; CHECK-NEXT:    [[DEST_GEP:%.*]] = getelementptr i64, ptr [[DEST]], i64 -1
-; CHECK-NEXT:    [[DEST_USE:%.*]] = load i8, ptr [[DEST_GEP]], align 1
 ; CHECK-NEXT:    [[SRC:%.*]] = alloca i8, align 4
-; CHECK-NEXT:    [[SRC_VAL:%.*]] = load i1, ptr [[SRC]], align 4
-; CHECK-NEXT:    store i1 [[SRC_VAL]], ptr [[DEST]], align 1
+; CHECK-NEXT:    [[DEST_GEP:%.*]] = getelementptr i64, ptr [[SRC]], i64 -1
+; CHECK-NEXT:    [[DEST_USE:%.*]] = load i8, ptr [[DEST_GEP]], align 1
 ; CHECK-NEXT:    ret i32 0
 ;
   %dest = alloca i1, align 1