[InstCombine] Perform memset -> load forwarding

[nikic]

Summary

InstCombine does some basic store to load forwarding. One case it currently misses is the case where the store is actually a memset. This patch adds support for this case. This is a minimal implementation that only handles a load at the memset base address, without an offset.

GVN is already capable of performing this optimization. Having it in InstCombine can help with phase ordering issues, similar to the existing store to load forwarding.

[spatel]

llvm/test/Transforms/InstCombine/load-store-forward.ll

282–283

Do we have a negative test like this except we're 1-bit short of covering the loaded amount?
Similarly, add a negative test where there's a store or some other interference between the memset and the load?

[nikic]

Add another test

[nikic]

llvm/test/Transforms/InstCombine/load-store-forward.ll

282–283

Do we have a negative test like this except we're 1-bit short of covering the loaded amount?

There was load_after_memset_0_too_small, but that's using a significantly larger load. I just added load_after_memset_0_too_small_by_one_bit to have a one bit size mismatch.

Similarly, add a negative test where there's a store or some other interference between the memset and the load?

This is covered by the load_after_memset_0_clobber test.

[spatel]

LGTM

llvm/lib/Analysis/Loads.cpp

if (StoreInst *SI = dyn_cast<StoreInst>(Inst)) {
 
     TypeSize StoreSize = DL.getTypeSizeInBits(Val->getType());
     TypeSize LoadSize = DL.getTypeSizeInBits(AccessTy);
     if (TypeSize::isKnownLE(LoadSize, StoreSize))
       if (auto *C = dyn_cast<Constant>(Val))
         return ConstantFoldLoadFromConst(C, AccessTy, DL);
   }
 
+  if (auto *MSI = dyn_cast<MemSetInst>(Inst)) {
+    // Don't forward from (non-atomic) memset to atomic load.
+    if (AtLeastAtomic)
+      return nullptr;
+
+    // Only handle constant memsets.
+    auto *Val = dyn_cast<ConstantInt>(MSI->getValue());
+    auto *Len = dyn_cast<ConstantInt>(MSI->getLength());
+    if (!Val || !Len)
+      return nullptr;
+
+    // TODO: Handle offsets.
+    Value *Dst = MSI->getDest();
+    if (!AreEquivalentAddressValues(Dst, Ptr))
+      return nullptr;
+
+    if (IsLoadCSE)
+      *IsLoadCSE = false;
+
+    // Make sure the read bytes are contained in the memset.
+    TypeSize LoadSize = DL.getTypeSizeInBits(AccessTy);
+    if (LoadSize.isScalable() ||
+        (Len->getValue() * 8).ult(LoadSize.getFixedSize()))
+      return nullptr;
+
+    APInt Splat = APInt::getSplat(LoadSize.getFixedSize(), Val->getValue());
+    ConstantInt *SplatC = ConstantInt::get(MSI->getContext(), Splat);
+    if (CastInst::isBitOrNoopPointerCastable(SplatC->getType(), AccessTy, DL))
+      return SplatC;
+
+    return nullptr;
+  }
+
   return nullptr;
 }
 
 Value *llvm::findAvailablePtrLoadStore(
     const MemoryLocation &Loc, Type *AccessTy, bool AtLeastAtomic,
     BasicBlock *ScanBB, BasicBlock::iterator &ScanFrom, unsigned MaxInstsToScan,
     AAResults *AA, bool *IsLoadCSE, unsigned *NumScanedInst) {
   if (MaxInstsToScan == 0)

llvm/test/Transforms/InstCombine/load-store-forward.ll

;
   store i8 1, ptr %a
   %v = load i1, ptr %a
   ret i1 %v
 }
 
 define i32 @load_after_memset_0(ptr %a) {
 ; CHECK-LABEL: @load_after_memset_0(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 0, i64 16, i1 false)
-; CHECK-NEXT:    [[V:%.*]] = load i32, ptr [[A]], align 4
-; CHECK-NEXT:    ret i32 [[V]]
+; CHECK-NEXT:    ret i32 0
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 0, i64 16, i1 false)
   %v = load i32, ptr %a
   ret i32 %v
 }
 
 define float @load_after_memset_0_float(ptr %a) {
 ; CHECK-LABEL: @load_after_memset_0_float(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 0, i64 16, i1 false)
-; CHECK-NEXT:    [[V:%.*]] = load float, ptr [[A]], align 4
-; CHECK-NEXT:    ret float [[V]]
+; CHECK-NEXT:    ret float 0.000000e+00
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 0, i64 16, i1 false)
   %v = load float, ptr %a
   ret float %v
 }
 
 define i27 @load_after_memset_0_non_byte_sized(ptr %a) {
 ; CHECK-LABEL: @load_after_memset_0_non_byte_sized(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 0, i64 16, i1 false)
-; CHECK-NEXT:    [[V:%.*]] = load i27, ptr [[A]], align 4
-; CHECK-NEXT:    ret i27 [[V]]
+; CHECK-NEXT:    ret i27 0
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 0, i64 16, i1 false)
   %v = load i27, ptr %a
   ret i27 %v
 }
 
 define <4 x i8> @load_after_memset_0_vec(ptr %a) {
 ; CHECK-LABEL: @load_after_memset_0_vec(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 0, i64 16, i1 false)
-; CHECK-NEXT:    [[V:%.*]] = load <4 x i8>, ptr [[A]], align 4
-; CHECK-NEXT:    ret <4 x i8> [[V]]
+; CHECK-NEXT:    ret <4 x i8> zeroinitializer
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 0, i64 16, i1 false)
   %v = load <4 x i8>, ptr %a
   ret <4 x i8> %v
 }
 
 define i32 @load_after_memset_1(ptr %a) {
 ; CHECK-LABEL: @load_after_memset_1(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 1, i64 16, i1 false)
-; CHECK-NEXT:    [[V:%.*]] = load i32, ptr [[A]], align 4
-; CHECK-NEXT:    ret i32 [[V]]
+; CHECK-NEXT:    ret i32 16843009
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 1, i64 16, i1 false)
   %v = load i32, ptr %a
   ret i32 %v
 }
 
 define float @load_after_memset_1_float(ptr %a) {
 ; CHECK-LABEL: @load_after_memset_1_float(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 1, i64 16, i1 false)
-; CHECK-NEXT:    [[V:%.*]] = load float, ptr [[A]], align 4
-; CHECK-NEXT:    ret float [[V]]
+; CHECK-NEXT:    ret float 0x3820202020000000
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 1, i64 16, i1 false)
   %v = load float, ptr %a
   ret float %v
 }
 
 define i27 @load_after_memset_1_non_byte_sized(ptr %a) {
 ; CHECK-LABEL: @load_after_memset_1_non_byte_sized(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 1, i64 16, i1 false)
-; CHECK-NEXT:    [[V:%.*]] = load i27, ptr [[A]], align 4
-; CHECK-NEXT:    ret i27 [[V]]
+; CHECK-NEXT:    ret i27 16843009
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 1, i64 16, i1 false)
   %v = load i27, ptr %a
   ret i27 %v
 }
 
 define <4 x i8> @load_after_memset_1_vec(ptr %a) {
 ; CHECK-LABEL: @load_after_memset_1_vec(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 1, i64 16, i1 false)
-; CHECK-NEXT:    [[V:%.*]] = load <4 x i8>, ptr [[A]], align 4
-; CHECK-NEXT:    ret <4 x i8> [[V]]
+; CHECK-NEXT:    ret <4 x i8> <i8 1, i8 1, i8 1, i8 1>
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 1, i64 16, i1 false)
   %v = load <4 x i8>, ptr %a
   ret <4 x i8> %v
 }
 
 define i32 @load_after_memset_unknown(ptr %a, i8 %byte) {
 ; CHECK-LABEL: @load_after_memset_unknown(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 [[BYTE:%.*]], i64 16, i1 false)
 ; CHECK-NEXT:    [[V:%.*]] = load i32, ptr [[A]], align 4
 ; CHECK-NEXT:    ret i32 [[V]]
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 %byte, i64 16, i1 false)
   %v = load i32, ptr %a
   ret i32 %v
 }
 
+; TODO: Handle load at offset.
 define i32 @load_after_memset_0_offset(ptr %a) {
 ; CHECK-LABEL: @load_after_memset_0_offset(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 0, i64 16, i1 false)
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[A]], i64 4
 ; CHECK-NEXT:    [[V:%.*]] = load i32, ptr [[GEP]], align 4
 ; CHECK-NEXT:    ret i32 [[V]]
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 0, i64 16, i1 false)

 ; CHECK-NEXT:    [[V:%.*]] = load i256, ptr [[A]], align 4
 ; CHECK-NEXT:    ret i256 [[V]]
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 0, i64 16, i1 false)
   %v = load i256, ptr %a
   ret i256 %v
 }
 
+define i129 @load_after_memset_0_too_small_by_one_bit(ptr %a) {
+; CHECK-LABEL: @load_after_memset_0_too_small_by_one_bit(
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 0, i64 16, i1 false)
+; CHECK-NEXT:    [[V:%.*]] = load i129, ptr [[A]], align 4
+; CHECK-NEXT:    ret i129 [[V]]
+;
+  call void @llvm.memset.p0.i64(ptr %a, i8 0, i64 16, i1 false)
+  %v = load i129, ptr %a
+  ret i129 %v
+}
+
 define i32 @load_after_memset_0_unknown_length(ptr %a, i64 %len) {
 ; CHECK-LABEL: @load_after_memset_0_unknown_length(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr align 1 [[A:%.*]], i8 0, i64 [[LEN:%.*]], i1 false)
 ; CHECK-NEXT:    [[V:%.*]] = load i32, ptr [[A]], align 4
 ; CHECK-NEXT:    ret i32 [[V]]
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 0, i64 %len, i1 false)
   %v = load i32, ptr %a