[InstCombine] Fix assert condition in `foldSelectShuffleOfSelectShuffle`

[n-omer]

Summary

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.

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

2008

clang-format this please :)

llvm/test/Transforms/InstCombine/foldSelectShuffleOfSelectShuffle.ll

3 ↗

(On Diff #468906)

Use the update_test_checks script to autogenerate the checks

7 ↗

(On Diff #468906)

remove local_unnamed_addr #0

[nikic]

llvm/test/Transforms/InstCombine/foldSelectShuffleOfSelectShuffle.ll

3 ↗

(On Diff #468906)

...and move this test next to wherever the other tests for this transform are.

[spatel]

llvm/test/Transforms/InstCombine/foldSelectShuffleOfSelectShuffle.ll

4 ↗

(On Diff #468906)

...and minimize the test :)

No loads or globals are needed to show the bug (if this is filed as a bug somewhere, add a reference).

define <4 x float> @identity_mask(<4 x float>%x, <4 x float> %y) {
  %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
}

[n-omer]

Address review comments

[spatel]

LGTM.

For reference, the bug was introduced with e239198cdbbf. Before that patch, the pattern would get reduced in multiple steps using demanded elements, so there's no difference in optimization on this kind of shuffle pattern AFAICT.

[RKSimon]

LGTM - cheers

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

static Instruction *foldSelectShuffleOfSelectShuffle(ShuffleVectorInst &Shuf) {
   // If the mask chooses from X (operand 0), it stays the same.
   // If the mask chooses from the earlier shuffle, the other mask value is
   // transferred to the combined select shuffle:
   // shuf X, (shuf X, Y, M1), M --> shuf X, Y, M'
   SmallVector<int, 16> NewMask(NumElts);
   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);
 }
 
 static Instruction *foldSelectShuffleWith1Binop(ShuffleVectorInst &Shuf) {
   assert(Shuf.isSelect() && "Must have select-equivalent shuffle");
 
   // Are we shuffling together some value and that same value after it has been
   // modified by a binop with a constant?

llvm/test/Transforms/InstCombine/shuffle_select.ll

 ; CHECK-NEXT:    [[S2:%.*]] = shufflevector <4 x i32> [[Y]], <4 x i32> [[X]], <4 x i32> <i32 0, i32 5, i32 6, i32 7>
 ; CHECK-NEXT:    ret <4 x i32> [[S2]]
 ;
   %s1 = shufflevector <4 x i32> %y, <4 x i32> %x, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
   call void @use_v4i32(<4 x i32> %s1)
   %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
+}