[InstCombine] Handle atomic memset in the same way as regular memset

[dneilson]

Summary

This change adds handling of the atomic memset intrinsic to the
code path that simplifies the regular memset. In practice this means
that we will now also expand a small constant-length atomic memset
into a single unordered atomic store.

[reames]

LGTM

[dneilson]

• Rebaseline prior to landing

lib/Transforms/InstCombine/InstCombineCalls.cpp

if (isa<AtomicMemTransferInst>(MI)) {
     S->setOrdering(AtomicOrdering::Unordered);
   }
 
   // Set the size of the copy to 0, it will be deleted on the next iteration.
   MI->setLength(Constant::getNullValue(MemOpLength->getType()));
   return MI;
 }
 
-Instruction *InstCombiner::SimplifyMemSet(MemSetInst *MI) {
+Instruction *InstCombiner::SimplifyAnyMemSet(AnyMemSetInst *MI) {
   unsigned Alignment = getKnownAlignment(MI->getDest(), DL, MI, &AC, &DT);
   if (MI->getDestAlignment() < Alignment) {
     MI->setDestAlignment(Alignment);
     return MI;
   }
 
   // Extract the length and alignment and fill if they are constant.
   ConstantInt *LenC = dyn_cast<ConstantInt>(MI->getLength());
if (Len <= 8 && isPowerOf2_32((uint32_t)Len)) {
     // Alignment 0 is identity for alignment 1 for memset, but not store.
     if (Alignment == 0) Alignment = 1;
 
     // Extract the fill value and store.
     uint64_t Fill = FillC->getZExtValue()*0x0101010101010101ULL;
     StoreInst *S = Builder.CreateStore(ConstantInt::get(ITy, Fill), Dest,
                                        MI->isVolatile());
     S->setAlignment(Alignment);
+    if (isa<AtomicMemSetInst>(MI))
+      S->setOrdering(AtomicOrdering::Unordered);
 
     // Set the size of the copy to 0, it will be deleted on the next iteration.
     MI->setLength(Constant::getNullValue(LenC->getType()));
     return MI;
   }
 
   return nullptr;
 }
if (AnyMemTransferInst *MTI = dyn_cast<AnyMemTransferInst>(MI)) {
         return eraseInstFromFunction(CI);
     }
 
     // If we can determine a pointer alignment that is bigger than currently
     // set, update the alignment.
     if (auto *MTI = dyn_cast<AnyMemTransferInst>(MI)) {
       if (Instruction *I = SimplifyAnyMemTransfer(MTI))
         return I;
-    } else if (MemSetInst *MSI = dyn_cast<MemSetInst>(MI)) {
-      if (Instruction *I = SimplifyMemSet(MSI))
+    } else if (auto *MSI = dyn_cast<AnyMemSetInst>(MI)) {
+      if (Instruction *I = SimplifyAnyMemSet(MSI))
         return I;
     }
 
     if (Changed) return II;
   }
 
   if (Instruction *I = SimplifyNVVMIntrinsic(II, *this))
     return I;

lib/Transforms/InstCombine/InstCombineInternal.h

private:
 
   Value *insertRangeTest(Value *V, const APInt &Lo, const APInt &Hi,
                          bool isSigned, bool Inside);
   Instruction *PromoteCastOfAllocation(BitCastInst &CI, AllocaInst &AI);
   Instruction *MatchBSwap(BinaryOperator &I);
   bool SimplifyStoreAtEndOfBlock(StoreInst &SI);
 
   Instruction *SimplifyAnyMemTransfer(AnyMemTransferInst *MI);
-  Instruction *SimplifyMemSet(MemSetInst *MI);
+  Instruction *SimplifyAnyMemSet(AnyMemSetInst *MI);
 
   Value *EvaluateInDifferentType(Value *V, Type *Ty, bool isSigned);
 
   /// Returns a value X such that Val = X * Scale, or null if none.
   ///
   /// If the multiplication is known not to overflow then NoSignedWrap is set.
   Value *Descale(Value *Val, APInt Scale, bool &NoSignedWrap);
 };
 
 } // end namespace llvm
 
 #undef DEBUG_TYPE
 
 #endif // LLVM_LIB_TRANSFORMS_INSTCOMBINE_INSTCOMBINEINTERNAL_H

test/Transforms/InstCombine/element-atomic-memintrins.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -instcombine -S < %s | FileCheck %s
 
 ;; ---- memset -----
 
 ; Ensure 0-length memset is removed
 define void @test_memset_zero_length(i8* %dest) {
 ; CHECK-LABEL: @test_memset_zero_length(
 ; CHECK-NEXT:    ret void
 ;
   call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 1 %dest, i8 1, i32 0, i32 1)
   ret void
 }
 
-; Placeholder test. This will chance once support for lowering atomic memsets is added to instcombine.
 define void @test_memset_to_store(i8* %dest) {
 ; CHECK-LABEL: @test_memset_to_store(
-; CHECK-NEXT:    call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 1 [[DEST:%.*]], i8 1, i32 1, i32 1)
-; CHECK-NEXT:    call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 1 [[DEST]], i8 1, i32 2, i32 1)
-; CHECK-NEXT:    call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 1 [[DEST]], i8 1, i32 4, i32 1)
-; CHECK-NEXT:    call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 1 [[DEST]], i8 1, i32 8, i32 1)
+; CHECK-NEXT:    store atomic i8 1, i8* [[DEST:%.*]] unordered, align 1
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[DEST]] to i16*
+; CHECK-NEXT:    store atomic i16 257, i16* [[TMP1]] unordered, align 1
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast i8* [[DEST]] to i32*
+; CHECK-NEXT:    store atomic i32 16843009, i32* [[TMP2]] unordered, align 1
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast i8* [[DEST]] to i64*
+; CHECK-NEXT:    store atomic i64 72340172838076673, i64* [[TMP3]] unordered, align 1
 ; CHECK-NEXT:    call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 1 [[DEST]], i8 1, i32 16, i32 1)
 ; CHECK-NEXT:    ret void
 ;
   call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 1 %dest, i8 1, i32 1, i32 1)
   call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 1 %dest, i8 1, i32 2, i32 1)
   call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 1 %dest, i8 1, i32 4, i32 1)
   call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 1 %dest, i8 1, i32 8, i32 1)
   call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 1 %dest, i8 1, i32 16, i32 1)
   ret void
 }
 
+define void @test_memset_to_store_2(i8* %dest) {
+; CHECK-LABEL: @test_memset_to_store_2(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[DEST:%.*]] to i16*
+; CHECK-NEXT:    store atomic i16 257, i16* [[TMP1]] unordered, align 2
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast i8* [[DEST]] to i32*
+; CHECK-NEXT:    store atomic i32 16843009, i32* [[TMP2]] unordered, align 2
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast i8* [[DEST]] to i64*
+; CHECK-NEXT:    store atomic i64 72340172838076673, i64* [[TMP3]] unordered, align 2
+; CHECK-NEXT:    call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 2 [[DEST]], i8 1, i32 16, i32 2)
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 2 %dest, i8 1, i32 2, i32 2)
+  call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 2 %dest, i8 1, i32 4, i32 2)
+  call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 2 %dest, i8 1, i32 8, i32 2)
+  call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 2 %dest, i8 1, i32 16, i32 2)
+  ret void
+}
+
+define void @test_memset_to_store_4(i8* %dest) {
+; CHECK-LABEL: @test_memset_to_store_4(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[DEST:%.*]] to i32*
+; CHECK-NEXT:    store atomic i32 16843009, i32* [[TMP1]] unordered, align 4
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast i8* [[DEST]] to i64*
+; CHECK-NEXT:    store atomic i64 72340172838076673, i64* [[TMP2]] unordered, align 4
+; CHECK-NEXT:    call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 4 [[DEST]], i8 1, i32 16, i32 4)
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 4 %dest, i8 1, i32 4, i32 4)
+  call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 4 %dest, i8 1, i32 8, i32 4)
+  call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 4 %dest, i8 1, i32 16, i32 4)
+  ret void
+}
+
+define void @test_memset_to_store_8(i8* %dest) {
+; CHECK-LABEL: @test_memset_to_store_8(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[DEST:%.*]] to i64*
+; CHECK-NEXT:    store atomic i64 72340172838076673, i64* [[TMP1]] unordered, align 8
+; CHECK-NEXT:    call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 8 [[DEST]], i8 1, i32 16, i32 8)
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 8 %dest, i8 1, i32 8, i32 8)
+  call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 8 %dest, i8 1, i32 16, i32 8)
+  ret void
+}
+
+define void @test_memset_to_store_16(i8* %dest) {
+; CHECK-LABEL: @test_memset_to_store_16(
+; CHECK-NEXT:    call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 16 [[DEST:%.*]], i8 1, i32 16, i32 16)
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* align 16 %dest, i8 1, i32 16, i32 16)
+  ret void
+}
+
 declare void @llvm.memset.element.unordered.atomic.p0i8.i32(i8* nocapture writeonly, i8, i32, i32) nounwind argmemonly
 
 
 ;; =========================================
 ;; ----- memmove ------
 
 
 @gconst = constant [32 x i8] c"0123456789012345678901234567890\00"