3 files changed, 44 insertions, 31 deletions
diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 7198203036b..9a91dcc341b 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -654,13 +654,12 @@ static ConstantSDNode *isConstOrConstSplat(SDValue N) {
   if (ConstantSDNode *CN = dyn_cast<ConstantSDNode>(N))
     return CN;
 
-  if (BuildVectorSDNode *BV = dyn_cast<BuildVectorSDNode>(N)) {
-    ConstantSDNode *CN = BV->getConstantSplatValue();
-
-    // BuildVectors can truncate their operands. Ignore that case here.
-    if (CN && CN->getValueType(0) == N.getValueType().getScalarType())
-      return CN;
-  }
+  if (BuildVectorSDNode *BV = dyn_cast<BuildVectorSDNode>(N))
+    if (SDValue Splat = BV->getConstantSplatValue())
+      if (auto *CN = dyn_cast<ConstantSDNode>(Splat))
+        // BuildVectors can truncate their operands. Ignore that case here.
+        if (CN->getValueType(0) == N.getValueType().getScalarType())
+          return CN;
 
   return nullptr;
 }
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 3a8a5f9601f..fb7d1b18d8b 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -6603,16 +6603,28 @@ bool BuildVectorSDNode::isConstantSplat(APInt &SplatValue,
   return true;
 }
 
-ConstantSDNode *BuildVectorSDNode::getConstantSplatValue() const {
-  SDValue Op0 = getOperand(0);
-  if (Op0.getOpcode() != ISD::Constant)
-    return nullptr;
+SDValue BuildVectorSDNode::getConstantSplatValue() const {
+  SDValue Splatted;
+  for (unsigned i = 0, e = getNumOperands(); i != e; ++i) {
+    SDValue Op = getOperand(i);
+    if (Op.getOpcode() == ISD::UNDEF)
+      continue;
+    if (Op.getOpcode() != ISD::Constant && Op.getOpcode() != ISD::ConstantFP)
+      return SDValue();
 
-  for (unsigned i = 1, e = getNumOperands(); i != e; ++i)
-    if (getOperand(i) != Op0)
-      return nullptr;
+    if (!Splatted)
+      Splatted = Op;
+    else if (Splatted != Op)
+      return SDValue();
+  }
+
+  if (!Splatted) {
+    assert(getOperand(0).getOpcode() == ISD::UNDEF &&
+           "Can only have a splat without a constant for all undefs.");
+    return getOperand(0);
+  }
 
-  return cast<ConstantSDNode>(Op0);
+  return Splatted;
 }
 
 bool BuildVectorSDNode::isConstant() const {
diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index ad91d4a87c1..1b3e42848bf 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -1152,14 +1152,15 @@ bool TargetLowering::isConstTrueVal(const SDNode *N) const {
 
   bool IsVec = false;
   const ConstantSDNode *CN = dyn_cast<ConstantSDNode>(N);
-  if (!CN) {
-    const BuildVectorSDNode *BV = dyn_cast<BuildVectorSDNode>(N);
-    if (!BV)
-      return false;
-
-    IsVec = true;
-    CN = BV->getConstantSplatValue();
-  }
+  if (!CN)
+    if (auto *BV = dyn_cast<BuildVectorSDNode>(N))
+      if (SDValue Splat = BV->getConstantSplatValue())
+        if (auto *SplatCN = dyn_cast<ConstantSDNode>(Splat)) {
+          IsVec = true;
+          CN = SplatCN;
+        }
+  if (!CN)
+    return false;
 
   switch (getBooleanContents(IsVec)) {
   case UndefinedBooleanContent:
@@ -1179,14 +1180,15 @@ bool TargetLowering::isConstFalseVal(const SDNode *N) const {
 
   bool IsVec = false;
   const ConstantSDNode *CN = dyn_cast<ConstantSDNode>(N);
-  if (!CN) {
-    const BuildVectorSDNode *BV = dyn_cast<BuildVectorSDNode>(N);
-    if (!BV)
-      return false;
-
-    IsVec = true;
-    CN = BV->getConstantSplatValue();
-  }
+  if (!CN)
+    if (auto *BV = dyn_cast<BuildVectorSDNode>(N))
+      if (SDValue Splat = BV->getConstantSplatValue())
+        if (auto *SplatCN = dyn_cast<ConstantSDNode>(Splat)) {
+          IsVec = true;
+          CN = SplatCN;
+        }
+  if (!CN)
+    return false;
 
   if (getBooleanContents(IsVec) == UndefinedBooleanContent)
     return !CN->getAPIntValue()[0];