[MemCpyOpt] Fix the invalid code modification for GEP

Relocate the GEP modification to a later stage of the function
performCallSlotOption, ensuring that the code remains unchanged
if the optimization fails.

Author: kaiyan96

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),

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
+}
+