[llvm] [MemCpyOpt] fix incorrect handling of lifetime markers (PR #143782)

[Jameson Nash via llvm-commits]

================
@@ -1366,56 +1366,68 @@ bool MemCpyOptPass::processMemSetMemCpyDependence(MemCpyInst *MemCpy,
 
 /// Determine whether the pointer V had only undefined content (due to Def) up
 /// to the given Size, either because it was freshly alloca'd or started its
-/// lifetime.
-static bool hasUndefContents(MemorySSA *MSSA, BatchAAResults &AA, Value *V,
-                             MemoryDef *Def, Value *Size) {
-  if (MSSA->isLiveOnEntryDef(Def))
-    return isa<AllocaInst>(getUnderlyingObject(V));
-
-  if (auto *II = dyn_cast_or_null<IntrinsicInst>(Def->getMemoryInst())) {
-    if (II->getIntrinsicID() == Intrinsic::lifetime_start) {
-      auto *LTSize = cast<ConstantInt>(II->getArgOperand(0));
-
-      if (auto *CSize = dyn_cast<ConstantInt>(Size)) {
-        if (AA.isMustAlias(V, II->getArgOperand(1)) &&
-            LTSize->getZExtValue() >= CSize->getZExtValue())
-          return true;
-      }
+/// lifetime by walking the MSSA graph.
+static bool hadUndefContentsBefore(MemorySSA *MSSA, BatchAAResults &BAA,
+                                   Value *V, MemoryAccess *Clobber,
+                                   MemoryLocation Loc, Value *Size) {
+  while (1) {
+    Clobber = MSSA->getWalker()->getClobberingMemoryAccess(Clobber, Loc, BAA);
+    MemoryDef *Def = dyn_cast<MemoryDef>(Clobber);
+    if (!Def)
+      return false;
+
+    if (MSSA->isLiveOnEntryDef(Def))
+      return isa<AllocaInst>(getUnderlyingObject(V));
+
+    if (auto *II = dyn_cast_or_null<IntrinsicInst>(Def->getMemoryInst())) {
+      if (II->getIntrinsicID() == Intrinsic::lifetime_start) {
+        auto *LTSize = cast<ConstantInt>(II->getArgOperand(0));
 
-      // If the lifetime.start covers a whole alloca (as it almost always
-      // does) and we're querying a pointer based on that alloca, then we know
-      // the memory is definitely undef, regardless of how exactly we alias.
-      // The size also doesn't matter, as an out-of-bounds access would be UB.
-      if (auto *Alloca = dyn_cast<AllocaInst>(getUnderlyingObject(V))) {
-        if (getUnderlyingObject(II->getArgOperand(1)) == Alloca) {
-          const DataLayout &DL = Alloca->getDataLayout();
-          if (std::optional<TypeSize> AllocaSize =
-                  Alloca->getAllocationSize(DL))
-            if (*AllocaSize == LTSize->getValue())
+        if (Size)
+          if (auto CSize = dyn_cast<ConstantInt>(Size))
+            if (BAA.isMustAlias(V, II->getArgOperand(1)) &&
+                LTSize->getZExtValue() >= CSize->getZExtValue())
               return true;
+
+        // If the lifetime.start covers a whole alloca (as it almost always
+        // does) and we're querying a pointer based on that alloca, then we know
+        // the memory is definitely undef, regardless of how exactly we alias.
+        // The size also doesn't matter, as an out-of-bounds access would be UB.
+        if (auto *Alloca = dyn_cast<AllocaInst>(getUnderlyingObject(V))) {
+          if (getUnderlyingObject(II->getArgOperand(1)) == Alloca) {
+            const DataLayout &DL = Alloca->getDataLayout();
+            if (std::optional<TypeSize> AllocaSize =
+                    Alloca->getAllocationSize(DL))
+              if (*AllocaSize == LTSize->getValue())
+                return true;
+          }
         }
+        Clobber = Def->getDefiningAccess();
+        continue;
+      } else if (II->getIntrinsicID() == Intrinsic::lifetime_end) {
+        Clobber = Def->getDefiningAccess();
+        continue;
----------------
vtjnash wrote:

I think your theory might be still correct about the proximal cause, it just only skips the limit specifically when the walk ended on a completely unrelated lifetime instruction. A local fix for that is:
```
diff --git a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp b/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
index 877600034f84..759b16206d5c 100644
--- a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
+++ b/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
@@ -1370,6 +1370,7 @@ bool MemCpyOptPass::processMemSetMemCpyDependence(MemCpyInst *MemCpy,
 static bool hadUndefContentsBefore(MemorySSA *MSSA, BatchAAResults &BAA,
                                    Value *V, MemoryAccess *Clobber,
                                    MemoryLocation Loc, Value *Size) {
+  Value *VBase = getUnderlyingObject(V);
   while (1) {
     Clobber = MSSA->getWalker()->getClobberingMemoryAccess(Clobber, Loc, BAA);
     MemoryDef *Def = dyn_cast<MemoryDef>(Clobber);
@@ -1377,12 +1378,18 @@ static bool hadUndefContentsBefore(MemorySSA *MSSA, BatchAAResults &BAA,
       return false;
 
     if (MSSA->isLiveOnEntryDef(Def))
-      return isa<AllocaInst>(getUnderlyingObject(V));
+      return isa<AllocaInst>(VBase);
 
     if (auto *II = dyn_cast_or_null<IntrinsicInst>(Def->getMemoryInst())) {
       if (II->getIntrinsicID() == Intrinsic::lifetime_start) {
         auto *LTSize = cast<ConstantInt>(II->getArgOperand(0));
 
+        // Check if the SSA Walk ended early due to heuristics or actually
+        // reached a lifetime instruction for this pointer.
+        Value *IIBase = getUnderlyingObject(II->getArgOperand(1));
+        if (VBase != IIBase)
+          return false;
+
         if (Size)
           if (auto CSize = dyn_cast<ConstantInt>(Size))
             if (BAA.isMustAlias(V, II->getArgOperand(1)) &&
@@ -1393,8 +1400,8 @@ static bool hadUndefContentsBefore(MemorySSA *MSSA, BatchAAResults &BAA,
         // does) and we're querying a pointer based on that alloca, then we know
         // the memory is definitely undef, regardless of how exactly we alias.
         // The size also doesn't matter, as an out-of-bounds access would be UB.
-        if (auto *Alloca = dyn_cast<AllocaInst>(getUnderlyingObject(V))) {
-          if (getUnderlyingObject(II->getArgOperand(1)) == Alloca) {
+        if (auto *Alloca = dyn_cast<AllocaInst>(VBase)) {
+          if (IIBase == Alloca) {
             const DataLayout &DL = Alloca->getDataLayout();
             if (std::optional<TypeSize> AllocaSize =
                     Alloca->getAllocationSize(DL))
@@ -1405,6 +1412,11 @@ static bool hadUndefContentsBefore(MemorySSA *MSSA, BatchAAResults &BAA,
         Clobber = Def->getDefiningAccess();
         continue;
       } else if (II->getIntrinsicID() == Intrinsic::lifetime_end) {
+        // Check if the SSA Walk ended early due to heuristics or actually
+        // reached a lifetime instruction for this pointer.
+        Value *IIBase = getUnderlyingObject(II->getArgOperand(1));
+        if (VBase != IIBase)
+          return false;
         Clobber = Def->getDefiningAccess();
         continue;
       }
```

It looks like the getClobberingMemoryAccess code also could be easily adjusted to expose the walk limit (it already does, just not accessible from the header file)

But since it seems you're implying this function shouldn't gain any benefit from looking past the lifetime, I could also just drop this loop.

The tests however look better to me, so I also am wondering if there is some way to improve MemorySSA to handle simple cases like this, where there is a noalias base pointer (such as an alloca or malloc) with only a few uses in the dominator tree before the current MemoryDef location, none of which were capturing?

https://github.com/llvm/llvm-project/pull/143782