[llvm] e1fd6d4 - [InstCombine] Fix assert condition in `foldSelectShuffleOfSelectShuffle`

[Nabeel Omer via llvm-commits]

Author: Nabeel Omer
Date: 2022-10-20T12:10:54Z
New Revision: e1fd6d49a3dfa59b4be85102a866a2068fa4760a

URL: https://github.com/llvm/llvm-project/commit/e1fd6d49a3dfa59b4be85102a866a2068fa4760a
DIFF: https://github.com/llvm/llvm-project/commit/e1fd6d49a3dfa59b4be85102a866a2068fa4760a.diff

LOG: [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

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
    llvm/test/Transforms/InstCombine/shuffle_select.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index 40acc2423a246..d50918629ba5c 100644
--- 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 357b8e1c9a1ae..044306f3b4811 100644
--- 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
+}