[X86][AVX512] Fix PR31515 - Do not flip vselect condition if it's not a vXi1 mask

[eladcohen]

Summary

r289653 added a case where

vselect <cond> <vector1> <all-zeros>

is transformed to:

vselect xor(cond, DAG.getConstant(1, DL, Cond.getValueType())) <all-zeros> <vector1>

This was not aimed to catch cases where Cond is not a vXi1 mask but it does (e.g. happens on KNL for select <16 x i1>). Moreover, when Cond type is VxiN (N > 1) then xor(cond, DAG.getConstant(1, DL, Cond.getValueType())) != NOT(cond).
This patch changes the above to xor with allones, and avoids entering the case for non-mask Conds.

[craig.topper]

lib/Target/X86/X86ISelLowering.cpp

28754

Why do we need getScalarSizeInBits() if we already know its i1?

[eladcohen]

lib/Target/X86/X86ISelLowering.cpp

28754

I thought it will be better to keep the code correct for any vector type just in case we decide one day to remove/change the

CondVT.getVectorElementType() == MVT::i1

condition.

[craig.topper]

LGTM

[eladcohen]

Thanks for the review!
Committed in r291745.

lib/Target/X86/X86ISelLowering.cpp

static SDValue combineVSelectWithAllOnesOrZeros(SDNode *N, SelectionDAG &DAG,
   EVT VT = LHS.getValueType();
   EVT CondVT = Cond.getValueType();
   SDLoc DL(N);
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
 
   if (N->getOpcode() != ISD::VSELECT)
     return SDValue();
 
+  assert(CondVT.isVector() && "Vector select expects a vector selector!");
+
   bool FValIsAllZeros = ISD::isBuildVectorAllZeros(LHS.getNode());
-  // Check if the first operand is all zeros.This situation only
-  // applies to avx512.
-  if (FValIsAllZeros  && Subtarget.hasAVX512() && Cond.hasOneUse()) {
+  // Check if the first operand is all zeros and Cond type is vXi1.
+  // This situation only applies to avx512.
+  if (FValIsAllZeros  && Subtarget.hasAVX512() && Cond.hasOneUse() &&
+      CondVT.getVectorElementType() == MVT::i1) {
       //Invert the cond to not(cond) : xor(op,allones)=not(op)
       SDValue CondNew = DAG.getNode(ISD::XOR, DL, Cond.getValueType(), Cond,
-        DAG.getConstant(1, DL, Cond.getValueType()));
+        DAG.getConstant(APInt::getAllOnesValue(CondVT.getScalarSizeInBits()),
+                        DL, CondVT));
       //Vselect cond, op1, op2 = Vselect not(cond), op2, op1
       return DAG.getNode(ISD::VSELECT, DL, VT, CondNew, RHS, LHS);
   }
-  assert(CondVT.isVector() && "Vector select expects a vector selector!");
 
   // To use the condition operand as a bitwise mask, it must have elements that
   // are the same size as the select elements. Ie, the condition operand must
   // have already been promoted from the IR select condition type <N x i1>.
   // Don't check if the types themselves are equal because that excludes
   // vector floating-point selects.
   if (CondVT.getScalarSizeInBits() != VT.getScalarSizeInBits())
     return SDValue();

test/CodeGen/X86/avx512-select.ll

 ; CHECK:       ## BB#0:
 ; CHECK-NEXT:    andl $1, %edi
 ; CHECK-NEXT:    kmovw %edi, %k1
 ; CHECK-NEXT:    vmovss %xmm1, %xmm0, %xmm0 {%k1}
 ; CHECK-NEXT:    retq
   %cond = select i1 %c, float %a, float %b
   ret float %cond
 }
+
+define <16 x i16> @pr31515(<16 x i1> %a, <16 x i1> %b, <16 x i16> %c) nounwind {
+; CHECK-LABEL: pr31515:
+; CHECK:       ## BB#0:
+; CHECK-NEXT:    vpmovsxbd %xmm1, %zmm1
+; CHECK-NEXT:    vpslld $31, %zmm1, %zmm1
+; CHECK-NEXT:    vpmovsxbd %xmm0, %zmm0
+; CHECK-NEXT:    vpslld $31, %zmm0, %zmm0
+; CHECK-NEXT:    vptestmd %zmm0, %zmm0, %k1
+; CHECK-NEXT:    vptestmd %zmm1, %zmm1, %k1 {%k1}
+; CHECK-NEXT:    vpternlogd $255, %zmm0, %zmm0, %zmm0 {%k1} {z}
+; CHECK-NEXT:    vpmovdw %zmm0, %ymm0
+; CHECK-NEXT:    vpandn %ymm2, %ymm0, %ymm0
+; CHECK-NEXT:    retq
+  %mask = and <16 x i1> %a, %b
+  %res = select <16 x i1> %mask, <16 x i16> zeroinitializer, <16 x i16> %c
+  ret <16 x i16> %res
+}
+