[RISCV] Lower inserts into identity shuffles as v{,f}slide1ups

[luke]

Summary

Inserting an element at index 0, into a shuffle that's an identity shuffle past the first
element, can be achieved with a single v{,f}slide1up. This patch detects such shuffles with
the isIdentityMask logic and lowers it to said vslide1up.

[frasercrmck]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

6189

This function seems to assert that the mask is not empty. Is it all possible we could end up triggering this assert with one-element vectors? Would we better be moving the VecVT.getVectorNumElements() > 1 check before calling isIdentityMask?

Could we try and add testing for the above case?

[reames]

Description wise, calling an slideup-by-one mask an identity mask feels very odd to me. It took me a bit reading the code to figure out what you meant by that, and I'd suggest rewording. I'd suggest something along the lines of "would be identity mask if shifted down by one position" or something along those lines.

Did you look at the alternative of adding a DAG combine to reverse the shuffle+insert form into the insert+shuffle form? That'd been what I'd been tentatively thinking about, but I hadn't explored it past idea phase.

llvm/test/CodeGen/RISCV/rvv/fixed-vector-shuffle-vslide1up.ll

318

You can go ahead and precommit these tests without further review.

[luke]

In D151736#4381994, @reames wrote:

Did you look at the alternative of adding a DAG combine to reverse the shuffle+insert form into the insert+shuffle form? That'd been what I'd been tentatively thinking about, but I hadn't explored it past idea phase.

Not yet, but I'd like to explore that option before returning to this patch. This combine feels very specific and doesn't seem to be catch any cases in the benchmarks I've been looking at.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

6189

Good catch, definitely meant to move the`VecVT.getVectorNumElements() > 1` first.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

SDValue RISCVTargetLowering::lowerINSERT_VECTOR_ELT(SDValue Op,
   auto [Mask, VL] = getDefaultVLOps(VecVT, ContainerVT, DL, DAG, Subtarget);
 
   SDValue ValInVec;
 
   if (IsLegalInsert) {
     unsigned Opc =
         VecVT.isFloatingPoint() ? RISCVISD::VFMV_S_F_VL : RISCVISD::VMV_S_X_VL;
     if (isNullConstant(Idx)) {
+
+      // Lower inserts at element 0 into what will become an identity
+      // shuffle as a vslide1up
+      // e.g. (insert_vec_elt (shuffle<undef,0,1,2> %v1, %v2), %val)
+      if (auto *Shuffle = dyn_cast<ShuffleVectorSDNode>(Op.getOperand(0))) {
+        ArrayRef<int> ShufMask = Shuffle->getMask();
+        if (ShuffleVectorInst::isIdentityMask(ShufMask.drop_front(1)) &&
+            VecVT.isFixedLengthVector() && VecVT.getVectorNumElements() > 1) {
+          // If the first index of the identity mask < mask size, then
+          // we're inserting into the first operand, otherwise we're
+          // inserting into the second.
+          Vec = ShufMask[1] < (int)ShufMask.size() ? Shuffle->getOperand(0)
+                                                   : Shuffle->getOperand(1);
+          Vec = convertToScalableVector(ContainerVT, Vec, DAG, Subtarget);
+
+          unsigned Opc = Val.getSimpleValueType().isFloatingPoint()
+                             ? RISCVISD::VFSLIDE1UP_VL
+                             : RISCVISD::VSLIDE1UP_VL;
+          SDValue Slideup =
+              DAG.getNode(Opc, DL, ContainerVT, DAG.getUNDEF(ContainerVT), Vec,
+                          Val, Mask, VL);
+          Slideup = convertFromScalableVector(VecVT, Slideup, DAG, Subtarget);
+          return Slideup;
+        }
+      }
+
       Vec = DAG.getNode(Opc, DL, ContainerVT, Vec, Val, VL);
       if (!VecVT.isFixedLengthVector())
         return Vec;
       return convertFromScalableVector(VecVT, Vec, DAG, Subtarget);
     }
     ValInVec = lowerScalarInsert(Val, VL, ContainerVT, DL, DAG, Subtarget);
   } else {
     // On RV32, i64-element vectors must be specially handled to place the

llvm/test/CodeGen/RISCV/rvv/fixed-vector-shuffle-vslide1up.ll

; CHECK-NEXT:    ret
   %v1 = shufflevector <4 x double> %v, <4 x double> %vb, <4 x i32> <i32 4, i32 5, i32 6, i32 0>
   ret <4 x double> %v1
 }
 
 define <2 x double> @vslide1up_v2f64_inverted(<2 x double> %v, double %b) {
 ; CHECK-LABEL: vslide1up_v2f64_inverted:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 2, e64, m1, ta, ma
-; CHECK-NEXT:    vrgather.vi v9, v8, 0
-; CHECK-NEXT:    vsetvli zero, zero, e64, m1, tu, ma
-; CHECK-NEXT:    vfmv.s.f v9, fa0
-; CHECK-NEXT:    vmv1r.v v8, v9
+; CHECK-NEXT:    vfslide1up.vf v9, v8, fa0
+; CHECK-NEXT:    vmv.v.v v8, v9
 ; CHECK-NEXT:    ret
   %v1 = shufflevector <2 x double> %v, <2 x double> poison, <2 x i32> <i32 0, i32 0>
   %v2 = insertelement <2 x double> %v1, double %b, i64 0
   ret <2 x double> %v2
 }
 
 define <4 x i8> @vslide1up_4xi8_inverted(<4 x i8> %v, i8 %b) {
 ; CHECK-LABEL: vslide1up_4xi8_inverted:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 4, e8, mf4, ta, ma
-; CHECK-NEXT:    vslideup.vi v9, v8, 1
-; CHECK-NEXT:    vsetvli zero, zero, e8, mf4, tu, ma
-; CHECK-NEXT:    vmv.s.x v9, a0
+; CHECK-NEXT:    vslide1up.vx v9, v8, a0
 ; CHECK-NEXT:    vmv1r.v v8, v9
 ; CHECK-NEXT:    ret
   %v1 = shufflevector <4 x i8> %v, <4 x i8> poison, <4 x i32> <i32 undef, i32 0, i32 1, i32 2>
   %v2 = insertelement <4 x i8> %v1, i8 %b, i64 0
   ret <4 x i8> %v2
 }
 
+define <2 x double> @vslide1up_v2f64_inverted_swapped(<2 x double> %v, double %b) {
+; CHECK-LABEL: vslide1up_v2f64_inverted_swapped:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetivli zero, 2, e64, m1, ta, ma
+; CHECK-NEXT:    vfslide1up.vf v9, v8, fa0
+; CHECK-NEXT:    vmv.v.v v8, v9
+; CHECK-NEXT:    ret
+  %v1 = shufflevector <2 x double> poison, <2 x double> %v, <2 x i32> <i32 2, i32 2>
+  %v2 = insertelement <2 x double> %v1, double %b, i64 0
+  ret <2 x double> %v2
+}
+
+define <4 x i8> @vslide1up_4xi8_inverted_swapped(<4 x i8> %v, i8 %b) {
+; CHECK-LABEL: vslide1up_4xi8_inverted_swapped:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetivli zero, 4, e8, mf4, ta, ma
+; CHECK-NEXT:    vslide1up.vx v9, v8, a0
+; CHECK-NEXT:    vmv1r.v v8, v9
+; CHECK-NEXT:    ret
+  %v1 = shufflevector <4 x i8> poison, <4 x i8> %v, <4 x i32> <i32 undef, i32 4, i32 5, i32 6>
+  %v2 = insertelement <4 x i8> %v1, i8 %b, i64 0
+  ret <4 x i8> %v2
+}
+
+define <4 x i64> @vslide1up_4xi64_inverted(<4 x i64> %v, i64 %b) {
+; RV32-LABEL: vslide1up_4xi64_inverted:
+; RV32:       # %bb.0:
+; RV32-NEXT:    vsetivli zero, 4, e64, m2, ta, ma
+; RV32-NEXT:    vslideup.vi v10, v8, 1
+; RV32-NEXT:    vsetivli zero, 2, e32, m2, tu, ma
+; RV32-NEXT:    vslide1down.vx v10, v10, a0
+; RV32-NEXT:    vslide1down.vx v10, v10, a1
+; RV32-NEXT:    vmv2r.v v8, v10
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: vslide1up_4xi64_inverted:
+; RV64:       # %bb.0:
+; RV64-NEXT:    vsetivli zero, 4, e64, m2, ta, ma
+; RV64-NEXT:    vslide1up.vx v10, v8, a0
+; RV64-NEXT:    vmv.v.v v8, v10
+; RV64-NEXT:    ret
+  %v1 = shufflevector <4 x i64> %v, <4 x i64> poison, <4 x i32> <i32 undef, i32 0, i32 1, i32 2>
+  %v2 = insertelement <4 x i64> %v1, i64 %b, i64 0
+  ret <4 x i64> %v2
+}
 
 ; The length of the shift is less than the suffix
 define <4 x i32> @vslide1up_4xi32_neg1(<4 x i32> %v, i32 %b) {
 ; CHECK-LABEL: vslide1up_4xi32_neg1:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 4, e32, m1, ta, ma
 ; CHECK-NEXT:    vmv.v.x v9, a0
 ; CHECK-NEXT:    vsetivli zero, 3, e32, m1, tu, ma
 ; CHECK-NEXT:    vslideup.vi v9, v8, 1
 ; CHECK-NEXT:    vmv1r.v v8, v9
 ; CHECK-NEXT:    ret
   %vb = insertelement <4 x i32> poison, i32 %b, i64 0
   %vb2 = insertelement <4 x i32> %vb, i32 %b, i64 3
   %v1 = shufflevector <4 x i32> %v, <4 x i32> %vb2, <4 x i32> <i32 4, i32 0, i32 1, i32 7>
   ret <4 x i32> %v1
 }