[MemCpyOpt] Use EarliestEscapeInfo (#110280)

Pass EarliestEscapeInfo to BatchAA in MemCpyOpt. This allows memcpy
elimination in cases where one of the involved pointers is captured
after the relevant memcpy/call.

Author: nikic

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

@@ -26,6 +26,7 @@ class AssumptionCache;
 class CallBase;
 class CallInst;
 class DominatorTree;
+class EarliestEscapeInfo;
 class Function;
 class Instruction;
 class LoadInst;
@@ -48,6 +49,7 @@ class MemCpyOptPass : public PassInfoMixin<MemCpyOptPass> {
   PostDominatorTree *PDT = nullptr;
   MemorySSA *MSSA = nullptr;
   MemorySSAUpdater *MSSAU = nullptr;
+  EarliestEscapeInfo *EEI = nullptr;
 
 public:
   MemCpyOptPass() = default;

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

@@ -284,6 +284,7 @@ static bool mayBeVisibleThroughUnwinding(Value *V, Instruction *Start,
 
 void MemCpyOptPass::eraseInstruction(Instruction *I) {
   MSSAU->removeMemoryAccess(I);
+  EEI->removeInstruction(I);
   I->eraseFromParent();
 }
 
@@ -638,7 +639,7 @@ bool MemCpyOptPass::processStoreOfLoad(StoreInst *SI, LoadInst *LI,
   if (!LI->isSimple() || !LI->hasOneUse() || LI->getParent() != SI->getParent())
     return false;
 
-  BatchAAResults BAA(*AA);
+  BatchAAResults BAA(*AA, EEI);
   auto *T = LI->getType();
   // Don't introduce calls to memcpy/memmove intrinsics out of thin air if
   // the corresponding libcalls are not available.
@@ -1147,14 +1148,14 @@ bool MemCpyOptPass::processMemCpyMemCpyDependence(MemCpyInst *M,
   IRBuilder<> Builder(M);
   auto *CopySource = MDep->getSource();
   Instruction *NewCopySource = nullptr;
-  auto CleanupOnRet = llvm::make_scope_exit([&NewCopySource] {
+  auto CleanupOnRet = llvm::make_scope_exit([&] {
     if (NewCopySource && NewCopySource->use_empty())
       // Safety: It's safe here because we will only allocate more instructions
       // after finishing all BatchAA queries, but we have to be careful if we
       // want to do something like this in another place. Then we'd probably
       // have to delay instruction removal until all transforms on an
       // instruction finished.
-      NewCopySource->eraseFromParent();
+      eraseInstruction(NewCopySource);
   });
   MaybeAlign CopySourceAlign = MDep->getSourceAlign();
   // We just need to calculate the actual size of the copy.
@@ -1751,7 +1752,7 @@ bool MemCpyOptPass::processMemCpy(MemCpyInst *M, BasicBlock::iterator &BBI) {
         return true;
       }
 
-  BatchAAResults BAA(*AA);
+  BatchAAResults BAA(*AA, EEI);
   // FIXME: Not using getClobberingMemoryAccess() here due to PR54682.
   MemoryAccess *AnyClobber = MA->getDefiningAccess();
   MemoryLocation DestLoc = MemoryLocation::getForDest(M);
@@ -1876,7 +1877,7 @@ bool MemCpyOptPass::processByValArgument(CallBase &CB, unsigned ArgNo) {
   if (!CallAccess)
     return false;
   MemCpyInst *MDep = nullptr;
-  BatchAAResults BAA(*AA);
+  BatchAAResults BAA(*AA, EEI);
   MemoryAccess *Clobber = MSSA->getWalker()->getClobberingMemoryAccess(
       CallAccess->getDefiningAccess(), Loc, BAA);
   if (auto *MD = dyn_cast<MemoryDef>(Clobber))
@@ -1949,7 +1950,7 @@ bool MemCpyOptPass::processByValArgument(CallBase &CB, unsigned ArgNo) {
 /// 4. The memcpy src is not modified during the call. (ModRef check shows no
 /// Mod.)
 bool MemCpyOptPass::processImmutArgument(CallBase &CB, unsigned ArgNo) {
-  BatchAAResults BAA(*AA);
+  BatchAAResults BAA(*AA, EEI);
   Value *ImmutArg = CB.getArgOperand(ArgNo);
 
   // 1. Ensure passed argument is immutable during call.
@@ -2117,6 +2118,8 @@ bool MemCpyOptPass::runImpl(Function &F, TargetLibraryInfo *TLI_,
   MSSA = MSSA_;
   MemorySSAUpdater MSSAU_(MSSA_);
   MSSAU = &MSSAU_;
+  EarliestEscapeInfo EEI_(*DT);
+  EEI = &EEI_;
 
   while (true) {
     if (!iterateOnFunction(F))

llvm/test/Transforms/MemCpyOpt/memcpy.ll

@@ -837,7 +837,7 @@ define void @memcpy_memcpy_escape_after1(ptr noalias %P, ptr noalias %Q) {
 ; CHECK-NEXT:    [[MEMTMP:%.*]] = alloca [32 x i8], align 16
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i32(ptr align 16 [[MEMTMP]], ptr align 16 [[P:%.*]], i32 32, i1 false)
 ; CHECK-NEXT:    call void @do_something()
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i32(ptr align 16 [[Q:%.*]], ptr align 16 [[MEMTMP]], i32 32, i1 false)
+; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i32(ptr align 16 [[Q:%.*]], ptr align 16 [[P]], i32 32, i1 false)
 ; CHECK-NEXT:    call void @capture(ptr [[MEMTMP]])
 ; CHECK-NEXT:    ret void
 ;
@@ -851,10 +851,8 @@ define void @memcpy_memcpy_escape_after1(ptr noalias %P, ptr noalias %Q) {
 
 define void @memcpy_memcpy_escape_after2(ptr noalias %P, ptr noalias %Q) {
 ; CHECK-LABEL: @memcpy_memcpy_escape_after2(
-; CHECK-NEXT:    [[MEMTMP:%.*]] = alloca [32 x i8], align 16
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i32(ptr align 16 [[MEMTMP]], ptr align 16 [[P:%.*]], i32 32, i1 false)
 ; CHECK-NEXT:    call void @do_something()
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i32(ptr align 16 [[Q:%.*]], ptr align 16 [[MEMTMP]], i32 32, i1 false)
+; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i32(ptr align 16 [[Q:%.*]], ptr align 16 [[P:%.*]], i32 32, i1 false)
 ; CHECK-NEXT:    call void @capture(ptr [[P]])
 ; CHECK-NEXT:    ret void
 ;
@@ -868,10 +866,8 @@ define void @memcpy_memcpy_escape_after2(ptr noalias %P, ptr noalias %Q) {
 
 define void @memcpy_byval_escape_after(ptr noalias %P) {
 ; CHECK-LABEL: @memcpy_byval_escape_after(
-; CHECK-NEXT:    [[A:%.*]] = alloca [8 x i8], align 1
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr align 4 [[A]], ptr align 4 [[P:%.*]], i64 8, i1 false)
 ; CHECK-NEXT:    call void @do_something()
-; CHECK-NEXT:    call void @test4a(ptr byval(i8) align 1 [[A]])
+; CHECK-NEXT:    call void @test4a(ptr byval(i8) align 1 [[P:%.*]])
 ; CHECK-NEXT:    call void @capture(ptr [[P]])
 ; CHECK-NEXT:    ret void
 ;
@@ -885,10 +881,8 @@ define void @memcpy_byval_escape_after(ptr noalias %P) {
 
 define void @memcpy_immut_escape_after(ptr align 4 noalias %val) {
 ; CHECK-LABEL: @memcpy_immut_escape_after(
-; CHECK-NEXT:    [[VAL1:%.*]] = alloca i8, align 4
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr align 4 [[VAL1]], ptr align 4 [[VAL:%.*]], i64 1, i1 false)
 ; CHECK-NEXT:    call void @do_something()
-; CHECK-NEXT:    call void @f(ptr noalias nocapture readonly align 4 [[VAL1]])
+; CHECK-NEXT:    call void @f(ptr noalias nocapture readonly align 4 [[VAL:%.*]])
 ; CHECK-NEXT:    call void @capture(ptr [[VAL]])
 ; CHECK-NEXT:    ret void
 ;