[DAGCombine] Allow scalable type dead store elimination

[dtemirbulatov]

Summary

Allow scalable type dead store elimination in DAGCombine.

[dtemirbulatov]

Missed to update sve-redundant-store.ll test after change, fixed.

[dtemirbulatov]

Fixed case where ST1 type size is bigger than in ST.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19926–19927

This looks really confusing. You can use TypeSize::isKnownLE, which also returns true when comparing sizes with different scalable properties, e.g. sizeof(v8i32) <= sizeof(nxv8i32).
You'd just need to add an additional check that the base pointer is the same.

Ideally we'd extend BaseIndexOffset to handle scalable sizes as well, but that's quite a bit of work and may not be worth it.

[dtemirbulatov]

Dropped (STBitSize <= ChainBitSize) checks, it is not required as STBase.contains() ignores those to handle.

[dtemirbulatov]

Addressed remakrs, using now TypeSize::isKnownLE.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19926–19927

For scalable vectors, can you avoid the call to STBase.contains() entirely and just check that the pointer is equal? Then you don't have to make the effort to compute the 'known' part of the STBitSize/ChainBitSize.

[dtemirbulatov]

Use TypeSize::isKnownLE only if base pointers are equal.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19915

nit: unintended change?

19926

I wonder if it's simpler to write this as follows:

if (ST->getMemoryVT().isScalableVector() ||
    ST1->getMemoryVT().isScalableVector()) {
  if (ST1->getBasePtr() == Ptr && TypeSize::isKnownLE(...))
    CombineTo(ST1, ST1->getChain());
} else {
  const BaseIndexOffset STBase = BaseIndexOffset::match(ST, DAG);
  const BaseIndexOffset ChainBase = BaseIndexOffset::match(ST1, DAG);
  if (STBase.contains(...))
    CombineTo(ST1, ST1->getChain());
}

That way the the logic between 'anything scalable' and everything else is clearly separated.

llvm/test/CodeGen/AArch64/sve-redundant-store.ll

21

This is missing some test cases, e.g.

• One where both stores are scalable, and same size.
• One where both stores are scalable, and second store is larger
• One where both stores are scalable, and first store is larger (negative test)

[dtemirbulatov]

Addressed comments.

[dtemirbulatov]

Fixed typo in test funtion name.

[sdesmalen]

Thanks for adding the tests. Just two more minor comments and then I'm happy to accept.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19931–19933

Can you combine the two if-statements with an &&? That saves a level of indentation.

19939

You can remove the } else { here, which saves a level of indentation.

[dtemirbulatov]

Combined two ifs

[dtemirbulatov]

Restoring the previous revision.

[sdesmalen]

LGTM!

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19931–19933

Sorry, I misread the indentation/curly braces here. The original version was fine, thanks for restoring it!

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (Ld->getBasePtr() == Ptr && ST->getMemoryVT() == Ld->getMemoryVT() &&
       // The store is dead, remove it.
       return Chain;
     }
   }
 
   // TODO: Can relax for unordered atomics (see D66309)
   if (StoreSDNode *ST1 = dyn_cast<StoreSDNode>(Chain)) {
     if (ST->isUnindexed() && ST->isSimple() &&
         ST1->isUnindexed() && ST1->isSimple()) {
       if (OptLevel != CodeGenOpt::None && ST1->getBasePtr() == Ptr &&
           ST1->getValue() == Value && ST->getMemoryVT() == ST1->getMemoryVT() &&
           ST->getAddressSpace() == ST1->getAddressSpace()) {
         // If this is a store followed by a store with the same value to the
         // same location, then the store is dead/noop.
         return Chain;
       }
 
       if (OptLevel != CodeGenOpt::None && ST1->hasOneUse() &&
           !ST1->getBasePtr().isUndef() &&
-          // BaseIndexOffset and the code below requires knowing the size
-          // of a vector, so bail out if MemoryVT is scalable.
-          !ST->getMemoryVT().isScalableVector() &&
-          !ST1->getMemoryVT().isScalableVector() &&
           ST->getAddressSpace() == ST1->getAddressSpace()) {
+        // If we consider two stores and one smaller in size is a scalable
+        // vector type and another one a bigger size store with a fixed type,
+        // then we could not allow the scalable store removal because we don't
+        // know its final size in the end.
+        if (ST->getMemoryVT().isScalableVector() ||
+            ST1->getMemoryVT().isScalableVector()) {
+          if (ST1->getBasePtr() == Ptr &&
+              TypeSize::isKnownLE(ST1->getMemoryVT().getStoreSize(),
+                                  ST->getMemoryVT().getStoreSize())) {
+            CombineTo(ST1, ST1->getChain());
+            return SDValue();
+          }
+        } else {
           const BaseIndexOffset STBase = BaseIndexOffset::match(ST, DAG);
           const BaseIndexOffset ChainBase = BaseIndexOffset::match(ST1, DAG);
-        unsigned STBitSize = ST->getMemoryVT().getFixedSizeInBits();
-        unsigned ChainBitSize = ST1->getMemoryVT().getFixedSizeInBits();
           // If this is a store who's preceding store to a subset of the current
           // location and no one other node is chained to that store we can
-        // effectively drop the store. Do not remove stores to undef as they may
-        // be used as data sinks.
-        if (STBase.contains(DAG, STBitSize, ChainBase, ChainBitSize)) {
+          // effectively drop the store. Do not remove stores to undef as they
+          // may be used as data sinks.
+          if (STBase.contains(DAG, ST->getMemoryVT().getFixedSizeInBits(),
+                              ChainBase,
+                              ST1->getMemoryVT().getFixedSizeInBits())) {
             CombineTo(ST1, ST1->getChain());
             return SDValue();
           }
         }
       }
     }
+  }
 
   // If this is an FP_ROUND or TRUNC followed by a store, fold this into a
   // truncating store.  We can do this even if this is already a truncstore.
   if ((Value.getOpcode() == ISD::FP_ROUND ||
        Value.getOpcode() == ISD::TRUNCATE) &&
       Value->hasOneUse() && ST->isUnindexed() &&
       TLI.canCombineTruncStore(Value.getOperand(0).getValueType(),
                                ST->getMemoryVT(), LegalOperations)) {

llvm/test/CodeGen/AArch64/sve-redundant-store.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -O2 -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
 
 ; #include <arm_sve.h>
 ; #include <stdint.h>
 ;
 ; void redundant_store(uint32_t *p, svint32_t v) {
 ;     *p = 1;
 ;     *(svint32_t *)p = v;
 ; }
-
-; Update me: Until dead store elimination is improved in DAGCombine, this will contain a redundant store.
-;
 define void @redundant_store(ptr nocapture %p, <vscale x 4 x i32> %v) {
 ; CHECK-LABEL: redundant_store:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w8, #1
 ; CHECK-NEXT:    ptrue p0.s
-; CHECK-NEXT:    str w8, [x0]
 ; CHECK-NEXT:    st1w { z0.s }, p0, [x0]
 ; CHECK-NEXT:    ret
   store i32 1, ptr %p, align 4
   store <vscale x 4 x i32> %v, <vscale x 4 x i32>* %p, align 16
   ret void
 }
+
+define void @two_scalable_same_size(ptr writeonly %ptr, <vscale x 4 x i32> %a, <vscale x 4 x i32> %b) {
+; CHECK-LABEL: two_scalable_same_size:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    st1w { z1.s }, p0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  store <vscale x 4 x i32> %a, ptr %ptr
+  store <vscale x 4 x i32> %b, ptr %ptr
+  ret void
+}
+
+; make sure that scalable store is present, becuase we don't know its final size.
+define void @keep_scalable_store(ptr writeonly %ptr, ptr %a, <vscale x 4 x i32> %b) {
+; CHECK-LABEL: keep_scalable_store:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ldp q2, q1, [x1]
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    st1w { z0.s }, p0, [x0]
+; CHECK-NEXT:    stp q2, q1, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %0 = load <8 x i32>, ptr %a
+  store <vscale x 4 x i32> %b, ptr %ptr
+  store <8 x i32> %0, ptr %ptr
+  ret void
+}
+
+define void @two_scalable_keep_stores(ptr writeonly %ptr, <vscale x 4 x i32> %a, <vscale x 4 x i64> %b) {
+; CHECK-LABEL: two_scalable_keep_stores:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    ptrue p1.s
+; CHECK-NEXT:    st1d { z2.d }, p0, [x0, #1, mul vl]
+; CHECK-NEXT:    st1d { z1.d }, p0, [x0]
+; CHECK-NEXT:    st1w { z0.s }, p1, [x0]
+; CHECK-NEXT:    ret
+entry:
+  store <vscale x 4 x i64> %b, ptr %ptr
+  store <vscale x 4 x i32> %a, ptr %ptr
+  ret void
+}
+
+define void @two_scalable_remove_store(ptr writeonly %ptr, <vscale x 4 x i32> %a, <vscale x 4 x i64> %b) {
+; CHECK-LABEL: two_scalable_remove_store:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    st1d { z2.d }, p0, [x0, #1, mul vl]
+; CHECK-NEXT:    st1d { z1.d }, p0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  store <vscale x 4 x i32> %a, ptr %ptr
+  store <vscale x 4 x i64> %b, ptr %ptr
+  ret void
+}