[X86] Add handling for select(icmp_uge(amt,BW),0,shift_logical(x,amt)) -> avx2 shift(x,amt)

We need to catch this otherwise pre-AVX512 targets will fold this to and(icmp_ult(amt,BW),shift_logical(x,amt))

Author: RKSimon

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -46190,20 +46190,32 @@ static SDValue combineSelect(SDNode *N, SelectionDAG &DAG,
   // to bitwidth-1 for unsigned shifts, effectively performing a maximum left
   // shift of bitwidth-1 positions. and returns zero for unsigned right shifts
   // exceeding bitwidth-1.
-  if (N->getOpcode() == ISD::VSELECT &&
-      (LHS.getOpcode() == ISD::SRL || LHS.getOpcode() == ISD::SHL) &&
-      supportedVectorVarShift(VT, Subtarget, LHS.getOpcode())) {
+  if (N->getOpcode() == ISD::VSELECT) {
     using namespace llvm::SDPatternMatch;
     // fold select(icmp_ult(amt,BW),shl(x,amt),0) -> avx2 psllv(x,amt)
     // fold select(icmp_ult(amt,BW),srl(x,amt),0) -> avx2 psrlv(x,amt)
-    if (ISD::isConstantSplatVectorAllZeros(RHS.getNode()) &&
+    if ((LHS.getOpcode() == ISD::SRL || LHS.getOpcode() == ISD::SHL) &&
+        supportedVectorVarShift(VT, Subtarget, LHS.getOpcode()) &&
+        ISD::isConstantSplatVectorAllZeros(RHS.getNode()) &&
         sd_match(Cond, m_SetCC(m_Specific(LHS.getOperand(1)),
                                m_SpecificInt(VT.getScalarSizeInBits()),
                                m_SpecificCondCode(ISD::SETULT)))) {
       return DAG.getNode(LHS.getOpcode() == ISD::SRL ? X86ISD::VSRLV
                                                      : X86ISD::VSHLV,
                          DL, VT, LHS.getOperand(0), LHS.getOperand(1));
     }
+    // fold select(icmp_uge(amt,BW),0,shl(x,amt)) -> avx2 psllv(x,amt)
+    // fold select(icmp_uge(amt,BW),0,srl(x,amt)) -> avx2 psrlv(x,amt)
+    if ((RHS.getOpcode() == ISD::SRL || RHS.getOpcode() == ISD::SHL) &&
+        supportedVectorVarShift(VT, Subtarget, RHS.getOpcode()) &&
+        ISD::isConstantSplatVectorAllZeros(LHS.getNode()) &&
+        sd_match(Cond, m_SetCC(m_Specific(RHS.getOperand(1)),
+                               m_SpecificInt(VT.getScalarSizeInBits()),
+                               m_SpecificCondCode(ISD::SETUGE)))) {
+      return DAG.getNode(RHS.getOpcode() == ISD::SRL ? X86ISD::VSRLV
+                                                     : X86ISD::VSHLV,
+                         DL, VT, RHS.getOperand(0), RHS.getOperand(1));
+    }
   }
 
   // Early exit check

llvm/test/CodeGen/X86/combine-shl.ll

@@ -1086,19 +1086,10 @@ define <4 x i32> @combine_vec_shl_commuted_clamped1(<4 x i32> %sh, <4 x i32> %am
 ; SSE41-NEXT:    pand %xmm2, %xmm0
 ; SSE41-NEXT:    retq
 ;
-; AVX2-LABEL: combine_vec_shl_commuted_clamped1:
-; AVX2:       # %bb.0:
-; AVX2-NEXT:    vpbroadcastd {{.*#+}} xmm2 = [31,31,31,31]
-; AVX2-NEXT:    vpsllvd %xmm1, %xmm0, %xmm0
-; AVX2-NEXT:    vpminud %xmm2, %xmm1, %xmm2
-; AVX2-NEXT:    vpcmpeqd %xmm2, %xmm1, %xmm1
-; AVX2-NEXT:    vpand %xmm0, %xmm1, %xmm0
-; AVX2-NEXT:    retq
-;
-; AVX512-LABEL: combine_vec_shl_commuted_clamped1:
-; AVX512:       # %bb.0:
-; AVX512-NEXT:    vpsllvd %xmm1, %xmm0, %xmm0
-; AVX512-NEXT:    retq
+; AVX-LABEL: combine_vec_shl_commuted_clamped1:
+; AVX:       # %bb.0:
+; AVX-NEXT:    vpsllvd %xmm1, %xmm0, %xmm0
+; AVX-NEXT:    retq
   %cmp.i = icmp uge <4 x i32> %amt, <i32 32, i32 32, i32 32, i32 32>
   %shl = shl <4 x i32> %sh, %amt
   %1 = select <4 x i1> %cmp.i, <4 x i32> zeroinitializer, <4 x i32> %shl

llvm/test/CodeGen/X86/combine-srl.ll

@@ -828,19 +828,10 @@ define <4 x i32> @combine_vec_lshr_commuted_clamped1(<4 x i32> %sh, <4 x i32> %a
 ; SSE41-NEXT:    pand %xmm2, %xmm0
 ; SSE41-NEXT:    retq
 ;
-; AVX2-LABEL: combine_vec_lshr_commuted_clamped1:
-; AVX2:       # %bb.0:
-; AVX2-NEXT:    vpbroadcastd {{.*#+}} xmm2 = [31,31,31,31]
-; AVX2-NEXT:    vpsrlvd %xmm1, %xmm0, %xmm0
-; AVX2-NEXT:    vpminud %xmm2, %xmm1, %xmm2
-; AVX2-NEXT:    vpcmpeqd %xmm2, %xmm1, %xmm1
-; AVX2-NEXT:    vpand %xmm0, %xmm1, %xmm0
-; AVX2-NEXT:    retq
-;
-; AVX512-LABEL: combine_vec_lshr_commuted_clamped1:
-; AVX512:       # %bb.0:
-; AVX512-NEXT:    vpsrlvd %xmm1, %xmm0, %xmm0
-; AVX512-NEXT:    retq
+; AVX-LABEL: combine_vec_lshr_commuted_clamped1:
+; AVX:       # %bb.0:
+; AVX-NEXT:    vpsrlvd %xmm1, %xmm0, %xmm0
+; AVX-NEXT:    retq
   %cmp.i = icmp uge <4 x i32> %amt, <i32 32, i32 32, i32 32, i32 32>
   %shr = lshr <4 x i32> %sh, %amt
   %1 = select <4 x i1> %cmp.i, <4 x i32> zeroinitializer, <4 x i32> %shr