[MemCpyOpt] Avoid infinite loop in processMemSetMemCpyDependence (PR54983)

[nikic]

Summary

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.

I'm a bit unsure on this patch -- 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). Though I can't say I particularly care about this transform either, so I'd be happy to go with that direction instead.

Fixes https://github.com/llvm/llvm-project/issues/54983.
Fixes https://github.com/llvm/llvm-project/issues/64886.

[nikic]

Rebase and handle undef/poison.

[fhahn]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

1300

I think the underlying issue here is that the builder folds this GEP to Dest is SrcSize is poison. This in turn leads to the new MemSet to must-alias the original memcpy, causing the cycle IIUC. Could we break the cycle here by bailing out if the pointer folds to Dest as more general solution and leave the cleanup of the cases where the size simplifies to 0 to instcombine?

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

35

might be good to use pointer arguments passed to the function here, instead of null/inttoptr so the poison is the only issue here

[nikic]

Avoid copy from null in test.

[nikic]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

1300

In this case, yes:

%1 = icmp ule i64 %len, poison
%2 = sub i64 %len, poison
%3 = select i1 %1, i64 0, i64 %2
%4 = icmp ule i64 %3, poison
%5 = sub i64 %3, poison
%6 = select i1 %4, i64 0, i64 %5
call void @llvm.memset.p0.i64(ptr align 1 inttoptr (i64 -1 to ptr), i8 0, i64 %6, i1 false)
call void @llvm.memcpy.p0.p0.i64(ptr inttoptr (i64 -1 to ptr), ptr null, i64 poison, i1 false)
ret void

But not in the other case:

%size = shl i64 0, 0
%1 = icmp ule i64 1, %size
%2 = sub i64 1, %size
%3 = select i1 %1, i64 0, i64 %2
%4 = getelementptr i8, ptr %p, i64 %size
%5 = icmp ule i64 %3, %size
%6 = sub i64 %3, %size
%7 = select i1 %5, i64 0, i64 %6
%8 = getelementptr i8, ptr %p, i64 %size
call void @llvm.memset.p0.i64(ptr align 1 %8, i8 0, i64 %7, i1 false)
call void @llvm.memcpy.p0.p0.i64(ptr %p, ptr %p2, i64 %size, i1 false)
ret void

Here the pointers aren't directly equal, but AA can determine they're equal.

So just checking that we get back the old Dest wouldn't fully avoid the issue.

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

35

I've replaced the null, but the inttoptr is important for the IRBuilder constant folding.

[davide]

This sounds reasonable to me. Thanks, assuming @fhahn has no other comments.

[asbirlea]

Proving the size is non-zero is a higher bar than proving the zero case and eliminating the memcpy in those cases. So this seems reasonable to me.

In cases where zero length cannot be proven by AA/InstSimplify in processMemCpy, is it still possible to end up in an infinite loop in processMemSetMemCpyDependence?

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #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"
 #include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/Analysis/PostDominators.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/BasicBlock.h"
bool MemCpyOptPass::processMemSetMemCpyDependence(MemCpyInst *MemCpy,
   Value *SizeDiff = Builder.CreateSub(DestSize, SrcSize);
   Value *MemsetLen = Builder.CreateSelect(
       Ule, ConstantInt::getNullValue(DestSize->getType()), SizeDiff);
   unsigned DestAS = Dest->getType()->getPointerAddressSpace();
   Instruction *NewMemSet = Builder.CreateMemSet(
       Builder.CreateGEP(
           Builder.getInt8Ty(),
           Builder.CreatePointerCast(Dest, Builder.getInt8PtrTy(DestAS)),
           SrcSize),
       MemSet->getOperand(1), MemsetLen, Alignment);
 
   assert(isa<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(MemCpy)) &&
          "MemCpy must be a MemoryDef");
   // The new memset is inserted before the memcpy, and it is known that the
   // memcpy's defining access is the memset about to be removed.
   auto *LastDef =
       cast<MemoryDef>(MSSAU->getMemorySSA()->getMemoryAccess(MemCpy));
bool MemCpyOptPass::performStackMoveOptzn(Instruction *Load, Instruction *Store,
   for (Instruction *I : NoAliasInstrs)
     I->setMetadata(LLVMContext::MD_noalias, nullptr);
 
   LLVM_DEBUG(dbgs() << "Stack Move: Performed staack-move optimization\n");
   NumStackMove++;
   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
 /// circumstances). This allows later passes to remove the first memcpy
 /// altogether.
 bool MemCpyOptPass::processMemCpy(MemCpyInst *M, BasicBlock::iterator &BBI) {
   // We can only optimize non-volatile memcpy's.
   if (M->isVolatile()) return false;
 
   // If the source and destination of the memcpy are the same, then zap it.
   if (M->getSource() == M->getDest()) {
     ++BBI;
     eraseInstruction(M);
     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())
       if (Value *ByteVal = isBytewiseValue(GV->getInitializer(),
                                            M->getModule()->getDataLayout())) {
         IRBuilder<> Builder(M);
         Instruction *NewM = Builder.CreateMemSet(
             M->getRawDest(), ByteVal, M->getLength(), M->getDestAlign(), false);

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

+; 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
+}