[MemCpyOpt] Preserve MemorySSA.

[fhahn]

Summary

This patch updates MemCpyOpt to preserve MemorySSA. It uses the
MemoryDef at the insertion point of the builder and inserts the new def
after that def.

In some cases, we just modify a memory instruction. In that case, get
the defining access, then remove the memory access and add a new one.
If the defining access is in a different block, insert a new def at the
beginning of the current block, otherwise after the defining access.

[fhahn]

Fix some bugs in the updating, add test cases.

There is one case where we turn a memmove into a memcpy, which currently is a no-op. We should probably remove the access and re-add it, but I am not sure how to best ensure we re-add the access at the correct position.

[asbirlea]

I was working on having MemCpyOpt preserve and use MSSA, in order to replace the use of MemDepAnalysis entirely. But I'm happy with having MSSA preserved first :-).

Some quick comments below.

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

427

assert (LastDef && "Must exist.");

437

Remove braces.

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

1 ↗

(On Diff #288764)

Typo in file name, rename file to preserve-memssa.

3 ↗

(On Diff #288764)

;REQUIRES: asserts

[fhahn]

In D86651#2251167, @asbirlea wrote:

I was working on having MemCpyOpt preserve and use MSSA, in order to replace the use of MemDepAnalysis entirely. But I'm happy with having MSSA preserved first :-).

I agree that MemCpyOpt is another good candidate the use MSSA and I think it should be possible to share some of the analysis/approach from DSE + MemorySSA to also remove the BB only limitation.

But preserving it is a relatively small step compared to that so we can take the initial step with DSE + MemorySSA :)

Some quick comments below.

Thanks, should be addressed!

[fhahn]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

427

Done, added an assert for both LastAcc and LastDef.

[fhahn]

Add MemorySSAAnalysis to preserved also for NPM.

[asbirlea]

A small note to test with EXPENSIVE_CHECKS as well, as there's extra verification there.

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

325

I'm not sure I fully understand how these two work.
LastDef is the last definition, last access, or the defining access for the new MemoryAccess created below. It should be always updated in the loop. Should it also be updated when BI is not a store or memset but the access is a MemoryDef?
LastAcc is the insertion point (perhaps rename this one?)

The insertion point for the memset in the BB, seems to be where BI left off, so then LastAcc is either the same as LastDef or one before, so it should be updated more often than just for non-store and non-memset?

If my understanding is incomplete, could you please explain?

Unrelated, below, in Ranges, the StartInst is also added and I think that's a store. So LastAcc could be initialized to MSSAU->getMemorySSA()->getMemoryAccess(StartInstI); outside the loop?

1094

Can you add a comment explaining this?
The new memset is inserted after the memcpy, but it is known that its defining access is the memset about to be removed which immediately precedes the memcpy.

[fhahn]

Add additional comments as suggested, thanks!

[fhahn]

Further simplify LastMemDef, MemInsertPoint management.

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

325

I'm not sure I fully understand how these two work.
LastDef is the last definition, last access, or the defining access for the new MemoryAccess created below. It should be always updated in the loop. Should it also be updated when BI is not a store or memset but the access is a MemoryDef?
LastAcc is the insertion point (perhaps rename this one?)

I renamed LastAcc to MemInsertPoint.

The insertion point for the memset in the BB, seems to be where BI left off, so then LastAcc is either the same as LastDef or one before, so it should be updated more often than just for non-store and non-memset?

Yes, MemInsertPoint (LastAcc) should be updated on each iteration. I updated the code. I also added comments explaining in more detail what both variables track. MemInsertPoint is either LastMemDef or the last memory user of before the insertion point for the new memsets.

Unrelated, below, in Ranges, the StartInst is also added and I think that's a store. So LastAcc could be initialized to MSSAU->getMemorySSA()->getMemoryAccess(StartInstI); outside the loop?

Yes I think it could be initialized with StartInstI, but it is set in each loop iteration now and should be much clearer, so I don't think it is necessary to additionally initialize it outside the loop

[asbirlea]

Diff looks good to me.
Same comment as before about testing with EXPENSIVE_CHECKS to catch potential missed updates.
Thank you for working on this!

llvm/include/llvm/Transforms/Scalar/MemCpyOptimizer.h

 class AssumptionCache;
 class CallInst;
 class DominatorTree;
 class Function;
 class Instruction;
 class MemCpyInst;
 class MemMoveInst;
 class MemoryDependenceResults;
+class MemorySSA;
+class MemorySSAUpdater;
 class MemSetInst;
 class StoreInst;
 class TargetLibraryInfo;
 class Value;
 
 class MemCpyOptPass : public PassInfoMixin<MemCpyOptPass> {
   MemoryDependenceResults *MD = nullptr;
   TargetLibraryInfo *TLI = nullptr;
   AliasAnalysis *AA = nullptr;
   AssumptionCache *AC = nullptr;
   DominatorTree *DT = nullptr;
+  MemorySSAUpdater *MSSAU = nullptr;
 
 public:
   MemCpyOptPass() = default;
 
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 
   // Glue for the old PM.
   bool runImpl(Function &F, MemoryDependenceResults *MD_,
                TargetLibraryInfo *TLI_, AliasAnalysis *AA_,
-               AssumptionCache *AC_, DominatorTree *DT_);
+               AssumptionCache *AC_, DominatorTree *DT_, MemorySSA *MSSA_);
 
 private:
   // Helper functions
   bool processStore(StoreInst *SI, BasicBlock::iterator &BBI);
   bool processMemSet(MemSetInst *SI, BasicBlock::iterator &BBI);
   bool processMemCpy(MemCpyInst *M, BasicBlock::iterator &BBI);
   bool processMemMove(MemMoveInst *M);
   bool performCallSlotOptzn(Instruction *cpy, Value *cpyDst, Value *cpySrc,

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/MemoryDependenceAnalysis.h"
 #include "llvm/Analysis/MemoryLocation.h"
+#include "llvm/Analysis/MemorySSA.h"
+#include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addPreserved<DominatorTreeWrapperPass>();
     AU.addPreserved<GlobalsAAWrapperPass>();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
     AU.addRequired<MemoryDependenceWrapperPass>();
     AU.addPreserved<MemoryDependenceWrapperPass>();
     AU.addRequired<AAResultsWrapperPass>();
     AU.addPreserved<AAResultsWrapperPass>();
+    AU.addPreserved<MemorySSAWrapperPass>();
   }
 };
 
 } // end anonymous namespace
 
 char MemCpyOptLegacyPass::ID = 0;
 
 /// The public interface to this file...
Instruction *MemCpyOptPass::tryMergingIntoMemset(Instruction *StartInst,
 
   // Okay, so we now have a single store that can be splatable.  Scan to find
   // all subsequent stores of the same value to offset from the same pointer.
   // Join these together into ranges, so we can decide whether contiguous blocks
   // are stored.
   MemsetRanges Ranges(DL);
 
   BasicBlock::iterator BI(StartInst);
+
+  // Keeps track of the last memory use or def before the insertion point for
+  // the new memset. The new MemoryDef for the inserted memsets will be inserted
+  // after MemInsertPoint. It points to either LastMemDef or to the last user
+  // before the insertion point of the memset, if there are any such users.
+  MemoryUseOrDef *MemInsertPoint = nullptr;
+  // Keeps track of the last MemoryDef between StartInst and the insertion point
+  // for the new memset. This will become the defining access of the inserted
+  // memsets.
+  MemoryDef *LastMemDef = nullptr;
   for (++BI; !BI->isTerminator(); ++BI) {
+    if (MSSAU) {
+      auto *CurrentAcc = cast_or_null<MemoryUseOrDef>(
+          MSSAU->getMemorySSA()->getMemoryAccess(&*BI));
+      if (CurrentAcc) {
+        MemInsertPoint = CurrentAcc;
+        if (auto *CurrentDef = dyn_cast<MemoryDef>(CurrentAcc))
+          LastMemDef = CurrentDef;
+      }
+    }
+
     if (!isa<StoreInst>(BI) && !isa<MemSetInst>(BI)) {
       // If the instruction is readnone, ignore it, otherwise bail out.  We
       // don't even allow readonly here because we don't want something like:
       // A[1] = 2; strlen(A); A[2] = 2; -> memcpy(A, ...); strlen(A).
       if (BI->mayWriteToMemory() || BI->mayReadFromMemory())
         break;
       continue;
     }
for (const MemsetRange &Range : Ranges) {
     StartPtr = Range.StartPtr;
 
     AMemSet = Builder.CreateMemSet(StartPtr, ByteVal, Range.End - Range.Start,
                                    MaybeAlign(Range.Alignment));
     LLVM_DEBUG(dbgs() << "Replace stores:\n"; for (Instruction *SI
                                                    : Range.TheStores) dbgs()
                                               << *SI << '\n';
                dbgs() << "With: " << *AMemSet << '\n');
-
     if (!Range.TheStores.empty())
       AMemSet->setDebugLoc(Range.TheStores[0]->getDebugLoc());
 
+    if (MSSAU) {
+      assert(LastMemDef && MemInsertPoint &&
+             "Both LastMemDef and MemInsertPoint need to be set");
+      auto *NewDef = cast<MemoryDef>(
+          MSSAU->createMemoryAccessAfter(AMemSet, LastMemDef, MemInsertPoint));
+      MSSAU->insertDef(NewDef, /*RenameUses=*/true);
+      LastMemDef = NewDef;
+      MemInsertPoint = NewDef;
+    }
+
     // Zap all the stores.
     for (Instruction *SI : Range.TheStores) {
+      if (MSSAU)
+        MSSAU->removeMemoryAccess(SI);
       MD->removeInstruction(SI);
       SI->eraseFromParent();
     }
+
     ++NumMemSetInfer;
   }
 
   return AMemSet;
 }
 
 // This method try to lift a store instruction before position P.
 // It will lift the store and its argument + that anything that
if (LI->isSimple() && LI->hasOneUse() &&
           else
             M = Builder.CreateMemCpy(
                 SI->getPointerOperand(), SI->getAlign(),
                 LI->getPointerOperand(), LI->getAlign(), Size);
 
           LLVM_DEBUG(dbgs() << "Promoting " << *LI << " to " << *SI << " => "
                             << *M << "\n");
 
+          if (MSSAU) {
+            assert(isa<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(P)));
+            auto *LastDef =
+                cast<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(P));
+            auto *NewAccess =
+                MSSAU->createMemoryAccessAfter(M, LastDef, LastDef);
+            MSSAU->insertDef(cast<MemoryDef>(NewAccess), /*RenameUses=*/true);
+            MSSAU->removeMemoryAccess(SI);
+            MSSAU->removeMemoryAccess(LI);
+          }
+
           MD->removeInstruction(SI);
           SI->eraseFromParent();
           MD->removeInstruction(LI);
           LI->eraseFromParent();
           ++NumMemCpyInstr;
 
           // Make sure we do not invalidate the iterator.
           BBI = M->getIterator();
if (LI->isSimple() && LI->hasOneUse() &&
 
       if (C) {
         bool changed = performCallSlotOptzn(
             LI, SI->getPointerOperand()->stripPointerCasts(),
             LI->getPointerOperand()->stripPointerCasts(),
             DL.getTypeStoreSize(SI->getOperand(0)->getType()),
             commonAlignment(SI->getAlign(), LI->getAlign()), C);
         if (changed) {
+          if (MSSAU) {
+            MSSAU->removeMemoryAccess(SI);
+            MSSAU->removeMemoryAccess(LI);
+          }
+
           MD->removeInstruction(SI);
           SI->eraseFromParent();
           MD->removeInstruction(LI);
           LI->eraseFromParent();
           ++NumMemCpyInstr;
           return true;
         }
       }
if (Value *ByteVal = isBytewiseValue(V, DL)) {
     if (T->isAggregateType()) {
       uint64_t Size = DL.getTypeStoreSize(T);
       IRBuilder<> Builder(SI);
       auto *M = Builder.CreateMemSet(SI->getPointerOperand(), ByteVal, Size,
                                      SI->getAlign());
 
       LLVM_DEBUG(dbgs() << "Promoting " << *SI << " to " << *M << "\n");
 
+      if (MSSAU) {
+        assert(isa<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(SI)));
+        auto *LastDef =
+            cast<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(SI));
+        auto *NewAccess = MSSAU->createMemoryAccessAfter(M, LastDef, LastDef);
+        MSSAU->insertDef(cast<MemoryDef>(NewAccess), /*RenameUses=*/true);
+        MSSAU->removeMemoryAccess(SI);
+      }
+
       MD->removeInstruction(SI);
       SI->eraseFromParent();
       NumMemSetInfer++;
 
       // Make sure we do not invalidate the iterator.
       BBI = M->getIterator();
       return true;
     }
bool MemCpyOptPass::processMemCpyMemCpyDependence(MemCpyInst *M,
 
   // If all checks passed, then we can transform M.
   LLVM_DEBUG(dbgs() << "MemCpyOptPass: Forwarding memcpy->memcpy src:\n"
                     << *MDep << '\n' << *M << '\n');
 
   // TODO: Is this worth it if we're creating a less aligned memcpy? For
   // example we could be moving from movaps -> movq on x86.
   IRBuilder<> Builder(M);
+  Instruction *NewM;
   if (UseMemMove)
-    Builder.CreateMemMove(M->getRawDest(), M->getDestAlign(),
+    NewM = Builder.CreateMemMove(M->getRawDest(), M->getDestAlign(),
                                  MDep->getRawSource(), MDep->getSourceAlign(),
                                  M->getLength(), M->isVolatile());
   else
-    Builder.CreateMemCpy(M->getRawDest(), M->getDestAlign(),
+    NewM = Builder.CreateMemCpy(M->getRawDest(), M->getDestAlign(),
                                 MDep->getRawSource(), MDep->getSourceAlign(),
                                 M->getLength(), M->isVolatile());
 
+  if (MSSAU) {
+    assert(isa<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(M)));
+    auto *LastDef = cast<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(M));
+    auto *NewAccess = MSSAU->createMemoryAccessAfter(NewM, LastDef, LastDef);
+    MSSAU->insertDef(cast<MemoryDef>(NewAccess), /*RenameUses=*/true);
+    MSSAU->removeMemoryAccess(M);
+  }
+
   // Remove the instruction we're replacing.
   MD->removeInstruction(M);
   M->eraseFromParent();
   ++NumMemCpyInstr;
   return true;
 }
 
 /// We've found that the (upward scanning) memory dependence of \p MemCpy is
if (DestSize->getType() != SrcSize->getType()) {
     else
       DestSize = Builder.CreateZExt(DestSize, SrcSize->getType());
   }
 
   Value *Ule = Builder.CreateICmpULE(DestSize, SrcSize);
   Value *SizeDiff = Builder.CreateSub(DestSize, SrcSize);
   Value *MemsetLen = Builder.CreateSelect(
       Ule, ConstantInt::getNullValue(DestSize->getType()), SizeDiff);
-  Builder.CreateMemSet(
+  Instruction *NewMemSet = Builder.CreateMemSet(
       Builder.CreateGEP(Dest->getType()->getPointerElementType(), Dest,
                         SrcSize),
       MemSet->getOperand(1), MemsetLen, MaybeAlign(Align));
 
+  if (MSSAU) {
+    assert(isa<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(MemCpy)) &&
+           "MemCpy must be a MemoryDef");
+    // The new memset is inserted after the memcpy, but it is known that its
+    // defining access is the memset about to be removed which immediately
+    // precedes the memcpy.
+    auto *LastDef =
+        cast<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(MemCpy));
+    auto *NewAccess = MSSAU->createMemoryAccessBefore(
+        NewMemSet, LastDef->getDefiningAccess(), LastDef);
+    MSSAU->insertDef(cast<MemoryDef>(NewAccess), /*RenameUses=*/true);
+    MSSAU->removeMemoryAccess(MemSet);
+  }
+
   MD->removeInstruction(MemSet);
   MemSet->eraseFromParent();
   return true;
 }
 
 /// Determine whether the instruction has undefined content for the given Size,
 /// either because it was freshly alloca'd or started its lifetime.
 static bool hasUndefContents(Instruction *I, ConstantInt *Size) {
MemDepResult DepInfo = MD->getPointerDependencyFrom(
         MemCpyLoc, true, MemSet->getIterator(), MemSet->getParent());
     if (DepInfo.isDef() && hasUndefContents(DepInfo.getInst(), CopySize))
       CopySize = MemSetSize;
     else
       return false;
   }
 
   IRBuilder<> Builder(MemCpy);
-  Builder.CreateMemSet(MemCpy->getRawDest(), MemSet->getOperand(1), CopySize,
-                       MaybeAlign(MemCpy->getDestAlignment()));
+  Instruction *NewM =
+      Builder.CreateMemSet(MemCpy->getRawDest(), MemSet->getOperand(1),
+                           CopySize, MaybeAlign(MemCpy->getDestAlignment()));
+  if (MSSAU) {
+    auto *LastDef =
+        cast<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(MemCpy));
+    auto *NewAccess = MSSAU->createMemoryAccessAfter(NewM, LastDef, LastDef);
+    MSSAU->insertDef(cast<MemoryDef>(NewAccess), /*RenameUses=*/true);
+  }
+
   return true;
 }
 
 /// Perform simplification of memcpy's.  If we have memcpy A
 /// which copies X to Y, and memcpy B which copies Y to Z, then we can rewrite
 /// B to be a memcpy from X to Z (or potentially a memmove, depending on
 /// circumstances). This allows later passes to remove the first memcpy
 /// altogether.
 bool MemCpyOptPass::processMemCpy(MemCpyInst *M, BasicBlock::iterator &BBI) {
   // We can only optimize non-volatile memcpy's.
   if (M->isVolatile()) return false;
 
   // If the source and destination of the memcpy are the same, then zap it.
   if (M->getSource() == M->getDest()) {
     ++BBI;
+    if (MSSAU)
+      MSSAU->removeMemoryAccess(M);
+
     MD->removeInstruction(M);
     M->eraseFromParent();
     return true;
   }
 
   // If copying from a constant, try to turn the memcpy into a memset.
   if (GlobalVariable *GV = dyn_cast<GlobalVariable>(M->getSource()))
     if (GV->isConstant() && GV->hasDefinitiveInitializer())
       if (Value *ByteVal = isBytewiseValue(GV->getInitializer(),
                                            M->getModule()->getDataLayout())) {
         IRBuilder<> Builder(M);
+        Instruction *NewM =
             Builder.CreateMemSet(M->getRawDest(), ByteVal, M->getLength(),
                                  MaybeAlign(M->getDestAlignment()), false);
+        if (MSSAU) {
+          auto *LastDef =
+              cast<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(M));
+          auto *NewAccess =
+              MSSAU->createMemoryAccessAfter(NewM, LastDef, LastDef);
+          MSSAU->insertDef(cast<MemoryDef>(NewAccess), /*RenameUses=*/true);
+          MSSAU->removeMemoryAccess(M);
+        }
+
         MD->removeInstruction(M);
         M->eraseFromParent();
         ++NumCpyToSet;
         return true;
       }
 
   MemDepResult DepInfo = MD->getDependency(M);
 
bool MemCpyOptPass::processMemCpy(MemCpyInst *M, BasicBlock::iterator &BBI) {
   if (DepInfo.isClobber()) {
     if (CallInst *C = dyn_cast<CallInst>(DepInfo.getInst())) {
       // FIXME: Can we pass in either of dest/src alignment here instead
       // of conservatively taking the minimum?
       Align Alignment = std::min(M->getDestAlign().valueOrOne(),
                                  M->getSourceAlign().valueOrOne());
       if (performCallSlotOptzn(M, M->getDest(), M->getSource(),
                                CopySize->getZExtValue(), Alignment, C)) {
+        if (MSSAU)
+          MSSAU->removeMemoryAccess(M);
+
         MD->removeInstruction(M);
         M->eraseFromParent();
         return true;
       }
     }
   }
 
   MemoryLocation SrcLoc = MemoryLocation::getForSource(M);
   MemDepResult SrcDepInfo = MD->getPointerDependencyFrom(
       SrcLoc, true, M->getIterator(), M->getParent());
 
   if (SrcDepInfo.isClobber()) {
     if (MemCpyInst *MDep = dyn_cast<MemCpyInst>(SrcDepInfo.getInst()))
       return processMemCpyMemCpyDependence(M, MDep);
   } else if (SrcDepInfo.isDef()) {
     if (hasUndefContents(SrcDepInfo.getInst(), CopySize)) {
+      if (MSSAU)
+        MSSAU->removeMemoryAccess(M);
+
       MD->removeInstruction(M);
       M->eraseFromParent();
       ++NumMemCpyInstr;
       return true;
     }
   }
 
   if (SrcDepInfo.isClobber())
     if (MemSetInst *MDep = dyn_cast<MemSetInst>(SrcDepInfo.getInst()))
       if (performMemCpyToMemSetOptzn(M, MDep)) {
+        if (MSSAU)
+          MSSAU->removeMemoryAccess(M);
         MD->removeInstruction(M);
         M->eraseFromParent();
         ++NumCpyToSet;
         return true;
       }
 
   return false;
 }
bool MemCpyOptPass::processMemMove(MemMoveInst *M) {
 
   // If not, then we know we can transform this.
   Type *ArgTys[3] = { M->getRawDest()->getType(),
                       M->getRawSource()->getType(),
                       M->getLength()->getType() };
   M->setCalledFunction(Intrinsic::getDeclaration(M->getModule(),
                                                  Intrinsic::memcpy, ArgTys));
 
+  // For MemorySSA nothing really changes (except that memcpy may imply stricter
+  // aliasing guarantees).
+
   // MemDep may have over conservative information about this instruction, just
   // conservatively flush it from the cache.
   MD->removeInstruction(M);
 
   ++NumMoveToCpy;
   return true;
 }
 

 }
 
 PreservedAnalyses MemCpyOptPass::run(Function &F, FunctionAnalysisManager &AM) {
   auto &MD = AM.getResult<MemoryDependenceAnalysis>(F);
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
   auto *AA = &AM.getResult<AAManager>(F);
   auto *AC = &AM.getResult<AssumptionAnalysis>(F);
   auto *DT = &AM.getResult<DominatorTreeAnalysis>(F);
+  auto *MSSA = AM.getCachedResult<MemorySSAAnalysis>(F);
 
-  bool MadeChange = runImpl(F, &MD, &TLI, AA, AC, DT);
+  bool MadeChange =
+      runImpl(F, &MD, &TLI, AA, AC, DT, MSSA ? &MSSA->getMSSA() : nullptr);
   if (!MadeChange)
     return PreservedAnalyses::all();
 
   PreservedAnalyses PA;
   PA.preserveSet<CFGAnalyses>();
   PA.preserve<GlobalsAA>();
   PA.preserve<MemoryDependenceAnalysis>();
+  if (MSSA)
+    PA.preserve<MemorySSAAnalysis>();
   return PA;
 }
 
 bool MemCpyOptPass::runImpl(Function &F, MemoryDependenceResults *MD_,
                             TargetLibraryInfo *TLI_, AliasAnalysis *AA_,
-                            AssumptionCache *AC_, DominatorTree *DT_) {
+                            AssumptionCache *AC_, DominatorTree *DT_,
+                            MemorySSA *MSSA_) {
   bool MadeChange = false;
   MD = MD_;
   TLI = TLI_;
   AA = AA_;
   AC = AC_;
   DT = DT_;
+  MemorySSAUpdater MSSAU_(MSSA_);
+  MSSAU = MSSA_ ? &MSSAU_ : nullptr;
   // If we don't have at least memset and memcpy, there is little point of doing
   // anything here.  These are required by a freestanding implementation, so if
   // even they are disabled, there is no point in trying hard.
   if (!TLI->has(LibFunc_memset) || !TLI->has(LibFunc_memcpy))
     return false;
 
   while (true) {
     if (!iterateOnFunction(F))
       break;
     MadeChange = true;
   }
 
+  if (MSSA_ && VerifyMemorySSA)
+    MSSA_->verifyMemorySSA();
+
   MD = nullptr;
   return MadeChange;
 }
 
 /// This is the main transformation entry point for a function.
 bool MemCpyOptLegacyPass::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;
 
   auto *MD = &getAnalysis<MemoryDependenceWrapperPass>().getMemDep();
   auto *TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
   auto *AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
   auto *AC = &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
   auto *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+  auto *MSSAWP = getAnalysisIfAvailable<MemorySSAWrapperPass>();
 
-  return Impl.runImpl(F, MD, TLI, AA, AC, DT);
+  return Impl.runImpl(F, MD, TLI, AA, AC, DT,
+                      MSSAWP ? &MSSAWP->getMSSA() : nullptr);
 }

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

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -aa-pipeline=basic-aa -passes='require<memoryssa>,memcpyopt' -verify-memoryssa -S %s | FileCheck %s
+
+; REQUIRES: asserts
+
+target datalayout = "e-m:o-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.15.0"
+
+%t = type <{ i8*, [4 x i8], i8*, i8*, i32, [8192 x i8] }>
+
+
+define i32 @test1(%t* %ptr) {
+; CHECK-LABEL: @test1(
+; CHECK-NEXT:  invoke.cont6:
+; CHECK-NEXT:    [[P_1:%.*]] = getelementptr inbounds [[T:%.*]], %t* [[PTR:%.*]], i64 0, i32 0
+; CHECK-NEXT:    [[P_1_C:%.*]] = bitcast i8** [[P_1]] to i8*
+; CHECK-NEXT:    [[P_2:%.*]] = getelementptr inbounds [[T]], %t* [[PTR]], i64 0, i32 4
+; CHECK-NEXT:    [[P_3:%.*]] = getelementptr inbounds [[T]], %t* [[PTR]], i64 0, i32 5, i64 0
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i8** [[P_1]] to i8*
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 8 [[TMP0]], i8 0, i64 20, i1 false)
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i32* [[P_2]] to i8*
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 8 [[TMP1]], i8 0, i64 8195, i1 false)
+; CHECK-NEXT:    ret i32 0
+;
+invoke.cont6:
+  %p.1 = getelementptr inbounds %t, %t* %ptr, i64 0, i32 0
+  %p.1.c = bitcast i8** %p.1 to i8*
+  call void @llvm.memset.p0i8.i64(i8* %p.1.c, i8 0, i64 20, i1 false)
+  store i8* null, i8** %p.1, align 8
+  %p.2 = getelementptr inbounds %t, %t* %ptr, i64 0, i32 4
+  store i32 0, i32* %p.2, align 8
+  %p.3 = getelementptr inbounds %t, %t* %ptr, i64 0, i32 5, i64 0
+  call void @llvm.memset.p0i8.i64(i8* %p.3, i8 0, i64 8191, i1 false)
+  ret i32 0
+}
+
+declare i8* @get_ptr()
+
+define void @test2(i8 *%in) {
+; CHECK-LABEL: @test2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CALL_I1_I:%.*]] = tail call i8* @get_ptr()
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, i8* [[CALL_I1_I]], i64 10
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 1 [[TMP0]], i8 0, i64 0, i1 false)
+; CHECK-NEXT:    tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* [[CALL_I1_I]], i8* [[IN:%.*]], i64 10, i1 false)
+; CHECK-NEXT:    ret void
+;
+entry:
+  %call.i1.i = tail call i8* @get_ptr()
+  tail call void @llvm.memset.p0i8.i64(i8* %call.i1.i, i8 0, i64 10, i1 false)
+  tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* %call.i1.i, i8* %in, i64 10, i1 false)
+  ret void
+}
+
+declare i8* @malloc(i64)
+
+define i32 @test3(i8* noalias %in) {
+; CHECK-LABEL: @test3(
+; CHECK-NEXT:    [[CALL_I_I_I:%.*]] = tail call i8* @malloc(i64 20)
+; CHECK-NEXT:    tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* [[CALL_I_I_I]], i8* [[IN:%.*]], i64 20, i1 false)
+; CHECK-NEXT:    ret i32 10
+;
+  %call.i.i.i = tail call i8* @malloc(i64 20)
+  tail call void @llvm.memmove.p0i8.p0i8.i64(i8* %call.i.i.i, i8* %in, i64 20, i1 false)
+  ret i32 10
+}
+
+define void @test4(i32 %n, i8* noalias %ptr.0, i8* noalias %ptr.1, i32* %ptr.2) unnamed_addr {
+; CHECK-LABEL: @test4(
+; CHECK-NEXT:    [[ELEM_I:%.*]] = getelementptr i8, i8* [[PTR_0:%.*]], i64 8
+; CHECK-NEXT:    store i32 [[N:%.*]], i32* [[PTR_2:%.*]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i8, i8* [[ELEM_I]], i64 10
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 1 [[TMP1]], i8 0, i64 0, i1 false)
+; CHECK-NEXT:    call void @llvm.memcpy.p0i8.p0i8.i64(i8* [[ELEM_I]], i8* [[PTR_1:%.*]], i64 10, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %elem.i = getelementptr i8, i8* %ptr.0, i64 8
+  call void @llvm.memset.p0i8.i64(i8* %elem.i, i8 0, i64 10, i1 false)
+  store i32 %n, i32* %ptr.2, align 8
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* %elem.i, i8* %ptr.1, i64 10, i1 false)
+  ret void
+}
+
+declare void @decompose(%t* nocapture)
+
+define void @test5(i32* %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:    [[TMP0:%.*]] = bitcast [128 x i8]* [[EARLY_DATA]] to i8*
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast %t* [[TMP]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 32, i8* [[TMP0]])
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* [[PTR:%.*]], align 8
+; CHECK-NEXT:    call fastcc void @decompose(%t* [[TMP]])
+; CHECK-NEXT:    call void @llvm.memcpy.p0i8.p0i8.i64(i8* [[TMP0]], i8* [[TMP1]], i64 32, i1 false)
+; CHECK-NEXT:    ret void
+;
+entry:
+  %early_data = alloca [128 x i8], align 8
+  %tmp = alloca %t, align 8
+  %0 = bitcast [128 x i8]* %early_data to i8*
+  %1 = bitcast %t* %tmp to i8*
+  call void @llvm.lifetime.start.p0i8(i64 32, i8* %0)
+  %2 = load i32, i32* %ptr, align 8
+  call fastcc void @decompose(%t* %tmp)
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* %0, i8* %1, i64 32, i1 false)
+  ret void
+}
+
+define i8 @test6(i8* %ptr, i8* noalias %ptr.1) {
+; CHECK-LABEL: @test6(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 24, i8* [[PTR:%.*]])
+; CHECK-NEXT:    [[TMP0:%.*]] = load i8, i8* [[PTR]], align 8
+; CHECK-NEXT:    call void @llvm.memcpy.p0i8.p0i8.i64(i8* [[PTR]], i8* [[PTR_1:%.*]], i64 24, i1 false)
+; CHECK-NEXT:    ret i8 [[TMP0]]
+;
+entry:
+  call void @llvm.lifetime.start.p0i8(i64 24, i8* %ptr)
+  %0 = load i8, i8* %ptr, align 8
+  call void @llvm.memmove.p0i8.p0i8.i64(i8* %ptr, i8* %ptr.1, i64 24, i1 false)
+  ret i8 %0
+}
+
+; Function Attrs: argmemonly nounwind willreturn
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture) #0
+
+; Function Attrs: argmemonly nounwind willreturn
+declare void @llvm.memcpy.p0i8.p0i8.i64(i8* noalias nocapture writeonly, i8* noalias nocapture readonly, i64, i1 immarg) #0
+
+; Function Attrs: argmemonly nounwind willreturn writeonly
+declare void @llvm.memset.p0i8.i64(i8* nocapture writeonly, i8, i64, i1 immarg) #1
+
+; Function Attrs: argmemonly nounwind willreturn
+declare void @llvm.memmove.p0i8.p0i8.i64(i8* nocapture, i8* nocapture readonly, i64, i1 immarg) #0
+
+attributes #0 = { argmemonly nounwind willreturn }
+attributes #1 = { argmemonly nounwind willreturn writeonly }