[SDAG] add demanded bits transform for bswap

[spatel]

Summary

A possible codegen regression for PowerPC is noted in D117406 because we don't recognize a pattern that demands a single byte from a bswap.

I was wondering if we have that fold in IR, and it has existed there since close to the beginning of LLVM:
https://github.com/llvm/llvm-project/blame/main/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp#L794
...so this patch copies that code as much as possible and adapts it for SDAG.

The test for PowerPC that would change in D117406 is over-reduced with undefs, so I recreated it for AArch64 and x86 by passing in pointer args and renamed the values to make the logic clearer.

[xbolva00]

powerpc test missing?

[craig.topper]

The test for PowerPC is over-reduced with undefs, so I made it valid again by passing in pointer args and renamed the values to make the logic clearer.

Is that part of a different patch or a pre-commit? There's no PowerPC test change in this patch.

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1828

Why can't we use getShiftAmountTy always? I fixed the problem that it only works for legal types a few months ago.

[spatel]

In D117508#3249337, @xbolva00 wrote:

powerpc test missing?

Sorry, the description wasn't clear (updated text). The PPC test is in D117406 but not affected here.

[spatel]

Just noticed that MSP430 can infinite loop because it doesn't have legal shifts. This needs a legal op check to be safe.

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1828

Ah, didn't realize that. I copied this from the case ISD::SIGN_EXTEND_INREG block (line 1876). I can update that as a cleanup.

[spatel]

Patch updated:

1. Removed unnecessary shift amount type-checks.
2. Added shift operation legality check (MSP430 tests fail without that).

[craig.topper]

LGTM

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

if (SimplifyDemandedBits(Src, DemandedSrcBits, DemandedElts, Known2, TLO,
                              Depth + 1))
       return true;
     Known.One = Known2.One.reverseBits();
     Known.Zero = Known2.Zero.reverseBits();
     break;
   }
   case ISD::BSWAP: {
     SDValue Src = Op.getOperand(0);
+
+    // If the only bits demanded come from one byte of the bswap result,
+    // just shift the input byte into position to eliminate the bswap.
+    unsigned NLZ = DemandedBits.countLeadingZeros();
+    unsigned NTZ = DemandedBits.countTrailingZeros();
+
+    // Round NTZ down to the next byte.  If we have 11 trailing zeros, then
+    // we need all the bits down to bit 8.  Likewise, round NLZ.  If we
+    // have 14 leading zeros, round to 8.
+    NLZ &= ~7;
+    NTZ &= ~7;
+    // If we need exactly one byte, we can do this transformation.
+    if (BitWidth - NLZ - NTZ == 8) {
+      unsigned ResultBit = NTZ;
+      unsigned InputBit = BitWidth - NTZ - 8;
+
+      // Replace this with either a left or right shift to get the byte into
+      // the right place.
+      unsigned ShiftOpcode = InputBit > ResultBit ? ISD::SRL : ISD::SHL;
+      if (!TLO.LegalOperations() || isOperationLegal(ShiftOpcode, VT)) {
+        EVT ShiftAmtTy = getShiftAmountTy(VT, DL);
+        unsigned ShiftAmount =
+            InputBit > ResultBit ? InputBit - ResultBit : ResultBit - InputBit;
+        SDValue ShAmt = TLO.DAG.getConstant(ShiftAmount, dl, ShiftAmtTy);
+        SDValue NewOp = TLO.DAG.getNode(ShiftOpcode, dl, VT, Src, ShAmt);
+        return TLO.CombineTo(Op, NewOp);
+      }
+    }
+
     APInt DemandedSrcBits = DemandedBits.byteSwap();
     if (SimplifyDemandedBits(Src, DemandedSrcBits, DemandedElts, Known2, TLO,
                              Depth + 1))
       return true;
     Known.One = Known2.One.byteSwap();
     Known.Zero = Known2.Zero.byteSwap();
     break;
   }

llvm/test/CodeGen/AArch64/bswap-known-bits.ll

; CHECK-NEXT:    ret
   %b = call i32 @llvm.bswap.i32(i32 %x)
   %r = trunc i32 %b to i8
   ret i8 %r
 }
 
 define i32 @demand_one_byte1(i32 %x) {
 ; CHECK-LABEL: demand_one_byte1:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    rev w8, w0
+; CHECK-NEXT:    lsr w8, w0, #8
 ; CHECK-NEXT:    and w0, w8, #0xff00
 ; CHECK-NEXT:    ret
   %b = call i32 @llvm.bswap.i32(i32 %x)
   %r = and i32 %b, 65280 ; 0x0000ff00
   ret i32 %r
 }
 
 define i32 @demand_one_byte2(i32 %x) {
 ; CHECK-LABEL: demand_one_byte2:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    rev w8, w0
+; CHECK-NEXT:    lsl w8, w0, #8
 ; CHECK-NEXT:    orr w0, w8, #0xff00ffff
 ; CHECK-NEXT:    ret
   %b = call i32 @llvm.bswap.i32(i32 %x)
   %r = or i32 %b, 4278255615 ; 0xff00ffff
   ret i32 %r
 }
 
 define i64 @demand_one_byte3(i64 %x) {
 ; CHECK-LABEL: demand_one_byte3:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    rev x8, x0
-; CHECK-NEXT:    lsr x0, x8, #56
+; CHECK-NEXT:    and x0, x0, #0xff
 ; CHECK-NEXT:    ret
   %b = call i64 @llvm.bswap.i64(i64 %x)
   %r = lshr i64 %b, 56
   ret i64 %r
 }
 
 define void @demand_one_loaded_byte(i64* %xp, i32* %yp) {
 ; CHECK-LABEL: demand_one_loaded_byte:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldr x8, [x0]
-; CHECK-NEXT:    lsr x8, x8, #8
-; CHECK-NEXT:    rev w8, w8
+; CHECK-NEXT:    ldrb w8, [x0, #4]
 ; CHECK-NEXT:    strb w8, [x1]
 ; CHECK-NEXT:    ret
   %x = load i64, i64* %xp, align 8
   %x_zzzz7654 = lshr i64 %x, 32
   %x_z7654zzz = shl nuw nsw i64 %x_zzzz7654, 24
   %x_4zzz = trunc i64 %x_z7654zzz to i32
   %y = load i32, i32* %yp, align 4
   %y_321z = and i32 %y, -256
   %x_zzz4 = call i32 @llvm.bswap.i32(i32 %x_4zzz)
   %r = or i32 %x_zzz4, %y_321z
   store i32 %r, i32* %yp, align 4
   ret void
 }

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

; X64-NEXT:    retq
   ret i32 %d
 }
 
 define void @demand_one_loaded_byte(i64* %xp, i32* %yp) {
 ; X86-LABEL: demand_one_loaded_byte:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %ecx
-; X86-NEXT:    movl (%ecx), %edx
-; X86-NEXT:    movl 4(%ecx), %ecx
-; X86-NEXT:    shldl $24, %edx, %ecx
-; X86-NEXT:    bswapl %ecx
+; X86-NEXT:    movb 4(%ecx), %cl
 ; X86-NEXT:    movb %cl, (%eax)
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: demand_one_loaded_byte:
 ; X64:       # %bb.0:
-; X64-NEXT:    movq (%rdi), %rax
-; X64-NEXT:    shrq $8, %rax
-; X64-NEXT:    bswapl %eax
+; X64-NEXT:    movb 4(%rdi), %al
 ; X64-NEXT:    movb %al, (%rsi)
 ; X64-NEXT:    retq
   %x = load i64, i64* %xp, align 8
   %x_zzzz7654 = lshr i64 %x, 32
   %x_z7654zzz = shl nuw nsw i64 %x_zzzz7654, 24
   %x_4zzz = trunc i64 %x_z7654zzz to i32
   %y = load i32, i32* %yp, align 4
   %y_321z = and i32 %y, -256
   %x_zzz4 = call i32 @llvm.bswap.i32(i32 %x_4zzz)
   %r = or i32 %x_zzz4, %y_321z
   store i32 %r, i32* %yp, align 4
   ret void
 }