[llvm] [memcpyopt] handle memcpy from memset in more cases (PR #140954)

[Jameson Nash via llvm-commits]

================
@@ -1418,51 +1437,52 @@ static bool hasUndefContents(MemorySSA *MSSA, BatchAAResults &AA, Value *V,
 bool MemCpyOptPass::performMemCpyToMemSetOptzn(MemCpyInst *MemCpy,
                                                MemSetInst *MemSet,
                                                BatchAAResults &BAA) {
-  // Make sure that memcpy(..., memset(...), ...), that is we are memsetting and
-  // memcpying from the same address. Otherwise it is hard to reason about.
-  if (!BAA.isMustAlias(MemSet->getRawDest(), MemCpy->getRawSource()))
-    return false;
-
   Value *MemSetSize = MemSet->getLength();
   Value *CopySize = MemCpy->getLength();
 
-  if (MemSetSize != CopySize) {
-    // Make sure the memcpy doesn't read any more than what the memset wrote.
-    // Don't worry about sizes larger than i64.
+  int64_t MOffset = 0;
+  const DataLayout &DL = MemCpy->getModule()->getDataLayout();
+  // We can only transforms memcpy's where the dest of one is the source of the
+  // other, or the memory transfer has a known offset from the memset.
+  if (MemCpy->getSource() != MemSet->getDest()) {
+    std::optional<int64_t> Offset =
+        MemCpy->getSource()->getPointerOffsetFrom(MemSet->getDest(), DL);
+    if (!Offset || *Offset < 0)
+      return false;
+    MOffset = *Offset;
+  }
 
-    // A known memset size is required.
+  MaybeAlign MDestAlign = MemCpy->getDestAlign();
+  if (MOffset != 0 || MemSetSize != CopySize) {
+    // Make sure the memcpy doesn't read any more than what the memset wrote,
+    // other than undef. Don't worry about sizes larger than i64. A known memset
+    // size is required.
     auto *CMemSetSize = dyn_cast<ConstantInt>(MemSetSize);
     if (!CMemSetSize)
       return false;
-
-    // A known memcpy size is also required.
+    // A known memcpy size is required.
     auto *CCopySize = dyn_cast<ConstantInt>(CopySize);
     if (!CCopySize)
       return false;
-    if (CCopySize->getZExtValue() > CMemSetSize->getZExtValue()) {
----------------
vtjnash wrote:

I believe this helper function should also be used in `performStackMoveOptzn`, I just didn't want to add additional improvements that re-use the logic to the current PR. I have a list of other TODOs to improve in this file (such as https://github.com/llvm/llvm-project/commit/827019c4bdfbb43b5c5c3a6094bb30e19c341ae8#diff-afb2b04eff951cb67d214bc6dfa62a087a72591451817dba7929b48c5fe635f9R1530) but they're blocked on getting this PR to progress

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