[X86] Prevent several calls to ISD::isConstantSplatVector from returning a narrower APInt than the original scalar type

[craig.topper]

Summary

ISD::isConstantSplatVector can shrink to the smallest splat width. But we don't check the size of the resulting APInt at all. This can cause us to misinterpret the results.

This patch just adds a flag to prevent the APInt from changing width.

Fixes PR34271.

[chandlerc]

LGTM, minor nits below, feel free to submit whenever.

include/llvm/CodeGen/SelectionDAGNodes.h

88

'returns' -> 'sets SplatValue to'

test/CodeGen/X86/pr34271.ll

7

The previous version of the generated code had a plausible instruction sequence that ultimately did the wrong thing... Is it possible to actually check the values being added here?

include/llvm/CodeGen/SelectionDAGNodes.h

 };
 
 namespace ISD {
 
   /// Node predicates
 
   /// If N is a BUILD_VECTOR node whose elements are all the same constant or
   /// undefined, return true and return the constant value in \p SplatValue.
-  bool isConstantSplatVector(const SDNode *N, APInt &SplatValue);
+  /// This returns the smallest possible splat unless AllowShrink is set to
+  /// false.
+  bool isConstantSplatVector(const SDNode *N, APInt &SplatValue,
+                             bool AllowShrink = true);
 
   /// Return true if the specified node is a BUILD_VECTOR where all of the
   /// elements are ~0 or undef.
   bool isBuildVectorAllOnes(const SDNode *N);
 
   /// Return true if the specified node is a BUILD_VECTOR where all of the
   /// elements are 0 or undef.
   bool isBuildVectorAllZeros(const SDNode *N);

lib/CodeGen/SelectionDAG/SelectionDAG.cpp

(void) Val2.convert(SelectionDAG::EVTToAPFloatSemantics(VT),
                       &losesInfo);
   return !losesInfo;
 }
 
 //===----------------------------------------------------------------------===//
 //                              ISD Namespace
 //===----------------------------------------------------------------------===//
 
-bool ISD::isConstantSplatVector(const SDNode *N, APInt &SplatVal) {
+bool ISD::isConstantSplatVector(const SDNode *N, APInt &SplatVal,
+                                bool AllowShrink) {
   auto *BV = dyn_cast<BuildVectorSDNode>(N);
   if (!BV)
     return false;
 
   APInt SplatUndef;
   unsigned SplatBitSize;
   bool HasUndefs;
-  EVT EltVT = N->getValueType(0).getVectorElementType();
-  return BV->isConstantSplat(SplatVal, SplatUndef, SplatBitSize, HasUndefs) &&
-         EltVT.getSizeInBits() >= SplatBitSize;
+  unsigned EltSize = N->getValueType(0).getVectorElementType().getSizeInBits();
+  unsigned MinSplatBits = AllowShrink ? 0 : EltSize;
+  return BV->isConstantSplat(SplatVal, SplatUndef, SplatBitSize, HasUndefs,
+                             MinSplatBits) &&
+         EltSize >= SplatBitSize;
 }
 
 // FIXME: AllOnes and AllZeros duplicate a lot of code. Could these be
 // specializations of the more general isConstantSplatVector()?
 
 bool ISD::isBuildVectorAllOnes(const SDNode *N) {
   // Look through a bit convert.
   while (N->getOpcode() == ISD::BITCAST)

lib/Target/X86/X86ISelLowering.cpp

static bool detectZextAbsDiff(const SDValue &Select, SDValue &Op0,
   // The first operand of SetCC is the first operand of the select, which is the
   // difference between the two input vectors.
   if (SetCC.getOperand(0) != SelectOp1)
     return false;
 
   // In SetLT case, The second operand of the comparison can be either 1 or 0.
   APInt SplatVal;
   if ((CC == ISD::SETLT) &&
-      !((ISD::isConstantSplatVector(SetCC.getOperand(1).getNode(), SplatVal) &&
-         SplatVal == 1) ||
+      !((ISD::isConstantSplatVector(SetCC.getOperand(1).getNode(), SplatVal,
+                                    /*AllowShrink*/false) &&
+         SplatVal.isOneValue()) ||
         (ISD::isBuildVectorAllZeros(SetCC.getOperand(1).getNode()))))
     return false;
 
   // In SetGT case, The second operand of the comparison can be either -1 or 0.
   if ((CC == ISD::SETGT) &&
       !(ISD::isBuildVectorAllZeros(SetCC.getOperand(1).getNode()) ||
         ISD::isBuildVectorAllOnes(SetCC.getOperand(1).getNode())))
     return false;
static SDValue combineAndMaskToShift(SDNode *N, SelectionDAG &DAG,
   SDValue Op1 = peekThroughBitcasts(N->getOperand(1));
   EVT VT0 = Op0.getValueType();
   EVT VT1 = Op1.getValueType();
 
   if (VT0 != VT1 || !VT0.isSimple() || !VT0.isInteger())
     return SDValue();
 
   APInt SplatVal;
-  if (!ISD::isConstantSplatVector(Op1.getNode(), SplatVal) ||
+  if (!ISD::isConstantSplatVector(Op1.getNode(), SplatVal,
+                                  /*AllowShrink*/false) ||
       !SplatVal.isMask())
     return SDValue();
 
   if (!SupportedVectorShiftWithImm(VT0.getSimpleVT(), Subtarget, ISD::SRL))
     return SDValue();
 
   unsigned EltBitWidth = VT0.getScalarSizeInBits();
   if (EltBitWidth != DAG.ComputeNumSignBits(Op0))
static SDValue detectUSatPattern(SDValue In, EVT VT) {
   if (In.getOpcode() != ISD::UMIN)
     return SDValue();
 
   //Saturation with truncation. We truncate from InVT to VT.
   assert(In.getScalarValueSizeInBits() > VT.getScalarSizeInBits() &&
     "Unexpected types for truncate operation");
 
   APInt C;
-  if (ISD::isConstantSplatVector(In.getOperand(1).getNode(), C)) {
+  if (ISD::isConstantSplatVector(In.getOperand(1).getNode(), C,
+                                 /*AllowShrink*/false)) {
     // C should be equal to UINT32_MAX / UINT16_MAX / UINT8_MAX according
     // the element size of the destination type.
     return C.isMask(VT.getScalarSizeInBits()) ? In.getOperand(0) :
       SDValue();
   }
   return SDValue();
 }
 
static SDValue combineIncDecVector(SDNode *N, SelectionDAG &DAG) {
   // Pseudo-legality check: getOnesVector() expects one of these types, so bail
   // out and wait for legalization if we have an unsupported vector length.
   EVT VT = N->getValueType(0);
   if (!VT.is128BitVector() && !VT.is256BitVector() && !VT.is512BitVector())
     return SDValue();
 
   SDNode *N1 = N->getOperand(1).getNode();
   APInt SplatVal;
-  if (!ISD::isConstantSplatVector(N1, SplatVal) || !SplatVal.isOneValue())
+  if (!ISD::isConstantSplatVector(N1, SplatVal, /*AllowShrink*/false) ||
+      !SplatVal.isOneValue())
     return SDValue();
 
   SDValue AllOnesVec = getOnesVector(VT, DAG, SDLoc(N));
   unsigned NewOpcode = N->getOpcode() == ISD::ADD ? ISD::SUB : ISD::ADD;
   return DAG.getNode(NewOpcode, SDLoc(N), VT, N->getOperand(0), AllOnesVec);
 }
 
 static SDValue combineAdd(SDNode *N, SelectionDAG &DAG,

test/CodeGen/X86/pr34271-1.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu -mattr=avx512vl,avx512bw | FileCheck %s
+
+define <16 x i16> @foo(<16 x i32> %i) {
+; CHECK-LABEL: foo:
+; CHECK:       # BB#0:
+; CHECK-NEXT:    vpminud {{.*}}(%rip){1to16}, %zmm0, %zmm0
+; CHECK-NEXT:    vpmovdw %zmm0, %ymm0
+; CHECK-NEXT:    retq
+  %x3 = icmp ult <16 x i32> %i, <i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009>
+  %x5 = select <16 x i1> %x3, <16 x i32> %i, <16 x i32> <i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009, i32 16843009>
+  %x6 = trunc <16 x i32> %x5 to <16 x i16>
+  ret <16 x i16> %x6
+}

test/CodeGen/X86/pr34271.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu | FileCheck %s
+
+define <4 x i32> @f(<4 x i32> %a) {
+; CHECK-LABEL: f:
+; CHECK:       # BB#0:
+; CHECK-NEXT:    paddd {{.*}}(%rip), %xmm0
+; CHECK-NEXT:    retq
+  %v = add nuw nsw <4 x i32> %a, <i32 16843009, i32 16843009, i32 16843009, i32 16843009>
+  ret <4 x i32> %v
+}