From 71051deff27928cff908ea794e09806eee662801 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Mon, 15 Jul 2024 09:41:11 +0200
Subject: [PATCH] [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 https://github.com/llvm/llvm-project/issues/98610.
---
 .../lib/Transforms/Scalar/MemCpyOptimizer.cpp |  13 ++-
 .../Transforms/MemCpyOpt/memcpy-zero-size.ll  |  15 +++
 .../MemCpyOpt/memset-memcpy-dbgloc.ll         |  13 ++-
 .../memset-memcpy-redundant-memset.ll         | 100 ++++++++++++++----
 llvm/test/Transforms/MemCpyOpt/opaque-ptr.ll  |   2 +-
 5 files changed, 115 insertions(+), 28 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp b/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
index 4063762c88a2e4..cee34f0a6da1f3 100644
--- a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
+++ b/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);
diff --git a/llvm/test/Transforms/MemCpyOpt/memcpy-zero-size.ll b/llvm/test/Transforms/MemCpyOpt/memcpy-zero-size.ll
index 33f7f530f6f995..e81044defbb3af 100644
--- a/llvm/test/Transforms/MemCpyOpt/memcpy-zero-size.ll
+++ b/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
+}
diff --git a/llvm/test/Transforms/MemCpyOpt/memset-memcpy-dbgloc.ll b/llvm/test/Transforms/MemCpyOpt/memset-memcpy-dbgloc.ll
index f8536aba2d19a9..3f577f09ada1eb 100644
--- a/llvm/test/Transforms/MemCpyOpt/memset-memcpy-dbgloc.ll
+++ b/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
diff --git a/llvm/test/Transforms/MemCpyOpt/memset-memcpy-redundant-memset.ll b/llvm/test/Transforms/MemCpyOpt/memset-memcpy-redundant-memset.ll
index 72a234091f79b0..68356397423507 100644
--- a/llvm/test/Transforms/MemCpyOpt/memset-memcpy-redundant-memset.ll
+++ b/llvm/test/Transforms/MemCpyOpt/memset-memcpy-redundant-memset.ll
@@ -7,7 +7,9 @@ target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
 
 define void @test_constant(i64 %src_size, ptr %dst, i64 %dst_size, i8 %c) {
 ; CHECK-LABEL: @test_constant(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ule i64 [[DST_SIZE:%.*]], [[SRC_SIZE:%.*]]
+; 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]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = sub i64 [[DST_SIZE]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i64 0, i64 [[TMP2]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr i8, ptr [[DST:%.*]], i64 [[SRC_SIZE]]
@@ -15,6 +17,8 @@ define void @test_constant(i64 %src_size, ptr %dst, i64 %dst_size, i8 %c) {
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST]], ptr @C, i64 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %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)
   call void @llvm.memcpy.p0.p0.i64(ptr %dst, ptr @C, i64 %src_size, i1 false)
   ret void
@@ -22,7 +26,9 @@ define void @test_constant(i64 %src_size, ptr %dst, i64 %dst_size, i8 %c) {
 
 define void @test(ptr %src, i64 %src_size, ptr noalias %dst, i64 %dst_size, i8 %c) {
 ; CHECK-LABEL: @test(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ule i64 [[DST_SIZE:%.*]], [[SRC_SIZE:%.*]]
+; 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]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = sub i64 [[DST_SIZE]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i64 0, i64 [[TMP2]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr i8, ptr [[DST:%.*]], i64 [[SRC_SIZE]]
@@ -30,6 +36,8 @@ define void @test(ptr %src, i64 %src_size, ptr noalias %dst, i64 %dst_size, i8 %
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST]], ptr [[SRC:%.*]], i64 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %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)
   call void @llvm.memcpy.p0.p0.i64(ptr %dst, ptr %src, i64 %src_size, i1 false)
   ret void
@@ -37,8 +45,10 @@ define void @test(ptr %src, i64 %src_size, ptr noalias %dst, i64 %dst_size, i8 %
 
 define void @test_different_types_i32_i64(ptr noalias %dst, ptr %src, i32 %dst_size, i64 %src_size, i8 %c) {
 ; CHECK-LABEL: @test_different_types_i32_i64(
+; CHECK-NEXT:    [[NON_ZERO:%.*]] = icmp ne i64 [[SRC_SIZE:%.*]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NON_ZERO]])
 ; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[DST_SIZE:%.*]] to i64
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp ule i64 [[TMP1]], [[SRC_SIZE:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ule i64 [[TMP1]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = sub i64 [[TMP1]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[TMP2]], i64 0, i64 [[TMP3]]
 ; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr i8, ptr [[DST:%.*]], i64 [[SRC_SIZE]]
@@ -46,6 +56,8 @@ define void @test_different_types_i32_i64(ptr noalias %dst, ptr %src, i32 %dst_s
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST]], ptr [[SRC:%.*]], i64 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %non.zero = icmp ne i64 %src_size, 0
+  call void @llvm.assume(i1 %non.zero)
   call void @llvm.memset.p0.i32(ptr %dst, i8 %c, i32 %dst_size, i1 false)
   call void @llvm.memcpy.p0.p0.i64(ptr %dst, ptr %src, i64 %src_size, i1 false)
   ret void
@@ -53,7 +65,9 @@ define void @test_different_types_i32_i64(ptr noalias %dst, ptr %src, i32 %dst_s
 
 define void @test_different_types_i128_i32(ptr noalias %dst, ptr %src, i128 %dst_size, i32 %src_size, i8 %c) {
 ; CHECK-LABEL: @test_different_types_i128_i32(
-; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[SRC_SIZE:%.*]] to i128
+; CHECK-NEXT:    [[NON_ZERO:%.*]] = icmp ne i32 [[SRC_SIZE:%.*]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NON_ZERO]])
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[SRC_SIZE]] to i128
 ; CHECK-NEXT:    [[TMP2:%.*]] = icmp ule i128 [[DST_SIZE:%.*]], [[TMP1]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = sub i128 [[DST_SIZE]], [[TMP1]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[TMP2]], i128 0, i128 [[TMP3]]
@@ -62,6 +76,8 @@ define void @test_different_types_i128_i32(ptr noalias %dst, ptr %src, i128 %dst
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i32(ptr [[DST]], ptr [[SRC:%.*]], i32 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %non.zero = icmp ne i32 %src_size, 0
+  call void @llvm.assume(i1 %non.zero)
   call void @llvm.memset.p0.i128(ptr %dst, i8 %c, i128 %dst_size, i1 false)
   call void @llvm.memcpy.p0.p0.i32(ptr %dst, ptr %src, i32 %src_size, i1 false)
   ret void
@@ -69,8 +85,10 @@ define void @test_different_types_i128_i32(ptr noalias %dst, ptr %src, i128 %dst
 
 define void @test_different_types_i32_i128(ptr noalias %dst, ptr %src, i32 %dst_size, i128 %src_size, i8 %c) {
 ; CHECK-LABEL: @test_different_types_i32_i128(
+; CHECK-NEXT:    [[NON_ZERO:%.*]] = icmp ne i128 [[SRC_SIZE:%.*]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NON_ZERO]])
 ; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[DST_SIZE:%.*]] to i128
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp ule i128 [[TMP1]], [[SRC_SIZE:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ule i128 [[TMP1]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = sub i128 [[TMP1]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[TMP2]], i128 0, i128 [[TMP3]]
 ; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr i8, ptr [[DST:%.*]], i128 [[SRC_SIZE]]
@@ -78,6 +96,8 @@ define void @test_different_types_i32_i128(ptr noalias %dst, ptr %src, i32 %dst_
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i128(ptr [[DST]], ptr [[SRC:%.*]], i128 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %non.zero = icmp ne i128 %src_size, 0
+  call void @llvm.assume(i1 %non.zero)
   call void @llvm.memset.p0.i32(ptr %dst, i8 %c, i32 %dst_size, i1 false)
   call void @llvm.memcpy.p0.p0.i128(ptr %dst, ptr %src, i128 %src_size, i1 false)
   ret void
@@ -85,7 +105,9 @@ define void @test_different_types_i32_i128(ptr noalias %dst, ptr %src, i32 %dst_
 
 define void @test_different_types_i64_i32(ptr noalias %dst, ptr %src, i64 %dst_size, i32 %src_size, i8 %c) {
 ; CHECK-LABEL: @test_different_types_i64_i32(
-; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[SRC_SIZE:%.*]] to i64
+; CHECK-NEXT:    [[NON_ZERO:%.*]] = icmp ne i32 [[SRC_SIZE:%.*]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NON_ZERO]])
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[SRC_SIZE]] to i64
 ; CHECK-NEXT:    [[TMP2:%.*]] = icmp ule i64 [[DST_SIZE:%.*]], [[TMP1]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = sub i64 [[DST_SIZE]], [[TMP1]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[TMP2]], i64 0, i64 [[TMP3]]
@@ -94,6 +116,8 @@ define void @test_different_types_i64_i32(ptr noalias %dst, ptr %src, i64 %dst_s
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i32(ptr [[DST]], ptr [[SRC:%.*]], i32 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %non.zero = icmp ne i32 %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)
   call void @llvm.memcpy.p0.p0.i32(ptr %dst, ptr %src, i32 %src_size, i1 false)
   ret void
@@ -146,14 +170,18 @@ define void @test_align_memcpy(ptr %src, ptr noalias %dst, i64 %dst_size) {
 
 define void @test_non_i8_dst_type(ptr %src, i64 %src_size, ptr noalias %dst_pi64, i64 %dst_size, i8 %c) {
 ; CHECK-LABEL: @test_non_i8_dst_type(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ule i64 [[DST_SIZE:%.*]], [[SRC_SIZE:%.*]]
+; 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]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = sub i64 [[DST_SIZE]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i64 0, i64 [[TMP2]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr i8, ptr [[DST_PI64:%.*]], i64 [[SRC_SIZE]]
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr align 1 [[TMP4]], i8 [[C:%.*]], i64 [[TMP3]], i1 false)
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST_PI64:%.*]], ptr [[SRC:%.*]], i64 [[SRC_SIZE]], i1 false)
+; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST_PI64]], ptr [[SRC:%.*]], i64 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %non.zero = icmp ne i64 %src_size, 0
+  call void @llvm.assume(i1 %non.zero)
   call void @llvm.memset.p0.i64(ptr %dst_pi64, i8 %c, i64 %dst_size, i1 false)
   call void @llvm.memcpy.p0.p0.i64(ptr %dst_pi64, ptr %src, i64 %src_size, i1 false)
   ret void
@@ -161,10 +189,14 @@ define void @test_non_i8_dst_type(ptr %src, i64 %src_size, ptr noalias %dst_pi64
 
 define void @test_different_dst(ptr noalias %dst2, ptr %src, i64 %src_size, ptr noalias %dst, i64 %dst_size) {
 ; CHECK-LABEL: @test_different_dst(
+; CHECK-NEXT:    [[NON_ZERO:%.*]] = icmp ne i64 [[SRC_SIZE:%.*]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NON_ZERO]])
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr [[DST:%.*]], i8 0, i64 [[DST_SIZE:%.*]], i1 false)
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST2:%.*]], ptr [[SRC:%.*]], i64 [[SRC_SIZE:%.*]], i1 false)
+; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST2:%.*]], ptr [[SRC:%.*]], i64 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %non.zero = icmp ne i64 %src_size, 0
+  call void @llvm.assume(i1 %non.zero)
   call void @llvm.memset.p0.i64(ptr %dst, i8 0, i64 %dst_size, i1 false)
   call void @llvm.memcpy.p0.p0.i64(ptr %dst2, ptr %src, i64 %src_size, i1 false)
   ret void
@@ -206,11 +238,15 @@ define void @test_intermediate_write(ptr %b) #0 {
 
 define void @test_throwing_call(ptr %src, i64 %src_size, ptr noalias %dst, i64 %dst_size, i8 %c) {
 ; CHECK-LABEL: @test_throwing_call(
+; CHECK-NEXT:    [[NON_ZERO:%.*]] = icmp ne i64 [[SRC_SIZE:%.*]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NON_ZERO]])
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr [[DST:%.*]], i8 [[C:%.*]], i64 [[DST_SIZE:%.*]], i1 false)
-; CHECK-NEXT:    call void @call() #[[ATTR2:[0-9]+]]
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST]], ptr [[SRC:%.*]], i64 [[SRC_SIZE:%.*]], i1 false)
+; CHECK-NEXT:    call void @call() #[[ATTR3:[0-9]+]]
+; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST]], ptr [[SRC:%.*]], i64 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %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)
   call void @call() readnone
   call void @llvm.memcpy.p0.p0.i64(ptr %dst, ptr %src, i64 %src_size, i1 false)
@@ -220,8 +256,10 @@ define void @test_throwing_call(ptr %src, i64 %src_size, ptr noalias %dst, i64 %
 define void @test_throwing_call_alloca(ptr %src, i64 %src_size, i64 %dst_size, i8 %c) {
 ; CHECK-LABEL: @test_throwing_call_alloca(
 ; CHECK-NEXT:    [[DST:%.*]] = alloca i8, align 1
-; CHECK-NEXT:    call void @call() #[[ATTR2]]
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ule i64 [[DST_SIZE:%.*]], [[SRC_SIZE:%.*]]
+; CHECK-NEXT:    [[NON_ZERO:%.*]] = icmp ne i64 [[SRC_SIZE:%.*]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NON_ZERO]])
+; CHECK-NEXT:    call void @call() #[[ATTR3]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ule i64 [[DST_SIZE:%.*]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = sub i64 [[DST_SIZE]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i64 0, i64 [[TMP2]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr i8, ptr [[DST]], i64 [[SRC_SIZE]]
@@ -230,6 +268,8 @@ define void @test_throwing_call_alloca(ptr %src, i64 %src_size, i64 %dst_size, i
 ; CHECK-NEXT:    ret void
 ;
   %dst = alloca i8
+  %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)
   call void @call() readnone
   call void @llvm.memcpy.p0.p0.i64(ptr %dst, ptr %src, i64 %src_size, i1 false)
@@ -240,10 +280,14 @@ define void @test_throwing_call_alloca(ptr %src, i64 %src_size, i64 %dst_size, i
 ; is not legal.
 define void @test_missing_noalias(ptr %src, i64 %src_size, ptr %dst, i64 %dst_size, i8 %c) {
 ; CHECK-LABEL: @test_missing_noalias(
+; CHECK-NEXT:    [[NON_ZERO:%.*]] = icmp ne i64 [[SRC_SIZE:%.*]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NON_ZERO]])
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr [[DST:%.*]], i8 [[C:%.*]], i64 [[DST_SIZE:%.*]], i1 false)
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST]], ptr [[SRC:%.*]], i64 [[SRC_SIZE:%.*]], i1 false)
+; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST]], ptr [[SRC:%.*]], i64 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %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)
   call void @llvm.memcpy.p0.p0.i64(ptr %dst, ptr %src, i64 %src_size, i1 false)
   ret void
@@ -261,9 +305,13 @@ define void @test_same_const_size(ptr noalias %src, ptr noalias %dst, i8 %c) {
 
 define void @test_same_dynamic_size(ptr noalias %src, ptr noalias %dst, i64 %size, i8 %c) {
 ; CHECK-LABEL: @test_same_dynamic_size(
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST:%.*]], ptr [[SRC:%.*]], i64 [[SIZE:%.*]], i1 false)
+; CHECK-NEXT:    [[NON_ZERO:%.*]] = icmp ne i64 [[SIZE:%.*]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NON_ZERO]])
+; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST:%.*]], ptr [[SRC:%.*]], i64 [[SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %non.zero = icmp ne i64 %size, 0
+  call void @llvm.assume(i1 %non.zero)
   call void @llvm.memset.p0.i64(ptr %dst, i8 %c, i64 %size, i1 false)
   call void @llvm.memcpy.p0.p0.i64(ptr %dst, ptr %src, i64 %size, i1 false)
   ret void
@@ -286,9 +334,11 @@ define void @test_must_alias_same_size(ptr noalias %src, ptr noalias %dst, i8 %c
 
 define void @test_must_alias_different_size(ptr noalias %src, i64 %src_size, ptr noalias %dst, i64 %dst_size, i8 %c) {
 ; CHECK-LABEL: @test_must_alias_different_size(
+; CHECK-NEXT:    [[NON_ZERO:%.*]] = icmp ne i64 [[SRC_SIZE:%.*]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NON_ZERO]])
 ; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr i8, ptr [[DST:%.*]], i64 16
 ; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr i8, ptr [[DST]], i64 16
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ule i64 [[DST_SIZE:%.*]], [[SRC_SIZE:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ule i64 [[DST_SIZE:%.*]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = sub i64 [[DST_SIZE]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i64 0, i64 [[TMP2]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr i8, ptr [[GEP2]], i64 [[SRC_SIZE]]
@@ -296,6 +346,8 @@ define void @test_must_alias_different_size(ptr noalias %src, i64 %src_size, ptr
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[GEP2]], ptr [[SRC:%.*]], i64 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %non.zero = icmp ne i64 %src_size, 0
+  call void @llvm.assume(i1 %non.zero)
   %gep1 = getelementptr i8, ptr %dst, i64 16
   call void @llvm.memset.p0.i64(ptr %gep1, i8 %c, i64 %dst_size, i1 false)
   %gep2 = getelementptr i8, ptr %dst, i64 16
@@ -305,14 +357,18 @@ define void @test_must_alias_different_size(ptr noalias %src, i64 %src_size, ptr
 
 define void @test_weird_element_type(ptr %src, i64 %src_size, ptr noalias %dst, i64 %dst_size, i8 %c) {
 ; CHECK-LABEL: @test_weird_element_type(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ule i64 [[DST_SIZE:%.*]], [[SRC_SIZE:%.*]]
+; 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]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = sub i64 [[DST_SIZE]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i64 0, i64 [[TMP2]]
-; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr i8, ptr [[DST:%.*]], i64 [[SRC_SIZE]]
-; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr align 1 [[TMP5]], i8 [[C:%.*]], i64 [[TMP3]], i1 false)
+; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr i8, ptr [[DST:%.*]], i64 [[SRC_SIZE]]
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr align 1 [[TMP4]], i8 [[C:%.*]], i64 [[TMP3]], i1 false)
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr [[DST]], ptr [[SRC:%.*]], i64 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %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)
   call void @llvm.memcpy.p0.p0.i64(ptr %dst, ptr %src, i64 %src_size, i1 false)
   ret void
@@ -320,7 +376,9 @@ define void @test_weird_element_type(ptr %src, i64 %src_size, ptr noalias %dst,
 
 define void @test_addrspace(ptr addrspace(1) %src, i64 %src_size, ptr addrspace(1) noalias %dst, i64 %dst_size, i8 %c) {
 ; CHECK-LABEL: @test_addrspace(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ule i64 [[DST_SIZE:%.*]], [[SRC_SIZE:%.*]]
+; 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]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = sub i64 [[DST_SIZE]], [[SRC_SIZE]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i64 0, i64 [[TMP2]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr i8, ptr addrspace(1) [[DST:%.*]], i64 [[SRC_SIZE]]
@@ -328,6 +386,8 @@ define void @test_addrspace(ptr addrspace(1) %src, i64 %src_size, ptr addrspace(
 ; CHECK-NEXT:    call void @llvm.memcpy.p1.p1.i64(ptr addrspace(1) [[DST]], ptr addrspace(1) [[SRC:%.*]], i64 [[SRC_SIZE]], i1 false)
 ; CHECK-NEXT:    ret void
 ;
+  %non.zero = icmp ne i64 %src_size, 0
+  call void @llvm.assume(i1 %non.zero)
   call void @llvm.memset.p1.i64(ptr addrspace(1) %dst, i8 %c, i64 %dst_size, i1 false)
   call void @llvm.memcpy.p1.p1.i64(ptr addrspace(1) %dst, ptr addrspace(1) %src, i64 %src_size, i1 false)
   ret void
diff --git a/llvm/test/Transforms/MemCpyOpt/opaque-ptr.ll b/llvm/test/Transforms/MemCpyOpt/opaque-ptr.ll
index dbd78604398f0c..a968195623780d 100644
--- a/llvm/test/Transforms/MemCpyOpt/opaque-ptr.ll
+++ b/llvm/test/Transforms/MemCpyOpt/opaque-ptr.ll
@@ -1,7 +1,7 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -passes=memcpyopt -S %s -verify-memoryssa | FileCheck %s
 
-define void @test_memset_memcpy(ptr %src, i64 %src_size, ptr noalias %dst, i64 %dst_size, i8 %c) {
+define void @test_memset_memcpy(ptr %src, i64 range(i64 1, 42) %src_size, ptr noalias %dst, i64 %dst_size, i8 %c) {
 ; CHECK-LABEL: @test_memset_memcpy(
 ; CHECK-NEXT:    [[TMP1:%.*]] = icmp ule i64 [[DST_SIZE:%.*]], [[SRC_SIZE:%.*]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = sub i64 [[DST_SIZE]], [[SRC_SIZE]]