[VP] Add widening for VP_STRIDED_LOAD and VP_STRIDED_STORE

[craig.topper]

LGTM

[frasercrmck]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

4452

This may well be copy/pasted from elsewhere (and possibly something I introduced...), but this assert looks fishy to me. Aren't we wanting to check that the widened Mask is the same length as the widened result type, like WidenVT below? This assert looks trivially true as it stands.

[loralb]

Changelog:

• Apply suggestions in the comments

[loralb]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

4452

Yep, correct; this should do what it is supposed to. Thanks!

[frasercrmck]

LGTM too

[loralb]

Changelog:

• Use poison instead of undef in tests

[loralb]

Changelog:

• Update tests

[simoll]

Please rebase for commit

[loralb]

Changelog:

• Rebase

[loralb]

thanks @simoll!

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h

private:
   SDValue WidenVecRes_BUILD_VECTOR(SDNode* N);
   SDValue WidenVecRes_CONCAT_VECTORS(SDNode* N);
   SDValue WidenVecRes_EXTEND_VECTOR_INREG(SDNode* N);
   SDValue WidenVecRes_EXTRACT_SUBVECTOR(SDNode* N);
   SDValue WidenVecRes_INSERT_SUBVECTOR(SDNode *N);
   SDValue WidenVecRes_INSERT_VECTOR_ELT(SDNode* N);
   SDValue WidenVecRes_LOAD(SDNode* N);
   SDValue WidenVecRes_VP_LOAD(VPLoadSDNode *N);
+  SDValue WidenVecRes_VP_STRIDED_LOAD(VPStridedLoadSDNode *N);
   SDValue WidenVecRes_MLOAD(MaskedLoadSDNode* N);
   SDValue WidenVecRes_MGATHER(MaskedGatherSDNode* N);
   SDValue WidenVecRes_VP_GATHER(VPGatherSDNode* N);
   SDValue WidenVecRes_ScalarOp(SDNode* N);
   SDValue WidenVecRes_Select(SDNode *N);
   SDValue WidenVSELECTMask(SDNode *N);
   SDValue WidenVecRes_SELECT_CC(SDNode* N);
   SDValue WidenVecRes_SETCC(SDNode* N);
private:
   SDValue WidenVecOp_BITCAST(SDNode *N);
   SDValue WidenVecOp_CONCAT_VECTORS(SDNode *N);
   SDValue WidenVecOp_EXTEND(SDNode *N);
   SDValue WidenVecOp_EXTRACT_VECTOR_ELT(SDNode *N);
   SDValue WidenVecOp_INSERT_SUBVECTOR(SDNode *N);
   SDValue WidenVecOp_EXTRACT_SUBVECTOR(SDNode *N);
   SDValue WidenVecOp_STORE(SDNode* N);
   SDValue WidenVecOp_VP_STORE(SDNode *N, unsigned OpNo);
+  SDValue WidenVecOp_VP_STRIDED_STORE(SDNode *N, unsigned OpNo);
   SDValue WidenVecOp_MSTORE(SDNode* N, unsigned OpNo);
   SDValue WidenVecOp_MGATHER(SDNode* N, unsigned OpNo);
   SDValue WidenVecOp_MSCATTER(SDNode* N, unsigned OpNo);
   SDValue WidenVecOp_VP_SCATTER(SDNode* N, unsigned OpNo);
   SDValue WidenVecOp_SETCC(SDNode* N);
   SDValue WidenVecOp_STRICT_FSETCC(SDNode* N);
   SDValue WidenVecOp_VSELECT(SDNode *N);
 

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

#endif
   case ISD::SETCC:             Res = WidenVecRes_SETCC(N); break;
   case ISD::UNDEF:             Res = WidenVecRes_UNDEF(N); break;
   case ISD::VECTOR_SHUFFLE:
     Res = WidenVecRes_VECTOR_SHUFFLE(cast<ShuffleVectorSDNode>(N));
     break;
   case ISD::VP_LOAD:
     Res = WidenVecRes_VP_LOAD(cast<VPLoadSDNode>(N));
     break;
+  case ISD::EXPERIMENTAL_VP_STRIDED_LOAD:
+    Res = WidenVecRes_VP_STRIDED_LOAD(cast<VPStridedLoadSDNode>(N));
+    break;
   case ISD::MLOAD:
     Res = WidenVecRes_MLOAD(cast<MaskedLoadSDNode>(N));
     break;
   case ISD::MGATHER:
     Res = WidenVecRes_MGATHER(cast<MaskedGatherSDNode>(N));
     break;
   case ISD::VP_GATHER:
     Res = WidenVecRes_VP_GATHER(cast<VPGatherSDNode>(N));
SDValue Res =
                     N->getBasePtr(), N->getOffset(), Mask, EVL,
                     N->getMemoryVT(), N->getMemOperand(), N->isExpandingLoad());
   // Legalize the chain result - switch anything that used the old chain to
   // use the new one.
   ReplaceValueWith(SDValue(N, 1), Res.getValue(1));
   return Res;
 }
 
+SDValue DAGTypeLegalizer::WidenVecRes_VP_STRIDED_LOAD(VPStridedLoadSDNode *N) {
+  SDLoc DL(N);
+
+  // The mask should be widened as well
+  SDValue Mask = N->getMask();
+  assert(getTypeAction(Mask.getValueType()) ==
+             TargetLowering::TypeWidenVector &&
+         "Unable to widen VP strided load");
+  Mask = GetWidenedVector(Mask);
+  assert(Mask.getValueType().getVectorElementCount() ==
+             TLI.getTypeToTransformTo(*DAG.getContext(), Mask.getValueType())
+                 .getVectorElementCount() &&
+         "Unable to widen VP strided load");
+
+  EVT WidenVT = TLI.getTypeToTransformTo(*DAG.getContext(), N->getValueType(0));
+  SDValue Res = DAG.getStridedLoadVP(
+      N->getAddressingMode(), N->getExtensionType(), WidenVT, DL, N->getChain(),
+      N->getBasePtr(), N->getOffset(), N->getStride(), Mask,
+      N->getVectorLength(), N->getMemoryVT(), N->getMemOperand(),
+      N->isExpandingLoad());
+
+  // Legalize the chain result - switch anything that used the old chain to
+  // use the new one.
+  ReplaceValueWith(SDValue(N, 1), Res.getValue(1));
+  return Res;
+}
+
 SDValue DAGTypeLegalizer::WidenVecRes_MLOAD(MaskedLoadSDNode *N) {
 
   EVT WidenVT = TLI.getTypeToTransformTo(*DAG.getContext(),N->getValueType(0));
   SDValue Mask = N->getMask();
   EVT MaskVT = Mask.getValueType();
   SDValue PassThru = GetWidenedVector(N->getPassThru());
   ISD::LoadExtType ExtType = N->getExtensionType();
   SDLoc dl(N);
#endif
 
   case ISD::BITCAST:            Res = WidenVecOp_BITCAST(N); break;
   case ISD::CONCAT_VECTORS:     Res = WidenVecOp_CONCAT_VECTORS(N); break;
   case ISD::INSERT_SUBVECTOR:   Res = WidenVecOp_INSERT_SUBVECTOR(N); break;
   case ISD::EXTRACT_SUBVECTOR:  Res = WidenVecOp_EXTRACT_SUBVECTOR(N); break;
   case ISD::EXTRACT_VECTOR_ELT: Res = WidenVecOp_EXTRACT_VECTOR_ELT(N); break;
   case ISD::STORE:              Res = WidenVecOp_STORE(N); break;
   case ISD::VP_STORE:           Res = WidenVecOp_VP_STORE(N, OpNo); break;
+  case ISD::EXPERIMENTAL_VP_STRIDED_STORE:
+    Res = WidenVecOp_VP_STRIDED_STORE(N, OpNo);
+    break;
   case ISD::MSTORE:             Res = WidenVecOp_MSTORE(N, OpNo); break;
   case ISD::MGATHER:            Res = WidenVecOp_MGATHER(N, OpNo); break;
   case ISD::MSCATTER:           Res = WidenVecOp_MSCATTER(N, OpNo); break;
   case ISD::VP_SCATTER:         Res = WidenVecOp_VP_SCATTER(N, OpNo); break;
   case ISD::SETCC:              Res = WidenVecOp_SETCC(N); break;
   case ISD::STRICT_FSETCC:
   case ISD::STRICT_FSETCCS:     Res = WidenVecOp_STRICT_FSETCC(N); break;
   case ISD::VSELECT:            Res = WidenVecOp_VSELECT(N); break;
assert(Mask.getValueType().getVectorElementCount() ==
          "Mask and data vectors should have the same number of elements");
   return DAG.getStoreVP(ST->getChain(), dl, StVal, ST->getBasePtr(),
                         ST->getOffset(), Mask, ST->getVectorLength(),
                         ST->getMemoryVT(), ST->getMemOperand(),
                         ST->getAddressingMode(), ST->isTruncatingStore(),
                         ST->isCompressingStore());
 }
 
+SDValue DAGTypeLegalizer::WidenVecOp_VP_STRIDED_STORE(SDNode *N,
+                                                      unsigned OpNo) {
+  assert((OpNo == 1 || OpNo == 4) &&
+         "Can widen only data or mask operand of vp_strided_store");
+  VPStridedStoreSDNode *SST = cast<VPStridedStoreSDNode>(N);
+  SDValue Mask = SST->getMask();
+  SDValue StVal = SST->getValue();
+  SDLoc DL(N);
+
+  if (OpNo == 1)
+    assert(getTypeAction(Mask.getValueType()) ==
+               TargetLowering::TypeWidenVector &&
+           "Unable to widen VP strided store");
+  else
+    assert(getTypeAction(StVal.getValueType()) ==
+               TargetLowering::TypeWidenVector &&
+           "Unable to widen VP strided store");
+
+  StVal = GetWidenedVector(StVal);
+  Mask = GetWidenedVector(Mask);
+
+  assert(StVal.getValueType().getVectorElementCount() ==
+             Mask.getValueType().getVectorElementCount() &&
+         "Data and mask vectors should have the same number of elements");
+
+  return DAG.getStridedStoreVP(
+      SST->getChain(), DL, StVal, SST->getBasePtr(), SST->getOffset(),
+      SST->getStride(), Mask, SST->getVectorLength(), SST->getMemoryVT(),
+      SST->getMemOperand(), SST->getAddressingMode(), SST->isTruncatingStore(),
+      SST->isCompressingStore());
+}
+
 SDValue DAGTypeLegalizer::WidenVecOp_MSTORE(SDNode *N, unsigned OpNo) {
   assert((OpNo == 1 || OpNo == 3) &&
          "Can widen only data or mask operand of mstore");
   MaskedStoreSDNode *MST = cast<MaskedStoreSDNode>(N);
   SDValue Mask = MST->getMask();
   EVT MaskVT = Mask.getValueType();
   SDValue StVal = MST->getValue();
   SDLoc dl(N);

llvm/test/CodeGen/RISCV/rvv/fixed-vectors-strided-vpload.ll

 ; CHECK-RV64-LABEL: strided_vpload_v8f64:
 ; CHECK-RV64:       # %bb.0:
 ; CHECK-RV64-NEXT:    vsetvli zero, a2, e64, m4, ta, mu
 ; CHECK-RV64-NEXT:    vlse64.v v8, (a0), a1, v0.t
 ; CHECK-RV64-NEXT:    ret
   %load = call <8 x double> @llvm.experimental.vp.strided.load.v8f64(double* %ptr, i32 %stride, <8 x i1> %m, i32 %evl)
   ret <8 x double> %load
 }
+
+; Widening
+define <3 x double> @strided_vpload_v3f64(double* %ptr, i32 signext %stride, <3 x i1> %mask, i32 zeroext %evl) {
+; CHECK-RV32-LABEL: strided_vpload_v3f64:
+; CHECK-RV32:       # %bb.0:
+; CHECK-RV32-NEXT:    vsetvli zero, a2, e64, m2, ta, mu
+; CHECK-RV32-NEXT:    vlse64.v v8, (a0), a1, v0.t
+; CHECK-RV32-NEXT:    ret
+;
+; CHECK-RV64-LABEL: strided_vpload_v3f64:
+; CHECK-RV64:       # %bb.0:
+; CHECK-RV64-NEXT:    vsetvli zero, a2, e64, m2, ta, mu
+; CHECK-RV64-NEXT:    vlse64.v v8, (a0), a1, v0.t
+; CHECK-RV64-NEXT:    ret
+  %v = call <3 x double> @llvm.experimental.vp.strided.load.v3f64(double* %ptr, i32 %stride, <3 x i1> %mask, i32 %evl)
+  ret <3 x double> %v
+}
+
+define <3 x double> @strided_vpload_v3f64_allones_mask(double* %ptr, i32 signext %stride, i32 zeroext %evl) {
+; CHECK-RV32-LABEL: strided_vpload_v3f64_allones_mask:
+; CHECK-RV32:       # %bb.0:
+; CHECK-RV32-NEXT:    vsetvli zero, a2, e64, m2, ta, mu
+; CHECK-RV32-NEXT:    vlse64.v v8, (a0), a1
+; CHECK-RV32-NEXT:    ret
+;
+; CHECK-RV64-LABEL: strided_vpload_v3f64_allones_mask:
+; CHECK-RV64:       # %bb.0:
+; CHECK-RV64-NEXT:    vsetvli zero, a2, e64, m2, ta, mu
+; CHECK-RV64-NEXT:    vlse64.v v8, (a0), a1
+; CHECK-RV64-NEXT:    ret
+  %one = insertelement <3 x i1> undef, i1 true, i32 0
+  %allones = shufflevector <3 x i1> %one, <3 x i1> poison, <3 x i32> zeroinitializer
+  %v = call <3 x double> @llvm.experimental.vp.strided.load.v3f64(double* %ptr, i32 %stride, <3 x i1> %allones, i32 %evl)
+  ret <3 x double> %v
+}
+
+declare <3 x double> @llvm.experimental.vp.strided.load.v3f64(double*, i32, <3 x i1>, i32)

llvm/test/CodeGen/RISCV/rvv/fixed-vectors-strided-vpstore.ll

 ; CHECK-RV64-NEXT:    vsetvli zero, a2, e8, mf8, ta, mu
 ; CHECK-RV64-NEXT:    vsse8.v v8, (a0), a1
 ; CHECK-RV64-NEXT:    ret
   %a = insertelement <2 x i1> undef, i1 true, i32 0
   %b = shufflevector <2 x i1> %a, <2 x i1> poison, <2 x i32> zeroinitializer
   call void @llvm.experimental.vp.strided.store.v2i8(<2 x i8> %val, i8* %ptr, i32 %stride, <2 x i1> %b, i32 %evl)
   ret void
 }
+
+; Widening
+define void @strided_vpstore_v3f32(<3 x float> %v, float *%ptr, i32 signext %stride, <3 x i1> %mask, i32 zeroext %evl) {
+; CHECK-RV32-LABEL: strided_vpstore_v3f32:
+; CHECK-RV32:       # %bb.0:
+; CHECK-RV32-NEXT:    vsetvli zero, a2, e32, m1, ta, mu
+; CHECK-RV32-NEXT:    vsse32.v v8, (a0), a1, v0.t
+; CHECK-RV32-NEXT:    ret
+;
+; CHECK-RV64-LABEL: strided_vpstore_v3f32:
+; CHECK-RV64:       # %bb.0:
+; CHECK-RV64-NEXT:    vsetvli zero, a2, e32, m1, ta, mu
+; CHECK-RV64-NEXT:    vsse32.v v8, (a0), a1, v0.t
+; CHECK-RV64-NEXT:    ret
+  call void @llvm.experimental.vp.strided.store.v3f32(<3 x float> %v, float* %ptr, i32 %stride, <3 x i1> %mask, i32 %evl)
+  ret void
+}
+
+define void @strided_vpstore_v3f32_allones_mask(<3 x float> %v, float *%ptr, i32 signext %stride, i32 zeroext %evl) {
+; CHECK-RV32-LABEL: strided_vpstore_v3f32_allones_mask:
+; CHECK-RV32:       # %bb.0:
+; CHECK-RV32-NEXT:    vsetvli zero, a2, e32, m1, ta, mu
+; CHECK-RV32-NEXT:    vsse32.v v8, (a0), a1
+; CHECK-RV32-NEXT:    ret
+;
+; CHECK-RV64-LABEL: strided_vpstore_v3f32_allones_mask:
+; CHECK-RV64:       # %bb.0:
+; CHECK-RV64-NEXT:    vsetvli zero, a2, e32, m1, ta, mu
+; CHECK-RV64-NEXT:    vsse32.v v8, (a0), a1
+; CHECK-RV64-NEXT:    ret
+  %one = insertelement <3 x i1> undef, i1 true, i32 0
+  %allones = shufflevector <3 x i1> %one, <3 x i1> poison, <3 x i32> zeroinitializer
+  call void @llvm.experimental.vp.strided.store.v3f32(<3 x float> %v, float* %ptr, i32 %stride, <3 x i1> %allones, i32 %evl)
+  ret void
+}
+
+declare void @llvm.experimental.vp.strided.store.v3f32(<3 x float>, float* , i32, <3 x i1>, i32)

llvm/test/CodeGen/RISCV/rvv/strided-vpload.ll

 ; CHECK-RV64-LABEL: strided_vpload_nxv8f64:
 ; CHECK-RV64:       # %bb.0:
 ; CHECK-RV64-NEXT:    vsetvli zero, a2, e64, m8, ta, mu
 ; CHECK-RV64-NEXT:    vlse64.v v8, (a0), a1, v0.t
 ; CHECK-RV64-NEXT:    ret
   %load = call <vscale x 8 x double> @llvm.experimental.vp.strided.load.nxv8f64(double* %ptr, i32 signext %stride, <vscale x 8 x i1> %m, i32 %evl)
   ret <vscale x 8 x double> %load
 }
+
+; Widening
+define <vscale x 3 x double> @strided_vpload_nxv3f64(double* %ptr, i32 signext %stride, <vscale x 3 x i1> %mask, i32 zeroext %evl) {
+; CHECK-RV32-LABEL: strided_vpload_nxv3f64:
+; CHECK-RV32:       # %bb.0:
+; CHECK-RV32-NEXT:    vsetvli zero, a2, e64, m4, ta, mu
+; CHECK-RV32-NEXT:    vlse64.v v8, (a0), a1, v0.t
+; CHECK-RV32-NEXT:    ret
+;
+; CHECK-RV64-LABEL: strided_vpload_nxv3f64:
+; CHECK-RV64:       # %bb.0:
+; CHECK-RV64-NEXT:    vsetvli zero, a2, e64, m4, ta, mu
+; CHECK-RV64-NEXT:    vlse64.v v8, (a0), a1, v0.t
+; CHECK-RV64-NEXT:    ret
+  %v = call <vscale x 3 x double> @llvm.experimental.vp.strided.load.nxv3f64(double* %ptr, i32 %stride, <vscale x 3 x i1> %mask, i32 %evl)
+  ret <vscale x 3 x double> %v
+}
+
+define <vscale x 3 x double> @strided_vpload_nxv3f64_allones_mask(double* %ptr, i32 signext %stride, i32 zeroext %evl) {
+; CHECK-RV32-LABEL: strided_vpload_nxv3f64_allones_mask:
+; CHECK-RV32:       # %bb.0:
+; CHECK-RV32-NEXT:    vsetvli zero, a2, e64, m4, ta, mu
+; CHECK-RV32-NEXT:    vlse64.v v8, (a0), a1
+; CHECK-RV32-NEXT:    ret
+;
+; CHECK-RV64-LABEL: strided_vpload_nxv3f64_allones_mask:
+; CHECK-RV64:       # %bb.0:
+; CHECK-RV64-NEXT:    vsetvli zero, a2, e64, m4, ta, mu
+; CHECK-RV64-NEXT:    vlse64.v v8, (a0), a1
+; CHECK-RV64-NEXT:    ret
+  %one = insertelement <vscale x 3 x i1> undef, i1 true, i32 0
+  %allones = shufflevector <vscale x 3 x i1> %one, <vscale x 3 x i1> poison, <vscale x 3 x i32> zeroinitializer
+  %v = call <vscale x 3 x double> @llvm.experimental.vp.strided.load.nxv3f64(double* %ptr, i32 %stride, <vscale x 3 x i1> %allones, i32 %evl)
+  ret <vscale x 3 x double> %v
+}
+
+declare <vscale x 3 x double> @llvm.experimental.vp.strided.load.nxv3f64(double*, i32, <vscale x 3 x i1>, i32)

llvm/test/CodeGen/RISCV/rvv/strided-vpstore.ll

 ; CHECK-RV64-NEXT:    vsetvli zero, a2, e8, mf8, ta, mu
 ; CHECK-RV64-NEXT:    vsse8.v v8, (a0), a1
 ; CHECK-RV64-NEXT:    ret
   %a = insertelement <vscale x 1 x i1> undef, i1 true, i32 0
   %b = shufflevector <vscale x 1 x i1> %a, <vscale x 1 x i1> poison, <vscale x 1 x i32> zeroinitializer
   call void @llvm.experimental.vp.strided.store.nxv1i8(<vscale x 1 x i8> %val, i8* %ptr, i32 %strided, <vscale x 1 x i1> %b, i32 %evl)
   ret void
 }
+
+; Widening
+define void @strided_vpstore_nxv3f32(<vscale x 3 x float> %v, float *%ptr, i32 signext %stride, <vscale x 3 x i1> %mask, i32 zeroext %evl) {
+; CHECK-RV32-LABEL: strided_vpstore_nxv3f32:
+; CHECK-RV32:       # %bb.0:
+; CHECK-RV32-NEXT:    vsetvli zero, a2, e32, m2, ta, mu
+; CHECK-RV32-NEXT:    vsse32.v v8, (a0), a1, v0.t
+; CHECK-RV32-NEXT:    ret
+;
+; CHECK-RV64-LABEL: strided_vpstore_nxv3f32:
+; CHECK-RV64:       # %bb.0:
+; CHECK-RV64-NEXT:    vsetvli zero, a2, e32, m2, ta, mu
+; CHECK-RV64-NEXT:    vsse32.v v8, (a0), a1, v0.t
+; CHECK-RV64-NEXT:    ret
+  call void @llvm.experimental.vp.strided.store.nxv3f32(<vscale x 3 x float> %v, float* %ptr, i32 %stride, <vscale x 3 x i1> %mask, i32 %evl)
+  ret void
+}
+
+define void @strided_vpstore_nxv3f32_allones_mask(<vscale x 3 x float> %v, float *%ptr, i32 signext %stride, i32 zeroext %evl) {
+; CHECK-RV32-LABEL: strided_vpstore_nxv3f32_allones_mask:
+; CHECK-RV32:       # %bb.0:
+; CHECK-RV32-NEXT:    vsetvli zero, a2, e32, m2, ta, mu
+; CHECK-RV32-NEXT:    vsse32.v v8, (a0), a1
+; CHECK-RV32-NEXT:    ret
+;
+; CHECK-RV64-LABEL: strided_vpstore_nxv3f32_allones_mask:
+; CHECK-RV64:       # %bb.0:
+; CHECK-RV64-NEXT:    vsetvli zero, a2, e32, m2, ta, mu
+; CHECK-RV64-NEXT:    vsse32.v v8, (a0), a1
+; CHECK-RV64-NEXT:    ret
+  %one = insertelement <vscale x 3 x i1> undef, i1 true, i32 0
+  %allones = shufflevector <vscale x 3 x i1> %one, <vscale x 3 x i1> poison, <vscale x 3 x i32> zeroinitializer
+  call void @llvm.experimental.vp.strided.store.nxv3f32(<vscale x 3 x float> %v, float* %ptr, i32 %stride, <vscale x 3 x i1> %allones, i32 %evl)
+  ret void
+}
+
+declare void @llvm.experimental.vp.strided.store.nxv3f32(<vscale x 3 x float>, float* , i32, <vscale x 3 x i1>, i32)