1 files changed, 18 insertions, 13 deletions

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 15baf14f2d9..c2cb0d8d779 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -1521,15 +1521,15 @@ SDValue SelectionDAG::getVectorShuffle(EVT VT, SDLoc dl, SDValue N1,
 
     // A splat should always show up as a build vector node.
     if (auto *BV = dyn_cast<BuildVectorSDNode>(V)) {
-      bool SplatHasUndefs;
-      SDValue Splat = BV->getSplatValue(SplatHasUndefs);
+      BitVector UndefElements;
+      SDValue Splat = BV->getSplatValue(&UndefElements);
       // If this is a splat of an undef, shuffling it is also undef.
       if (Splat && Splat.getOpcode() == ISD::UNDEF)
         return getUNDEF(VT);
 
       // We only have a splat which can skip shuffles if there is a splatted
       // value and no undef lanes rearranged by the shuffle.
-      if (Splat && !SplatHasUndefs) {
+      if (Splat && UndefElements.none()) {
         // Splat of <x, x, ..., x>, return <x, x, ..., x>, provided that the
         // number of elements match or the value splatted is a zero constant.
         if (V.getValueType().getVectorNumElements() ==
@@ -6635,17 +6635,22 @@ bool BuildVectorSDNode::isConstantSplat(APInt &SplatValue,
   return true;
 }
 
-SDValue BuildVectorSDNode::getSplatValue(bool &HasUndefElements) const {
-  HasUndefElements = false;
+SDValue BuildVectorSDNode::getSplatValue(BitVector *UndefElements) const {
+  if (UndefElements) {
+    UndefElements->clear();
+    UndefElements->resize(getNumOperands());
+  }
   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)
+    if (Op.getOpcode() == ISD::UNDEF) {
+      if (UndefElements)
+        (*UndefElements)[i] = true;
+    } else if (!Splatted) {
       Splatted = Op;
-    else if (Splatted != Op)
+    } else if (Splatted != Op) {
       return SDValue();
+    }
   }
 
   if (!Splatted) {
@@ -6658,15 +6663,15 @@ SDValue BuildVectorSDNode::getSplatValue(bool &HasUndefElements) const {
 }
 
 ConstantSDNode *
-BuildVectorSDNode::getConstantSplatNode(bool &HasUndefElements) const {
+BuildVectorSDNode::getConstantSplatNode(BitVector *UndefElements) const {
   return dyn_cast_or_null<ConstantSDNode>(
-      getSplatValue(HasUndefElements).getNode());
+      getSplatValue(UndefElements).getNode());
 }
 
 ConstantFPSDNode *
-BuildVectorSDNode::getConstantFPSplatNode(bool &HasUndefElements) const {
+BuildVectorSDNode::getConstantFPSplatNode(BitVector *UndefElements) const {
   return dyn_cast_or_null<ConstantFPSDNode>(
-      getSplatValue(HasUndefElements).getNode());
+      getSplatValue(UndefElements).getNode());
 }
 
 bool BuildVectorSDNode::isConstant() const {