[X86] Enforce strict pre-legalization to combine in scalarizeExtEltFP

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -45842,7 +45842,8 @@ static SDValue combineExtractWithShuffle(SDNode *N, SelectionDAG &DAG,
 /// Extracting a scalar FP value from vector element 0 is free, so extract each
 /// operand first, then perform the math as a scalar op.
 static SDValue scalarizeExtEltFP(SDNode *ExtElt, SelectionDAG &DAG,
-                                 const X86Subtarget &Subtarget) {
+                                 const X86Subtarget &Subtarget,
+                                 TargetLowering::DAGCombinerInfo &DCI) {
   assert(ExtElt->getOpcode() == ISD::EXTRACT_VECTOR_ELT && "Expected extract");
   SDValue Vec = ExtElt->getOperand(0);
   SDValue Index = ExtElt->getOperand(1);
@@ -45877,13 +45878,14 @@ static SDValue scalarizeExtEltFP(SDNode *ExtElt, SelectionDAG &DAG,
   // Vector FP selects don't fit the pattern of FP math ops (because the
   // condition has a different type and we have to change the opcode), so deal
   // with those here.
-  // FIXME: This is restricted to pre type legalization by ensuring the setcc
-  // has i1 elements. If we loosen this we need to convert vector bool to a
-  // scalar bool.
-  if (Vec.getOpcode() == ISD::VSELECT &&
+  // FIXME: This is restricted to pre type legalization. If we loosen this we
+  // need to convert vector bool to a scalar bool.
+  if (DCI.getDAGCombineLevel() < llvm::AfterLegalizeTypes &&
+      Vec.getOpcode() == ISD::VSELECT &&
       Vec.getOperand(0).getOpcode() == ISD::SETCC &&
-      Vec.getOperand(0).getValueType().getScalarType() == MVT::i1 &&
       Vec.getOperand(0).getOperand(0).getValueType() == VecVT) {
+    assert(Vec.getOperand(0).getValueType().getScalarType() == MVT::i1 &&
+           "Unexpected cond type for combine");
     // ext (sel Cond, X, Y), 0 --> sel (ext Cond, 0), (ext X, 0), (ext Y, 0)
     SDLoc DL(ExtElt);
     SDValue Ext0 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL,
@@ -46242,7 +46244,7 @@ static SDValue combineExtractVectorElt(SDNode *N, SelectionDAG &DAG,
   if (SDValue V = combineArithReduction(N, DAG, Subtarget))
     return V;
 
-  if (SDValue V = scalarizeExtEltFP(N, DAG, Subtarget))
+  if (SDValue V = scalarizeExtEltFP(N, DAG, Subtarget, DCI))
     return V;
 
   if (CIdx)

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

@@ -0,0 +1,38 @@
+; RUN: llc < %s | FileCheck %s
+
+define void @extvselectsetcc_illegal(i1 %cond, <8 x float> %vec, ptr %ptr1, ptr %ptr2) #0 {
+; CHECK-LABEL: extvselectsetcc_illegal:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vxorps	%xmm1, %xmm1, %xmm1
+; CHECK-NEXT:	 vcmpnltps	%ymm1, %ymm0, %k1
+; CHECK-NEXT:	 vbroadcastss	.LCPI0_0(%rip), %xmm0   # xmm0 = [NaN,NaN,NaN,NaN]
+; CHECK-NEXT:	 vinsertf32x4	$0, %xmm0, %ymm0, %ymm0 {%k1} {z}
+; CHECK-NEXT:	 vxorps	%xmm1, %xmm1, %xmm1
+; CHECK-NEXT:	 vmulss	%xmm1, %xmm0, %xmm0
+; CHECK-NEXT:	 vmulss	.LCPI0_0(%rip), %xmm0, %xmm2
+; CHECK-NEXT:	 vbroadcastss	%xmm2, %ymm2
+; CHECK-NEXT:	 testb	$1, %dil
+; CHECK-NEXT:	 cmoveq	%rdx, %rsi
+; CHECK-NEXT:	 vmovups	%ymm2, (%rsi)
+; CHECK-NEXT:	 vmulss	%xmm1, %xmm0, %xmm0
+; CHECK-NEXT:	 vbroadcastss	%xmm0, %ymm0
+; CHECK-NEXT:	 vmovups	%ymm0, (%rdx)
+; CHECK-NEXT:	 vzeroupper
+; CHECK-NEXT:	 retq
+  %cmp = fcmp olt <8 x float> %vec, zeroinitializer
+  %sel1 = select <8 x i1> %cmp, <8 x float> zeroinitializer, <8 x float> 
+  <float 0x7FF8000000000000, float 0x7FF8000000000000, float 0x7FF8000000000000, float 0x7FF8000000000000, 
+  float 0x7FF8000000000000, float 0x7FF8000000000000, float 0x7FF8000000000000, float 0x7FF8000000000000>
+  %fmul1 = fmul <8 x float> zeroinitializer, %sel1
+  %shuffle = shufflevector <8 x float> %fmul1, <8 x float> zeroinitializer, <8 x i32> zeroinitializer
+  %fmul2 = fmul <8 x float> %shuffle, 
+  <float 0x7FF8000000000000, float 0x7FF8000000000000, float 0x7FF8000000000000, float 0x7FF8000000000000, 
+  float 0x7FF8000000000000, float 0x7FF8000000000000, float 0x7FF8000000000000, float 0x7FF8000000000000>  
+  %sel2 = select i1 %cond, ptr %ptr1, ptr %ptr2
+  store <8 x float> %fmul2, ptr %sel2, align 4
+  %fmul3 = fmul <8 x float> %shuffle, zeroinitializer
+  store <8 x float> %fmul3, ptr %ptr2, align 4
+  ret void
+}
+
+attributes #0 = { "target-cpu"="skylake-avx512" }