[InstCombine] Add a combine for a shuffle of similar bitcasts

[sanwou01]

Summary

Some intrinsics wrapper code has the habit of ignoring the type of the
elements in vectors, thinking of vector registers as a "bag of bits". As
a consequence, some operations are shared between vectors of different
types are shared. For example, functions that rearrange elements in a
vector can be shared between vectors of int32 and float.

This can result in bitcasts in awkward places that prevent the backend
from recognizing some instructions. For AArch64 in particular, it
inhibits the selection of dup from a general purpose register (GPR), and
mov from GPR to a vector lane.

This patch adds a pattern in InstCombine to move the bitcasts past the
shufflevector if this is possible. Sometimes this even allows
InstCombine to remove the bitcast entirely, as in the included tests.

Alternatively this could be done with a few extra patterns in the
AArch64 backend, but InstCombine seems like a better place for this.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

2297

One of the original bitcasts needs to be one-use, else we'll increase instruction count.

2299

Can we do anything for the case where we have element count mismatches?

2301

[lebedev.ri]

Whoops, phab did something weird

[sanwou01]

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

2297

Good point! I can bail out if neither is one-use.

2299

Yes, I'm being fairly conservative here. I think the following example would be legal to transform similarly. Is that what you had in mind?

%0 = bitcast <4 x i32> %a to <4 x float>
%1 = bitcast <2 x i32> %b to <2 x float>
%2 = shufflevector <4 x float> %0, <2 x float> %1, <2 x i32> <i32 3, i32 5>

2301

Nice, thanks!

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

2299

I actually don't have anything particular in mind, just asking.
That being said, i think both operands of a shufflevector must have the same type
(including vector element count), so i'm not sure that example works?

[spatel]

We intentionally do not create new shuffle masks in instcombine because we can't guarantee that codegen can lower arbitrary masks efficiently, but this patch seems fine since it just re-uses the existing mask.
If there is motivation to handle casts of different-sized elements (and therefore requires a new mask), you might look at building on VectorCombine::foldBitcastShuf(). We use the cost model there to avoid creating unsupported shuffles.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

2299

It should be ok to handle a length-changing shuffle, but we need to add type checks to make it safe.
Note that as-is this patch doesn't have the right combination - this crashes:

define <4 x double> @bc_shuf(<4 x i32> %x, <4 x i32> %y) {
  %xb = bitcast <4 x i32> %x to <2 x double>
  %yb = bitcast <4 x i32> %y to <2 x double>
  %r = shufflevector <2 x double> %xb, <2 x double> %yb, <4 x i32> <i32 1, i32 3, i32 0, i32 2>
  ret <4 x double> %r
}

So definitely add some negative tests to make sure we don't break things.

[sanwou01]

In D97397#2585516, @spatel wrote:

We intentionally do not create new shuffle masks in instcombine because we can't guarantee that codegen can lower arbitrary masks efficiently, but this patch seems fine since it just re-uses the existing mask.
If there is motivation to handle casts of different-sized elements (and therefore requires a new mask), you might look at building on VectorCombine::foldBitcastShuf(). We use the cost model there to avoid creating unsupported shuffles.

I don't have a motivation to handle bitcasts that change the element size (or scalar to vector, for that matter), so I'd prefer to keep this simple.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

2299

Good catch, I've added tests for length-changing shuffles (which as you point out we can handle without changing the mask) and some negative tests for bitcasts that change the element size.

[sanwou01]

Address comments.

[lebedev.ri]

I'd like to see some test improvements:

1. There are no tests with extra uses on bitcasts
2. There are no tests with something like

%xb = bitcast <2 x half> %x to <2 x i16>
%yb = bitcast <2 x bfloat> %y to <2 x i16>
%r = shufflevector <2 x i16> %xb, <2 x i16> %yb, <4 x i16> <i16 3, i16 2, i16 1, i16 0>

I suspect this will miscompile?

3. Some tests have unneeded stuff. They should only contain 2 bitcasts and a shufflevector (and a ret), nothing more.

[sanwou01]

In D97397#2594000, @lebedev.ri wrote:

I'd like to see some test improvements:

1. There are no tests with extra uses on bitcasts
2. There are no tests with something like

%xb = bitcast <2 x half> %x to <2 x i16>
%yb = bitcast <2 x bfloat> %y to <2 x i16>
%r = shufflevector <2 x i16> %xb, <2 x i16> %yb, <4 x i16> <i16 3, i16 2, i16 1, i16 0>

I suspect this will miscompile?

3. Some tests have unneeded stuff. They should only contain 2 bitcasts and a shufflevector (and a ret), nothing more.

Thanks for the review, @lebedev.ri !

1. Fair enough, I'll add some tests for that.
2. I'll add this as a test, but I don't think it miscompiles: I do check that the input types of the two bitcasts are identical.
3. Can do. The first two tests were derived from my motivating examples, but I can simplify them some more as you point out, thanks.

[sanwou01]

Add a few more tests and some further test reduction.

[sanwou01]

@spatel @lebedev.ri thanks for the comments so far. Any other comments, or is this okay as is?

[lebedev.ri]

LGTM

[spatel]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

if (auto *V = SimplifyShuffleVectorInst(LHS, RHS, SVI.getShuffleMask(),
     return replaceInstUsesWith(SVI, V);
 
   // Bail out for scalable vectors
   if (isa<ScalableVectorType>(LHS->getType()))
     return nullptr;
 
   unsigned VWidth = cast<FixedVectorType>(SVI.getType())->getNumElements();
   unsigned LHSWidth = cast<FixedVectorType>(LHS->getType())->getNumElements();
+
+  // shuffle (bitcast X), (bitcast Y), Mask --> bitcast (shuffle X, Y, Mask)
+  //
+  // if X and Y are of the same (vector) type, and the element size is not
+  // changed by the bitcasts, we can distribute the bitcasts through the
+  // shuffle, hopefully reducing the number of instructions. We make sure that
+  // at least one bitcast only has one use, so we don't *increase* the number of
+  // instructions here.
+  Value *X, *Y;
+  if (match(LHS, m_BitCast(m_Value(X))) && match(RHS, m_BitCast(m_Value(Y))) &&
+      X->getType()->isVectorTy() && X->getType() == Y->getType() &&
+      X->getType()->getScalarSizeInBits() ==
+          SVI.getType()->getScalarSizeInBits() &&
+      (LHS->hasOneUse() || RHS->hasOneUse())) {
+    Value *V = Builder.CreateShuffleVector(X, Y, SVI.getShuffleMask(),
+                                           SVI.getName() + ".uncasted");
+    return new BitCastInst(V, SVI.getType());
+  }
+
   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.
   //       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) {
     // Try to create a scaled mask constant.
     auto *XType = cast<FixedVectorType>(X->getType());
     unsigned XNumElts = XType->getNumElements();
     SmallVector<int, 16> ScaledMask;
     if (XNumElts >= VWidth) {
       assert(XNumElts % VWidth == 0 && "Unexpected vector bitcast");

llvm/test/Transforms/InstCombine/shuffle-cast-dist.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 define <2 x float> @vtrn1(<2 x i32> %v)
 ; CHECK-LABEL: @vtrn1(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[VB1:%.*]] = bitcast <2 x i32> [[V:%.*]] to <2 x float>
-; CHECK-NEXT:    [[VB2:%.*]] = bitcast <2 x i32> [[V]] to <2 x float>
-; CHECK-NEXT:    [[R:%.*]] = shufflevector <2 x float> [[VB1]], <2 x float> [[VB2]], <2 x i32> <i32 0, i32 2>
+; CHECK-NEXT:    [[R_UNCASTED:%.*]] = shufflevector <2 x i32> [[V:%.*]], <2 x i32> undef, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[R:%.*]] = bitcast <2 x i32> [[R_UNCASTED]] to <2 x float>
 ; CHECK-NEXT:    ret <2 x float> [[R]]
 ;
 {
 entry:
   %vb1 = bitcast <2 x i32> %v to <2 x float>
   %vb2 = bitcast <2 x i32> %v to <2 x float>
   %r = shufflevector <2 x float> %vb1, <2 x float> %vb2, <2 x i32> <i32 0, i32 2>
   ret <2 x float> %r
 }
 
 define <2 x float> @vtrn2(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @vtrn2(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[XB:%.*]] = bitcast <2 x i32> [[X:%.*]] to <2 x float>
-; CHECK-NEXT:    [[YB:%.*]] = bitcast <2 x i32> [[Y:%.*]] to <2 x float>
-; CHECK-NEXT:    [[R:%.*]] = shufflevector <2 x float> [[XB]], <2 x float> [[YB]], <2 x i32> <i32 1, i32 3>
+; CHECK-NEXT:    [[R_UNCASTED:%.*]] = shufflevector <2 x i32> [[X:%.*]], <2 x i32> [[Y:%.*]], <2 x i32> <i32 1, i32 3>
+; CHECK-NEXT:    [[R:%.*]] = bitcast <2 x i32> [[R_UNCASTED]] to <2 x float>
 ; CHECK-NEXT:    ret <2 x float> [[R]]
 ;
 entry:
   %xb = bitcast <2 x i32> %x to <2 x float>
   %yb = bitcast <2 x i32> %y to <2 x float>
   %r = shufflevector <2 x float> %xb, <2 x float> %yb, <2 x i32> <i32 1, i32 3>
   ret <2 x float> %r
 }
 
 
 define <4 x float> @bc_shuf_lenchange(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @bc_shuf_lenchange(
-; CHECK-NEXT:    [[XB:%.*]] = bitcast <2 x i32> [[X:%.*]] to <2 x float>
-; CHECK-NEXT:    [[YB:%.*]] = bitcast <2 x i32> [[Y:%.*]] to <2 x float>
-; CHECK-NEXT:    [[R:%.*]] = shufflevector <2 x float> [[XB]], <2 x float> [[YB]], <4 x i32> <i32 3, i32 2, i32 1, i32 0>
+; CHECK-NEXT:    [[R_UNCASTED:%.*]] = shufflevector <2 x i32> [[X:%.*]], <2 x i32> [[Y:%.*]], <4 x i32> <i32 3, i32 2, i32 1, i32 0>
+; CHECK-NEXT:    [[R:%.*]] = bitcast <4 x i32> [[R_UNCASTED]] to <4 x float>
 ; CHECK-NEXT:    ret <4 x float> [[R]]
 ;
   %xb = bitcast <2 x i32> %x to <2 x float>
   %yb = bitcast <2 x i32> %y to <2 x float>
   %r = shufflevector <2 x float> %xb, <2 x float> %yb, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
   ret <4 x float> %r
 }
 
;
   %xb = bitcast <4 x i32> %x to <4 x float>
   %r = shufflevector <4 x float> %xb, <4 x float> %xb, <2 x i32> <i32 0, i32 4>
   ret <2 x float> %r
 }
 
 define <4 x float> @bc_shuf_y_hasoneuse(<4 x i32> %x, <4 x i32> %y){
 ; CHECK-LABEL: @bc_shuf_y_hasoneuse(
 ; CHECK-NEXT:    [[XB:%.*]] = bitcast <4 x i32> [[X:%.*]] to <4 x float>
-; CHECK-NEXT:    [[YB:%.*]] = bitcast <4 x i32> [[Y:%.*]] to <4 x float>
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x float> [[XB]], <4 x float> [[YB]], <4 x i32> <i32 0, i32 1, i32 4, i32 5>
-; CHECK-NEXT:    [[R:%.*]] = fadd <4 x float> [[SHUF]], [[XB]]
+; CHECK-NEXT:    [[SHUF_UNCASTED:%.*]] = shufflevector <4 x i32> [[X]], <4 x i32> [[Y:%.*]], <4 x i32> <i32 0, i32 1, i32 4, i32 5>
+; CHECK-NEXT:    [[SHUF:%.*]] = bitcast <4 x i32> [[SHUF_UNCASTED]] to <4 x float>
+; CHECK-NEXT:    [[R:%.*]] = fadd <4 x float> [[XB]], [[SHUF]]
 ; CHECK-NEXT:    ret <4 x float> [[R]]
 ;
   %xb = bitcast <4 x i32> %x to <4 x float>
   %yb = bitcast <4 x i32> %y to <4 x float>
   %shuf = shufflevector <4 x float> %xb, <4 x float> %yb, <4 x i32> <i32 0, i32 1, i32 4, i32 5>
   %r = fadd <4 x float> %xb, %shuf
   ret <4 x float> %r
 }