[MemCpyOpt] Fix MemorySSA preservation

[nikic]

Summary

moveUp() moves instructions, so we should move the corresponding memory accesses as well. We should also move the store instruction itself: Even though we'll end up removing it later, this gives us a correct MemoryDef to replace.

Noticed this discrepancy while looking at D26739.

[asbirlea]

Thank you for fixing this!

Could you clang-format the patch?

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

554

Is it possible for the cast here to fail?
P is known to modify the Load location a the callsite. Is it possible this is true for specific cases not modeled as MemoryDefs in MemorySSA (e.g. when using a non-standard aa pipeline, and basicaa is disabled, certain accesses are seen as "can access/modify" but are ignored in the MemorySSA modeling - see MemorySSA.cpp:1766).

[nikic]

Be more conservative when fetching the MemoryDef for P.

[nikic]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

554

I wasn't able to come up with a case that fails based on some quick experiments -- anything I tried would get foiled by the above safety checks without basicaa. However, it doesn't cost us anything to be conservative here, so I changed this to use a dyn_cast_or_null.

[nikic]

Revert previous change.

[nikic]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

554

Actually, thinking about this again, just using dyn_cast_or_null here isn't safe: Just because we can't get the MemoryDef for P does not mean that we don't have to move accesses in MemorySSA, we just no longer know where they need to be moved to. With that in mind, I think it is safer to leave the cast<> and let is assert if it turns out that we can run into this issue after all.

[asbirlea]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

554

Agreed, we'd still need an insertion point.
I think the logic would be something like this:

MemoryAccess *InsertPoint = nullptr;
Instruction *LastToMove = ToLift[ToList.size()-1];
for (auto &Ins : make_range(P->getIterator(), --LastToMove->getIterator())) {
  if (InsertPoint = MSSAU->getMemorySSA()->getMemoryAccess(&Ins))
    break;
}

////// below

// No need to move accesses is InsertPoint == nullptr
if (MSSAU && InsertPoint) {
  if (MemoryUseOrDef *MA = MSSAU->getMemorySSA()->getMemoryAccess(I))
    MSSAU->moveBefore(MA, InsertPoint);
}

[nikic]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

554

This doesn't look quite right to me either. Imagine we have this kind of sequence:

P
ToLift[1]
MA
ToLift[0] = SI

While there is no memory access between P and ToLift[1], it is still necessary to move the ToLift[0] access above MA.

I think we would have to do something like this to handle it:

MemoryAccess *InsertPoint =
    MSSAU ? MSSAU->getMemorySSA()->getMemoryAccess(P);
// Below
if (MSSAU) {
  if (MemoryUseOrDef *MA = MSSAU->getMemorySSA()->getMemoryAccess(I)) {
    if (InsertPoint)
      MSSAU->moveBefore(MA, InsertPoint);
    else
      InsertPoint = MA;
  }
}

[nikic]

Handle P not being a MemoryAccess.

[asbirlea]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

554

You're right, there may be accesses that are interleaved with ToLift entries.
But not all entries in ToList can have an access based on the check below (or is that not the case?).
For example:

P
ToLift[1]      -> getMemAccess = nullptr
MA
ToLift[0] = SI  -> needs to be moved before MA, MA needs to be found

In your example, ToLift[0] would be moved before ToLift[1], and that doesn't look right.

A simpler approach may be:

• get first access before P (if P has an access get the iterator for memory accesses and do a decrement; if it doesn't walk instructions up to LI, which has an access for sure, or pass this info into the method from the call site which does this exact walk).
• in the loop below do:

assert (MemInsertPoint && "Exists");
if (MemoryUseOrDef *MA = MSSAU->getMemorySSA()->getMemoryAccess(I))
  MSSAU->moveAfter(MA, MemInsertPoint);
  MemInsertPoint = MA;
}

[nikic]

Next try...

[nikic]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

554

You're right, my previous code didn't really make sense, as the relevant memory access isn't part of ToLift. I've now updated it according to your suggestion, hopefully correctly. This is pretty ugly, but at least should go away once we determine this based on MemorySSA.

[asbirlea]

LGTM.

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

bool MemCpyOptPass::moveUp(StoreInst *SI, Instruction *P, const LoadInst *LI) {
   // Keep track of the arguments of all instruction we plan to lift
   // so we can make sure to lift them as well if appropriate.
   DenseSet<Instruction*> Args;
   if (auto *Ptr = dyn_cast<Instruction>(SI->getPointerOperand()))
     if (Ptr->getParent() == SI->getParent())
       Args.insert(Ptr);
 
   // Instruction to lift before P.
-  SmallVector<Instruction*, 8> ToLift;
+  SmallVector<Instruction *, 8> ToLift{SI};
 
   // Memory locations of lifted instructions.
   SmallVector<MemoryLocation, 8> MemLocs{StoreLoc};
 
   // Lifted calls.
   SmallVector<const CallBase *, 8> Calls;
 
   const MemoryLocation LoadLoc = MemoryLocation::get(LI);
for (unsigned k = 0, e = C->getNumOperands(); k != e; ++k)
         if (A->getParent() == SI->getParent()) {
           // Cannot hoist user of P above P
           if(A == P) return false;
           Args.insert(A);
         }
       }
   }
 
-  // We made it, we need to lift
+  // Find MSSA insertion point. Normally P will always have a corresponding
+  // memory access before which we can insert. However, with non-standard AA
+  // pipelines, there may be a mismatch between AA and MSSA, in which case we
+  // will scan for a memory access before P. In either case, we know for sure
+  // that at least the load will have a memory access.
+  // TODO: Simplify this once P will be determined by MSSA, in which case the
+  // discrepancy can no longer occur.
+  MemoryUseOrDef *MemInsertPoint = nullptr;
+  if (MSSAU) {
+    if (MemoryUseOrDef *MA = MSSAU->getMemorySSA()->getMemoryAccess(P)) {
+      MemInsertPoint = cast<MemoryUseOrDef>(--MA->getIterator());
+    } else {
+      const Instruction *ConstP = P;
+      for (const Instruction &I : make_range(++ConstP->getReverseIterator(),
+                                             ++LI->getReverseIterator())) {
+        if (MemoryUseOrDef *MA = MSSAU->getMemorySSA()->getMemoryAccess(&I)) {
+          MemInsertPoint = MA;
+          break;
+        }
+      }
+    }
+  }
+
+  // We made it, we need to lift.
   for (auto *I : llvm::reverse(ToLift)) {
     LLVM_DEBUG(dbgs() << "Lifting " << *I << " before " << *P << "\n");
     I->moveBefore(P);
+    if (MSSAU) {
+      assert(MemInsertPoint && "Must have found insert point");
+      if (MemoryUseOrDef *MA = MSSAU->getMemorySSA()->getMemoryAccess(I)) {
+        MSSAU->moveAfter(MA, MemInsertPoint);
+        MemInsertPoint = MA;
+      }
+    }
   }
 
   return true;
 }
 
 bool MemCpyOptPass::processStore(StoreInst *SI, BasicBlock::iterator &BBI) {
   if (!SI->isSimple()) return false;
 
if (LI->isSimple() && LI->hasOneUse() &&
             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));
+                cast<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(SI));
             auto *NewAccess =
                 MSSAU->createMemoryAccessAfter(M, LastDef, LastDef);
             MSSAU->insertDef(cast<MemoryDef>(NewAccess), /*RenameUses=*/true);
           }
 
           eraseInstruction(SI);
           eraseInstruction(LI);
           ++NumMemCpyInstr;

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

entry:
   %2 = getelementptr inbounds [4 x i32], [4 x i32]* %ptr, i64 0, i32 2
   store i32 0, i32* %2, align 1
   %3 = getelementptr inbounds [4 x i32], [4 x i32]* %ptr, i64 0, i32 3
   store i32 0, i32* %3, align 1
   call void @clobber()
   ret void
 }
 
+define void @test8(%t* noalias %src, %t* %dst) {
+; CHECK-LABEL: @test8(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast %t* [[SRC:%.*]] to i8*
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast %t* [[DST:%.*]] to i8*
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast %t* [[SRC]] to i8*
+; CHECK-NEXT:    call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 1 [[TMP2]], i8* align 1 [[TMP3]], i64 8224, i1 false)
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 1 [[TMP1]], i8 0, i64 8224, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %1 = load %t, %t* %src
+  store %t zeroinitializer, %t* %src
+  store %t %1, %t* %dst
+  ret void
+}
+
 declare void @clobber()
 
 ; 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 }