3 files changed, 47 insertions, 13 deletions
diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 7198203036b..5d2b02b4bff 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -655,10 +655,14 @@ static ConstantSDNode *isConstOrConstSplat(SDValue N) {
     return CN;
 
   if (BuildVectorSDNode *BV = dyn_cast<BuildVectorSDNode>(N)) {
-    ConstantSDNode *CN = BV->getConstantSplatValue();
+    bool HasUndefElements;
+    ConstantSDNode *CN = BV->getConstantSplatNode(HasUndefElements);
 
     // BuildVectors can truncate their operands. Ignore that case here.
-    if (CN && CN->getValueType(0) == N.getValueType().getScalarType())
+    // FIXME: We blindly ignore splats which include undef which is overly
+    // pessimistic.
+    if (CN && !HasUndefElements &&
+        CN->getValueType(0) == N.getValueType().getScalarType())
       return CN;
   }
 
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 3a8a5f9601f..e41d459f39d 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -1506,8 +1506,10 @@ SDValue SelectionDAG::getVectorShuffle(EVT VT, SDLoc dl, SDValue N1,
     return N1;
 
   // Shuffling a constant splat doesn't change the result.
+  bool SplatHasUndefs;
   if (N2Undef && N1.getOpcode() == ISD::BUILD_VECTOR)
-    if (cast<BuildVectorSDNode>(N1)->getConstantSplatValue())
+    if (cast<BuildVectorSDNode>(N1)->getConstantSplatNode(SplatHasUndefs) &&
+        !SplatHasUndefs)
       return N1;
 
   FoldingSetNodeID ID;
@@ -6603,16 +6605,38 @@ bool BuildVectorSDNode::isConstantSplat(APInt &SplatValue,
   return true;
 }
 
-ConstantSDNode *BuildVectorSDNode::getConstantSplatValue() const {
-  SDValue Op0 = getOperand(0);
-  if (Op0.getOpcode() != ISD::Constant)
-    return nullptr;
+SDValue BuildVectorSDNode::getSplatValue(bool &HasUndefElements) const {
+  HasUndefElements = false;
+  SDValue Splatted;
+  for (unsigned i = 0, e = getNumOperands(); i != e; ++i) {
+    SDValue Op = getOperand(i);
+    if (Op.getOpcode() == ISD::UNDEF)
+      HasUndefElements = true;
+    else if (!Splatted)
+      Splatted = Op;
+    else if (Splatted != Op)
+      return SDValue();
+  }
 
-  for (unsigned i = 1, e = getNumOperands(); i != e; ++i)
-    if (getOperand(i) != Op0)
-      return nullptr;
+  if (!Splatted) {
+    assert(getOperand(0).getOpcode() == ISD::UNDEF &&
+           "Can only have a splat without a constant for all undefs.");
+    return getOperand(0);
+  }
+
+  return Splatted;
+}
+
+ConstantSDNode *
+BuildVectorSDNode::getConstantSplatNode(bool &HasUndefElements) const {
+  return dyn_cast_or_null<ConstantSDNode>(
+      getSplatValue(HasUndefElements).getNode());
+}
 
-  return cast<ConstantSDNode>(Op0);
+ConstantFPSDNode *
+BuildVectorSDNode::getConstantFPSplatNode(bool &HasUndefElements) const {
+  return dyn_cast_or_null<ConstantFPSDNode>(
+      getSplatValue(HasUndefElements).getNode());
 }
 
 bool BuildVectorSDNode::isConstant() const {
diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index ad91d4a87c1..74de034777d 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -1158,7 +1158,10 @@ bool TargetLowering::isConstTrueVal(const SDNode *N) const {
       return false;
 
     IsVec = true;
-    CN = BV->getConstantSplatValue();
+    bool HasUndefElements;
+    CN = BV->getConstantSplatNode(HasUndefElements);
+    if (HasUndefElements)
+      return false; // Can't blindly collapse the undef values.
   }
 
   switch (getBooleanContents(IsVec)) {
@@ -1185,7 +1188,10 @@ bool TargetLowering::isConstFalseVal(const SDNode *N) const {
       return false;
 
     IsVec = true;
-    CN = BV->getConstantSplatValue();
+    bool HasUndefElements;
+    CN = BV->getConstantSplatNode(HasUndefElements);
+    if (HasUndefElements)
+      return false; // Can't blindly collapse the undef values.
   }
 
   if (getBooleanContents(IsVec) == UndefinedBooleanContent)