[llvm] [MemCpyOpt] Fix the invalid code modification for GEP (PR #68479)

[Kai Yan via llvm-commits]

https://github.com/kaiyan96 updated https://github.com/llvm/llvm-project/pull/68479

>From 40084e54f8274a241f286c0a297f6df994041743 Mon Sep 17 00:00:00 2001
From: aklkaiyan <aklkaiyan at tencent.com>
Date: Sat, 7 Oct 2023 10:53:27 +0800
Subject: [PATCH] [MemCpyOpt] Fix the invalid code modification for GEP

Apply a rollback for the GEP modification in the performCallSlotOption function
after the optimization has failed.
---
 .../lib/Transforms/Scalar/MemCpyOptimizer.cpp |  8 ++-
 .../MemCpyOpt/memcpy-gep-modification.ll      | 60 +++++++++++++++++++
 2 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/MemCpyOpt/memcpy-gep-modification.ll

diff --git a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp b/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
index 783ef57805610b9..688bcfa57158975 100644
--- a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
+++ b/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
@@ -1043,12 +1043,13 @@ bool MemCpyOptPass::performCallSlotOptzn(Instruction *cpyLoad,
 
   // Since we're changing the parameter to the callsite, we need to make sure
   // that what would be the new parameter dominates the callsite.
+  bool NeedMoveGEP = false;
   if (!DT->dominates(cpyDest, C)) {
     // Support moving a constant index GEP before the call.
     auto *GEP = dyn_cast<GetElementPtrInst>(cpyDest);
     if (GEP && GEP->hasAllConstantIndices() &&
         DT->dominates(GEP->getPointerOperand(), C))
-      GEP->moveBefore(C);
+      NeedMoveGEP = true;
     else
       return false;
   }
@@ -1101,6 +1102,11 @@ bool MemCpyOptPass::performCallSlotOptzn(Instruction *cpyLoad,
     cast<AllocaInst>(cpyDest)->setAlignment(srcAlign);
   }
 
+  if (NeedMoveGEP) {
+    auto *GEP = dyn_cast<GetElementPtrInst>(cpyDest);
+    GEP->moveBefore(C);
+  }
+
   if (SkippedLifetimeStart) {
     SkippedLifetimeStart->moveBefore(C);
     MSSAU->moveBefore(MSSA->getMemoryAccess(SkippedLifetimeStart),
diff --git a/llvm/test/Transforms/MemCpyOpt/memcpy-gep-modification.ll b/llvm/test/Transforms/MemCpyOpt/memcpy-gep-modification.ll
new file mode 100644
index 000000000000000..ba6faf3578c5027
--- /dev/null
+++ b/llvm/test/Transforms/MemCpyOpt/memcpy-gep-modification.ll
@@ -0,0 +1,60 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -passes=memcpyopt < %s -verify-memoryssa | FileCheck %s
+
+%struct.MaskedType = type { i8, i8 }
+
+declare void @llvm.lifetime.start.p0(i64, ptr nocapture) #0
+declare void @llvm.lifetime.end.p0(i64, ptr nocapture) #0
+declare void @MaskedFunction1(ptr, ptr addrspace(1))
+declare void @MaskedFunction2(ptr, ptr)
+
+define i8 @test_gep_not_modified(ptr %in0, ptr %in1) {
+; CHECK-LABEL: @test_gep_not_modified(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[FUNCALLOC:%.*]] = alloca [[STRUCT_MASKEDTYPE:%.*]], align 4
+; CHECK-NEXT:    [[PTRALLOC:%.*]] = alloca i8, align 1
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 4, ptr [[PTRALLOC]])
+; CHECK-NEXT:    [[ADDRSPACECAST:%.*]] = addrspacecast ptr [[PTRALLOC]] to ptr addrspace(1)
+; CHECK-NEXT:    call void @MaskedFunction1(ptr [[IN1:%.*]], ptr addrspace(1) [[ADDRSPACECAST]])
+; CHECK-NEXT:    [[LOAD1:%.*]] = load i8, ptr [[PTRALLOC]], align 1
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 4, ptr [[PTRALLOC]])
+; CHECK-NEXT:    [[GETELEMPTR1:%.*]] = getelementptr inbounds [[STRUCT_MASKEDTYPE]], ptr [[FUNCALLOC]], i32 0, i32 1
+; CHECK-NEXT:    store i8 [[LOAD1]], ptr [[GETELEMPTR1]], align 1
+; CHECK-NEXT:    ret i8 0
+;
+entry:
+  %funcAlloc = alloca %struct.MaskedType, align 4
+  %ptrAlloc = alloca i8, align 1
+  call void @llvm.lifetime.start.p0(i64 4, ptr %ptrAlloc) #0
+  %addrspaceCast = addrspacecast ptr %ptrAlloc to ptr addrspace(1)
+  call void @MaskedFunction1(ptr %in1, ptr addrspace(1) %addrspaceCast)
+  %load1 = load i8, ptr %ptrAlloc, align 1
+  call void @llvm.lifetime.end.p0(i64 4, ptr %ptrAlloc) #0
+  %getElemPtr1 = getelementptr inbounds %struct.MaskedType, ptr %funcAlloc, i32 0, i32 1
+  store i8 %load1, ptr %getElemPtr1, align 1
+  ret i8 0
+}
+
+define i8 @test_gep_modified(ptr %in0, ptr %in1) {
+; CHECK-LABEL: @test_gep_modified(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[FUNCALLOC:%.*]] = alloca [[STRUCT_MASKEDTYPE:%.*]], align 4
+; CHECK-NEXT:    [[PTRALLOC:%.*]] = alloca i8, align 1
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 4, ptr [[PTRALLOC]])
+; CHECK-NEXT:    [[GETELEMPTR1:%.*]] = getelementptr inbounds [[STRUCT_MASKEDTYPE]], ptr [[FUNCALLOC]], i32 0, i32 1
+; CHECK-NEXT:    call void @MaskedFunction2(ptr [[IN1:%.*]], ptr [[GETELEMPTR1]])
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 4, ptr [[PTRALLOC]])
+; CHECK-NEXT:    ret i8 0
+;
+entry:
+  %funcAlloc = alloca %struct.MaskedType, align 4
+  %ptrAlloc = alloca i8, align 1
+  call void @llvm.lifetime.start.p0(i64 4, ptr %ptrAlloc) #0
+  call void @MaskedFunction2(ptr %in1, ptr %ptrAlloc)
+  %load1 = load i8, ptr %ptrAlloc, align 1
+  call void @llvm.lifetime.end.p0(i64 4, ptr %ptrAlloc) #0
+  %getElemPtr1 = getelementptr inbounds %struct.MaskedType, ptr %funcAlloc, i32 0, i32 1
+  store i8 %load1, ptr %getElemPtr1, align 1
+  ret i8 0
+}
+