[RISCV][VP] Add RVV codegen for [nX]vXi1 vp.select

[victor-eds]

Summary

Expand [nX]vXi1 vp.select the same way as [nX]vXi1 vselect.

[frasercrmck]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp

1204

nit: Implement

1211

nit: EVL would be a more appropriate name, I think.

1216

nit: VP_XOR

1221

The following expansion is only correct under certain conditions: when the vselect condition and operand types are identically-sized integer vectors and when the mask is -1/0. See the original ExpandVSELECT for that.

Since we don't have an in-tree test case for anything other than expanding [nx]vXi1 VSELECT maybe it's sufficient for now to resort to UnrollVectorOp whenever the operand types aren't i1 vectors? UnrollVectorOp will no doubt fail for VP_SELECT right now, but we wouldn't be making anything worse.

Not sure: any thoughts @craig.topper? @RKSimon?

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp

1204

This code only works if Mask's VT is the same as Op1/Op2 VT but that isn't guaranteed in general. We should call UnrollVectorOp for that just like ExpandVSELECT.

1221

I think I'm fine with calling Unroll if it isn't an nxvXi1 vector.

1222

Why a lamdba? Can't we just do

SDValue Ones = DAG.getAllOnesConstant(DL, VT);
SDValue NotMask = DAG.getNode(ISD::VP_XOR, DL, VT, Mask, Ones, Ones, Length);

Arguably that should be

DAG.getNode(ISD::VP_XOR, DL, VT, Mask, Ones, Mask, Length);

so the XOR has the same mask as the other ops.

[victor-eds]

Unroll op if the operands are not i1 vectors

[craig.topper]

LGTM

[victor-eds]

Rebase

[victor-eds]

@craig.topper Can you please land this one for me? I tried to fix the pipeline by rebasing (hence the last diff), but it appears to be broken in main. I'll ask for commit access after this patch has been landed so that I don't need to ask you for it again. Thanks!

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp

class VectorLegalizer {
   SDValue ExpandZERO_EXTEND_VECTOR_INREG(SDNode *Node);
 
   /// Expand bswap of vectors into a shuffle if legal.
   SDValue ExpandBSWAP(SDNode *Node);
 
   /// Implement vselect in terms of XOR, AND, OR when blend is not
   /// supported by the target.
   SDValue ExpandVSELECT(SDNode *Node);
+  SDValue ExpandVP_SELECT(SDNode *Node);
   SDValue ExpandSELECT(SDNode *Node);
   std::pair<SDValue, SDValue> ExpandLoad(SDNode *N);
   SDValue ExpandStore(SDNode *N);
   SDValue ExpandFNEG(SDNode *Node);
   void ExpandFSUB(SDNode *Node, SmallVectorImpl<SDValue> &Results);
   void ExpandSETCC(SDNode *Node, SmallVectorImpl<SDValue> &Results);
   void ExpandBITREVERSE(SDNode *Node, SmallVectorImpl<SDValue> &Results);
   void ExpandUADDSUBO(SDNode *Node, SmallVectorImpl<SDValue> &Results);
#include "llvm/IR/ConstrainedOps.def"
   case ISD::BITREVERSE:
   case ISD::CTLZ:
   case ISD::CTTZ:
   case ISD::CTLZ_ZERO_UNDEF:
   case ISD::CTTZ_ZERO_UNDEF:
   case ISD::CTPOP:
   case ISD::SELECT:
   case ISD::VSELECT:
+  case ISD::VP_SELECT:
   case ISD::SELECT_CC:
   case ISD::ZERO_EXTEND:
   case ISD::ANY_EXTEND:
   case ISD::TRUNCATE:
   case ISD::SIGN_EXTEND:
   case ISD::FP_TO_SINT:
   case ISD::FP_TO_UINT:
   case ISD::FNEG:
case ISD::ZERO_EXTEND_VECTOR_INREG:
     Results.push_back(ExpandZERO_EXTEND_VECTOR_INREG(Node));
     return;
   case ISD::BSWAP:
     Results.push_back(ExpandBSWAP(Node));
     return;
   case ISD::VSELECT:
     Results.push_back(ExpandVSELECT(Node));
     return;
+  case ISD::VP_SELECT:
+    Results.push_back(ExpandVP_SELECT(Node));
+    return;
   case ISD::SELECT:
     Results.push_back(ExpandSELECT(Node));
     return;
   case ISD::FP_TO_UINT:
     ExpandFP_TO_UINT(Node, Results);
     return;
   case ISD::UINT_TO_FP:
     ExpandUINT_TO_FLOAT(Node, Results);
SDValue VectorLegalizer::ExpandVSELECT(SDNode *Node) {
   SDValue NotMask = DAG.getNOT(DL, Mask, VT);
 
   Op1 = DAG.getNode(ISD::AND, DL, VT, Op1, Mask);
   Op2 = DAG.getNode(ISD::AND, DL, VT, Op2, NotMask);
   SDValue Val = DAG.getNode(ISD::OR, DL, VT, Op1, Op2);
   return DAG.getNode(ISD::BITCAST, DL, Node->getValueType(0), Val);
 }
 
+SDValue VectorLegalizer::ExpandVP_SELECT(SDNode *Node) {
+  // Implement VP_SELECT in terms of VP_XOR, VP_AND and VP_OR on platforms which
+  // do not support it natively.
+  SDLoc DL(Node);
+
+  SDValue Mask = Node->getOperand(0);
+  SDValue Op1 = Node->getOperand(1);
+  SDValue Op2 = Node->getOperand(2);
+  SDValue EVL = Node->getOperand(3);
+
+  EVT VT = Mask.getValueType();
+
+  // If we can't even use the basic vector operations of
+  // VP_AND,VP_OR,VP_XOR, we will have to scalarize the op.
+  if (TLI.getOperationAction(ISD::VP_AND, VT) == TargetLowering::Expand ||
+      TLI.getOperationAction(ISD::VP_XOR, VT) == TargetLowering::Expand ||
+      TLI.getOperationAction(ISD::VP_OR, VT) == TargetLowering::Expand)
+    return DAG.UnrollVectorOp(Node);
+
+  // This operation also isn't safe when the operands aren't also booleans.
+  if (Op1.getValueType().getVectorElementType() != MVT::i1)
+    return DAG.UnrollVectorOp(Node);
+
+  SDValue Ones = DAG.getAllOnesConstant(DL, VT);
+  SDValue NotMask = DAG.getNode(ISD::VP_XOR, DL, VT, Mask, Ones, Mask, EVL);
+
+  Op1 = DAG.getNode(ISD::VP_AND, DL, VT, Op1, Mask, Mask, EVL);
+  Op2 = DAG.getNode(ISD::VP_AND, DL, VT, Op2, NotMask, Mask, EVL);
+  return DAG.getNode(ISD::VP_OR, DL, VT, Op1, Op2, Mask, EVL);
+}
+
 void VectorLegalizer::ExpandFP_TO_UINT(SDNode *Node,
                                        SmallVectorImpl<SDValue> &Results) {
   // Attempt to expand using TargetLowering.
   SDValue Result, Chain;
   if (TLI.expandFP_TO_UINT(Node, Result, Chain, DAG)) {
     Results.push_back(Result);
     if (Node->isStrictFPOpcode())
       Results.push_back(Chain);

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

for (MVT VT : BoolVecVTs) {
       setOperationAction(ISD::EXTRACT_SUBVECTOR, VT, Custom);
 
       setOperationAction(ISD::INSERT_VECTOR_ELT, VT, Custom);
       setOperationAction(ISD::EXTRACT_VECTOR_ELT, VT, Custom);
 
       setOperationAction(ISD::SELECT, VT, Custom);
       setOperationAction(ISD::SELECT_CC, VT, Expand);
       setOperationAction(ISD::VSELECT, VT, Expand);
+      setOperationAction(ISD::VP_SELECT, VT, Expand);
 
       setOperationAction(ISD::VP_AND, VT, Custom);
       setOperationAction(ISD::VP_OR, VT, Custom);
       setOperationAction(ISD::VP_XOR, VT, Custom);
 
       setOperationAction(ISD::VECREDUCE_AND, VT, Custom);
       setOperationAction(ISD::VECREDUCE_OR, VT, Custom);
       setOperationAction(ISD::VECREDUCE_XOR, VT, Custom);

llvm/test/CodeGen/RISCV/rvv/fixed-vectors-vselect-vp.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=riscv32 -mattr=+d,+experimental-zfh,+experimental-v -target-abi=ilp32d -riscv-v-vector-bits-min=128 \
 ; RUN:   -verify-machineinstrs < %s | FileCheck %s
 ; RUN: llc -mtriple=riscv64 -mattr=+d,+experimental-zfh,+experimental-v -target-abi=lp64d -riscv-v-vector-bits-min=128 \
 ; RUN:   -verify-machineinstrs < %s | FileCheck %s
 
+declare <1 x i1> @llvm.vp.select.v1i1(<1 x i1>, <1 x i1>, <1 x i1>, i32)
+
+define <1 x i1> @select_v1i1(<1 x i1> %a, <1 x i1> %b, <1 x i1> %c, i32 zeroext %evl) {
+; CHECK-LABEL: select_v1i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e8, mf8, ta, mu
+; CHECK-NEXT:    vmandn.mm v9, v9, v0
+; CHECK-NEXT:    vmand.mm v8, v8, v0
+; CHECK-NEXT:    vmor.mm v0, v8, v9
+; CHECK-NEXT:    ret
+  %v = call <1 x i1> @llvm.vp.select.v1i1(<1 x i1> %a, <1 x i1> %b, <1 x i1> %c, i32 %evl)
+  ret <1 x i1> %v
+}
+
+declare <2 x i1> @llvm.vp.select.v2i1(<2 x i1>, <2 x i1>, <2 x i1>, i32)
+
+define <2 x i1> @select_v2i1(<2 x i1> %a, <2 x i1> %b, <2 x i1> %c, i32 zeroext %evl) {
+; CHECK-LABEL: select_v2i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e8, mf8, ta, mu
+; CHECK-NEXT:    vmandn.mm v9, v9, v0
+; CHECK-NEXT:    vmand.mm v8, v8, v0
+; CHECK-NEXT:    vmor.mm v0, v8, v9
+; CHECK-NEXT:    ret
+  %v = call <2 x i1> @llvm.vp.select.v2i1(<2 x i1> %a, <2 x i1> %b, <2 x i1> %c, i32 %evl)
+  ret <2 x i1> %v
+}
+
+declare <4 x i1> @llvm.vp.select.v4i1(<4 x i1>, <4 x i1>, <4 x i1>, i32)
+
+define <4 x i1> @select_v4i1(<4 x i1> %a, <4 x i1> %b, <4 x i1> %c, i32 zeroext %evl) {
+; CHECK-LABEL: select_v4i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e8, mf4, ta, mu
+; CHECK-NEXT:    vmandn.mm v9, v9, v0
+; CHECK-NEXT:    vmand.mm v8, v8, v0
+; CHECK-NEXT:    vmor.mm v0, v8, v9
+; CHECK-NEXT:    ret
+  %v = call <4 x i1> @llvm.vp.select.v4i1(<4 x i1> %a, <4 x i1> %b, <4 x i1> %c, i32 %evl)
+  ret <4 x i1> %v
+}
+
+declare <8 x i1> @llvm.vp.select.v8i1(<8 x i1>, <8 x i1>, <8 x i1>, i32)
+
+define <8 x i1> @select_v8i1(<8 x i1> %a, <8 x i1> %b, <8 x i1> %c, i32 zeroext %evl) {
+; CHECK-LABEL: select_v8i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e8, mf2, ta, mu
+; CHECK-NEXT:    vmandn.mm v9, v9, v0
+; CHECK-NEXT:    vmand.mm v8, v8, v0
+; CHECK-NEXT:    vmor.mm v0, v8, v9
+; CHECK-NEXT:    ret
+  %v = call <8 x i1> @llvm.vp.select.v8i1(<8 x i1> %a, <8 x i1> %b, <8 x i1> %c, i32 %evl)
+  ret <8 x i1> %v
+}
+
+declare <16 x i1> @llvm.vp.select.v16i1(<16 x i1>, <16 x i1>, <16 x i1>, i32)
+
+define <16 x i1> @select_v16i1(<16 x i1> %a, <16 x i1> %b, <16 x i1> %c, i32 zeroext %evl) {
+; CHECK-LABEL: select_v16i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e8, m1, ta, mu
+; CHECK-NEXT:    vmandn.mm v9, v9, v0
+; CHECK-NEXT:    vmand.mm v8, v8, v0
+; CHECK-NEXT:    vmor.mm v0, v8, v9
+; CHECK-NEXT:    ret
+  %v = call <16 x i1> @llvm.vp.select.v16i1(<16 x i1> %a, <16 x i1> %b, <16 x i1> %c, i32 %evl)
+  ret <16 x i1> %v
+}
+
 declare <2 x i8> @llvm.vp.select.v2i8(<2 x i1>, <2 x i8>, <2 x i8>, i32)
 
 define <2 x i8> @select_v2i8(<2 x i1> %a, <2 x i8> %b, <2 x i8> %c, i32 zeroext %evl) {
 ; CHECK-LABEL: select_v2i8:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli zero, a0, e8, mf8, ta, mu
 ; CHECK-NEXT:    vmerge.vvm v8, v9, v8, v0
 ; CHECK-NEXT:    ret

llvm/test/CodeGen/RISCV/rvv/vselect-vp.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=riscv32 -mattr=+d,+experimental-zfh,+experimental-v -target-abi=ilp32d \
 ; RUN:     -verify-machineinstrs < %s | FileCheck %s
 ; RUN: llc -mtriple=riscv64 -mattr=+d,+experimental-zfh,+experimental-v -target-abi=lp64d \
 ; RUN:     -verify-machineinstrs < %s | FileCheck %s
 
+declare <vscale x 1 x i1> @llvm.vp.select.nxv1i1(<vscale x 1 x i1>, <vscale x 1 x i1>, <vscale x 1 x i1>, i32)
+
+define <vscale x 1 x i1> @select_nxv1i1(<vscale x 1 x i1> %a, <vscale x 1 x i1> %b, <vscale x 1 x i1> %c, i32 zeroext %evl) {
+; CHECK-LABEL: select_nxv1i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e8, mf8, ta, mu
+; CHECK-NEXT:    vmandn.mm v9, v9, v0
+; CHECK-NEXT:    vmand.mm v8, v8, v0
+; CHECK-NEXT:    vmor.mm v0, v8, v9
+; CHECK-NEXT:    ret
+  %v = call <vscale x 1 x i1> @llvm.vp.select.nxv1i1(<vscale x 1 x i1> %a, <vscale x 1 x i1> %b, <vscale x 1 x i1> %c, i32 %evl)
+  ret <vscale x 1 x i1> %v
+}
+
+declare <vscale x 2 x i1> @llvm.vp.select.nxv2i1(<vscale x 2 x i1>, <vscale x 2 x i1>, <vscale x 2 x i1>, i32)
+
+define <vscale x 2 x i1> @select_nxv2i1(<vscale x 2 x i1> %a, <vscale x 2 x i1> %b, <vscale x 2 x i1> %c, i32 zeroext %evl) {
+; CHECK-LABEL: select_nxv2i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e8, mf4, ta, mu
+; CHECK-NEXT:    vmandn.mm v9, v9, v0
+; CHECK-NEXT:    vmand.mm v8, v8, v0
+; CHECK-NEXT:    vmor.mm v0, v8, v9
+; CHECK-NEXT:    ret
+  %v = call <vscale x 2 x i1> @llvm.vp.select.nxv2i1(<vscale x 2 x i1> %a, <vscale x 2 x i1> %b, <vscale x 2 x i1> %c, i32 %evl)
+  ret <vscale x 2 x i1> %v
+}
+
+declare <vscale x 4 x i1> @llvm.vp.select.nxv4i1(<vscale x 4 x i1>, <vscale x 4 x i1>, <vscale x 4 x i1>, i32)
+
+define <vscale x 4 x i1> @select_nxv4i1(<vscale x 4 x i1> %a, <vscale x 4 x i1> %b, <vscale x 4 x i1> %c, i32 zeroext %evl) {
+; CHECK-LABEL: select_nxv4i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e8, mf2, ta, mu
+; CHECK-NEXT:    vmandn.mm v9, v9, v0
+; CHECK-NEXT:    vmand.mm v8, v8, v0
+; CHECK-NEXT:    vmor.mm v0, v8, v9
+; CHECK-NEXT:    ret
+  %v = call <vscale x 4 x i1> @llvm.vp.select.nxv4i1(<vscale x 4 x i1> %a, <vscale x 4 x i1> %b, <vscale x 4 x i1> %c, i32 %evl)
+  ret <vscale x 4 x i1> %v
+}
+
+declare <vscale x 8 x i1> @llvm.vp.select.nxv8i1(<vscale x 8 x i1>, <vscale x 8 x i1>, <vscale x 8 x i1>, i32)
+
+define <vscale x 8 x i1> @select_nxv8i1(<vscale x 8 x i1> %a, <vscale x 8 x i1> %b, <vscale x 8 x i1> %c, i32 zeroext %evl) {
+; CHECK-LABEL: select_nxv8i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e8, m1, ta, mu
+; CHECK-NEXT:    vmandn.mm v9, v9, v0
+; CHECK-NEXT:    vmand.mm v8, v8, v0
+; CHECK-NEXT:    vmor.mm v0, v8, v9
+; CHECK-NEXT:    ret
+  %v = call <vscale x 8 x i1> @llvm.vp.select.nxv8i1(<vscale x 8 x i1> %a, <vscale x 8 x i1> %b, <vscale x 8 x i1> %c, i32 %evl)
+  ret <vscale x 8 x i1> %v
+}
+
+declare <vscale x 16 x i1> @llvm.vp.select.nxv16i1(<vscale x 16 x i1>, <vscale x 16 x i1>, <vscale x 16 x i1>, i32)
+
+define <vscale x 16 x i1> @select_nxv16i1(<vscale x 16 x i1> %a, <vscale x 16 x i1> %b, <vscale x 16 x i1> %c, i32 zeroext %evl) {
+; CHECK-LABEL: select_nxv16i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e8, m2, ta, mu
+; CHECK-NEXT:    vmandn.mm v9, v9, v0
+; CHECK-NEXT:    vmand.mm v8, v8, v0
+; CHECK-NEXT:    vmor.mm v0, v8, v9
+; CHECK-NEXT:    ret
+  %v = call <vscale x 16 x i1> @llvm.vp.select.nxv16i1(<vscale x 16 x i1> %a, <vscale x 16 x i1> %b, <vscale x 16 x i1> %c, i32 %evl)
+  ret <vscale x 16 x i1> %v
+}
+
+declare <vscale x 32 x i1> @llvm.vp.select.nxv32i1(<vscale x 32 x i1>, <vscale x 32 x i1>, <vscale x 32 x i1>, i32)
+
+define <vscale x 32 x i1> @select_nxv32i1(<vscale x 32 x i1> %a, <vscale x 32 x i1> %b, <vscale x 32 x i1> %c, i32 zeroext %evl) {
+; CHECK-LABEL: select_nxv32i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e8, m4, ta, mu
+; CHECK-NEXT:    vmandn.mm v9, v9, v0
+; CHECK-NEXT:    vmand.mm v8, v8, v0
+; CHECK-NEXT:    vmor.mm v0, v8, v9
+; CHECK-NEXT:    ret
+  %v = call <vscale x 32 x i1> @llvm.vp.select.nxv32i1(<vscale x 32 x i1> %a, <vscale x 32 x i1> %b, <vscale x 32 x i1> %c, i32 %evl)
+  ret <vscale x 32 x i1> %v
+}
+
+declare <vscale x 64 x i1> @llvm.vp.select.nxv64i1(<vscale x 64 x i1>, <vscale x 64 x i1>, <vscale x 64 x i1>, i32)
+
+define <vscale x 64 x i1> @select_nxv64i1(<vscale x 64 x i1> %a, <vscale x 64 x i1> %b, <vscale x 64 x i1> %c, i32 zeroext %evl) {
+; CHECK-LABEL: select_nxv64i1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetvli zero, a0, e8, m8, ta, mu
+; CHECK-NEXT:    vmandn.mm v9, v9, v0
+; CHECK-NEXT:    vmand.mm v8, v8, v0
+; CHECK-NEXT:    vmor.mm v0, v8, v9
+; CHECK-NEXT:    ret
+  %v = call <vscale x 64 x i1> @llvm.vp.select.nxv64i1(<vscale x 64 x i1> %a, <vscale x 64 x i1> %b, <vscale x 64 x i1> %c, i32 %evl)
+  ret <vscale x 64 x i1> %v
+}
+
 declare <vscale x 1 x i8> @llvm.vp.select.nxv1i8(<vscale x 1 x i1>, <vscale x 1 x i8>, <vscale x 1 x i8>, i32)
 
 define <vscale x 1 x i8> @select_nxv1i8(<vscale x 1 x i1> %a, <vscale x 1 x i8> %b, <vscale x 1 x i8> %c, i32 zeroext %evl) {
 ; CHECK-LABEL: select_nxv1i8:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli zero, a0, e8, mf8, ta, mu
 ; CHECK-NEXT:    vmerge.vvm v8, v9, v8, v0
 ; CHECK-NEXT:    ret