[MemCpyOpt] Avoid infinite loop in processMemSetMemCpyDependence (PR54983)

This adds an additional transform to drop zero-size memcpys, also in
the case where the size is only zero after instruction simplification.
The motivation is the case from PR54983 where the size is non-trivially
zero, and processMemSetMemCpyDependence() keeps trying to reduce the
memset size by zero bytes.

This fix it's not really principled. It only works on the premise that
if InstSimplify doesn't realize the size is zero, then AA also won't.

The principled approach would be to instead add a isKnownNonZero()
guard to the processMemSetMemCpyDependence() transform, but I
suspect that would render that optimization mostly useless (at least
it breaks all the existing test coverage -- worth noting that the
constant size case is also handled by DSE, so I think this transform
is primarily about the dynamic size case).

Fixes #54983.
Fixes #64886.

Differential Revision: https://reviews.llvm.org/D124078

Author: nikic

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

@@ -22,6 +22,7 @@
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/GlobalsModRef.h"
+#include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/Analysis/MemorySSA.h"
@@ -1626,6 +1627,16 @@ bool MemCpyOptPass::performStackMoveOptzn(Instruction *Load, Instruction *Store,
   return true;
 }
 
+static bool isZeroSize(Value *Size) {
+  if (auto *I = dyn_cast<Instruction>(Size))
+    if (auto *Res = simplifyInstruction(I, I->getModule()->getDataLayout()))
+      Size = Res;
+  // Treat undef/poison size like zero.
+  if (auto *C = dyn_cast<Constant>(Size))
+    return isa<UndefValue>(C) || C->isNullValue();
+  return false;
+}
+
 /// 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
@@ -1642,6 +1653,14 @@ bool MemCpyOptPass::processMemCpy(MemCpyInst *M, BasicBlock::iterator &BBI) {
     return true;
   }
 
+  // If the size is zero, remove the memcpy. This also prevents infinite loops
+  // in processMemSetMemCpyDependence, which is a no-op for zero-length memcpys.
+  if (isZeroSize(M->getLength())) {
+    ++BBI;
+    eraseInstruction(M);
+    return true;
+  }
+
   // If copying from a constant, try to turn the memcpy into a memset.
   if (auto *GV = dyn_cast<GlobalVariable>(M->getSource()))
     if (GV->isConstant() && GV->hasDefinitiveInitializer())

llvm/test/Transforms/MemCpyOpt/memcpy-zero-size.ll

@@ -0,0 +1,36 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -passes=memcpyopt < %s | FileCheck %s
+
+declare void @llvm.memset.p0.i64(ptr, i8, i64, i1)
+declare void @llvm.memcpy.p0.p0.i64(ptr, ptr, i64, i1)
+
+define void @zero_size(ptr %p, ptr %p2) {
+; CHECK-LABEL: @zero_size(
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.memcpy.p0.p0.i64(ptr %p, ptr %p2, i64 0, i1 false)
+  ret void
+}
+
+; The memcpy size is zero in a way that is non-trivial, but understood by AA.
+define void @pr54983(ptr %p, ptr noalias %p2) {
+; CHECK-LABEL: @pr54983(
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr [[P:%.*]], i8 0, i64 1, i1 false)
+; CHECK-NEXT:    [[SIZE:%.*]] = shl i64 0, 0
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.memset.p0.i64(ptr %p, i8 0, i64 1, i1 false)
+  %size = shl i64 0, 0
+  call void @llvm.memcpy.p0.p0.i64(ptr %p, ptr %p2, i64 %size, i1 false)
+  ret void
+}
+
+define void @pr64886(i64 %len, ptr noalias %p) {
+; CHECK-LABEL: @pr64886(
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr inttoptr (i64 -1 to ptr), i8 0, i64 [[LEN:%.*]], i1 false)
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.memset.p0.i64(ptr inttoptr (i64 -1 to ptr), i8 0, i64 %len, i1 false)
+  call void @llvm.memcpy.p0.p0.i64(ptr inttoptr (i64 -1 to ptr), ptr %p, i64 poison, i1 false)
+  ret void
+}