[llvm] [SelectionDAG] Allow vselect in foldBinOpIntoSelect (PR #143283)

[via llvm-commits]

https://github.com/AZero13 updated https://github.com/llvm/llvm-project/pull/143283

>From 04979113b4196ffc791bde0b0bcc7d25e3f5e43d Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Sat, 7 Jun 2025 13:03:10 -0400
Subject: [PATCH] [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.
---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp    | 16 ++++++++++------
 .../CodeGen/RISCV/rvv/narrow-shift-extend.ll     |  3 +--
 llvm/test/CodeGen/X86/extract-vselect-setcc.ll   | 10 +++++-----
 3 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index aba3c0f80a024..e4970a79371f8 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/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);
diff --git a/llvm/test/CodeGen/RISCV/rvv/narrow-shift-extend.ll b/llvm/test/CodeGen/RISCV/rvv/narrow-shift-extend.ll
index 557e809f4f9a8..5c451ebb6b267 100644
--- a/llvm/test/CodeGen/RISCV/rvv/narrow-shift-extend.ll
+++ b/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
diff --git a/llvm/test/CodeGen/X86/extract-vselect-setcc.ll b/llvm/test/CodeGen/X86/extract-vselect-setcc.ll
index 96c8e773d5edd..eff130b25dfab 100644
--- a/llvm/test/CodeGen/X86/extract-vselect-setcc.ll
+++ b/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