[RISCV] Support STEP_VECTOR with a step greater than one

[frasercrmck]

Summary

DAGCombiner was recently taught how to combine STEP_VECTOR nodes,
meaning the step value is no longer guaranteed to be one by the time it
reaches the backend for lowering.

This patch supports such cases on RISC-V by lowering to other step
values to a multiply following the vid.v instruction. It includes a
small optimization for common cases where the multiply can be expressed
as a shift left.

[frasercrmck]

• rebase

[junparser]

D100812 has checked in. FYI

[frasercrmck]

• rebase

[rogfer01]

I hit this one recently, too!

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

3812

I used DAG.getSplatVector here (I guess makes little difference) but I didn't consider the power-of-two case.

I'm curious that we cannot combine a case like this

  t13: nxv1i64 = RISCVISD::VID_VL t12, Register:i64 $x0
  t14: nxv1i64 = splat_vector Constant:i64<2>
t15: nxv1i64 = mul t13, t14

into

  t13: nxv1i64 = RISCVISD::VID_VL t12, Register:i64 $x0
  t16: nxv1i64 = splat_vector Constant:i64<1>
t15: nxv1i64 = shl t13, t16

[frasercrmck]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

3812

I don't think we can use getSplatVector because at this stage for i64 vectors on RV32, we can't introduce an illegal i64 type and so it would be nxvXi64 = splat_vector i32 which is ill-formed.

Regarding the second part, I think with splat_vector that would be possible at this stage -- it relies on getNode(ISD::MUL) folding constant arithmetic as we're post-combine -- but as I mentioned we can't use that node.

Perhaps if we conditionally used splat_vector and splat_vector_parts it would work, but splat_vector_parts isn't yet hooked up to the optimization machinery. I think it's better to do this awkward MUL/SHL lowering ourselves so that rv32 and rv64 behave the same way, but I can appreciate the DRY side of the argument.

[rogfer01]

LGTM. Thanks @frasercrmck !

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

3812

RIght, thanks for the explanation I had forgotten about rv32.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

SDValue RISCVTargetLowering::lowerEXTRACT_SUBVECTOR(SDValue Op,
   Slidedown = DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, SubVecVT, Slidedown,
                           DAG.getConstant(0, DL, XLenVT));
 
   // We might have bitcast from a mask type: cast back to the original type if
   // required.
   return DAG.getBitcast(Op.getSimpleValueType(), Slidedown);
 }
 
-// Implement step_vector to the vid instruction.
+// Lower step_vector to the vid instruction. Any non-identity step value must
+// be accounted for my manual expansion.
 SDValue RISCVTargetLowering::lowerSTEP_VECTOR(SDValue Op,
                                               SelectionDAG &DAG) const {
   SDLoc DL(Op);
-  assert(Op.getConstantOperandAPInt(0) == 1 && "Unexpected step value");
   MVT VT = Op.getSimpleValueType();
+  MVT XLenVT = Subtarget.getXLenVT();
   SDValue Mask, VL;
   std::tie(Mask, VL) = getDefaultScalableVLOps(VT, DL, DAG, Subtarget);
-  return DAG.getNode(RISCVISD::VID_VL, DL, VT, Mask, VL);
+  SDValue StepVec = DAG.getNode(RISCVISD::VID_VL, DL, VT, Mask, VL);
+  uint64_t StepValImm = Op.getConstantOperandVal(0);
+  if (StepValImm != 1) {
+    assert(Op.getOperand(0).getValueType() == XLenVT &&
+           "Unexpected step value type");
+    if (isPowerOf2_64(StepValImm)) {
+      SDValue StepVal =
+          DAG.getNode(RISCVISD::VMV_V_X_VL, DL, VT,
+                      DAG.getConstant(Log2_64(StepValImm), DL, XLenVT));
+      StepVec = DAG.getNode(ISD::SHL, DL, VT, StepVec, StepVal);
+    } else {
+      SDValue StepVal =
+          DAG.getNode(RISCVISD::VMV_V_X_VL, DL, VT, Op.getOperand(0));
+      StepVec = DAG.getNode(ISD::MUL, DL, VT, StepVec, StepVal);
+    }
+  }
+  return StepVec;
 }
 
 // Implement vector_reverse using vrgather.vv with indices determined by
 // subtracting the id of each element from (VLMAX-1). This will convert
 // the indices like so:
 // (0, 1,..., VLMAX-2, VLMAX-1) -> (VLMAX-1, VLMAX-2,..., 1, 0).
 // TODO: This code assumes VLMAX <= 65536 for LMUL=8 SEW=16.
 SDValue RISCVTargetLowering::lowerVECTOR_REVERSE(SDValue Op,

llvm/test/CodeGen/RISCV/rvv/stepvector.ll

 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e8,m1,ta,mu
 ; CHECK-NEXT:    vid.v v8
 ; CHECK-NEXT:    ret
   %v = call <vscale x 8 x i8> @llvm.experimental.stepvector.nxv8i8()
   ret <vscale x 8 x i8> %v
 }
 
+define <vscale x 8 x i8> @add_stepvector_nxv8i8() {
+; CHECK-LABEL: add_stepvector_nxv8i8:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e8,m1,ta,mu
+; CHECK-NEXT:    vid.v v25
+; CHECK-NEXT:    vsll.vi v8, v25, 1
+; CHECK-NEXT:    ret
+entry:
+  %0 = call <vscale x 8 x i8> @llvm.experimental.stepvector.nxv8i8()
+  %1 = call <vscale x 8 x i8> @llvm.experimental.stepvector.nxv8i8()
+  %2 = add <vscale x 8 x i8> %0, %1
+  ret <vscale x 8 x i8> %2
+}
+
+define <vscale x 8 x i8> @mul_stepvector_nxv8i8() {
+; CHECK-LABEL: mul_stepvector_nxv8i8:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e8,m1,ta,mu
+; CHECK-NEXT:    vid.v v25
+; CHECK-NEXT:    addi a0, zero, 3
+; CHECK-NEXT:    vmul.vx v8, v25, a0
+; CHECK-NEXT:    ret
+entry:
+  %0 = insertelement <vscale x 8 x i8> poison, i8 3, i32 0
+  %1 = shufflevector <vscale x 8 x i8> %0, <vscale x 8 x i8> poison, <vscale x 8 x i32> zeroinitializer
+  %2 = call <vscale x 8 x i8> @llvm.experimental.stepvector.nxv8i8()
+  %3 = mul <vscale x 8 x i8> %2, %1
+  ret <vscale x 8 x i8> %3
+}
+
+define <vscale x 8 x i8> @shl_stepvector_nxv8i8() {
+; CHECK-LABEL: shl_stepvector_nxv8i8:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e8,m1,ta,mu
+; CHECK-NEXT:    vid.v v25
+; CHECK-NEXT:    vsll.vi v8, v25, 2
+; CHECK-NEXT:    ret
+entry:
+  %0 = insertelement <vscale x 8 x i8> poison, i8 2, i32 0
+  %1 = shufflevector <vscale x 8 x i8> %0, <vscale x 8 x i8> poison, <vscale x 8 x i32> zeroinitializer
+  %2 = call <vscale x 8 x i8> @llvm.experimental.stepvector.nxv8i8()
+  %3 = shl <vscale x 8 x i8> %2, %1
+  ret <vscale x 8 x i8> %3
+}
+
 declare <vscale x 16 x i8> @llvm.experimental.stepvector.nxv16i8()
 
 define <vscale x 16 x i8> @stepvector_nxv16i8() {
 ; CHECK-LABEL: stepvector_nxv16i8:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e8,m2,ta,mu
 ; CHECK-NEXT:    vid.v v8
 ; CHECK-NEXT:    ret

 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e16,m4,ta,mu
 ; CHECK-NEXT:    vid.v v8
 ; CHECK-NEXT:    ret
   %v = call <vscale x 16 x i16> @llvm.experimental.stepvector.nxv16i16()
   ret <vscale x 16 x i16> %v
 }
 
+define <vscale x 16 x i16> @add_stepvector_nxv16i16() {
+; CHECK-LABEL: add_stepvector_nxv16i16:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e16,m4,ta,mu
+; CHECK-NEXT:    vid.v v28
+; CHECK-NEXT:    vsll.vi v8, v28, 1
+; CHECK-NEXT:    ret
+entry:
+  %0 = call <vscale x 16 x i16> @llvm.experimental.stepvector.nxv16i16()
+  %1 = call <vscale x 16 x i16> @llvm.experimental.stepvector.nxv16i16()
+  %2 = add <vscale x 16 x i16> %0, %1
+  ret <vscale x 16 x i16> %2
+}
+
+define <vscale x 16 x i16> @mul_stepvector_nxv16i16() {
+; CHECK-LABEL: mul_stepvector_nxv16i16:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e16,m4,ta,mu
+; CHECK-NEXT:    vid.v v28
+; CHECK-NEXT:    addi a0, zero, 3
+; CHECK-NEXT:    vmul.vx v8, v28, a0
+; CHECK-NEXT:    ret
+entry:
+  %0 = insertelement <vscale x 16 x i16> poison, i16 3, i32 0
+  %1 = shufflevector <vscale x 16 x i16> %0, <vscale x 16 x i16> poison, <vscale x 16 x i32> zeroinitializer
+  %2 = call <vscale x 16 x i16> @llvm.experimental.stepvector.nxv16i16()
+  %3 = mul <vscale x 16 x i16> %2, %1
+  ret <vscale x 16 x i16> %3
+}
+
+define <vscale x 16 x i16> @shl_stepvector_nxv16i16() {
+; CHECK-LABEL: shl_stepvector_nxv16i16:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e16,m4,ta,mu
+; CHECK-NEXT:    vid.v v28
+; CHECK-NEXT:    vsll.vi v8, v28, 2
+; CHECK-NEXT:    ret
+entry:
+  %0 = insertelement <vscale x 16 x i16> poison, i16 2, i32 0
+  %1 = shufflevector <vscale x 16 x i16> %0, <vscale x 16 x i16> poison, <vscale x 16 x i32> zeroinitializer
+  %2 = call <vscale x 16 x i16> @llvm.experimental.stepvector.nxv16i16()
+  %3 = shl <vscale x 16 x i16> %2, %1
+  ret <vscale x 16 x i16> %3
+}
+
 declare <vscale x 32 x i16> @llvm.experimental.stepvector.nxv32i16()
 
 define <vscale x 32 x i16> @stepvector_nxv32i16() {
 ; CHECK-LABEL: stepvector_nxv32i16:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e16,m8,ta,mu
 ; CHECK-NEXT:    vid.v v8
 ; CHECK-NEXT:    ret

 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e32,m8,ta,mu
 ; CHECK-NEXT:    vid.v v8
 ; CHECK-NEXT:    ret
   %v = call <vscale x 16 x i32> @llvm.experimental.stepvector.nxv16i32()
   ret <vscale x 16 x i32> %v
 }
 
+define <vscale x 16 x i32> @add_stepvector_nxv16i32() {
+; CHECK-LABEL: add_stepvector_nxv16i32:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e32,m8,ta,mu
+; CHECK-NEXT:    vid.v v8
+; CHECK-NEXT:    vsll.vi v8, v8, 1
+; CHECK-NEXT:    ret
+entry:
+  %0 = call <vscale x 16 x i32> @llvm.experimental.stepvector.nxv16i32()
+  %1 = call <vscale x 16 x i32> @llvm.experimental.stepvector.nxv16i32()
+  %2 = add <vscale x 16 x i32> %0, %1
+  ret <vscale x 16 x i32> %2
+}
+
+define <vscale x 16 x i32> @mul_stepvector_nxv16i32() {
+; CHECK-LABEL: mul_stepvector_nxv16i32:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e32,m8,ta,mu
+; CHECK-NEXT:    vid.v v8
+; CHECK-NEXT:    addi a0, zero, 3
+; CHECK-NEXT:    vmul.vx v8, v8, a0
+; CHECK-NEXT:    ret
+entry:
+  %0 = insertelement <vscale x 16 x i32> poison, i32 3, i32 0
+  %1 = shufflevector <vscale x 16 x i32> %0, <vscale x 16 x i32> poison, <vscale x 16 x i32> zeroinitializer
+  %2 = call <vscale x 16 x i32> @llvm.experimental.stepvector.nxv16i32()
+  %3 = mul <vscale x 16 x i32> %2, %1
+  ret <vscale x 16 x i32> %3
+}
+
+define <vscale x 16 x i32> @shl_stepvector_nxv16i32() {
+; CHECK-LABEL: shl_stepvector_nxv16i32:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e32,m8,ta,mu
+; CHECK-NEXT:    vid.v v8
+; CHECK-NEXT:    vsll.vi v8, v8, 2
+; CHECK-NEXT:    ret
+entry:
+  %0 = insertelement <vscale x 16 x i32> poison, i32 2, i32 0
+  %1 = shufflevector <vscale x 16 x i32> %0, <vscale x 16 x i32> poison, <vscale x 16 x i32> zeroinitializer
+  %2 = call <vscale x 16 x i32> @llvm.experimental.stepvector.nxv16i32()
+  %3 = shl <vscale x 16 x i32> %2, %1
+  ret <vscale x 16 x i32> %3
+}
+
 declare <vscale x 1 x i64> @llvm.experimental.stepvector.nxv1i64()
 
 define <vscale x 1 x i64> @stepvector_nxv1i64() {
 ; CHECK-LABEL: stepvector_nxv1i64:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e64,m1,ta,mu
 ; CHECK-NEXT:    vid.v v8
 ; CHECK-NEXT:    ret

 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e64,m8,ta,mu
 ; CHECK-NEXT:    vid.v v8
 ; CHECK-NEXT:    ret
   %v = call <vscale x 8 x i64> @llvm.experimental.stepvector.nxv8i64()
   ret <vscale x 8 x i64> %v
 }
 
+define <vscale x 8 x i64> @add_stepvector_nxv8i64() {
+; CHECK-LABEL: add_stepvector_nxv8i64:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e64,m8,ta,mu
+; CHECK-NEXT:    vid.v v8
+; CHECK-NEXT:    vsll.vi v8, v8, 1
+; CHECK-NEXT:    ret
+entry:
+  %0 = call <vscale x 8 x i64> @llvm.experimental.stepvector.nxv8i64()
+  %1 = call <vscale x 8 x i64> @llvm.experimental.stepvector.nxv8i64()
+  %2 = add <vscale x 8 x i64> %0, %1
+  ret <vscale x 8 x i64> %2
+}
+
+define <vscale x 8 x i64> @mul_stepvector_nxv8i64() {
+; CHECK-LABEL: mul_stepvector_nxv8i64:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e64,m8,ta,mu
+; CHECK-NEXT:    vid.v v8
+; CHECK-NEXT:    addi a0, zero, 3
+; CHECK-NEXT:    vmul.vx v8, v8, a0
+; CHECK-NEXT:    ret
+entry:
+  %0 = insertelement <vscale x 8 x i64> poison, i64 3, i32 0
+  %1 = shufflevector <vscale x 8 x i64> %0, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
+  %2 = call <vscale x 8 x i64> @llvm.experimental.stepvector.nxv8i64()
+  %3 = mul <vscale x 8 x i64> %2, %1
+  ret <vscale x 8 x i64> %3
+}
+
+define <vscale x 8 x i64> @shl_stepvector_nxv8i64() {
+; CHECK-LABEL: shl_stepvector_nxv8i64:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetvli a0, zero, e64,m8,ta,mu
+; CHECK-NEXT:    vid.v v8
+; CHECK-NEXT:    vsll.vi v8, v8, 2
+; CHECK-NEXT:    ret
+entry:
+  %0 = insertelement <vscale x 8 x i64> poison, i64 2, i32 0
+  %1 = shufflevector <vscale x 8 x i64> %0, <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer
+  %2 = call <vscale x 8 x i64> @llvm.experimental.stepvector.nxv8i64()
+  %3 = shl <vscale x 8 x i64> %2, %1
+  ret <vscale x 8 x i64> %3
+}
+
 declare <vscale x 16 x i64> @llvm.experimental.stepvector.nxv16i64()
 
 define <vscale x 16 x i64> @stepvector_nxv16i64() {
 ; CHECK-LABEL: stepvector_nxv16i64:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    csrr a0, vlenb
 ; CHECK-NEXT:    vsetvli a1, zero, e64,m8,ta,mu
 ; CHECK-NEXT:    vid.v v8
 ; CHECK-NEXT:    vadd.vx v16, v8, a0
 ; CHECK-NEXT:    ret
   %v = call <vscale x 16 x i64> @llvm.experimental.stepvector.nxv16i64()
   ret <vscale x 16 x i64> %v
 }
+
+define <vscale x 16 x i64> @add_stepvector_nxv16i64() {
+; RV32-LABEL: add_stepvector_nxv16i64:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    csrr a0, vlenb
+; RV32-NEXT:    slli a0, a0, 1
+; RV32-NEXT:    vsetvli a1, zero, e64,m8,ta,mu
+; RV32-NEXT:    vmv.v.x v8, a0
+; RV32-NEXT:    addi a0, zero, 32
+; RV32-NEXT:    vsll.vx v8, v8, a0
+; RV32-NEXT:    vsrl.vx v16, v8, a0
+; RV32-NEXT:    vid.v v8
+; RV32-NEXT:    vsll.vi v8, v8, 1
+; RV32-NEXT:    vadd.vv v16, v8, v16
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: add_stepvector_nxv16i64:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    csrr a0, vlenb
+; RV64-NEXT:    slli a0, a0, 1
+; RV64-NEXT:    vsetvli a1, zero, e64,m8,ta,mu
+; RV64-NEXT:    vid.v v8
+; RV64-NEXT:    vsll.vi v8, v8, 1
+; RV64-NEXT:    vadd.vx v16, v8, a0
+; RV64-NEXT:    ret
+entry:
+  %0 = call <vscale x 16 x i64> @llvm.experimental.stepvector.nxv16i64()
+  %1 = call <vscale x 16 x i64> @llvm.experimental.stepvector.nxv16i64()
+  %2 = add <vscale x 16 x i64> %0, %1
+  ret <vscale x 16 x i64> %2
+}
+
+define <vscale x 16 x i64> @mul_stepvector_nxv16i64() {
+; RV32-LABEL: mul_stepvector_nxv16i64:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    vsetvli a0, zero, e64,m8,ta,mu
+; RV32-NEXT:    vid.v v8
+; RV32-NEXT:    addi a0, zero, 3
+; RV32-NEXT:    vmul.vx v8, v8, a0
+; RV32-NEXT:    csrr a0, vlenb
+; RV32-NEXT:    srli a0, a0, 3
+; RV32-NEXT:    addi a1, zero, 24
+; RV32-NEXT:    mul a0, a0, a1
+; RV32-NEXT:    vmv.v.x v16, a0
+; RV32-NEXT:    addi a0, zero, 32
+; RV32-NEXT:    vsll.vx v16, v16, a0
+; RV32-NEXT:    vsrl.vx v16, v16, a0
+; RV32-NEXT:    vadd.vv v16, v8, v16
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: mul_stepvector_nxv16i64:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    vsetvli a0, zero, e64,m8,ta,mu
+; RV64-NEXT:    vid.v v8
+; RV64-NEXT:    addi a0, zero, 3
+; RV64-NEXT:    vmul.vx v8, v8, a0
+; RV64-NEXT:    csrr a0, vlenb
+; RV64-NEXT:    srli a0, a0, 3
+; RV64-NEXT:    addi a1, zero, 24
+; RV64-NEXT:    mul a0, a0, a1
+; RV64-NEXT:    vadd.vx v16, v8, a0
+; RV64-NEXT:    ret
+entry:
+  %0 = insertelement <vscale x 16 x i64> poison, i64 3, i32 0
+  %1 = shufflevector <vscale x 16 x i64> %0, <vscale x 16 x i64> poison, <vscale x 16 x i32> zeroinitializer
+  %2 = call <vscale x 16 x i64> @llvm.experimental.stepvector.nxv16i64()
+  %3 = mul <vscale x 16 x i64> %2, %1
+  ret <vscale x 16 x i64> %3
+}
+
+define <vscale x 16 x i64> @shl_stepvector_nxv16i64() {
+; RV32-LABEL: shl_stepvector_nxv16i64:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    csrr a0, vlenb
+; RV32-NEXT:    slli a0, a0, 2
+; RV32-NEXT:    vsetvli a1, zero, e64,m8,ta,mu
+; RV32-NEXT:    vmv.v.x v8, a0
+; RV32-NEXT:    addi a0, zero, 32
+; RV32-NEXT:    vsll.vx v8, v8, a0
+; RV32-NEXT:    vsrl.vx v16, v8, a0
+; RV32-NEXT:    vid.v v8
+; RV32-NEXT:    vsll.vi v8, v8, 2
+; RV32-NEXT:    vadd.vv v16, v8, v16
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: shl_stepvector_nxv16i64:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    csrr a0, vlenb
+; RV64-NEXT:    slli a0, a0, 2
+; RV64-NEXT:    vsetvli a1, zero, e64,m8,ta,mu
+; RV64-NEXT:    vid.v v8
+; RV64-NEXT:    vsll.vi v8, v8, 2
+; RV64-NEXT:    vadd.vx v16, v8, a0
+; RV64-NEXT:    ret
+entry:
+  %0 = insertelement <vscale x 16 x i64> poison, i64 2, i32 0
+  %1 = shufflevector <vscale x 16 x i64> %0, <vscale x 16 x i64> poison, <vscale x 16 x i32> zeroinitializer
+  %2 = call <vscale x 16 x i64> @llvm.experimental.stepvector.nxv16i64()
+  %3 = shl <vscale x 16 x i64> %2, %1
+  ret <vscale x 16 x i64> %3
+}