[MemCpyOpt] Fix infinite loop in memset+memcpy fold (#98638)

For the case where the memcpy size is zero, this transform is a complex
no-op. This can lead to an infinite loop when the size is zero in a way
that BasicAA understands, because it can still understand that dst and
dst + src_size are MustAlias.

I've tried to mitigate this before using the isZeroSize() check, but we
can hit cases where InstSimplify doesn't understand that the size is
zero, but BasicAA does.

As such, this bites the bullet and adds an explicit isKnownNonZero()
check to guard against no-op transforms.

Fixes #98610.

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

@@ -1296,6 +1296,15 @@ bool MemCpyOptPass::processMemSetMemCpyDependence(MemCpyInst *MemCpy,
   if (!BAA.isMustAlias(MemSet->getDest(), MemCpy->getDest()))
     return false;
 
+  // Don't perform the transform if src_size may be zero. In that case, the
+  // transform is essentially a complex no-op and may lead to an infinite
+  // loop if BasicAA is smart enough to understand that dst and dst + src_size
+  // are still MustAlias after the transform.
+  Value *SrcSize = MemCpy->getLength();
+  if (!isKnownNonZero(SrcSize,
+                      SimplifyQuery(MemCpy->getDataLayout(), DT, AC, MemCpy)))
+    return false;
+
   // Check that src and dst of the memcpy aren't the same. While memcpy
   // operands cannot partially overlap, exact equality is allowed.
   if (isModSet(BAA.getModRefInfo(MemCpy, MemoryLocation::getForSource(MemCpy))))
@@ -1312,7 +1321,6 @@ bool MemCpyOptPass::processMemSetMemCpyDependence(MemCpyInst *MemCpy,
   // Use the same i8* dest as the memcpy, killing the memset dest if different.
   Value *Dest = MemCpy->getRawDest();
   Value *DestSize = MemSet->getLength();
-  Value *SrcSize = MemCpy->getLength();
 
   if (mayBeVisibleThroughUnwinding(Dest, MemSet, MemCpy))
     return false;
@@ -1726,8 +1734,7 @@ 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 the size is zero, remove the memcpy.
   if (isZeroSize(M->getLength())) {
     ++BBI;
     eraseInstruction(M);

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

@@ -34,3 +34,18 @@ define void @pr64886(i64 %len, ptr noalias %p) {
   call void @llvm.memcpy.p0.p0.i64(ptr inttoptr (i64 -1 to ptr), ptr %p, i64 poison, i1 false)
   ret void
 }
+
+define void @pr98610(ptr %p, ptr noalias %p2) {
+; CHECK-LABEL: @pr98610(
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr [[P:%.*]], i8 0, i64 1, i1 false)
+; CHECK-NEXT:    [[ZERO_EXT:%.*]] = zext i32 0 to i64
+; CHECK-NEXT:    [[MUL:%.*]] = mul i64 [[ZERO_EXT]], 1
+; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[P]], ptr [[P2:%.*]], i64 [[MUL]], i1 false)
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.memset.p0.i64(ptr %p, i8 0, i64 1, i1 false)
+  %zero.ext = zext i32 0 to i64
+  %mul = mul i64 %zero.ext, 1
+  call void @llvm.memcpy.p0.p0.i64(ptr %p, ptr %p2, i64 %mul, i1 false)
+  ret void
+}

llvm/test/Transforms/MemCpyOpt/memset-memcpy-dbgloc.ll

@@ -8,15 +8,20 @@ declare void @llvm.memset.p0.i64(ptr nocapture, i8, i64, i1)
 declare void @llvm.memcpy.p0.p0.i64(ptr nocapture, ptr nocapture readonly, i64, i1)
 
 define void @test_constant(i64 %src_size, ptr %dst, i64 %dst_size, i8 %c) !dbg !5 {
-; CHECK-LABEL: @test_constant(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ule i64 [[DST_SIZE:%.*]], [[SRC_SIZE:%.*]], !dbg [[DBG11:![0-9]+]]
+; CHECK-LABEL: define void @test_constant(
+; CHECK-SAME: i64 [[SRC_SIZE:%.*]], ptr [[DST:%.*]], i64 [[DST_SIZE:%.*]], i8 [[C:%.*]]) !dbg [[DBG5:![0-9]+]] {
+; CHECK-NEXT:    [[NON_ZERO:%.*]] = icmp ne i64 [[SRC_SIZE]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NON_ZERO]])
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ule i64 [[DST_SIZE]], [[SRC_SIZE]], !dbg [[DBG11:![0-9]+]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = sub i64 [[DST_SIZE]], [[SRC_SIZE]], !dbg [[DBG11]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i64 0, i64 [[TMP2]], !dbg [[DBG11]]
-; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr i8, ptr [[DST:%.*]], i64 [[SRC_SIZE]], !dbg [[DBG11]]
-; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr align 1 [[TMP4]], i8 [[C:%.*]], i64 [[TMP3]], i1 false), !dbg [[DBG11]]
+; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr i8, ptr [[DST]], i64 [[SRC_SIZE]], !dbg [[DBG11]]
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr align 1 [[TMP4]], i8 [[C]], i64 [[TMP3]], i1 false), !dbg [[DBG11]]
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST]], ptr @C, i64 [[SRC_SIZE]], i1 false), !dbg [[DBG12:![0-9]+]]
 ; CHECK-NEXT:    ret void, !dbg [[DBG13:![0-9]+]]
 ;
+  %non.zero = icmp ne i64 %src_size, 0
+  call void @llvm.assume(i1 %non.zero)
   call void @llvm.memset.p0.i64(ptr %dst, i8 %c, i64 %dst_size, i1 false), !dbg !11
   call void @llvm.memcpy.p0.p0.i64(ptr %dst, ptr @C, i64 %src_size, i1 false), !dbg !12
   ret void, !dbg !13