[X86] Adding a check against i64 inputs in combineScalarToVector

[mike.dvoretsky]

Summary

This patch adds checks to prevent creating scalar_to_vector DAG nodes with i64 inputs when folding scalar masking, fixing PR37879.

[craig.topper]

llvm/lib/Target/X86/X86ISelLowering.cpp

39127 ↗

(On Diff #152081)

You should probably just check for i8 and i32. I'm not sure anything prevents illegal types from getting here. Which might get legalized into i64 or i16 and be just as broken.

39142 ↗

(On Diff #152081)

Same here

[mike.dvoretsky]

Changed the type checks to ensure that the input has a valid type, rather than guarding against specific invalid types.

[craig.topper]

LGTM

llvm/trunk/lib/Target/X86/X86ISelLowering.cpp

static SDValue combineScalarToVector(SDNode *N, SelectionDAG &DAG) {
   // This occurs frequently in our masked scalar intrinsic code and our
   // floating point select lowering with AVX512.
   // TODO: SimplifyDemandedBits instead?
   if (VT == MVT::v1i1 && Src.getOpcode() == ISD::AND && Src.hasOneUse())
     if (auto *C = dyn_cast<ConstantSDNode>(Src.getOperand(1)))
       if (C->getAPIntValue().isOneValue()) {
         SDValue Mask = Src.getOperand(0);
         if (Mask.getOpcode() == ISD::TRUNCATE &&
-            Mask.getOperand(0).getValueType() != MVT::i16)
+            (Mask.getOperand(0).getValueType() == MVT::i8 ||
+             Mask.getOperand(0).getValueType() == MVT::i32))
           Mask = Mask.getOperand(0);
         return DAG.getNode(ISD::SCALAR_TO_VECTOR, SDLoc(N), MVT::v1i1, Mask);
       }
 
   // The result of AND may also be truncated. This occurs in code for lowered
   // masked scalar intrinsics.
   if (VT == MVT::v1i1 && Src.getOpcode() == ISD::TRUNCATE && Src.hasOneUse() &&
       Src.getOperand(0).getOpcode() == ISD::AND &&
       Src.getOperand(0).hasOneUse())
     if (auto *C = dyn_cast<ConstantSDNode>(Src.getOperand(0).getOperand(1)))
       if (C->getAPIntValue().isOneValue()) {
         SDValue Mask = Src.getOperand(0).getOperand(0);
         if (Mask.getOpcode() == ISD::TRUNCATE &&
-            Mask.getOperand(0).getValueType() != MVT::i16)
+            (Mask.getOperand(0).getValueType() == MVT::i8 ||
+             Mask.getOperand(0).getValueType() == MVT::i32))
           Mask = Mask.getOperand(0);
-        // Check if the initial value is an i16. scalar_to_vector fails to
-        // select for that type, so the combine should be aborted.
-        if (Mask.getValueType() == MVT::i16)
+        // Check if the initial value is of a legal type for scalar_to_vector.
+        if (Mask.getValueType() != MVT::i8 && Mask.getValueType() != MVT::i32)
           return SDValue();
         return DAG.getNode(ISD::SCALAR_TO_VECTOR, SDLoc(N), MVT::v1i1, Mask);
       }
 
   return SDValue();
 }
 
 // Simplify PMULDQ and PMULUDQ operations.

llvm/trunk/test/CodeGen/X86/pr37879.ll

-; XFAIL: *
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -O3 < %s -mtriple=x86_64-apple-darwin -mattr=+avx512bw | FileCheck %s
 
 define double @foo(i32** nocapture readonly) #0 {
+; CHECK-LABEL: foo:
+; CHECK:       ## %bb.0:
+; CHECK-NEXT:    movq (%rax), %rax
+; CHECK-NEXT:    vcvtsi2sdq %rax, %xmm0, %xmm1
+; CHECK-NEXT:    vmovsd {{.*#+}} xmm0 = mem[0],zero
+; CHECK-NEXT:    kmovd %eax, %k1
+; CHECK-NEXT:    vmovsd %xmm1, %xmm0, %xmm0 {%k1}
+; CHECK-NEXT:    retq
   %2 = load i64, i64* undef, align 8
   %3 = and i64 %2, 1
   %4 = icmp eq i64 %3, 0
   %5 = sitofp i64 %2 to double
   %6 = select i1 %4, double 1.000000e+00, double %5
   ret double %6
 }