[DAG] Attempt to fold bswap(shl(x,c)) -> zext(bswap(trunc(shl(x,c-bw/2))))

[RKSimon]

Summary

If the shl is at least half the bitwidth (i.e. the lower half of the bswap source is zero), then we can reduce the shift and perform the bswap at half the bitwidth and just zero extend.

I've currently allowed any shift value >= bw/2, but we could limit this to modulo16 so that the shl is always folded away? I'll probably enforce that limit for the InstCombine variant of this fold for PR53867, but I was wondering whether we should be more relaxed in DAG.

Based off PR51391 + PR53867

[RKSimon]

rebase

[spatel]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

9594

Is this intentionally different than the usual "N0.hasOneUse()"?

9599

If I'm seeing it correctly, we need to either tighten this clause by removing !LegalOperations or add the "BW % 16" requirement.

As-is (and we probably want a test either way), we could end up worse. For example on AArch64:

define i32 @test_bswap32_shift17(i32 %a0) {
  %s = shl i32 %a0, 17
  %b = call i32 @llvm.bswap.i32(i32 %s)
  ret i32 %b
}

is:

	lsl	w8, w0, #17
	rev	w0, w8

but this patch would make it:

	lsl	w8, w0, #1
	rev	w8, w8
	lsr	w0, w8, #16

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

9594

I'll change it - I'm just used to writing it that way to handle binops with duplicated ops.

9599

Yeah, I figured allowing any shift amount was probably too far - I'll reduce it to (ShAmt %16) == 0 - its what I've been doing on the InstCombine variant that I'm still working on.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

9602

Should we go through getShiftAmountTy?

[RKSimon]

address feedback

[spatel]

LGTM

[nathanchance]

This patch causes a failed assertion when building the arm64 Linux kernel.

A reduced C reproducer:

struct mlx5_ifc_cmd_hca_cap_bits {
  char gid_table_size[16]
};
enum { MLX5_CAP_GENERAL, MLX5_CAP_NUM };
struct mlx5_hca_cap {
  int cur[0]
};
struct mlx5_core_dev {
  struct {
    struct mlx5_hca_cap *hca[MLX5_CAP_NUM]
  } caps
};
short mlx5_query_hca_port_props_1;
mlx5_query_hca_port___trans_tmp_1;
mlx5_query_hca_port() {
  struct mlx5_core_dev *dev = to_mdev();
  int rep = mlx5_query_hca_port___trans_tmp_1 = __builtin_constant_p(0);
  gid_tbl_len(
      (mlx5_query_hca_port___trans_tmp_1
           ? *dev->caps.hca[MLX5_CAP_GENERAL]->cur << 24 |
                 (*dev->caps.hca[MLX5_CAP_GENERAL]->cur & 16711680) >> 8 |
                 (*dev->caps.hca[MLX5_CAP_GENERAL]->cur & 4278190080) >> 24
           : 0) &
      (1 << sizeof((struct mlx5_ifc_cmd_hca_cap_bits *)0)->gid_table_size) - 1);
  mlx5_query_hca_port_props_1 = rep;
}

$ clang --target=aarch64-linux-gnu -O2 -c -o /dev/null main.i
...
clang: /home/nathan/cbl/src/llvm-project/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp:979: void (anonymous namespace)::SelectionDAGLegalize::LegalizeOp(llvm::SDNode *): Assertion `(TLI.getTypeAction(*DAG.getContext(), Op.getValueType()) == TargetLowering::TypeLegal || Op.getOpcode() == ISD::TargetConstant || Op.getOpcode() == ISD::Register) && "Unexpected illegal type!"' failed.
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 main.i
1.      <eof> parser at end of file
2.      Code generation
3.      Running pass 'Function Pass Manager' on module 'main.i'.
4.      Running pass 'AArch64 Instruction Selection' on function '@mlx5_query_hca_port'
...

A reduced LLVM IR reproducer:

target triple = "aarch64-unknown-linux-gnu"

@mlx5_query_hca_port_props_1 = external global i16

define i32 @mlx5_query_hca_port(i16* %mlx5_query_hca_port_props_1) {
entry:
  %0 = load i32, i32* null, align 4
  %1 = and i32 %0, -65536
  %2 = tail call i32 @llvm.bswap.i32(i32 %1)
  %and16 = zext i32 %2 to i64
  %call17 = tail call i32 bitcast (i32 (...)* @gid_tbl_len to i32 (i64)*)(i64 %and16)
  store i16 0, i16* %mlx5_query_hca_port_props_1, align 4
  ret i32 0
}

declare i32 @to_mdev(...)

declare i32 @gid_tbl_len(...)

; Function Attrs: nofree nosync nounwind readnone speculatable willreturn
declare i32 @llvm.bswap.i32(i32) #0

attributes #0 = { nofree nosync nounwind readnone speculatable willreturn }

$ llc reduced.ll
llc: /home/nathan/cbl/src/llvm-project/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp:979: void (anonymous namespace)::SelectionDAGLegalize::LegalizeOp(llvm::SDNode *): Assertion `(TLI.getTypeAction(*DAG.getContext(), Op.getValueType()) == TargetLowering::TypeLegal || Op.getOpcode() == ISD::TargetConstant || Op.getOpcode() == ISD::Register) && "Unexpected illegal type!"' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.      Program arguments: llc reduced.ll
1.      Running pass 'Function Pass Manager' on module 'reduced.ll'.
2.      Running pass 'AArch64 Instruction Selection' on function '@mlx5_query_hca_port'
...

[RKSimon]

@nathanchance Thanks for the repro, I'll get this fixed over the weekend

[RKSimon]

@nathanchance Please can you confirm if rGfadd20f80d69 fixes your build?

[nathanchance]

In D120192#3348461, @RKSimon wrote:

@nathanchance Please can you confirm if rGfadd20f80d69 fixes your build?

Yes, from my preliminary tests, everything appears okay with that fix! Thank you for the quick turnaround.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitBSWAP(SDNode *N) {
   // isn't supported, it will be expanded to bswap followed by a manual reversal
   // of bits in each byte. By placing bswaps before bitreverse, we can remove
   // the two bswaps if the bitreverse gets expanded.
   if (N0.getOpcode() == ISD::BITREVERSE && N0.hasOneUse()) {
     SDValue BSwap = DAG.getNode(ISD::BSWAP, DL, VT, N0.getOperand(0));
     return DAG.getNode(ISD::BITREVERSE, DL, VT, BSwap);
   }
 
+  // fold (bswap shl(x,c)) -> (zext(bswap(trunc(shl(x,sub(c,bw/2))))))
+  // iff x >= bw/2 (i.e. lower half is known zero)
+  unsigned BW = VT.getScalarSizeInBits();
+  if (BW >= 32 && N0.getOpcode() == ISD::SHL && N->isOnlyUserOf(N0.getNode())) {
+    auto *ShAmt = dyn_cast<ConstantSDNode>(N0.getOperand(1));
+    EVT HalfVT = EVT::getIntegerVT(*DAG.getContext(), BW / 2);
+    if (ShAmt && ShAmt->getAPIntValue().uge(BW / 2) &&
+        TLI.isTruncateFree(VT, HalfVT) &&
+        (!LegalOperations || hasOperation(ISD::BSWAP, HalfVT))) {
+      SDValue Res = N0.getOperand(0);
+      if (uint64_t NewShAmt = (ShAmt->getZExtValue() - (BW / 2))) {
+        EVT ShiftAmtVT = N0.getOperand(1).getValueType();
+        Res = DAG.getNode(ISD::SHL, DL, VT, Res,
+                          DAG.getConstant(NewShAmt, DL, ShiftAmtVT));
+      }
+      Res = DAG.getZExtOrTrunc(Res, DL, HalfVT);
+      Res = DAG.getNode(ISD::BSWAP, DL, HalfVT, Res);
+      return DAG.getZExtOrTrunc(Res, DL, VT);
+    }
+  }
+
   return SDValue();
 }
 
 SDValue DAGCombiner::visitBITREVERSE(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   EVT VT = N->getValueType(0);
 
   // fold (bitreverse c1) -> c2

llvm/test/CodeGen/AArch64/load-combine-big-endian.ll

; CHECK-NEXT:    ret
   ret i32 %tmp48
 }
 ; i8* p; // p is 2 byte aligned
 ; (i32) p[0] | ((i32) p[1] << 8)
 define i32 @zext_load_i32_by_i8(i32* %arg) {
 ; CHECK-LABEL: zext_load_i32_by_i8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldrh w8, [x0]
-; CHECK-NEXT:    lsl w8, w8, #16
-; CHECK-NEXT:    rev w0, w8
+; CHECK-NEXT:    rev w8, w8
+; CHECK-NEXT:    lsr w0, w8, #16
 ; CHECK-NEXT:    ret
   %tmp = bitcast i32* %arg to i8*
   %tmp1 = getelementptr inbounds i8, i8* %tmp, i32 0
   %tmp2 = load i8, i8* %tmp1, align 2
   %tmp3 = zext i8 %tmp2 to i32
   %tmp4 = getelementptr inbounds i8, i8* %tmp, i32 1
   %tmp5 = load i8, i8* %tmp4, align 1
   %tmp6 = zext i8 %tmp5 to i32

llvm/test/CodeGen/AArch64/load-combine.ll

; CHECK-NEXT:    ret
   ret i32 %tmp8
 }
 ; i8* p; // p is 2 byte aligned
 ; (i32) p[1] | ((i32) p[0] << 8)
 define i32 @zext_load_i32_by_i8_bswap(i32* %arg) {
 ; CHECK-LABEL: zext_load_i32_by_i8_bswap:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldrh w8, [x0]
-; CHECK-NEXT:    lsl w8, w8, #16
-; CHECK-NEXT:    rev w0, w8
+; CHECK-NEXT:    rev w8, w8
+; CHECK-NEXT:    lsr w0, w8, #16
 ; CHECK-NEXT:    ret
 
   %tmp = bitcast i32* %arg to i8*
   %tmp1 = getelementptr inbounds i8, i8* %tmp, i32 1
   %tmp2 = load i8, i8* %tmp1, align 1
   %tmp3 = zext i8 %tmp2 to i32
   %tmp4 = getelementptr inbounds i8, i8* %tmp, i32 0
   %tmp5 = load i8, i8* %tmp4, align 2

llvm/test/CodeGen/X86/combine-bswap.ll

; X64-NEXT:    retq
   %r = or i32 %x_zzz4, %y_321z
   store i32 %r, i32* %yp, align 4
   ret void
 }
 
 define i64 @test_bswap64_shift48_zext(i16 %a0) {
 ; X86-LABEL: test_bswap64_shift48_zext:
 ; X86:       # %bb.0:
-; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; X86-NEXT:    shll $16, %eax
-; X86-NEXT:    bswapl %eax
+; X86-NEXT:    movzwl {{[0-9]+}}(%esp), %eax
+; X86-NEXT:    rolw $8, %ax
+; X86-NEXT:    movzwl %ax, %eax
 ; X86-NEXT:    xorl %edx, %edx
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: test_bswap64_shift48_zext:
 ; X64:       # %bb.0:
-; X64-NEXT:    movl %edi, %eax
-; X64-NEXT:    shlq $48, %rax
-; X64-NEXT:    bswapq %rax
+; X64-NEXT:    rolw $8, %di
+; X64-NEXT:    movzwl %di, %eax
 ; X64-NEXT:    retq
   %z = zext i16 %a0 to i64
   %s = shl i64 %z, 48
   %b = call i64 @llvm.bswap.i64(i64 %s)
   ret i64 %b
 }
 
 define i64 @test_bswap64_shift48(i64 %a0) {
 ; X86-LABEL: test_bswap64_shift48:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; X86-NEXT:    shll $16, %eax
-; X86-NEXT:    bswapl %eax
+; X86-NEXT:    rolw $8, %ax
+; X86-NEXT:    movzwl %ax, %eax
 ; X86-NEXT:    xorl %edx, %edx
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: test_bswap64_shift48:
 ; X64:       # %bb.0:
-; X64-NEXT:    movq %rdi, %rax
-; X64-NEXT:    shlq $48, %rax
-; X64-NEXT:    bswapq %rax
+; X64-NEXT:    rolw $8, %di
+; X64-NEXT:    movzwl %di, %eax
 ; X64-NEXT:    retq
   %s = shl i64 %a0, 48
   %b = call i64 @llvm.bswap.i64(i64 %s)
   ret i64 %b
 }
 
 define i32 @test_bswap32_shift17(i32 %a0) {
 ; X86-LABEL: test_bswap32_shift17:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; X86-NEXT:    shll $17, %eax
-; X86-NEXT:    bswapl %eax
+; X86-NEXT:    addl %eax, %eax
+; X86-NEXT:    rolw $8, %ax
+; X86-NEXT:    movzwl %ax, %eax
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: test_bswap32_shift17:
 ; X64:       # %bb.0:
-; X64-NEXT:    movl %edi, %eax
-; X64-NEXT:    shll $17, %eax
-; X64-NEXT:    bswapl %eax
+; X64-NEXT:    addl %edi, %edi
+; X64-NEXT:    rolw $8, %di
+; X64-NEXT:    movzwl %di, %eax
 ; X64-NEXT:    retq
   %s = shl i32 %a0, 17
   %b = call i32 @llvm.bswap.i32(i32 %s)
   ret i32 %b
 }
 
-; negative test
+; negative test (x86 can fold once bswap64 is split to 2 x bswap32)
 define i64 @test_bswap64_shift17(i64 %a0) {
 ; X86-LABEL: test_bswap64_shift17:
 ; X86:       # %bb.0:
-; X86-NEXT:    movl {{[0-9]+}}(%esp), %edx
+; X86-NEXT:    movl {{[0-9]+}}(%esp), %ecx
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; X86-NEXT:    shldl $17, %edx, %eax
-; X86-NEXT:    shll $17, %edx
+; X86-NEXT:    shldl $17, %ecx, %eax
 ; X86-NEXT:    bswapl %eax
-; X86-NEXT:    bswapl %edx
+; X86-NEXT:    addl %ecx, %ecx
+; X86-NEXT:    rolw $8, %cx
+; X86-NEXT:    movzwl %cx, %edx
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: test_bswap64_shift17:
 ; X64:       # %bb.0:
 ; X64-NEXT:    movq %rdi, %rax
 ; X64-NEXT:    shlq $17, %rax
 ; X64-NEXT:    bswapq %rax
 ; X64-NEXT:    retq

llvm/test/CodeGen/X86/load-combine.ll

; CHECK64-NEXT:    retq
   %tmp7 = shl nuw nsw i32 %tmp6, 24
   %tmp8 = or i32 %tmp7, %tmp30
   ret i32 %tmp8
 }
 
 ; i8* p;
 ; (i32) p[1] | ((i32) p[0] << 8)
 define i32 @zext_load_i32_by_i8_bswap(i32* %arg) {
-; CHECK-LABEL: zext_load_i32_by_i8_bswap:
-; CHECK:       # %bb.0:
-; CHECK-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; CHECK-NEXT:    movzwl (%eax), %eax
-; CHECK-NEXT:    shll $16, %eax
-; CHECK-NEXT:    bswapl %eax
-; CHECK-NEXT:    retl
+; BSWAP-LABEL: zext_load_i32_by_i8_bswap:
+; BSWAP:       # %bb.0:
+; BSWAP-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; BSWAP-NEXT:    movzwl (%eax), %eax
+; BSWAP-NEXT:    rolw $8, %ax
+; BSWAP-NEXT:    movzwl %ax, %eax
+; BSWAP-NEXT:    retl
 ;
-; CHECK64-LABEL: zext_load_i32_by_i8_bswap:
-; CHECK64:       # %bb.0:
-; CHECK64-NEXT:    movzwl (%rdi), %eax
-; CHECK64-NEXT:    shll $16, %eax
-; CHECK64-NEXT:    bswapl %eax
-; CHECK64-NEXT:    retq
+; MOVBE-LABEL: zext_load_i32_by_i8_bswap:
+; MOVBE:       # %bb.0:
+; MOVBE-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; MOVBE-NEXT:    movbew (%eax), %ax
+; MOVBE-NEXT:    movzwl %ax, %eax
+; MOVBE-NEXT:    retl
+;
+; BSWAP64-LABEL: zext_load_i32_by_i8_bswap:
+; BSWAP64:       # %bb.0:
+; BSWAP64-NEXT:    movzwl (%rdi), %eax
+; BSWAP64-NEXT:    rolw $8, %ax
+; BSWAP64-NEXT:    movzwl %ax, %eax
+; BSWAP64-NEXT:    retq
+;
+; MOVBE64-LABEL: zext_load_i32_by_i8_bswap:
+; MOVBE64:       # %bb.0:
+; MOVBE64-NEXT:    movbew (%rdi), %ax
+; MOVBE64-NEXT:    movzwl %ax, %eax
+; MOVBE64-NEXT:    retq
   %tmp = bitcast i32* %arg to i8*
   %tmp1 = getelementptr inbounds i8, i8* %tmp, i32 1
   %tmp2 = load i8, i8* %tmp1, align 1
   %tmp3 = zext i8 %tmp2 to i32
   %tmp4 = getelementptr inbounds i8, i8* %tmp, i32 0
   %tmp5 = load i8, i8* %tmp4, align 1
   %tmp6 = zext i8 %tmp5 to i32
   %tmp7 = shl nuw nsw i32 %tmp6, 8