[DAGCombiner] Simplify shuffle mask index if the referenced input element is UNDEF

[RKSimon]

Summary

Fixes one of the cases in PR34041.

[zvi]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15130

Do the affected tests cover both the 'if' and the 'else'?

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15130

Not properly, I'll add some extra tests

[RKSimon]

Updated with second test that simplifies the mask based on both shuffle inputs (note: further improvements are possible on this shuffle in future patches).

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

static SDValue simplifyShuffleOperands(ShuffleVectorSDNode *SVN, SDValue N0,
   SDValue S0 = simplifyShuffleOperandRecursively(N0UsedElements, N0, DAG);
   SDValue S1 = simplifyShuffleOperandRecursively(N1UsedElements, N1, DAG);
   if (S0 == N0 && S1 == N1)
     return SDValue();
 
   return DAG.getVectorShuffle(VT, SDLoc(SVN), S0, S1, SVN->getMask());
 }
 
+static SDValue simplifyShuffleMask(ShuffleVectorSDNode *SVN, SDValue N0,
+                                   SDValue N1, SelectionDAG &DAG) {
+  // TODO - handle cases other than BUILD_VECTOR.
+  auto *BV0 = dyn_cast<BuildVectorSDNode>(N0);
+  auto *BV1 = dyn_cast<BuildVectorSDNode>(N1);
+  if (!BV0 && !BV1)
+    return SDValue();
+
+  EVT VT = SVN->getValueType(0);
+  unsigned NumElts = VT.getVectorNumElements();
+
+  bool Changed = false;
+  SmallVector<int, 8> NewMask;
+  for (unsigned i = 0; i != NumElts; ++i) {
+    int Idx = SVN->getMaskElt(i);
+    if (BV0 && 0 <= Idx && Idx < (int)NumElts &&
+        BV0->getOperand(Idx).isUndef()) {
+      Changed = true;
+      Idx = -1;
+    } else if (BV1 && Idx > (int)NumElts &&
+               BV1->getOperand(Idx - NumElts).isUndef()) {
+      Changed = true;
+      Idx = -1;
+    }
+    NewMask.push_back(Idx);
+  }
+  if (Changed)
+    return DAG.getVectorShuffle(VT, SDLoc(SVN), N0, N1, NewMask);
+
+  return SDValue();
+}
+
 // Tries to turn a shuffle of two CONCAT_VECTORS into a single concat,
 // or turn a shuffle of a single concat into simpler shuffle then concat.
 static SDValue partitionShuffleOfConcats(SDNode *N, SelectionDAG &DAG) {
   EVT VT = N->getValueType(0);
   unsigned NumElts = VT.getVectorNumElements();
 
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
for (unsigned i = 0; i != NumElts; ++i) {
         Changed = true;
       }
       NewMask.push_back(Idx);
     }
     if (Changed)
       return DAG.getVectorShuffle(VT, SDLoc(N), N0, N1, NewMask);
   }
 
+  // Simplify shuffle mask if a referenced element is UNDEF.
+  if (SDValue V = simplifyShuffleMask(SVN, N0, N1, DAG))
+    return V;
+
   // A shuffle of a single vector that is a splat can always be folded.
   if (auto *N0Shuf = dyn_cast<ShuffleVectorSDNode>(N0))
     if (N1->isUndef() && N0Shuf->isSplat())
       return combineShuffleOfSplat(SVN->getMask(), N0Shuf, DAG);
 
   // If it is a splat, check if the argument vector is another splat or a
   // build_vector.
   if (SVN->isSplat() && SVN->getSplatIndex() < (int)NumElts) {

test/CodeGen/X86/avx-vbroadcast.ll

; X64-NEXT:    retq
   ret <4 x double> %6
 }
 
 ; PR34041
 define <4 x double> @broadcast_shuffle_1000(double* %p) {
 ; X32-LABEL: broadcast_shuffle_1000:
 ; X32:       ## BB#0:
 ; X32-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; X32-NEXT:    vmovddup {{.*#+}} xmm0 = mem[0,0]
-; X32-NEXT:    vinsertf128 $1, %xmm0, %ymm0, %ymm0
+; X32-NEXT:    vbroadcastsd (%eax), %ymm0
 ; X32-NEXT:    retl
 ;
 ; X64-LABEL: broadcast_shuffle_1000:
 ; X64:       ## BB#0:
-; X64-NEXT:    vmovddup {{.*#+}} xmm0 = mem[0,0]
-; X64-NEXT:    vinsertf128 $1, %xmm0, %ymm0, %ymm0
+; X64-NEXT:    vbroadcastsd (%rdi), %ymm0
 ; X64-NEXT:    retq
   %1 = load double, double* %p
   %2 = insertelement <2 x double> undef, double %1, i32 0
   %3 = shufflevector <2 x double> %2, <2 x double> undef, <4 x i32> <i32 1, i32 0, i32 0, i32 0>
   ret <4 x double> %3
 }
 
 ;