[InstCombine] recognize bswap disguised as shufflevector

[spatel]

Summary

bitcast <N x i8> (shuf X, undef, <N, N-1,...0>) to i{N*8}--> bswap (bitcast X to i{N*8})

In PR43146:
https://bugs.llvm.org/show_bug.cgi?id=43146
...we have a more complicated case where SLP is making a mess of bswap. This patch won't do anything for that currently, but we need to improve bswap recognition in instcombine, SLP, and/or a standalone pass to avoid that problem.

This is limited using the data-layout so we don't try to do this transform with actual vector types. The backend does not appear to have folds to convert in either direction, so we don't want to mess up something that is actually better lowered as a shuffle.

On x86, we're trading something like this:

vmovd	%edi, %xmm0
vpshufb	LCPI0_0(%rip), %xmm0, %xmm0 ## xmm0 = xmm0[3,2,1,0,u,u,u,u,u,u,u,u,u,u,u,u]
vmovd	%xmm0, %eax

For:

movl	%edi, %eax
bswapl	%eax

[lebedev.ri]

Looks good.
The same applies to llvm.bitreverse.*, i believe, if scalar type is i1?

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

2397–2412

Precommit please

[spatel]

In D66965#1651426, @lebedev.ri wrote:

Looks good.
The same applies to llvm.bitreverse.*, i believe, if scalar type is i1?

Yes. That one is probably rarer and might want slightly different constraints. I'll add a TODO for now.

[spatel]

Patch updated:

1. Rebased after cleanup - rL370399.
2. Added TODO comment about bitreverse.

[lebedev.ri]

Some more thoughts, but i think this looks ok.

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

2424–2426

Should Shuf->isReverse() be checked last, it is likely most costly check here?

llvm/test/Transforms/InstCombine/bswap.ll

284–286

I think you also want a test where the shuffle mask is widening (to a legal type)
We currently won't accept it, but we might one day, and i'm not sure what that code would do.

[spatel]

Patch updated:

1. Optimization: rearranged comparisons for better perf.
2. Added widening shuffle negative test.

[lebedev.ri]

No further comments from me here, thanks.

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

if (SrcVTy->getNumElements() == 1) {
       if (auto *InsElt = dyn_cast<InsertElementInst>(Src))
         return new BitCastInst(InsElt->getOperand(1), DestTy);
     }
   }
 
   if (auto *Shuf = dyn_cast<ShuffleVectorInst>(Src)) {
     // Okay, we have (bitcast (shuffle ..)).  Check to see if this is
     // a bitcast to a vector with the same # elts.
     Value *ShufOp0 = Shuf->getOperand(0);
     Value *ShufOp1 = Shuf->getOperand(1);
     unsigned NumShufElts = Shuf->getType()->getVectorNumElements();
     unsigned NumSrcVecElts = ShufOp0->getType()->getVectorNumElements();
     if (Shuf->hasOneUse() && DestTy->isVectorTy() &&
         DestTy->getVectorNumElements() == NumShufElts &&
         NumShufElts == NumSrcVecElts) {
       BitCastInst *Tmp;
       // If either of the operands is a cast from CI.getType(), then
       // evaluating the shuffle in the casted destination's type will allow
       // us to eliminate at least one cast.
       if (((Tmp = dyn_cast<BitCastInst>(ShufOp0)) &&
            Tmp->getOperand(0)->getType() == DestTy) ||
           ((Tmp = dyn_cast<BitCastInst>(ShufOp1)) &&
            Tmp->getOperand(0)->getType() == DestTy)) {
         Value *LHS = Builder.CreateBitCast(ShufOp0, DestTy);
         Value *RHS = Builder.CreateBitCast(ShufOp1, DestTy);
         // Return a new shuffle vector.  Use the same element ID's, as we
         // know the vector types match #elts.
         return new ShuffleVectorInst(LHS, RHS, Shuf->getOperand(2));
       }
     }
+
+    // A bitcasted-to-scalar and byte-reversing shuffle is better recognized as
+    // a byte-swap:
+    // bitcast <N x i8> (shuf X, undef, <N, N-1,...0>) --> bswap (bitcast X)
+    // TODO: We should match the related pattern for bitreverse.
+    if (DestTy->isIntegerTy() &&
+        DL.isLegalInteger(DestTy->getScalarSizeInBits()) &&
+        SrcTy->getScalarSizeInBits() == 8 && NumShufElts % 2 == 0 &&
+        Shuf->hasOneUse() && Shuf->isReverse()) {
+      assert(ShufOp0->getType() == SrcTy && "Unexpected shuffle mask");
+      assert(isa<UndefValue>(ShufOp1) && "Unexpected shuffle op");
+      Function *Bswap =
+          Intrinsic::getDeclaration(CI.getModule(), Intrinsic::bswap, DestTy);
+      Value *ScalarX = Builder.CreateBitCast(ShufOp0, DestTy);
+      return IntrinsicInst::Create(Bswap, { ScalarX });
+    }
   }
 
   // Handle the A->B->A cast, and there is an intervening PHI node.
   if (PHINode *PN = dyn_cast<PHINode>(Src))
     if (Instruction *I = optimizeBitCastFromPhi(CI, PN))
       return I;
 
   if (Instruction *I = canonicalizeBitCastExtElt(CI, *this))

llvm/test/Transforms/InstCombine/bswap.ll

;
   %shl1 = and i32 %and2, 65280
   %or = or i32 %and1, %shl1
   %conv = trunc i32 %or to i16
   ret i16 %conv
 }
 
 define i32 @shuf_4bytes(<4 x i8> %x) {
 ; CHECK-LABEL: @shuf_4bytes(
-; CHECK-NEXT:    [[BSWAP:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> undef, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
-; CHECK-NEXT:    [[CAST:%.*]] = bitcast <4 x i8> [[BSWAP]] to i32
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast <4 x i8> [[X:%.*]] to i32
+; CHECK-NEXT:    [[CAST:%.*]] = call i32 @llvm.bswap.i32(i32 [[TMP1]])
 ; CHECK-NEXT:    ret i32 [[CAST]]
 ;
   %bswap = shufflevector <4 x i8> %x, <4 x i8> undef, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
   %cast = bitcast <4 x i8> %bswap to i32
   ret i32 %cast
 }
 
 define i32 @shuf_load_4bytes(<4 x i8>* %p) {
 ; CHECK-LABEL: @shuf_load_4bytes(
-; CHECK-NEXT:    [[X:%.*]] = load <4 x i8>, <4 x i8>* [[P:%.*]], align 4
-; CHECK-NEXT:    [[BSWAP:%.*]] = shufflevector <4 x i8> [[X]], <4 x i8> undef, <4 x i32> <i32 3, i32 2, i32 undef, i32 0>
-; CHECK-NEXT:    [[CAST:%.*]] = bitcast <4 x i8> [[BSWAP]] to i32
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast <4 x i8>* [[P:%.*]] to i32*
+; CHECK-NEXT:    [[X1:%.*]] = load i32, i32* [[TMP1]], align 4
+; CHECK-NEXT:    [[CAST:%.*]] = call i32 @llvm.bswap.i32(i32 [[X1]])
 ; CHECK-NEXT:    ret i32 [[CAST]]
 ;
   %x = load <4 x i8>, <4 x i8>* %p
   %bswap = shufflevector <4 x i8> %x, <4 x i8> undef, <4 x i32> <i32 3, i32 2, i32 undef, i32 0>
   %cast = bitcast <4 x i8> %bswap to i32
   ret i32 %cast
 }
 
 define i32 @shuf_bitcast_twice_4bytes(i32 %x) {
 ; CHECK-LABEL: @shuf_bitcast_twice_4bytes(
-; CHECK-NEXT:    [[CAST1:%.*]] = bitcast i32 [[X:%.*]] to <4 x i8>
-; CHECK-NEXT:    [[BSWAP:%.*]] = shufflevector <4 x i8> [[CAST1]], <4 x i8> undef, <4 x i32> <i32 undef, i32 2, i32 1, i32 0>
-; CHECK-NEXT:    [[CAST2:%.*]] = bitcast <4 x i8> [[BSWAP]] to i32
+; CHECK-NEXT:    [[CAST2:%.*]] = call i32 @llvm.bswap.i32(i32 [[X:%.*]])
 ; CHECK-NEXT:    ret i32 [[CAST2]]
 ;
   %cast1 = bitcast i32 %x to <4 x i8>
   %bswap = shufflevector <4 x i8> %cast1, <4 x i8> undef, <4 x i32> <i32 undef, i32 2, i32 1, i32 0>
   %cast2 = bitcast <4 x i8> %bswap to i32
   ret i32 %cast2
 }
 
+; Negative test - extra use
 declare void @use(<4 x i8>)
 
 define i32 @shuf_4bytes_extra_use(<4 x i8> %x) {
 ; CHECK-LABEL: @shuf_4bytes_extra_use(
 ; CHECK-NEXT:    [[BSWAP:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> undef, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
 ; CHECK-NEXT:    call void @use(<4 x i8> [[BSWAP]])
 ; CHECK-NEXT:    [[CAST:%.*]] = bitcast <4 x i8> [[BSWAP]] to i32
 ; CHECK-NEXT:    ret i32 [[CAST]]
 ;
   %bswap = shufflevector <4 x i8> %x, <4 x i8> undef, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
   call void @use(<4 x i8> %bswap)
   %cast = bitcast <4 x i8> %bswap to i32
   ret i32 %cast
 }
 
+; Negative test - scalar type is not in the data layout
+
 define i128 @shuf_16bytes(<16 x i8> %x) {
 ; CHECK-LABEL: @shuf_16bytes(
 ; CHECK-NEXT:    [[BSWAP:%.*]] = shufflevector <16 x i8> [[X:%.*]], <16 x i8> undef, <16 x i32> <i32 15, i32 14, i32 13, i32 12, i32 11, i32 10, i32 9, i32 8, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0>
 ; CHECK-NEXT:    [[CAST:%.*]] = bitcast <16 x i8> [[BSWAP]] to i128
 ; CHECK-NEXT:    ret i128 [[CAST]]
 ;
   %bswap = shufflevector <16 x i8> %x, <16 x i8> undef, <16 x i32> <i32 15, i32 14, i32 13, i32 12, i32 11, i32 10, i32 9, i32 8, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0>
   %cast = bitcast <16 x i8> %bswap to i128
   ret i128 %cast
 }
 
+; Negative test - don't touch widening shuffles (for now)
+
 define i32 @shuf_2bytes_widening(<2 x i8> %x) {
 ; CHECK-LABEL: @shuf_2bytes_widening(
 ; CHECK-NEXT:    [[BSWAP:%.*]] = shufflevector <2 x i8> [[X:%.*]], <2 x i8> undef, <4 x i32> <i32 1, i32 0, i32 undef, i32 undef>
 ; CHECK-NEXT:    [[CAST:%.*]] = bitcast <4 x i8> [[BSWAP]] to i32
 ; CHECK-NEXT:    ret i32 [[CAST]]
 ;
   %bswap = shufflevector <2 x i8> %x, <2 x i8> undef, <4 x i32> <i32 1, i32 0, i32 undef, i32 undef>
   %cast = bitcast <4 x i8> %bswap to i32
   ret i32 %cast
 }