[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 5bde6fa98da40ea6fded2a8a790cf313fdd979c3 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

move test to the success zone
---
 .../lib/Transforms/Scalar/MemCpyOptimizer.cpp | 22 +++++----
 llvm/test/Transforms/MemCpyOpt/stack-move.ll  | 47 ++++++++-----------
 2 files changed, 32 insertions(+), 37 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..dee630f470d0053 100644
--- a/llvm/test/Transforms/MemCpyOpt/stack-move.ll
+++ b/llvm/test/Transforms/MemCpyOpt/stack-move.ll
@@ -45,6 +45,25 @@ define void @basic_memcpy() {
   ret void
 }
 
+define i32 @use_not_dominated_by_src_alloca() {
+; CHECK-LABEL: define i32 @use_not_dominated_by_src_alloca() {
+; CHECK-NEXT:    [[SRC:%.*]] = alloca i8, align 4
+; 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
+  ; Replacing the use of dest with src causes no domination uses.
+  %dest.gep = getelementptr i64, ptr %dest, i64 -1
+  %dest.use = load i8, ptr %dest.gep, align 1
+  %src = alloca i8, align 4
+  %src.val = load i1, ptr %src, align 4
+
+  store i1 %src.val, ptr %dest, align 1
+
+  ret i32 0
+}
+
 define void @basic_memmove() {
 ; CHECK-LABEL: define void @basic_memmove() {
 ; CHECK-NEXT:    [[SRC:%.*]] = alloca [[STRUCT_FOO:%.*]], align 4
@@ -346,7 +365,6 @@ define void @avoid_memory_use_last_user_crash() {
   ret void
 }
 
-; TODO: if the last user is terminator, we won't insert lifetime.end.
 ; For multi-bb patch, we will insert it for next immediate post dominator block.
 define void @terminator_lastuse() personality i32 0 {
 ; CHECK-LABEL: define void @terminator_lastuse() personality i32 0 {
@@ -718,7 +736,6 @@ unr2:
 
 }
 
-; TODO: merge allocas for multi basicblock loop case.
 define void @multi_bb_loop(i32 %n) {
 ; CHECK-LABEL: define void @multi_bb_loop
 ; CHECK-SAME: (i32 [[N:%.*]]) {
@@ -981,32 +998,6 @@ bb2:
 
 ; Optimization failures follow:
 
-; TODO: we can merge those alloca if we move src alloca to the start of the BB.
-; Tests that a the optimization isn't performed,
-; 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:    ret i32 0
-;
-  %dest = alloca i1, align 1
-  ; Replacing the use of dest with src causes no domination uses.
-  %dest.gep = getelementptr i64, ptr %dest, i64 -1
-  %dest.use = load i8, ptr %dest.gep, align 1
-  %src = alloca i8, align 4
-  %src.val = load i1, ptr %src, align 4
-
-  store i1 %src.val, ptr %dest, align 1
-
-  ret i32 0
-}
-
-
 ; Tests that a memcpy that doesn't completely overwrite a stack value is a use
 ; for the purposes of liveness analysis, not a definition.
 define void @incomplete_memcpy() {