[SelectionDAG] Allow vselect in selection folds

It seems vselect was also meant to be an option given the comment and the fact vectors are allowed and the kind is checked too.

Author: AZero13

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -2490,7 +2490,8 @@ SDValue DAGCombiner::foldBinOpIntoSelect(SDNode *BO) {
   unsigned SelOpNo = 0;
   SDValue Sel = BO->getOperand(0);
   auto BinOpcode = BO->getOpcode();
-  if (Sel.getOpcode() != ISD::SELECT || !Sel.hasOneUse()) {
+  if ((Sel.getOpcode() != ISD::SELECT && Sel.getOpcode() != ISD::VSELECT) ||
+      !Sel.hasOneUse()) {
     SelOpNo = 1;
     Sel = BO->getOperand(1);
 
@@ -2506,7 +2507,8 @@ SDValue DAGCombiner::foldBinOpIntoSelect(SDNode *BO) {
     }
   }
 
-  if (Sel.getOpcode() != ISD::SELECT || !Sel.hasOneUse())
+  if ((Sel.getOpcode() != ISD::SELECT && Sel.getOpcode() != ISD::VSELECT) ||
+      !Sel.hasOneUse())
     return SDValue();
 
   SDValue CT = Sel.getOperand(1);
@@ -10017,7 +10019,8 @@ SDValue DAGCombiner::visitShiftByConstant(SDNode *N) {
                             BinOpLHSVal.getOpcode() == ISD::SRL) &&
                            isa<ConstantSDNode>(BinOpLHSVal.getOperand(1));
   bool IsCopyOrSelect = BinOpLHSVal.getOpcode() == ISD::CopyFromReg ||
-                        BinOpLHSVal.getOpcode() == ISD::SELECT;
+                        BinOpLHSVal.getOpcode() == ISD::SELECT ||
+                        BinOpLHSVal.getOpcode() == ISD::VSELECT;
 
   if (!IsShiftByConstant && !IsCopyOrSelect)
     return SDValue();
@@ -13435,7 +13438,7 @@ static SDValue tryToFoldExtendOfConstant(SDNode *N, const SDLoc &DL,
   // fold (sext (select cond, c1, c2)) -> (select cond, sext c1, sext c2)
   // fold (zext (select cond, c1, c2)) -> (select cond, zext c1, zext c2)
   // fold (aext (select cond, c1, c2)) -> (select cond, sext c1, sext c2)
-  if (N0->getOpcode() == ISD::SELECT) {
+  if (N0->getOpcode() == ISD::SELECT || N0->getOpcode() == ISD::VSELECT) {
     SDValue Op1 = N0->getOperand(1);
     SDValue Op2 = N0->getOperand(2);
     if (isa<ConstantSDNode>(Op1) && isa<ConstantSDNode>(Op2) &&
@@ -17763,10 +17766,11 @@ SDValue DAGCombiner::visitFMUL(SDNode *N) {
   // fold (fmul X, (select (fcmp X > 0.0), -1.0, 1.0)) -> (fneg (fabs X))
   // fold (fmul X, (select (fcmp X > 0.0), 1.0, -1.0)) -> (fabs X)
   if (Flags.hasNoNaNs() && Flags.hasNoSignedZeros() &&
-      (N0.getOpcode() == ISD::SELECT || N1.getOpcode() == ISD::SELECT) &&
+      (N0.getOpcode() == ISD::SELECT || N0.getOpcode() == ISD::VSELECT ||
+       N1.getOpcode() == ISD::SELECT || N1.getOpcode() == ISD::VSELECT) &&
       TLI.isOperationLegal(ISD::FABS, VT)) {
     SDValue Select = N0, X = N1;
-    if (Select.getOpcode() != ISD::SELECT)
+    if (Select.getOpcode() != ISD::SELECT && Select.getOpcode() != ISD::VSELECT)
       std::swap(Select, X);
 
     SDValue Cond = Select.getOperand(0);

llvm/test/CodeGen/RISCV/rvv/narrow-shift-extend.ll

@@ -144,8 +144,7 @@ define <vscale x 4 x i32> @test_vloxei7(ptr %ptr, <vscale x 4 x i1> %offset, i64
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    vsetvli zero, a1, e64, m4, ta, ma
 ; CHECK-NEXT:    vmv.v.i v8, 0
-; CHECK-NEXT:    vmerge.vim v8, v8, 1, v0
-; CHECK-NEXT:    vsll.vi v12, v8, 2
+; CHECK-NEXT:    vmerge.vim v12, v8, 4, v0
 ; CHECK-NEXT:    vsetvli zero, zero, e32, m2, ta, ma
 ; CHECK-NEXT:    vloxei64.v v8, (a0), v12
 ; CHECK-NEXT:    ret

llvm/test/CodeGen/X86/extract-vselect-setcc.ll

@@ -5,11 +5,11 @@ define void @PR117684(i1 %cond, <8 x float> %vec, ptr %ptr1, ptr %ptr2) #0 {
 ; CHECK-LABEL: PR117684:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vxorps %xmm1, %xmm1, %xmm1
-; CHECK-NEXT:    vcmpnltss %xmm1, %xmm0, %k1
-; CHECK-NEXT:    vbroadcastss {{.*#+}} xmm0 = [NaN,NaN,NaN,NaN]
-; CHECK-NEXT:    vinsertf32x4 $0, %xmm0, %ymm0, %ymm0 {%k1} {z}
-; CHECK-NEXT:    vmulss %xmm1, %xmm0, %xmm0
-; CHECK-NEXT:    vmulss {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm2
+; CHECK-NEXT:    vmovss {{.*#+}} xmm2 = [NaN,0.0E+0,0.0E+0,0.0E+0]
+; CHECK-NEXT:    vcmpltss %xmm1, %xmm0, %k1
+; CHECK-NEXT:    vmovaps %xmm2, %xmm0
+; CHECK-NEXT:    vmovss %xmm1, %xmm0, %xmm0 {%k1}
+; CHECK-NEXT:    vmulss %xmm2, %xmm0, %xmm2
 ; CHECK-NEXT:    vbroadcastss %xmm2, %ymm2
 ; CHECK-NEXT:    testb $1, %dil
 ; CHECK-NEXT:    cmoveq %rdx, %rsi