[memcpyopt] allow more memcpy-to-memcpy optimziation

Allow the memcpy-to-memcpy optimization even when the sizes are not
identical. For example, it might have been generated as a small slice of
a larger struct (currently only for zero offset however), or might be
only storing to part of an oversized alloca.

Author: vtjnash

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

@@ -23,6 +23,7 @@
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/GlobalsModRef.h"
+#include "llvm/Analysis/InstSimplifyFolder.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/MemoryLocation.h"
@@ -1431,6 +1432,28 @@ static bool overreadUndefContents(MemorySSA *MSSA, MemCpyInst *MemCpy,
   return false;
 }
 
+// If only the MemSrc instruction is known, a similar but slightly weaker
+// analysis can apply
+static bool anyOverreadUndefContents(MemorySSA *MSSA, Instruction *Store,
+                                     BatchAAResults &BAA) {
+  MemoryLocation Loc;
+  Value *Ptr;
+  if (auto SI = dyn_cast<StoreInst>(Store)) {
+    Loc = MemoryLocation::get(SI);
+    Ptr = SI->getPointerOperand();
+  } else if (auto MI = dyn_cast<MemCpyInst>(Store)) {
+    Loc = MemoryLocation::getForDest(MI);
+    Ptr = MI->getDest();
+  } else {
+    llvm_unreachable("performStackMoveOptzn must have a known store kind");
+  }
+  MemoryAccess *MemAccess = MSSA->getMemoryAccess(Store)->getDefiningAccess();
+  if (hadUndefContentsBefore(MSSA, BAA, Ptr, MemAccess, Loc, nullptr))
+    return true;
+  return false;
+}
+
+
 /// Transform memcpy to memset when its source was just memset.
 /// In other words, turn:
 /// \code
@@ -1524,30 +1547,49 @@ bool MemCpyOptPass::performStackMoveOptzn(Instruction *Load, Instruction *Store,
     return false;
   }
 
-  // Check that copy is full with static size.
-  const DataLayout &DL = DestAlloca->getDataLayout();
-  std::optional<TypeSize> SrcSize = SrcAlloca->getAllocationSize(DL);
-  if (!SrcSize || Size != *SrcSize) {
-    LLVM_DEBUG(dbgs() << "Stack Move: Source alloca size mismatch\n");
-    return false;
-  }
-  std::optional<TypeSize> DestSize = DestAlloca->getAllocationSize(DL);
-  if (!DestSize || Size != *DestSize) {
-    LLVM_DEBUG(dbgs() << "Stack Move: Destination alloca size mismatch\n");
-    return false;
-  }
-
   if (!SrcAlloca->isStaticAlloca() || !DestAlloca->isStaticAlloca())
     return false;
 
+  Type *SrcType = SrcAlloca->getAllocatedType();
+  Type *DestType = DestAlloca->getAllocatedType();
+  // If they don't have common type, then they will need to be converted to a
+  // common size at runtime
+  const auto &DL = SrcAlloca->getDataLayout();
+  TypeSize SrcSize = DL.getTypeAllocSize(SrcType);
+  TypeSize DestSize = DL.getTypeAllocSize(DestType);
+  if (SrcType != DestType)
+    if (SrcSize != DestSize)
+      if (!SrcSize.isFixed() || !DestSize.isFixed())
+        return false;
+
+  // Check that copy is full with dest size, either because it wrote every byte,
+  // or it was fresh.
+  std::optional<TypeSize> FullSize = DestAlloca->getAllocationSize(DL);
+  if (!FullSize || Size != *FullSize)
+    if (!anyOverreadUndefContents(MSSA, Store, BAA)) {
+      LLVM_DEBUG(dbgs() << "Stack Move: Destination alloca size mismatch\n");
+      return false;
+    }
+
+  // Check if it will be legal to combine allocas without breaking dominator.
+  // TODO: Try to hoist the arguments (recursively) instead of giving up
+  // immediately.
+  bool MoveSrc = !DT->dominates(SrcAlloca, DestAlloca);
+  if (MoveSrc) {
+    if (!DT->dominates(SrcAlloca->getArraySize(), DestAlloca))
+      return false;
+  } else {
+    if (!DT->dominates(DestAlloca->getArraySize(), SrcAlloca))
+      return false;
+  }
+
   // Check that src and dest are never captured, unescaped allocas. Also
   // find the nearest common dominator and postdominator for all users in
   // order to shrink wrap the lifetimes, and instructions with noalias metadata
   // to remove them.
 
   SmallVector<Instruction *, 4> LifetimeMarkers;
   SmallSet<Instruction *, 4> AAMetadataInstrs;
-  bool SrcNotDom = false;
 
   auto CaptureTrackingWithModRef =
       [&](Instruction *AI,
@@ -1561,10 +1603,6 @@ bool MemCpyOptPass::performStackMoveOptzn(Instruction *Load, Instruction *Store,
       Instruction *I = Worklist.pop_back_val();
       for (const Use &U : I->uses()) {
         auto *UI = cast<Instruction>(U.getUser());
-        // 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(
@@ -1678,15 +1716,43 @@ bool MemCpyOptPass::performStackMoveOptzn(Instruction *Load, Instruction *Store,
   if (!CaptureTrackingWithModRef(SrcAlloca, SrcModRefCallback))
     return false;
 
-  // We can do the transformation. First, move the SrcAlloca to the start of the
-  // BB.
-  if (SrcNotDom)
-    SrcAlloca->moveBefore(*SrcAlloca->getParent(),
-                          SrcAlloca->getParent()->getFirstInsertionPt());
+  // We can now do the transformation. First move the Src if it was after Dest.
+  if (MoveSrc)
+    SrcAlloca->moveBefore(DestAlloca->getIterator());
+
   // Align the allocas appropriately.
   SrcAlloca->setAlignment(
       std::max(SrcAlloca->getAlign(), DestAlloca->getAlign()));
 
+  // Size the allocas appropriately.
+  Value *SrcArraySize = SrcAlloca->getArraySize();
+  Value *DestArraySize = DestAlloca->getArraySize();
+  IRBuilder<InstSimplifyFolder> Builder(SrcAlloca->getContext(),
+                                        InstSimplifyFolder(DL));
+  Builder.SetInsertPoint(SrcAlloca);
+  Type *Int32Ty = Builder.getInt32Ty();
+  if (SrcType != DestType && SrcSize != DestSize) {
+    SrcAlloca->setAllocatedType(Type::getInt8Ty(Load->getContext()));
+    if (SrcArraySize->getType() != Int32Ty)
+      SrcArraySize = Builder.CreateZExtOrTrunc(SrcArraySize, Int32Ty);
+    if (DestArraySize->getType() != Int32Ty)
+      DestArraySize = Builder.CreateZExtOrTrunc(DestArraySize, Int32Ty);
+    SrcArraySize = Builder.CreateMul(
+        SrcArraySize, ConstantInt::get(Int32Ty, SrcSize.getFixedValue()), "",
+        true, true);
+    DestArraySize = Builder.CreateMul(
+        DestArraySize, ConstantInt::get(Int32Ty, DestSize.getFixedValue()), "",
+        true, true);
+  }
+  if (SrcArraySize != DestArraySize) {
+    if (SrcArraySize->getType() != DestArraySize->getType()) {
+      SrcArraySize = Builder.CreateZExtOrTrunc(SrcArraySize, Int32Ty);
+      DestArraySize = Builder.CreateZExtOrTrunc(DestArraySize, Int32Ty);
+    }
+    SrcAlloca->setOperand(0, Builder.CreateBinaryIntrinsic(
+                                 Intrinsic::umax, SrcArraySize, DestArraySize));
+  }
+
   // Merge the two allocas.
   DestAlloca->replaceAllUsesWith(SrcAlloca);
   eraseInstruction(DestAlloca);

llvm/test/Transforms/MemCpyOpt/preserve-memssa.ll

@@ -76,12 +76,9 @@ declare void @decompose(ptr nocapture)
 define void @test5(ptr %ptr) {
 ; CHECK-LABEL: @test5(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[EARLY_DATA:%.*]] = alloca [128 x i8], align 8
-; CHECK-NEXT:    [[TMP:%.*]] = alloca [[T:%.*]], align 8
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 32, ptr [[EARLY_DATA]])
+; CHECK-NEXT:    [[TMP:%.*]] = alloca i8, i32 8224, align 8
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i32, ptr [[PTR:%.*]], align 8
 ; CHECK-NEXT:    call fastcc void @decompose(ptr [[TMP]])
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[EARLY_DATA]], ptr [[TMP]], i64 32, i1 false)
 ; CHECK-NEXT:    ret void
 ;
 entry:

llvm/test/Transforms/MemCpyOpt/stack-move.ll

@@ -1029,16 +1029,10 @@ bb2:
 ; for the purposes of liveness analysis, not a definition.
 define void @incomplete_memcpy() {
 ; CHECK-LABEL: define void @incomplete_memcpy() {
-; CHECK-NEXT:    [[SRC:%.*]] = alloca [[STRUCT_FOO:%.*]], align 4
-; CHECK-NEXT:    [[DEST:%.*]] = alloca [[STRUCT_FOO]], align 4
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 12, ptr captures(none) [[SRC]])
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 12, ptr captures(none) [[DEST]])
-; CHECK-NEXT:    store [[STRUCT_FOO]] { i32 10, i32 20, i32 30 }, ptr [[SRC]], align 4
-; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @use_nocapture(ptr noundef captures(none) [[SRC]])
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr align 4 [[DEST]], ptr align 4 [[SRC]], i64 11, i1 false)
+; CHECK-NEXT:    [[DEST:%.*]] = alloca [[STRUCT_FOO:%.*]], align 4
+; CHECK-NEXT:    store [[STRUCT_FOO]] { i32 10, i32 20, i32 30 }, ptr [[DEST]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @use_nocapture(ptr noundef captures(none) [[DEST]])
 ; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @use_nocapture(ptr noundef captures(none) [[DEST]])
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 12, ptr captures(none) [[SRC]])
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 12, ptr captures(none) [[DEST]])
 ; CHECK-NEXT:    ret void
 ;
   %src = alloca %struct.Foo, align 4
@@ -1058,17 +1052,10 @@ define void @incomplete_memcpy() {
 ; for the purposes of liveness analysis, not a definition.
 define void @incomplete_store() {
 ; CHECK-LABEL: define void @incomplete_store() {
-; CHECK-NEXT:    [[SRC:%.*]] = alloca [[STRUCT_FOO:%.*]], align 4
-; CHECK-NEXT:    [[DEST:%.*]] = alloca [[STRUCT_FOO]], align 4
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 12, ptr captures(none) [[SRC]])
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 12, ptr captures(none) [[DEST]])
-; CHECK-NEXT:    store [[STRUCT_FOO]] { i32 10, i32 20, i32 30 }, ptr [[SRC]], align 4
-; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @use_nocapture(ptr noundef captures(none) [[SRC]])
-; CHECK-NEXT:    [[TMP2:%.*]] = load i32, ptr [[SRC]], align 4
-; CHECK-NEXT:    store i32 [[TMP2]], ptr [[DEST]], align 4
+; CHECK-NEXT:    [[DEST:%.*]] = alloca [[STRUCT_FOO:%.*]], align 4
+; CHECK-NEXT:    store [[STRUCT_FOO]] { i32 10, i32 20, i32 30 }, ptr [[DEST]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @use_nocapture(ptr noundef captures(none) [[DEST]])
 ; CHECK-NEXT:    [[TMP3:%.*]] = call i32 @use_nocapture(ptr noundef captures(none) [[DEST]])
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 12, ptr captures(none) [[SRC]])
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 12, ptr captures(none) [[DEST]])
 ; CHECK-NEXT:    ret void
 ;
   %src = alloca %struct.Foo, align 4
@@ -1182,15 +1169,9 @@ define void @dynamically_sized_memcpy(i64 %size) {
 define void @mismatched_alloca_size() {
 ; CHECK-LABEL: define void @mismatched_alloca_size() {
 ; CHECK-NEXT:    [[SRC:%.*]] = alloca i8, i64 24, align 4
-; CHECK-NEXT:    [[DEST:%.*]] = alloca i8, i64 12, align 4
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 24, ptr captures(none) [[SRC]])
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 12, ptr captures(none) [[DEST]])
 ; CHECK-NEXT:    store [[STRUCT_FOO:%.*]] { i32 10, i32 20, i32 30 }, ptr [[SRC]], align 4
 ; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @use_nocapture(ptr captures(none) [[SRC]])
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr align 4 [[DEST]], ptr align 4 [[SRC]], i64 12, i1 false)
-; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @use_nocapture(ptr captures(none) [[DEST]])
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 24, ptr captures(none) [[SRC]])
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 12, ptr captures(none) [[DEST]])
+; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @use_nocapture(ptr captures(none) [[SRC]])
 ; CHECK-NEXT:    ret void
 ;
   %src = alloca i8, i64 24, align 4