[InstCombine] Fix assert condition in `foldSelectShuffleOfSelectShuffle`

Bug introduced in e239198cdbbf.

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

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index 40acc24..d509186 100644 (file)
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -2004,7 +2004,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);
 }
 

diff --git a/llvm/test/Transforms/InstCombine/shuffle_select.ll b/llvm/test/Transforms/InstCombine/shuffle_select.ll
index 357b8e1..044306f 100644 (file)
--- a/llvm/test/Transforms/InstCombine/shuffle_select.ll
+++ b/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
+}