[MemCpyOpt] handle memcpy from memset for non-constant sizes (#143727)

Allows forwarding memset to memcpy for mismatching unknown sizes if overread has undef contents. In that case we can refine the undef bytes to the memset value. Refs #140954 which laid some of the groundwork for this.

Author: vtjnash

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

@@ -1440,7 +1440,7 @@ bool MemCpyOptPass::performMemCpyToMemSetOptzn(MemCpyInst *MemCpy,
  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.
+ // other, or they have a known offset.
  if (MemCpy->getSource() != MemSet->getDest()) {
  std::optional<int64_t> Offset =
  MemCpy->getSource()->getPointerOffsetFrom(MemSet->getDest(), DL);
@@ -1451,28 +1451,28 @@ bool MemCpyOptPass::performMemCpyToMemSetOptzn(MemCpyInst *MemCpy,
 
  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.
+ // other than undef. Don't worry about sizes larger than i64.
  auto *CMemSetSize = dyn_cast<ConstantInt>(MemSetSize);
- if (!CMemSetSize)
- return false;
-
- // A known memcpy size is also required.
  auto *CCopySize = dyn_cast<ConstantInt>(CopySize);
- if (!CCopySize)
- return false;
- if (CCopySize->getZExtValue() + MOffset > CMemSetSize->getZExtValue()) {
+ if (!CMemSetSize || !CCopySize ||
+ CCopySize->getZExtValue() + MOffset > CMemSetSize->getZExtValue()) {
  if (!overreadUndefContents(MSSA, MemCpy, MemSet, BAA))
  return false;
- // Clip the memcpy to the bounds of the memset
- if (MOffset == 0)
- CopySize = MemSetSize;
- else
- CopySize =
- ConstantInt::get(CopySize->getType(),
- CMemSetSize->getZExtValue() <= (uint64_t)MOffset
- ? 0
- : CMemSetSize->getZExtValue() - MOffset);
+
+ if (CMemSetSize && CCopySize) {
+ // If both have constant sizes and offsets, clip the memcpy to the
+ // bounds of the memset if applicable.
+ assert(CCopySize->getZExtValue() + MOffset >
+ CMemSetSize->getZExtValue());
+ if (MOffset == 0)
+ CopySize = MemSetSize;
+ else
+ CopySize =
+ ConstantInt::get(CopySize->getType(),
+ CMemSetSize->getZExtValue() <= (uint64_t)MOffset
+ ? 0
+ : CMemSetSize->getZExtValue() - MOffset);
+ }
  }
  }
 

llvm/test/Transforms/MemCpyOpt/variable-sized-memset-memcpy.ll

@@ -19,12 +19,12 @@ define void @test(ptr %src, i8 %c, i64 %size) {
 }
 
 ; Differing sizes, but would be UB if size1 < size2 since the memcpy would reference outside of the first alloca
-define void @negative_test(ptr %src, i8 %c, i64 %size1, i64 %size2) {
-; CHECK-LABEL: @negative_test(
+define void @dynsize_test(ptr %src, i8 %c, i64 %size1, i64 %size2) {
+; CHECK-LABEL: @dynsize_test(
 ; CHECK-NEXT: [[DST1:%.*]] = alloca i8, i64 [[SIZE1:%.*]], align 1
 ; CHECK-NEXT: [[DST2:%.*]] = alloca i8, i64 [[SIZE2:%.*]], align 1
 ; CHECK-NEXT: call void @llvm.memset.p0.i64(ptr align 8 [[DST1]], i8 [[C:%.*]], i64 [[SIZE1]], i1 false)
-; CHECK-NEXT: call void @llvm.memcpy.p0.p0.i64(ptr align 8 [[DST2]], ptr align 8 [[DST1]], i64 [[SIZE2]], i1 false)
+; CHECK-NEXT: call void @llvm.memset.p0.i64(ptr align 8 [[DST2]], i8 [[C]], i64 [[SIZE2]], i1 false)
 ; CHECK-NEXT: ret void
 ;
  %dst1 = alloca i8, i64 %size1