[InstCombine] try to reduce shuffle with bitcasted operand

[spatel]

Summary

shuf (bitcast X), undef, Mask --> bitcast X'

The 'inverse shuffles' test (shuf_bitcast_operand) is a pattern that would emerge in the motivating examples from PR35454:
https://bugs.llvm.org/show_bug.cgi?id=35454
(if we proceed with D76727)

We can deal with this class of patterns in generic instcombine because we are not creating any new shuffles, just a bitcast.

This provides an opportunity to exercise Alive2's vector support:
http://volta.cs.utah.edu:8080/z/mwDUZf

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1927

Its seems a pity to create a speculative ConstantVector like this - can we easily refactor SimplifyShuffleVectorInst to take ArrayRef<int> ?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1927

The diffs we want are part of D72467.
@efriedma - is that patch held up on anything? We may not have settled on the scalable vector future, but changing the fixed vector implementation would make things easier for existing code today.

[efriedma]

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1927

I'll go over it one more time, but I don't think it's held up on anything in particular. We didn't completely reach consensus on scalable vector shuffles, but the result is almost certainly compatible with that patch. I'll try to land it today.

[lebedev.ri]

This seems reasonable to me, unless @RKSimon has concerns.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1920–1922

Same comment about reserve+range-based-for, although i hope that will get obsoleted by D72467?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1920–1922

Yes, if we convert the mask from a constant operand to a special constant, then this and other code gets simplified away.

[spatel]

D72467 is not actually a parent, but adding that here because this would conflict (force another rebase of that patch).

[spatel]

Patch updated - no logic changes, but rebased/reduced with new shuffle instruction.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

for (int i = 0, e = Mask.size(); i != e; ++i) {
     }
   }
   return new ShuffleVectorInst(X, Y, ConstantVector::get(NewMask));
 }
 
 Instruction *InstCombiner::visitShuffleVectorInst(ShuffleVectorInst &SVI) {
   Value *LHS = SVI.getOperand(0);
   Value *RHS = SVI.getOperand(1);
-  if (auto *V =
-          SimplifyShuffleVectorInst(LHS, RHS, SVI.getShuffleMask(),
-                                    SVI.getType(), SQ.getWithInstruction(&SVI)))
+  SimplifyQuery ShufQuery = SQ.getWithInstruction(&SVI);
+  if (auto *V = SimplifyShuffleVectorInst(LHS, RHS, SVI.getShuffleMask(),
+                                          SVI.getType(), ShufQuery))
     return replaceInstUsesWith(SVI, V);
 
   // shuffle x, x, mask --> shuffle x, undef, mask'
   unsigned VWidth = SVI.getType()->getVectorNumElements();
   unsigned LHSWidth = LHS->getType()->getVectorNumElements();
   ArrayRef<int> Mask = SVI.getShuffleMask();
   Type *Int32Ty = Type::getInt32Ty(SVI.getContext());
+
+  // Peek through a bitcasted shuffle operand by scaling the mask. If the
+  // simulated shuffle can simplify, then this shuffle is unnecessary:
+  // shuf (bitcast X), undef, Mask --> bitcast X'
+  // TODO: This could be extended to allow length-changing shuffles and/or casts
+  //       to narrower elements. The transform might also be obsoleted if we
+  //       allowed canonicalization of bitcasted shuffles.
+  Value *X;
+  if (match(LHS, m_BitCast(m_Value(X))) && match(RHS, m_Undef()) &&
+      X->getType()->isVectorTy() && VWidth == LHSWidth &&
+      X->getType()->getVectorNumElements() >= VWidth) {
+    // Create the scaled mask constant.
+    Type *XType = X->getType();
+    unsigned XNumElts = XType->getVectorNumElements();
+    assert(XNumElts % VWidth == 0 && "Unexpected vector bitcast");
+    unsigned ScaleFactor = XNumElts / VWidth;
+    SmallVector<int, 16> ScaledMask;
+    scaleShuffleMask(ScaleFactor, Mask, ScaledMask);
+
+    // If the shuffled source vector simplifies, cast that value to this
+    // shuffle's type.
+    if (auto *V = SimplifyShuffleVectorInst(X, UndefValue::get(XType),
+                                            ScaledMask, XType, ShufQuery))
+      return BitCastInst::Create(Instruction::BitCast, V, SVI.getType());
+  }
+
   if (LHS == RHS) {
     assert(!isa<UndefValue>(RHS) && "Shuffle with 2 undef ops not simplified?");
     // Remap any references to RHS to use LHS.
     SmallVector<int, 16> Elts;
     for (unsigned i = 0; i != VWidth; ++i) {
       // Propagate undef elements or force mask to LHS.
       if (Mask[i] < 0)
         Elts.push_back(UndefMaskElem);

llvm/test/Transforms/InstCombine/shufflevec-bitcast.ll

;
   %v = shufflevector <16 x i8> %w, <16 x i8> undef, <4 x i32> <i32 12, i32 13, i32 14, i32 15>
   %f = bitcast <4 x i8> %v to float
   %i = bitcast <4 x i8> %v to i32
   store i32 %i, i32* %o1, align 4
   store float %f, float* %o2, align 4
   ret void
 }
 
+; Shuffle-of-bitcast-splat --> splat-bitcast
+
 define <4 x i16> @splat_bitcast_operand(<8 x i8> %x) {
 ; CHECK-LABEL: @splat_bitcast_operand(
-; CHECK-NEXT:    [[S1:%.*]] = shufflevector <8 x i8> [[X:%.*]], <8 x i8> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 undef, i32 undef>
-; CHECK-NEXT:    [[BC:%.*]] = bitcast <8 x i8> [[S1]] to <4 x i16>
-; CHECK-NEXT:    [[S2:%.*]] = shufflevector <4 x i16> [[BC]], <4 x i16> undef, <4 x i32> <i32 0, i32 2, i32 1, i32 0>
+; CHECK-NEXT:    [[S1:%.*]] = shufflevector <8 x i8> [[X:%.*]], <8 x i8> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
+; CHECK-NEXT:    [[S2:%.*]] = bitcast <8 x i8> [[S1]] to <4 x i16>
 ; CHECK-NEXT:    ret <4 x i16> [[S2]]
 ;
   %s1 = shufflevector <8 x i8> %x, <8 x i8> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
   %bc = bitcast <8 x i8> %s1 to <4 x i16>
   %s2 = shufflevector <4 x i16> %bc, <4 x i16> undef, <4 x i32> <i32 0, i32 2, i32 1, i32 0>
   ret <4 x i16> %s2
 }
 
+; Shuffle-of-bitcast-splat --> splat-bitcast
+
 define <4 x i16> @splat_bitcast_operand_uses(<8 x i8> %x) {
 ; CHECK-LABEL: @splat_bitcast_operand_uses(
 ; CHECK-NEXT:    [[S1:%.*]] = shufflevector <8 x i8> [[X:%.*]], <8 x i8> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
 ; CHECK-NEXT:    [[BC:%.*]] = bitcast <8 x i8> [[S1]] to <4 x i16>
 ; CHECK-NEXT:    call void @use(<4 x i16> [[BC]])
-; CHECK-NEXT:    [[S2:%.*]] = shufflevector <4 x i16> [[BC]], <4 x i16> undef, <4 x i32> <i32 0, i32 2, i32 1, i32 0>
+; CHECK-NEXT:    [[S2:%.*]] = bitcast <8 x i8> [[S1]] to <4 x i16>
 ; CHECK-NEXT:    ret <4 x i16> [[S2]]
 ;
   %s1 = shufflevector <8 x i8> %x, <8 x i8> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
   %bc = bitcast <8 x i8> %s1 to <4 x i16>
   call void @use(<4 x i16> %bc)
   %s2 = shufflevector <4 x i16> %bc, <4 x i16> undef, <4 x i32> <i32 0, i32 2, i32 1, i32 0>
   ret <4 x i16> %s2
 }
 
+; Shuffle-of-bitcast-splat --> splat-bitcast
+
 define <4 x i32> @splat_bitcast_operand_same_size_src_elt(<4 x float> %x) {
 ; CHECK-LABEL: @splat_bitcast_operand_same_size_src_elt(
-; CHECK-NEXT:    [[S1:%.*]] = shufflevector <4 x float> [[X:%.*]], <4 x float> undef, <4 x i32> <i32 2, i32 2, i32 2, i32 undef>
-; CHECK-NEXT:    [[BC:%.*]] = bitcast <4 x float> [[S1]] to <4 x i32>
-; CHECK-NEXT:    [[S2:%.*]] = shufflevector <4 x i32> [[BC]], <4 x i32> undef, <4 x i32> <i32 0, i32 2, i32 1, i32 0>
+; CHECK-NEXT:    [[S1:%.*]] = shufflevector <4 x float> [[X:%.*]], <4 x float> undef, <4 x i32> <i32 2, i32 2, i32 2, i32 2>
+; CHECK-NEXT:    [[S2:%.*]] = bitcast <4 x float> [[S1]] to <4 x i32>
 ; CHECK-NEXT:    ret <4 x i32> [[S2]]
 ;
   %s1 = shufflevector <4 x float> %x, <4 x float> undef, <4 x i32> <i32 2, i32 2, i32 2, i32 2>
   %bc = bitcast <4 x float> %s1 to <4 x i32>
   %s2 = shufflevector <4 x i32> %bc, <4 x i32> undef, <4 x i32> <i32 0, i32 2, i32 1, i32 0>
   ret <4 x i32> %s2
 }
 
 ; Scaled mask is inverse of first mask.
 
 define <4 x i32> @shuf_bitcast_operand(<16 x i8> %x) {
 ; CHECK-LABEL: @shuf_bitcast_operand(
-; CHECK-NEXT:    [[S1:%.*]] = shufflevector <16 x i8> [[X:%.*]], <16 x i8> undef, <16 x i32> <i32 12, i32 13, i32 14, i32 15, i32 8, i32 9, i32 10, i32 11, i32 4, i32 5, i32 6, i32 7, i32 0, i32 1, i32 2, i32 3>
-; CHECK-NEXT:    [[BC:%.*]] = bitcast <16 x i8> [[S1]] to <4 x i32>
-; CHECK-NEXT:    [[S2:%.*]] = shufflevector <4 x i32> [[BC]], <4 x i32> undef, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
+; CHECK-NEXT:    [[S2:%.*]] = bitcast <16 x i8> [[X:%.*]] to <4 x i32>
 ; CHECK-NEXT:    ret <4 x i32> [[S2]]
 ;
   %s1 = shufflevector <16 x i8> %x, <16 x i8> undef, <16 x i32> <i32 12, i32 13, i32 14, i32 15, i32 8, i32 9, i32 10, i32 11, i32 4, i32 5, i32 6, i32 7, i32 0, i32 1, i32 2, i32 3>
   %bc = bitcast <16 x i8> %s1 to <4 x i32>
   %s2 = shufflevector <4 x i32> %bc, <4 x i32> undef, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
   ret <4 x i32> %s2
 }
 
+; TODO: Could allow fold for length-changing shuffles.
+
 define <5 x i16> @splat_bitcast_operand_change_type(<8 x i8> %x) {
 ; CHECK-LABEL: @splat_bitcast_operand_change_type(
 ; CHECK-NEXT:    [[S1:%.*]] = shufflevector <8 x i8> [[X:%.*]], <8 x i8> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
 ; CHECK-NEXT:    [[BC:%.*]] = bitcast <8 x i8> [[S1]] to <4 x i16>
 ; CHECK-NEXT:    [[S2:%.*]] = shufflevector <4 x i16> [[BC]], <4 x i16> undef, <5 x i32> <i32 0, i32 2, i32 1, i32 0, i32 3>
 ; CHECK-NEXT:    ret <5 x i16> [[S2]]
 ;
   %s1 = shufflevector <8 x i8> %x, <8 x i8> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
   %bc = bitcast <8 x i8> %s1 to <4 x i16>
   %s2 = shufflevector <4 x i16> %bc, <4 x i16> undef, <5 x i32> <i32 0, i32 2, i32 1, i32 0, i32 3>
   ret <5 x i16> %s2
 }
 
+; TODO: Could allow fold for cast to narrow element.
+
 define <4 x i16> @splat_bitcast_operand_wider_src_elt(<2 x i32> %x) {
 ; CHECK-LABEL: @splat_bitcast_operand_wider_src_elt(
 ; CHECK-NEXT:    [[S1:%.*]] = shufflevector <2 x i32> [[X:%.*]], <2 x i32> undef, <2 x i32> <i32 1, i32 undef>
 ; CHECK-NEXT:    [[BC:%.*]] = bitcast <2 x i32> [[S1]] to <4 x i16>
 ; CHECK-NEXT:    [[S2:%.*]] = shufflevector <4 x i16> [[BC]], <4 x i16> undef, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
 ; CHECK-NEXT:    ret <4 x i16> [[S2]]
 ;
   %s1 = shufflevector <2 x i32> %x, <2 x i32> undef, <2 x i32> <i32 1, i32 1>
   %bc = bitcast <2 x i32> %s1 to <4 x i16>
   %s2 = shufflevector <4 x i16> %bc, <4 x i16> undef, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
   ret <4 x i16> %s2
 }