[DAGCombine] Allow DAGCombine to remove dead masked stores

[dtemirbulatov]

Summary

Remove a dead masked store if another one has the same base pointer and mask.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11167

Is this check necessary? I would expect that it doesn't run the DAGCombiner when OptLevel == CodeGenOpt::None. Does it still run the DAGCombiner with -O0?

11168

A store returns a Chain, which always has one use, so this test can be removed.

11169

This seems like an odd case to ever happen, and you did not add a specific test for this case. Why did you add this check?

11170

Is this check necessary? I would expect that if both pointers are equal, they use the same address space?

11172

I think this can be isKnownEQ, because for the predicates to match, the sizes must match too.

llvm/test/CodeGen/AArch64/sve-dead-masked-store.ll

12

Another case that can be handled is where the second store has an 'all true' mask, then it doesn't really matter what the mask of the first store is, e.g.

%alltrue.ins = insertelement <vscale x 4 x i1> poison, i1 true, i32 0
%alltrue = shufflevector  <vscale x 4 x i1> %splat.ins,  <vscale x 4 x i1> poison,  <vscale x 4 x i32> zeroinitializer
call void @llvm.masked.store.nxv4i32(<vscale x 4 x i32> %val, ptr %a, i32 4, <vscale x 4 x i1> %mask)
call void @llvm.masked.store.nxv4i32(<vscale x 4 x i32> %val, ptr %a, i32 4, <vscale x 4 x i1> %alltrue)

[dtemirbulatov]

Resolved comments.

[dtemirbulatov]

I found an error in my implementation: we could not remove the store if the chained store is a fixed type and the store we consider to remove is a scalable type, since we don't know scalable type size in the runtime. fixed.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11121–11124

Can this be replaced by:

ElementCount::isKnownLE(MST1->getMemoryVT().getStoreSize(),
                        MST->getMemoryVT().getStoreSize())

?

[dtemirbulatov]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11121–11124

hmm, I could image somethis like:

%alltrue.ins = insertelement <vscale x 4 x i1> poison, i1 true, i32 0
%alltrue = shufflevector  <vscale x 4 x i1> %alltrue.ins,  <vscale x 4 x i1> poison,  <vscale x 4 x i32> zeroinitializer
call void @llvm.masked.store.nxv4i32(<vscale x 4 x i32> %val, ptr %a, i32 4, <vscale x 4 x i1> %mask)
call void @llvm.masked.store.nxv4i16(<vscale x 4 x i16> %val1, ptr %a, i32 4, <vscale x 4 x i1> %alltrue)

that ends up with st1h, probably TypeSize::KnownLE suits better here.

[dtemirbulatov]

Replaced to TypeSize::isKnownLE() usage for all constant true case. Added few tests.

[sdesmalen]

Thanks for adding the tests @dtemirbulatov! I've just got two more nits on the tests, then I'm happy to accept.

llvm/test/CodeGen/AArch64/sve-dead-masked-store.ll

37–38

Can you change this to:

call void @llvm.masked.store.nxv4i16(<vscale x 4 x i16> %val, ptr %a, i32 4, <vscale x 4 x i1> %mask)
call void @llvm.masked.store.nxv4i32(<vscale x 4 x i32> %val1, ptr %a, i32 4, <vscale x 4 x i1> %mask) ; with same value for the mask

such that it's more clear that the first store is eliminated?

Now it code-generates the store of nxv4i64 into two stores, because nxv4i64 is not a legal type and needs splitting.

52

For this test, can you use the same %mask value?

[dtemirbulatov]

Adjust according to comments, fixed error with the same mask but different type size by forbid to delete the store.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11119

fixed error with the same mask but different type size by forbid to delete the store

I'm not sure if this is the reason you made this change, but the test that previously removed the redundant store:

define void @dead_masked_store_same_mask_bigger_type(<vscale x 4 x i16> %val, <vscale x 4 x i32> %val1, ptr %a, <vscale x 4 x i1> %mask) {
; CHECK-LABEL: dead_masked_store_same_mask_bigger_type:
; CHECK:       // %bb.0:
; CHECK-NEXT:    st1w { z1.s }, p0, [x0]
; CHECK-NEXT:    ret
  call void @llvm.masked.store.nxv4i16(<vscale x 4 x i16> %val, ptr %a, i32 4, <vscale x 4 x i1> %mask)
  call void @llvm.masked.store.nxv4i32(<vscale x 4 x i32> %val1, ptr %a, i32 4, <vscale x 4 x i1> %mask)
  ret void
}

Is now doing this:

define void @dead_masked_store_same_mask_bigger_type(<vscale x 4 x i16> %val, <vscale x 4 x i32> %val1, ptr %a, <vscale x 4 x i1> %mask) {
; CHECK-LABEL: dead_masked_store_same_mask_bigger_type:
; CHECK:       // %bb.0:
; CHECK-NEXT:    st1h { z0.s }, p0, [x0]
; CHECK-NEXT:    st1w { z1.s }, p0, [x0]
; CHECK-NEXT:    ret
  call void @llvm.masked.store.nxv4i16(<vscale x 4 x i16> %val, ptr %a, i32 4, <vscale x 4 x i1> %mask)
  call void @llvm.masked.store.nxv4i32(<vscale x 4 x i32> %val1, ptr %a, i32 4, <vscale x 4 x i1> %mask)
  ret void
}

The second store is overwriting all the i16 elements that were stored in the first llvm.masked.store, with i32 elements from the second llvm.masked.store, so that means the first store can be removed right?

[dtemirbulatov]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11119

The mask is the same for both stores, but because we types are different starting then offsets of starting and ending affected elements different too. I think we have to keep both store in those cases.

[dtemirbulatov]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11119

The mask is the same for both stores, but because types are different starting then offsets of starting and ending affected elements different too. I think we have to keep both store in those cases.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

11119

You're right, your current code is correct. The elements wouldn't necessarily line-up in memory, e.g.

call void @llvm.masked.store.v4i16(<4 x i16> <a, b, c, d>, ptr %a, i32 4, <4 x i1> <1, 0, 0, 1>)  ; would store `a` and `d` resulting in   { a__d____ }
call void @llvm.masked.store.v4i32(<4 x i32> <e, f, g, h>, ptr %a, i32 4, <4 x i1> <1, 0, 0, 1>)   ; would store `e` and `h` resulting in  { ee_d__hh }

so the first call can indeed not be removed.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitMSTORE(SDNode *N) {
   if (ISD::isConstantSplatVectorAllZeros(Mask.getNode()))
     return Chain;
 
   // If this is a masked load with an all ones mask, we can use a unmasked load.
   // FIXME: Can we do this for indexed, compressing, or truncating stores?
   if (ISD::isConstantSplatVectorAllOnes(Mask.getNode()) && MST->isUnindexed() &&
       !MST->isCompressingStore() && !MST->isTruncatingStore())
     return DAG.getStore(MST->getChain(), SDLoc(N), MST->getValue(),
                         MST->getBasePtr(), MST->getPointerInfo(),
                         MST->getOriginalAlign(), MachineMemOperand::MOStore,
                         MST->getAAInfo());
 
   // Try transforming N to an indexed store.
   if (CombineToPreIndexedLoadStore(N) || CombineToPostIndexedLoadStore(N))
     return SDValue(N, 0);
 
   if (MST->isTruncatingStore() && MST->isUnindexed() &&
       Value.getValueType().isInteger() &&
       (!isa<ConstantSDNode>(Value) ||
        !cast<ConstantSDNode>(Value)->isOpaque())) {
     APInt TruncDemandedBits =
         APInt::getLowBitsSet(Value.getScalarValueSizeInBits(),
if ((Value.getOpcode() == ISD::TRUNCATE) && Value->hasOneUse() &&
     auto Mask = TLI.promoteTargetBoolean(DAG, MST->getMask(),
                                          Value.getOperand(0).getValueType());
     return DAG.getMaskedStore(Chain, SDLoc(N), Value.getOperand(0), Ptr,
                               MST->getOffset(), Mask, MST->getMemoryVT(),
                               MST->getMemOperand(), MST->getAddressingMode(),
                               /*IsTruncating=*/true);
   }
 
+  // Remove a masked store if base pointers and masks are equal.
+  if (MaskedStoreSDNode *MST1 = dyn_cast<MaskedStoreSDNode>(Chain)) {
+    if (MST->isUnindexed() && MST->isSimple() && MST1->isUnindexed() &&
+        MST1->isSimple() && OptLevel != CodeGenOpt::None &&
+        MST1->getBasePtr() == Ptr && MST1->hasOneUse() &&
+        !MST1->getBasePtr().isUndef() &&
+        MST->getAddressSpace() == MST1->getAddressSpace() &&
+        Mask == MST1->getMask() &&
+        TypeSize::isKnownLE(MST1->getMemoryVT().getStoreSize(),
+                            MST->getMemoryVT().getStoreSize())) {
+      CombineTo(MST1, MST1->getChain());
+      if (N->getOpcode() != ISD::DELETED_NODE)
+        AddToWorklist(N);
+      return SDValue(N, 0);
+    }
+  }
+
   return SDValue();
 }
 
 SDValue DAGCombiner::visitVPGATHER(SDNode *N) {
   VPGatherSDNode *MGT = cast<VPGatherSDNode>(N);
   SDValue Mask = MGT->getMask();
   SDValue Chain = MGT->getChain();
   SDValue Index = MGT->getIndex();

llvm/test/CodeGen/AArch64/sve-dead-masked-store.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64--linux-gnu -mattr=+sve < %s | FileCheck %s
+
+define void @dead_masked_store(<vscale x 4 x i32> %val, ptr %a, <vscale x 4 x i1> %mask) nounwind {
+; CHECK-LABEL: dead_masked_store:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    st1w { z0.s }, p0, [x0]
+; CHECK-NEXT:    ret
+  call void @llvm.masked.store.nxv4i32(<vscale x 4 x i32> %val, ptr %a, i32 4, <vscale x 4 x i1> %mask)
+  call void @llvm.masked.store.nxv4i32(<vscale x 4 x i32> %val, ptr %a, i32 4, <vscale x 4 x i1> %mask)
+  ret void
+}
+declare void @llvm.masked.store.nxv4i32(<vscale x 4 x i32>, <vscale x 4 x i32>*, i32, <vscale x 4 x i1>)

llvm/test/CodeGen/X86/masked_store.ll

 ; SSE4-NEXT:    testb $8, %al
 ; SSE4-NEXT:    je LBB30_16
 ; SSE4-NEXT:  LBB30_15: ## %cond.store14
 ; SSE4-NEXT:    extractps $3, %xmm1, 12(%rdi)
 ; SSE4-NEXT:    retq
 ;
 ; AVX1OR2-LABEL: PR11210:
 ; AVX1OR2:       ## %bb.0:
-; AVX1OR2-NEXT:    vmaskmovps %xmm0, %xmm2, (%rdi)
 ; AVX1OR2-NEXT:    vmaskmovps %xmm1, %xmm2, (%rdi)
 ; AVX1OR2-NEXT:    retq
 ;
 ; AVX512F-LABEL: PR11210:
 ; AVX512F:       ## %bb.0:
 ; AVX512F-NEXT:    ## kill: def $xmm2 killed $xmm2 def $zmm2
 ; AVX512F-NEXT:    ## kill: def $xmm1 killed $xmm1 def $zmm1
-; AVX512F-NEXT:    ## kill: def $xmm0 killed $xmm0 def $zmm0
-; AVX512F-NEXT:    vpxor %xmm3, %xmm3, %xmm3
-; AVX512F-NEXT:    vpcmpgtd %zmm2, %zmm3, %k0
+; AVX512F-NEXT:    vpxor %xmm0, %xmm0, %xmm0
+; AVX512F-NEXT:    vpcmpgtd %zmm2, %zmm0, %k0
 ; AVX512F-NEXT:    kshiftlw $12, %k0, %k0
 ; AVX512F-NEXT:    kshiftrw $12, %k0, %k1
-; AVX512F-NEXT:    vmovups %zmm0, (%rdi) {%k1}
 ; AVX512F-NEXT:    vmovups %zmm1, (%rdi) {%k1}
 ; AVX512F-NEXT:    vzeroupper
 ; AVX512F-NEXT:    retq
 ;
 ; AVX512VLDQ-LABEL: PR11210:
 ; AVX512VLDQ:       ## %bb.0:
 ; AVX512VLDQ-NEXT:    vpmovd2m %xmm2, %k1
-; AVX512VLDQ-NEXT:    vmovups %xmm0, (%rdi) {%k1}
 ; AVX512VLDQ-NEXT:    vmovups %xmm1, (%rdi) {%k1}
 ; AVX512VLDQ-NEXT:    retq
 ;
 ; AVX512VLBW-LABEL: PR11210:
 ; AVX512VLBW:       ## %bb.0:
-; AVX512VLBW-NEXT:    vpxor %xmm3, %xmm3, %xmm3
-; AVX512VLBW-NEXT:    vpcmpgtd %xmm2, %xmm3, %k1
-; AVX512VLBW-NEXT:    vmovups %xmm0, (%rdi) {%k1}
+; AVX512VLBW-NEXT:    vpxor %xmm0, %xmm0, %xmm0
+; AVX512VLBW-NEXT:    vpcmpgtd %xmm2, %xmm0, %k1
 ; AVX512VLBW-NEXT:    vmovups %xmm1, (%rdi) {%k1}
 ; AVX512VLBW-NEXT:    retq
 ;
 ; X86-AVX512-LABEL: PR11210:
 ; X86-AVX512:       ## %bb.0:
 ; X86-AVX512-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-AVX512-NEXT:    vpmovd2m %xmm2, %k1
-; X86-AVX512-NEXT:    vmovups %xmm0, (%eax) {%k1}
 ; X86-AVX512-NEXT:    vmovups %xmm1, (%eax) {%k1}
 ; X86-AVX512-NEXT:    retl
   %bc = bitcast <2 x i64> %mask to <4 x i32>
   %trunc = icmp slt <4 x i32> %bc, zeroinitializer
   call void @llvm.masked.store.v4f32.p0(<4 x float> %x, ptr %ptr, i32 1, <4 x i1> %trunc)
   call void @llvm.masked.store.v4f32.p0(<4 x float> %y, ptr %ptr, i32 1, <4 x i1> %trunc)
   ret void
 }