[InstCombine] Fix assert condition in foldSelectShuffleOfSelectShuffle

Bug introduced in e239198.

The assert() is making an assumption that the resulting shuffle mask
will always select elements from both vectors, this is untrue in the
case of two shuffles being folded if the former shuffle has a mask with
undef elements in it. In such a case folding the shuffles might result
in a mask which only selects from one of the vectors because the other
elements (in the mask) are undef.

Differential Revision: https://reviews.llvm.org/D136256

Author: omern1

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

@@ -2009,7 +2009,10 @@ static Instruction *foldSelectShuffleOfSelectShuffle(ShuffleVectorInst &Shuf) {
   for (unsigned i = 0; i != NumElts; ++i)
     NewMask[i] = Mask[i] < (signed)NumElts ? Mask[i] : Mask1[i];
 
-  assert(ShuffleVectorInst::isSelectMask(NewMask) && "Unexpected shuffle mask");
+  // A select mask with undef elements might look like an identity mask.
+  assert((ShuffleVectorInst::isSelectMask(NewMask) ||
+          ShuffleVectorInst::isIdentityMask(NewMask)) &&
+         "Unexpected shuffle mask");
   return new ShuffleVectorInst(X, Y, NewMask);
 }
 

llvm/test/Transforms/InstCombine/shuffle_select.ll

@@ -1621,3 +1621,13 @@ define <4 x i32> @sel_common_op_extra_use(<4 x i32> %x, <4 x i32> %y) {
   %s2 = shufflevector <4 x i32> %s1, <4 x i32> %x, <4 x i32> <i32 0, i32 1, i32 6, i32 7>
   ret <4 x i32> %s2
 }
+
+define <4 x float> @identity_mask(<4 x float>%x, <4 x float> %y) {
+; CHECK-LABEL: @identity_mask(
+; CHECK-NEXT:    [[S2:%.*]] = shufflevector <4 x float> [[X:%.*]], <4 x float> poison, <4 x i32> <i32 0, i32 undef, i32 2, i32 3>
+; CHECK-NEXT:    ret <4 x float> [[S2]]
+;
+  %s1 = shufflevector <4 x float> %x, <4 x float> %y, <4 x i32> <i32 0, i32 5, i32 undef, i32 undef>
+  %s2 = shufflevector <4 x float> %s1, <4 x float> %x, <4 x i32> <i32 0, i32 undef, i32 6, i32 7>
+  ret <4 x float> %s2
+}