[InstCombine] try to narrow shifted bswap-of-zext

[spatel]

Summary

This is the IR counterpart to 370ebc9d9a573d6 which provided a bswap narrowing fix for issue #53867.

Here we can be more general (although I'm not sure yet what would happen for illegal types in codegen - too rare to worry about?):
https://alive2.llvm.org/ce/z/3-CPfo

This assumes that we have moved the shift after the bswap as proposed in D122010.

[RKSimon]

LGTM

[nathanchance]

This patch causes a crash while building the Linux kernel for arm64. A simplifier C reproducer:

$ cat csio_lnode.i
struct csio_lnode {
  unsigned nport_id
};
csio_ln_vnp_read_cbfn_mbp_0;
csio_ln_vnp_read_cbfn_rsp_0() {
  struct csio_lnode *ln = csio_ln_vnp_read_cbfn_mbp_0;
  int nport_id;
  memcpy(&nport_id, csio_ln_vnp_read_cbfn_rsp_0, 3);
  ln->nport_id =
      __builtin_constant_p(0) ? nport_id << 24 | (nport_id & 83040) >> 8 : 0;
  ln->nport_id = ln->nport_id >> 8;
}

$ clang --target=aarch64-linux-gnu -O2 -c -o /dev/null csio_lnode.i
...
bswap must be an even number of bytes
  %3 = tail call i24 @llvm.bswap.i24(i24 %2)
in function csio_ln_vnp_read_cbfn_rsp_0
fatal error: error in backend: Broken function found, compilation aborted!
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace, preprocessed source, and associated run script.
Stack dump:
0.      Program arguments: clang --target=aarch64-linux-gnu -O2 -c -o /dev/null csio_lnode.i
1.      <eof> parser at end of file
2.      Code generation
3.      Running pass 'Function Pass Manager' on module 'csio_lnode.i'.
4.      Running pass 'Module Verifier' on function '@csio_ln_vnp_read_cbfn_rsp_0'
...

[nathanchance]

I have reverted this in 4e0008dcbe9fce99b9727e8bbeb129efc7bf2d80 to hopefully keep our CI from going completely red. I am happy to test a new revision as needed.

[spatel]

In D122166#3401171, @nathanchance wrote:

I have reverted this in 4e0008dcbe9fce99b9727e8bbeb129efc7bf2d80 to hopefully keep our CI from going completely red. I am happy to test a new revision as needed.

Thanks for the revert! I missed a check on the source type size, so we tried to create a bogus bswap of i24. I'll add a regression test based on the provided source file and update the patch.

[DavidSpickett]

We (Linaro) also saw the same crash when building for Arm.

[spatel]

In D122166#3405357, @DavidSpickett wrote:

We (Linaro) also saw the same crash when building for Arm.

Thanks - I got an automated mail from your tester a few minutes ago. Please let me know if you see any problems with the updated version of the patch ( 0fcff69bcb3d ).

[DavidSpickett]

Oh good to know, I wasn't sure if we were sending reports externally.

Thanks for the fix, I'll keep you posted. (and you can always reply to the email if you have questions)

[DavidSpickett]

Oh and you might get some more mails as it runs through different configurations. Again I'll check if they're new or not.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

if (match(Op1, m_APInt(C))) {
     // lshr i32 (mul nuw X, Pow2+1), 16 --> and X, Pow2-1
     // TODO: Generalize to allow more than just half-width shifts?
     const APInt *MulC;
     if (match(Op0, m_NUWMul(m_Value(X), m_APInt(MulC))) &&
         ShAmtC * 2 == BitWidth && (*MulC - 1).isPowerOf2() &&
         MulC->logBase2() == ShAmtC)
       return BinaryOperator::CreateAnd(X, ConstantInt::get(Ty, *MulC - 2));
 
+    // Try to narrow a bswap:
+    // (bswap (zext X)) >> C --> zext (bswap X >> C')
+    // In the case where the shift amount equals the bitwidth difference, the
+    // shift is eliminated.
+    if (match(Op0, m_OneUse(m_Intrinsic<Intrinsic::bswap>(
+                       m_OneUse(m_ZExt(m_Value(X))))))) {
+      // TODO: If the shift amount is less than the zext, we could shift left.
+      unsigned SrcWidth = X->getType()->getScalarSizeInBits();
+      if (ShAmtC >= BitWidth - SrcWidth) {
+        Value *NarrowSwap = Builder.CreateUnaryIntrinsic(Intrinsic::bswap, X);
+        Value *NewShift =
+            Builder.CreateLShr(NarrowSwap, ShAmtC - (BitWidth - SrcWidth));
+        return new ZExtInst(NewShift, Ty);
+      }
+    }
+
     // If the shifted-out value is known-zero, then this is an exact shift.
     if (!I.isExact() &&
         MaskedValueIsZero(Op0, APInt::getLowBitsSet(BitWidth, ShAmtC), 0, &I)) {
       I.setIsExact();
       return &I;
     }
   }
 

llvm/test/Transforms/InstCombine/lshr.ll

 ;
   %x.shifted = lshr i32 %x, %y
   %cmp = icmp sge i32 %x.shifted, %x
   ret i1 %cmp
 }
 
 define i32 @narrow_bswap(i16 %x) {
 ; CHECK-LABEL: @narrow_bswap(
-; CHECK-NEXT:    [[Z:%.*]] = zext i16 [[X:%.*]] to i32
-; CHECK-NEXT:    [[B:%.*]] = call i32 @llvm.bswap.i32(i32 [[Z]])
-; CHECK-NEXT:    [[S:%.*]] = lshr exact i32 [[B]], 16
+; CHECK-NEXT:    [[TMP1:%.*]] = call i16 @llvm.bswap.i16(i16 [[X:%.*]])
+; CHECK-NEXT:    [[S:%.*]] = zext i16 [[TMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %z = zext i16 %x to i32
   %b = call i32 @llvm.bswap.i32(i32 %z)
   %s = lshr i32 %b, 16
   ret i32 %s
 }
 
 define i128 @narrow_bswap_extra_wide(i16 %x) {
 ; CHECK-LABEL: @narrow_bswap_extra_wide(
-; CHECK-NEXT:    [[Z:%.*]] = zext i16 [[X:%.*]] to i128
-; CHECK-NEXT:    [[B:%.*]] = call i128 @llvm.bswap.i128(i128 [[Z]])
-; CHECK-NEXT:    [[S:%.*]] = lshr exact i128 [[B]], 112
+; CHECK-NEXT:    [[TMP1:%.*]] = call i16 @llvm.bswap.i16(i16 [[X:%.*]])
+; CHECK-NEXT:    [[S:%.*]] = zext i16 [[TMP1]] to i128
 ; CHECK-NEXT:    ret i128 [[S]]
 ;
   %z = zext i16 %x to i128
   %b = call i128 @llvm.bswap.i128(i128 %z)
   %s = lshr i128 %b, 112
   ret i128 %s
 }
 
+; TODO: The bswap can be narrowed followed by shl.
+
 define i32 @narrow_bswap_undershift(i16 %x) {
 ; CHECK-LABEL: @narrow_bswap_undershift(
 ; CHECK-NEXT:    [[Z:%.*]] = zext i16 [[X:%.*]] to i32
 ; CHECK-NEXT:    [[B:%.*]] = call i32 @llvm.bswap.i32(i32 [[Z]])
 ; CHECK-NEXT:    [[S:%.*]] = lshr exact i32 [[B]], 8
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %z = zext i16 %x to i32
   %b = call i32 @llvm.bswap.i32(i32 %z)
   %s = lshr i32 %b, 8
   ret i32 %s
 }
 
 define <2 x i64> @narrow_bswap_splat(<2 x i16> %x) {
 ; CHECK-LABEL: @narrow_bswap_splat(
-; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i16> [[X:%.*]] to <2 x i64>
-; CHECK-NEXT:    [[B:%.*]] = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> [[Z]])
-; CHECK-NEXT:    [[S:%.*]] = lshr exact <2 x i64> [[B]], <i64 48, i64 48>
+; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i16> @llvm.bswap.v2i16(<2 x i16> [[X:%.*]])
+; CHECK-NEXT:    [[S:%.*]] = zext <2 x i16> [[TMP1]] to <2 x i64>
 ; CHECK-NEXT:    ret <2 x i64> [[S]]
 ;
   %z = zext <2 x i16> %x to <2 x i64>
   %b = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> %z)
   %s = lshr <2 x i64> %b, <i64 48, i64 48>
   ret <2 x i64> %s
 }
 
+; TODO: poison/undef in the shift amount is ok to propagate.
+
 define <2 x i64> @narrow_bswap_splat_poison_elt(<2 x i16> %x) {
 ; CHECK-LABEL: @narrow_bswap_splat_poison_elt(
 ; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i16> [[X:%.*]] to <2 x i64>
 ; CHECK-NEXT:    [[B:%.*]] = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> [[Z]])
 ; CHECK-NEXT:    [[S:%.*]] = lshr <2 x i64> [[B]], <i64 48, i64 poison>
 ; CHECK-NEXT:    ret <2 x i64> [[S]]
 ;
   %z = zext <2 x i16> %x to <2 x i64>
   %b = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> %z)
   %s = lshr <2 x i64> %b, <i64 48, i64 poison>
   ret <2 x i64> %s
 }
 
 define <2 x i64> @narrow_bswap_overshift(<2 x i32> %x) {
 ; CHECK-LABEL: @narrow_bswap_overshift(
-; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i32> [[X:%.*]] to <2 x i64>
-; CHECK-NEXT:    [[B:%.*]] = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> [[Z]])
-; CHECK-NEXT:    [[S:%.*]] = lshr <2 x i64> [[B]], <i64 48, i64 48>
+; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i32> @llvm.bswap.v2i32(<2 x i32> [[X:%.*]])
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr <2 x i32> [[TMP1]], <i32 16, i32 16>
+; CHECK-NEXT:    [[S:%.*]] = zext <2 x i32> [[TMP2]] to <2 x i64>
 ; CHECK-NEXT:    ret <2 x i64> [[S]]
 ;
   %z = zext <2 x i32> %x to <2 x i64>
   %b = call <2 x i64> @llvm.bswap.v2i64(<2 x i64> %z)
   %s = lshr <2 x i64> %b, <i64 48, i64 48>
   ret <2 x i64> %s
 }
 
 define i128 @narrow_bswap_overshift2(i96 %x) {
 ; CHECK-LABEL: @narrow_bswap_overshift2(
-; CHECK-NEXT:    [[Z:%.*]] = zext i96 [[X:%.*]] to i128
-; CHECK-NEXT:    [[B:%.*]] = call i128 @llvm.bswap.i128(i128 [[Z]])
-; CHECK-NEXT:    [[S:%.*]] = lshr i128 [[B]], 61
+; CHECK-NEXT:    [[TMP1:%.*]] = call i96 @llvm.bswap.i96(i96 [[X:%.*]])
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr i96 [[TMP1]], 29
+; CHECK-NEXT:    [[S:%.*]] = zext i96 [[TMP2]] to i128
 ; CHECK-NEXT:    ret i128 [[S]]
 ;
   %z = zext i96 %x to i128
   %b = call i128 @llvm.bswap.i128(i128 %z)
   %s = lshr i128 %b, 61
   ret i128 %s
 }