[RISCV] Use scalar stores for splats of zero to memory up to XLen

[reames]

Summary

The direct motivation here is to undo an unprofitable vectorization performed by SLP, but the transform seems generally useful as well. If we are storing a zero to memory, we can use a single scalar store (from X0) for all power of two sizes up to XLen.

Note: We can extend this transform in a bunch of ways. I'm deliberately starting narrow to focus on the two questions highlighted just above.

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

12176

Did you intend to leave FLAGIT in here?

[reames]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

12176

Nope. :) That's my usual easy grep flag for local work, will remove.

[reames]

Noticed while looking at something else that DAG combine store merging will produce the <2 x i64> vector stores this is splitting. The particular test case came from memset lowering. Oddly, this seeming conflict appears to work out, but it feels like we should really fix the store merging preference before handling the two stores part of this.

(following is removed from review description as it no longer applies, saved only for my later reference)

One concern here is that splitting the store into two instructions could introduce a store-to-load forwarding stall. I think this is worth doing, but what do others think? Should we restrict this to a maximum of XLEN sized operations which can be done in a single store?

This transform is potentially increasing the number of memory operations (from 1 to 2). The same restriction mentioned above would resolve this as well.

[craig.topper]

This seems reasonable to me.

[luke]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

12189

Do we want to check if the memory access is fast? Currently we always report unaligned scalar accesses as slow, but that sounds like it's at odds with this combine

bool RISCVTargetLowering::allowsMisalignedMemoryAccesses(
    EVT VT, unsigned AddrSpace, Align Alignment, MachineMemOperand::Flags Flags,
    unsigned *Fast) const {
  if (!VT.isVector()) {
    if (Fast)
      *Fast = 0;
    return Subtarget.enableUnalignedScalarMem();
  }

[asb]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

12189

Luke and I chatted about this earlier today, and the specific case would be where a vector access that is aligned according to the natural alignment of its elements is converted to a wider scalar access that is misaligned. Though now reminding myself of the how +unaligned-scalar-mem is used (https://reviews.llvm.org/D126085), this is probably not a concern, as it's only enabled if unaligned scalar mem is performant enough to be worth using.

[luke]

[reames]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

12189

Honestly, this is whole area is a bit of a mess. What exactly the fast flag means (it's a unsigned so it has many possible states) is very poorly defined. It's also extremely target specific, but interacts with generic codegen, so the best kind of fun.

At the moment, we're interpreting it as basically a hint that misaligned accesses should be aggressively formed - as opposed to simply used when natural during lowering.

In this case, I think moving forward without Fast is the right practical answer, but this is a bit of a code smell.

Thanks for pointing this out, I'd not considered it. I'm going to be looking at a couple other aspects of mem-op lowering, and I'll see if I can clean this up.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

case ISD::STORE: {
     if (DCI.isAfterLegalizeDAG())
       if (SDValue V = performMemPairCombine(N, DCI))
         return V;
 
     if (N->getOpcode() != ISD::STORE)
       break;
 
     auto *Store = cast<StoreSDNode>(N);
+    EVT MemVT = Store->getMemoryVT();
     SDValue Val = Store->getValue();
+
+    // Using vector to store zeros requires e.g.:
+    //   vsetivli   zero, 2, e64, m1, ta, ma
+    //   vmv.v.i    v8, 0
+    //   vse64.v    v8, (a0)
+    // If sufficiently aligned, we can use at most one scalar store to zero
+    // initialize any power-of-two size up to XLen bits.
+    if (DCI.isBeforeLegalize() && !Store->isTruncatingStore() &&
+        !Store->isIndexed() && ISD::isBuildVectorAllZeros(Val.getNode()) &&
+        MemVT.getVectorElementType().bitsLE(Subtarget.getXLenVT()) &&
+        isPowerOf2_64(MemVT.getSizeInBits()) &&
+        MemVT.getSizeInBits() <= Subtarget.getXLen()) {
+      assert(!MemVT.isScalableVector());
+      auto NewVT = MVT::getIntegerVT(MemVT.getSizeInBits());
+      if (allowsMemoryAccessForAlignment(*DAG.getContext(), DAG.getDataLayout(),
+                                         NewVT, *Store->getMemOperand())) {
+        SDLoc DL(N);
+        SDValue Chain = Store->getChain();
+        auto NewV = DAG.getConstant(0, DL, NewVT);
+        return DAG.getStore(Chain, DL, NewV, Store->getBasePtr(),
+                            Store->getPointerInfo(), Store->getOriginalAlign(),
+                            Store->getMemOperand()->getFlags());
+      }
+    }
+
     // Combine store of vmv.x.s/vfmv.f.s to vse with VL of 1.
     // vfmv.f.s is represented as extract element from 0. Match it late to avoid
     // any illegal types.
     if (Val.getOpcode() == RISCVISD::VMV_X_S ||
         (DCI.isAfterLegalizeDAG() &&
          Val.getOpcode() == ISD::EXTRACT_VECTOR_ELT &&
          isNullConstant(Val.getOperand(1)))) {
       SDValue Src = Val.getOperand(0);
       MVT VecVT = Src.getSimpleValueType();
-      EVT MemVT = Store->getMemoryVT();
       // VecVT should be scalable and memory VT should match the element type.
       if (VecVT.isScalableVector() &&
           MemVT == VecVT.getVectorElementType()) {
         SDLoc DL(N);
         MVT MaskVT = getMaskTypeFor(VecVT);
         return DAG.getStoreVP(
             Store->getChain(), DL, Src, Store->getBasePtr(), Store->getOffset(),
             DAG.getConstant(1, DL, MaskVT),

llvm/test/CodeGen/RISCV/rvv/fixed-vectors-int-splat.ll

; LMULMAX1-RV64-NEXT:    ret
   %b = shufflevector <4 x i64> %a, <4 x i64> poison, <4 x i32> zeroinitializer
   store <4 x i64> %b, ptr %x
   ret void
 }
 
 define void @splat_zero_v2i16(ptr %p) {
 ; CHECK-LABEL: splat_zero_v2i16:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    vsetivli zero, 2, e16, mf4, ta, ma
-; CHECK-NEXT:    vmv.v.i v8, 0
-; CHECK-NEXT:    vse16.v v8, (a0)
+; CHECK-NEXT:    sw zero, 0(a0)
 ; CHECK-NEXT:    ret
   store <2 x i16> zeroinitializer, ptr %p
   ret void
 }
 
 define void @splat_zero_v2i16_unaligned(ptr %p) {
 ; CHECK-LABEL: splat_zero_v2i16_unaligned:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 2, e16, mf4, ta, ma
 ; CHECK-NEXT:    vmv.v.i v8, 0
 ; CHECK-NEXT:    vsetivli zero, 4, e8, mf4, ta, ma
 ; CHECK-NEXT:    vse8.v v8, (a0)
 ; CHECK-NEXT:    ret
   store <2 x i16> zeroinitializer, ptr %p, align 1
   ret void
 }
 
 define void @splat_zero_v4i16(ptr %p) {
-; CHECK-LABEL: splat_zero_v4i16:
-; CHECK:       # %bb.0:
-; CHECK-NEXT:    vsetivli zero, 4, e16, mf2, ta, ma
-; CHECK-NEXT:    vmv.v.i v8, 0
-; CHECK-NEXT:    vse16.v v8, (a0)
-; CHECK-NEXT:    ret
+; LMULMAX8-RV32-LABEL: splat_zero_v4i16:
+; LMULMAX8-RV32:       # %bb.0:
+; LMULMAX8-RV32-NEXT:    vsetivli zero, 4, e16, mf2, ta, ma
+; LMULMAX8-RV32-NEXT:    vmv.v.i v8, 0
+; LMULMAX8-RV32-NEXT:    vse16.v v8, (a0)
+; LMULMAX8-RV32-NEXT:    ret
+;
+; LMULMAX2-RV32-LABEL: splat_zero_v4i16:
+; LMULMAX2-RV32:       # %bb.0:
+; LMULMAX2-RV32-NEXT:    vsetivli zero, 4, e16, mf2, ta, ma
+; LMULMAX2-RV32-NEXT:    vmv.v.i v8, 0
+; LMULMAX2-RV32-NEXT:    vse16.v v8, (a0)
+; LMULMAX2-RV32-NEXT:    ret
+;
+; LMULMAX1-RV32-LABEL: splat_zero_v4i16:
+; LMULMAX1-RV32:       # %bb.0:
+; LMULMAX1-RV32-NEXT:    vsetivli zero, 4, e16, mf2, ta, ma
+; LMULMAX1-RV32-NEXT:    vmv.v.i v8, 0
+; LMULMAX1-RV32-NEXT:    vse16.v v8, (a0)
+; LMULMAX1-RV32-NEXT:    ret
+;
+; LMULMAX8-RV64-LABEL: splat_zero_v4i16:
+; LMULMAX8-RV64:       # %bb.0:
+; LMULMAX8-RV64-NEXT:    sd zero, 0(a0)
+; LMULMAX8-RV64-NEXT:    ret
+;
+; LMULMAX2-RV64-LABEL: splat_zero_v4i16:
+; LMULMAX2-RV64:       # %bb.0:
+; LMULMAX2-RV64-NEXT:    sd zero, 0(a0)
+; LMULMAX2-RV64-NEXT:    ret
+;
+; LMULMAX1-RV64-LABEL: splat_zero_v4i16:
+; LMULMAX1-RV64:       # %bb.0:
+; LMULMAX1-RV64-NEXT:    sd zero, 0(a0)
+; LMULMAX1-RV64-NEXT:    ret
   store <4 x i16> zeroinitializer, ptr %p
   ret void
 }
 
 define void @splat_zero_v2i32(ptr %p) {
-; CHECK-LABEL: splat_zero_v2i32:
-; CHECK:       # %bb.0:
-; CHECK-NEXT:    vsetivli zero, 2, e32, mf2, ta, ma
-; CHECK-NEXT:    vmv.v.i v8, 0
-; CHECK-NEXT:    vse32.v v8, (a0)
-; CHECK-NEXT:    ret
+; LMULMAX8-RV32-LABEL: splat_zero_v2i32:
+; LMULMAX8-RV32:       # %bb.0:
+; LMULMAX8-RV32-NEXT:    vsetivli zero, 2, e32, mf2, ta, ma
+; LMULMAX8-RV32-NEXT:    vmv.v.i v8, 0
+; LMULMAX8-RV32-NEXT:    vse32.v v8, (a0)
+; LMULMAX8-RV32-NEXT:    ret
+;
+; LMULMAX2-RV32-LABEL: splat_zero_v2i32:
+; LMULMAX2-RV32:       # %bb.0:
+; LMULMAX2-RV32-NEXT:    vsetivli zero, 2, e32, mf2, ta, ma
+; LMULMAX2-RV32-NEXT:    vmv.v.i v8, 0
+; LMULMAX2-RV32-NEXT:    vse32.v v8, (a0)
+; LMULMAX2-RV32-NEXT:    ret
+;
+; LMULMAX1-RV32-LABEL: splat_zero_v2i32:
+; LMULMAX1-RV32:       # %bb.0:
+; LMULMAX1-RV32-NEXT:    vsetivli zero, 2, e32, mf2, ta, ma
+; LMULMAX1-RV32-NEXT:    vmv.v.i v8, 0
+; LMULMAX1-RV32-NEXT:    vse32.v v8, (a0)
+; LMULMAX1-RV32-NEXT:    ret
+;
+; LMULMAX8-RV64-LABEL: splat_zero_v2i32:
+; LMULMAX8-RV64:       # %bb.0:
+; LMULMAX8-RV64-NEXT:    sd zero, 0(a0)
+; LMULMAX8-RV64-NEXT:    ret
+;
+; LMULMAX2-RV64-LABEL: splat_zero_v2i32:
+; LMULMAX2-RV64:       # %bb.0:
+; LMULMAX2-RV64-NEXT:    sd zero, 0(a0)
+; LMULMAX2-RV64-NEXT:    ret
+;
+; LMULMAX1-RV64-LABEL: splat_zero_v2i32:
+; LMULMAX1-RV64:       # %bb.0:
+; LMULMAX1-RV64-NEXT:    sd zero, 0(a0)
+; LMULMAX1-RV64-NEXT:    ret
   store <2 x i32> zeroinitializer, ptr %p
   ret void
 }
 
 ; Not a power of two and requires more than two scalar stores.
 define void @splat_zero_v7i16(ptr %p) {
 ; LMULMAX8-RV32-LABEL: splat_zero_v7i16:
 ; LMULMAX8-RV32:       # %bb.0:

llvm/test/CodeGen/RISCV/rvv/fixed-vectors-mask-splat.ll

 ; CHECK-NEXT:    ret
   store <1 x i1> <i1 1>, ptr %x
   ret void
 }
 
 define void @splat_zeros_v2i1(ptr %x) {
 ; CHECK-LABEL: splat_zeros_v2i1:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    vsetivli zero, 2, e8, mf8, ta, ma
-; CHECK-NEXT:    vmclr.m v0
-; CHECK-NEXT:    vmv.v.i v8, 0
-; CHECK-NEXT:    vmerge.vim v8, v8, 1, v0
-; CHECK-NEXT:    vsetivli zero, 8, e8, mf2, ta, ma
-; CHECK-NEXT:    vmv.v.i v9, 0
-; CHECK-NEXT:    vsetivli zero, 2, e8, mf2, tu, ma
-; CHECK-NEXT:    vslideup.vi v9, v8, 0
-; CHECK-NEXT:    vsetivli zero, 8, e8, mf2, ta, ma
-; CHECK-NEXT:    vmsne.vi v8, v9, 0
-; CHECK-NEXT:    vsm.v v8, (a0)
+; CHECK-NEXT:    sb zero, 0(a0)
 ; CHECK-NEXT:    ret
   store <2 x i1> zeroinitializer, ptr %x
   ret void
 }
 
 define void @splat_v1i1(ptr %x, i1 %y) {
 ; CHECK-LABEL: splat_v1i1:
 ; CHECK:       # %bb.0:
; CHECK-NEXT:    ret
   %b = shufflevector <4 x i1> %a, <4 x i1> poison, <4 x i32> zeroinitializer
   store <4 x i1> %b, ptr %x
   ret void
 }
 
 define void @splat_zeros_v8i1(ptr %x) {
 ; CHECK-LABEL: splat_zeros_v8i1:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    vsetivli zero, 8, e8, mf2, ta, ma
-; CHECK-NEXT:    vmclr.m v8
-; CHECK-NEXT:    vsm.v v8, (a0)
+; CHECK-NEXT:    sb zero, 0(a0)
 ; CHECK-NEXT:    ret
   store <8 x i1> zeroinitializer, ptr %x
   ret void
 }
 
 define void @splat_v8i1(ptr %x, i1 %y) {
 ; CHECK-LABEL: splat_v8i1:
 ; CHECK:       # %bb.0:

 ; CHECK-NEXT:    ret
   %a = insertelement <16 x i1> poison, i1 %y, i32 0
   %b = shufflevector <16 x i1> %a, <16 x i1> poison, <16 x i32> zeroinitializer
   store <16 x i1> %b, ptr %x
   ret void
 }
 
 define void @splat_zeros_v32i1(ptr %x) {
-; LMULMAX2-LABEL: splat_zeros_v32i1:
-; LMULMAX2:       # %bb.0:
-; LMULMAX2-NEXT:    li a1, 32
-; LMULMAX2-NEXT:    vsetvli zero, a1, e8, m2, ta, ma
-; LMULMAX2-NEXT:    vmclr.m v8
-; LMULMAX2-NEXT:    vsm.v v8, (a0)
-; LMULMAX2-NEXT:    ret
-;
-; LMULMAX1-RV32-LABEL: splat_zeros_v32i1:
-; LMULMAX1-RV32:       # %bb.0:
-; LMULMAX1-RV32-NEXT:    vsetivli zero, 16, e8, m1, ta, ma
-; LMULMAX1-RV32-NEXT:    vmclr.m v8
-; LMULMAX1-RV32-NEXT:    vsm.v v8, (a0)
-; LMULMAX1-RV32-NEXT:    addi a0, a0, 2
-; LMULMAX1-RV32-NEXT:    vsm.v v8, (a0)
-; LMULMAX1-RV32-NEXT:    ret
-;
-; LMULMAX1-RV64-LABEL: splat_zeros_v32i1:
-; LMULMAX1-RV64:       # %bb.0:
-; LMULMAX1-RV64-NEXT:    vsetivli zero, 16, e8, m1, ta, ma
-; LMULMAX1-RV64-NEXT:    vmclr.m v8
-; LMULMAX1-RV64-NEXT:    vsm.v v8, (a0)
-; LMULMAX1-RV64-NEXT:    addi a0, a0, 2
-; LMULMAX1-RV64-NEXT:    vsm.v v8, (a0)
-; LMULMAX1-RV64-NEXT:    ret
+; CHECK-LABEL: splat_zeros_v32i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    sw zero, 0(a0)
+; CHECK-NEXT:    ret
   store <32 x i1> zeroinitializer, ptr %x
   ret void
 }
 
 define void @splat_v32i1(ptr %x, i1 %y) {
 ; LMULMAX2-LABEL: splat_v32i1:
 ; LMULMAX2:       # %bb.0:
 ; LMULMAX2-NEXT:    andi a1, a1, 1