[InstCombine] try to canonicalize logical shift after bswap

[spatel]

Summary

When shifting by a byte-multiple:
bswap (shl X, C) --> lshr (bswap X), C
bswap (lshr X, C) --> shl (bswap X), C

This is an IR implementation of a transform suggested in D120648. The "swaps cancel" test models the motivating optimization from that proposal.

Alive2 checks (as noted in the other review, we could use knownbits to handle shift-by-variable-amount, but that can be an enhancement patch):
https://alive2.llvm.org/ce/z/pXUaRf
https://alive2.llvm.org/ce/z/ZnaMLf

[RKSimon]

I haven't had chance to look at https://github.com/llvm/llvm-project/issues/53867 yet but this has a lot in common with it (and feel free to grab it if you have the bandwidth) - please can you add test coverage for that as well?

[spatel]

In D122010#3392573, @RKSimon wrote:

I haven't had chance to look at https://github.com/llvm/llvm-project/issues/53867 yet but this has a lot in common with it (and feel free to grab it if you have the bandwidth) - please can you add test coverage for that as well?

Sure - if we canonicalize in this direction, then that problem won't be solved by this patch, but I think it becomes a simpler fold (we don't have to worry about the trailing trunc/demanded bits):
https://alive2.llvm.org/ce/z/zC5Gez

[RKSimon]

LGTM with one (optional) minor

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1360

Worth using KnownBits to allow non-uniform shift amounts?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1360

I'll put a TODO comment on it. I haven't seen a motivating case for variable or non-splat shifts yet.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

if (Instruction *NewMinMax = factorizeMinMaxTree(II))
        return NewMinMax;
 
     break;
   }
   case Intrinsic::bswap: {
     Value *IIOperand = II->getArgOperand(0);
     Value *X = nullptr;
 
+    // Try to canonicalize bswap-of-logical-shift-by-8-bit-multiple as
+    // inverse-shift-of-bswap:
+    // bswap (shl X, C) --> lshr (bswap X), C
+    // bswap (lshr X, C) --> shl (bswap X), C
+    // TODO: Use knownbits to allow variable shift and non-splat vector match.
+    BinaryOperator *BO;
+    if (match(IIOperand, m_OneUse(m_BinOp(BO)))) {
+      const APInt *C;
+      if (match(BO, m_LogicalShift(m_Value(X), m_APIntAllowUndef(C))) &&
+          (*C & 7) == 0) {
+        Value *NewSwap = Builder.CreateUnaryIntrinsic(Intrinsic::bswap, X);
+        BinaryOperator::BinaryOps InverseShift =
+            BO->getOpcode() == Instruction::Shl ? Instruction::LShr
+                                                : Instruction::Shl;
+        return BinaryOperator::Create(InverseShift, NewSwap, BO->getOperand(1));
+      }
+    }
+
     KnownBits Known = computeKnownBits(IIOperand, 0, II);
     uint64_t LZ = alignDown(Known.countMinLeadingZeros(), 8);
     uint64_t TZ = alignDown(Known.countMinTrailingZeros(), 8);
     unsigned BW = Known.getBitWidth();
 
     // bswap(x) -> shift(x) if x has exactly one "active byte"
     if (BW - LZ - TZ == 8) {
       assert(LZ != TZ && "active byte cannot be in the middle");

llvm/test/Transforms/InstCombine/bswap-fold.ll

 ;
   %t2 = call i32 @llvm.bswap.i32( i32 %a )
   %t4 = and i32 %t2, 255
   ret i32 %t4
 }
 
 define i32 @lshr8_i32(i32 %x) {
 ; CHECK-LABEL: @lshr8_i32(
-; CHECK-NEXT:    [[S:%.*]] = lshr i32 [[X:%.*]], 8
-; CHECK-NEXT:    [[R:%.*]] = call i32 @llvm.bswap.i32(i32 [[S]])
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.bswap.i32(i32 [[X:%.*]])
+; CHECK-NEXT:    [[R:%.*]] = shl i32 [[TMP1]], 8
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = lshr i32 %x, 8
   %r = call i32 @llvm.bswap.i32(i32 %s)
   ret i32 %r
 }
 
 define <2 x i32> @lshr16_v2i32(<2 x i32> %x) {
 ; CHECK-LABEL: @lshr16_v2i32(
-; CHECK-NEXT:    [[S:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 16, i32 16>
-; CHECK-NEXT:    [[R:%.*]] = call <2 x i32> @llvm.bswap.v2i32(<2 x i32> [[S]])
+; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i32> @llvm.bswap.v2i32(<2 x i32> [[X:%.*]])
+; CHECK-NEXT:    [[R:%.*]] = shl <2 x i32> [[TMP1]], <i32 16, i32 16>
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %s = lshr <2 x i32> %x, <i32 16, i32 16>
   %r = call <2 x i32> @llvm.bswap.v2i32(<2 x i32> %s)
   ret <2 x i32> %r
 }
 
 define i32 @lshr24_i32(i32 %x) {
 ; CHECK-LABEL: @lshr24_i32(
-; CHECK-NEXT:    [[S:%.*]] = and i32 [[X:%.*]], -16777216
-; CHECK-NEXT:    ret i32 [[S]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], -16777216
+; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %s = lshr i32 %x, 24
   %r = call i32 @llvm.bswap.i32(i32 %s)
   ret i32 %r
 }
 
+; negative test - need shift-by-8-bit-multiple
+
 define i32 @lshr12_i32(i32 %x) {
 ; CHECK-LABEL: @lshr12_i32(
 ; CHECK-NEXT:    [[S:%.*]] = lshr i32 [[X:%.*]], 12
 ; CHECK-NEXT:    [[R:%.*]] = call i32 @llvm.bswap.i32(i32 [[S]])
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = lshr i32 %x, 12
   %r = call i32 @llvm.bswap.i32(i32 %s)
   ret i32 %r
 }
 
+; negative test - uses
+
 define i32 @lshr8_i32_use(i32 %x, i32* %p) {
 ; CHECK-LABEL: @lshr8_i32_use(
 ; CHECK-NEXT:    [[S:%.*]] = lshr i32 [[X:%.*]], 12
 ; CHECK-NEXT:    store i32 [[S]], i32* [[P:%.*]], align 4
 ; CHECK-NEXT:    [[R:%.*]] = call i32 @llvm.bswap.i32(i32 [[S]])
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = lshr i32 %x, 12
   store i32 %s, i32* %p
   %r = call i32 @llvm.bswap.i32(i32 %s)
   ret i32 %r
 }
 
 define i64 @shl16_i64(i64 %x) {
 ; CHECK-LABEL: @shl16_i64(
-; CHECK-NEXT:    [[S:%.*]] = shl i64 [[X:%.*]], 16
-; CHECK-NEXT:    [[R:%.*]] = call i64 @llvm.bswap.i64(i64 [[S]])
+; CHECK-NEXT:    [[TMP1:%.*]] = call i64 @llvm.bswap.i64(i64 [[X:%.*]])
+; CHECK-NEXT:    [[R:%.*]] = lshr i64 [[TMP1]], 16
 ; CHECK-NEXT:    ret i64 [[R]]
 ;
   %s = shl i64 %x, 16
   %r = call i64 @llvm.bswap.i64(i64 %s)
   ret i64 %r
 }
 
+; poison vector element propagates
+
 define <2 x i64> @shl16_v2i64(<2 x i64> %x) {
 ; CHECK-LABEL: @shl16_v2i64(
-; CHECK-NEXT:    [[S:%.*]] = shl <2 x i64> [[X:%.*]], <i64 poison, i64 24>
-; CHECK-NEXT:    [[R:%.*]] = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> [[S]])
+; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> [[X:%.*]])
+; CHECK-NEXT:    [[R:%.*]] = lshr <2 x i64> [[TMP1]], <i64 poison, i64 24>
 ; CHECK-NEXT:    ret <2 x i64> [[R]]
 ;
   %s = shl <2 x i64> %x, <i64 poison, i64 24>
   %r = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> %s)
   ret <2 x i64> %r
 }
 
 define i64 @shl56_i64(i64 %x) {
 ; CHECK-LABEL: @shl56_i64(
-; CHECK-NEXT:    [[S:%.*]] = and i64 [[X:%.*]], 255
-; CHECK-NEXT:    ret i64 [[S]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and i64 [[X:%.*]], 255
+; CHECK-NEXT:    ret i64 [[TMP1]]
 ;
   %s = shl i64 %x, 56
   %r = call i64 @llvm.bswap.i64(i64 %s)
   ret i64 %r
 }
 
+; negative test - need shift-by-8-bit-multiple
+
 define i64 @shl42_i64(i64 %x) {
 ; CHECK-LABEL: @shl42_i64(
 ; CHECK-NEXT:    [[S:%.*]] = shl i64 [[X:%.*]], 42
 ; CHECK-NEXT:    [[R:%.*]] = call i64 @llvm.bswap.i64(i64 [[S]])
 ; CHECK-NEXT:    ret i64 [[R]]
 ;
   %s = shl i64 %x, 42
   %r = call i64 @llvm.bswap.i64(i64 %s)
   ret i64 %r
 }
 
+; negative test - uses
+
 define i32 @shl8_i32_use(i32 %x, i32* %p) {
 ; CHECK-LABEL: @shl8_i32_use(
 ; CHECK-NEXT:    [[S:%.*]] = shl i32 [[X:%.*]], 8
 ; CHECK-NEXT:    store i32 [[S]], i32* [[P:%.*]], align 4
 ; CHECK-NEXT:    [[R:%.*]] = call i32 @llvm.bswap.i32(i32 [[S]])
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %s = shl i32 %x, 8
   store i32 %s, i32* %p
   %r = call i32 @llvm.bswap.i32(i32 %s)
   ret i32 %r
 }
 
+; swaps cancel
+
 define i64 @swap_shl16_i64(i64 %x) {
 ; CHECK-LABEL: @swap_shl16_i64(
-; CHECK-NEXT:    [[B:%.*]] = call i64 @llvm.bswap.i64(i64 [[X:%.*]])
-; CHECK-NEXT:    [[S:%.*]] = shl i64 [[B]], 16
-; CHECK-NEXT:    [[R:%.*]] = call i64 @llvm.bswap.i64(i64 [[S]])
+; CHECK-NEXT:    [[R:%.*]] = lshr i64 [[X:%.*]], 16
 ; CHECK-NEXT:    ret i64 [[R]]
 ;
   %b = call i64 @llvm.bswap.i64(i64 %x)
   %s = shl i64 %b, 16
   %r = call i64 @llvm.bswap.i64(i64 %s)
   ret i64 %r
 }
 

 ; CHECK-NEXT:    [[TMP3:%.*]] = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> [[TMP2]])
 ; CHECK-NEXT:    ret <2 x i64> [[TMP3]]
 ;
   %2 = shl <2 x i64> %0, <i64 56, i64 55>  ; second elem has 9 active high bits
   %3 = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> %2)
   ret <2 x i64> %3
 }
 
-; negative test
+; TODO: This should fold to 'and'.
 define <2 x i64> @bs_active_high_undef(<2 x i64> %0) {
 ; CHECK-LABEL: @bs_active_high_undef(
-; CHECK-NEXT:    [[TMP2:%.*]] = shl <2 x i64> [[TMP0:%.*]], <i64 56, i64 undef>
-; CHECK-NEXT:    [[TMP3:%.*]] = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> [[TMP2]])
+; CHECK-NEXT:    [[TMP2:%.*]] = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> [[TMP0:%.*]])
+; CHECK-NEXT:    [[TMP3:%.*]] = lshr <2 x i64> [[TMP2]], <i64 56, i64 undef>
 ; CHECK-NEXT:    ret <2 x i64> [[TMP3]]
 ;
   %2 = shl <2 x i64> %0, <i64 56, i64 undef>
   %3 = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> %2)
   ret <2 x i64> %3
 }
 
 define i64 @bs_active_high8_multiuse(i64 %0) {