[DAGCombine] Improve ReduceLoadWidth for SRL

[samparker]

Summary

If the SRL node is only used by an AND, we may be able to set the ExtVT to the width of the mask, making the AND redundant. To support this, another check has been added in isLegalNarrowLoad which queries whether the load is valid.

I've had to change a couple of X86 codegen tests, but I really don't know if they're correct. @RKSimon, if you could check these that would be great.

[niravd]

All of the current tests look good to me. I think you should add some tests for non-byte shifts.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

3793 ↗

(On Diff #127344)

What if ShAmt isn't a multiple of 8?

[samparker]

Good point, I will add some tests for that.

[samparker]

Added some tests to handle non-byte shifts and masks, which highlighted that those cases weren't handled! This fix means that I've reverted one of the x86 test changes too.

[RKSimon]

test/CodeGen/X86/h-registers-1.ll

18 ↗

(On Diff #127380)

These checks aren't testing the zero-extension. I've updated the test at rL321003 - please can you rebase and then regenerate the test with the update_llc_test_checks.py script to check the diff?

[samparker]

Rebased and update the X86 test, plus I've added a couple of other ARM tests.

[RKSimon]

x86 codegen looks OK, one minor but no other comments from me.

test/CodeGen/X86/h-registers-1.ll

76 ↗

(On Diff #127485)

For clarity please can you put this back on the previous line with the other arguments?

[niravd]

LGTM modulo rksimon's comment.

[samparker]

Will do, Thanks for the review!

[samparker]

I kindly got a reproducer from @nemanjai so now I've, hopefully, sorted out the issue that was causing big endian stage two builders to fail.

[samparker]

The trunc_i64_mask_srl test has been added. This example showed that I had a typo in my patch, I was checking the user of N and not 'N0'! So I've also renamed some variables to reduce the chance of possible confusion again.

[nemanjai]

LGTM. I have tested this on PPC64BE bootstrap and the previous failures are gone. Also, can you please add the test case I sent you as a part of this patch? Clearly, we had missing coverage there.

[samparker]

Thanks. Yes, I will do.

[gchatelet]

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

3841

I stumbled upon this code while debugging a regression in D81379.
AFAICT ShAmt (ShiftAmount?) can be a non power of two, especially when using the sanitizers.
For instance, ShAmt can be 48, it is a round number of byte but not a valid alignment 48/8==6.

I'm not quite sure what is the expected fix so I'm commenting here to get some feedback.
It seems to me that it should be the natural alignment for ExtVT possibly constrained by the original LoadN alignment.
Something along those lines:

if (ShAmt) {
  assert(ShAmt % 8 == 0 && "ShAmt is byte offset");
  const unsigned ByteShAmt = ShAmt / 8;
  const Align LDSTAlign = LDST->getAlign();
  const Align NarrowAlign = commonAlignment(LDSTAlign, ByteShAmt);
  if (!TLI.allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), MemVT,
                              LDST->getAddressSpace(), NarrowAlign.value(),
                              LDST->getMemOperand()->getFlags()))
    return false;
}

Would that make sense?

[samparker]

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

3841

I'm not sure about the assert - as I remember, this function should just return false for anything that wouldn't be a legal load, whereas it's perfectly legal to have a shift of, say, 48 somewhere in the DAG. Other than that, looks like a good fix to me.

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

bool DAGCombiner::isLegalNarrowLoad(LoadSDNode *LoadN, ISD::LoadExtType ExtType,
   // For the transform to be legal, the load must produce only two values
   // (the value loaded and the chain).  Don't transform a pre-increment
   // load, for example, which produces an extra value.  Otherwise the
   // transformation is not equivalent, and the downstream logic to replace
   // uses gets things wrong.
   if (LoadN->getNumValues() > 2)
     return false;
 
+ // Only allow byte offsets.
+  if (ShAmt % 8)
+    return false;
+
+  // Ensure that this isn't going to produce an unsupported unaligned access.
+  if (ShAmt && !TLI.allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(),
+                                       ExtVT, LoadN->getAddressSpace(),
+                                       ShAmt / 8))
+    return false;
+
+
   // If the load that we're shrinking is an extload and we're not just
   // discarding the extension we can't simply shrink the load. Bail.
   // TODO: It would be possible to merge the extensions in some cases.
   if (LoadN->getExtensionType() != ISD::NON_EXTLOAD &&
       LoadN->getMemoryVT().getSizeInBits() < ExtVT.getSizeInBits() + ShAmt)
     return false;
 
   if (!TLI.shouldReduceLoadWidth(LoadN, ExtType, ExtVT))
if (Opc == ISD::SIGN_EXTEND_INREG) {
 
     unsigned ActiveBits = AndC->getAPIntValue().countTrailingOnes();
     ExtType = ISD::ZEXTLOAD;
     ExtVT = EVT::getIntegerVT(*DAG.getContext(), ActiveBits);
   }
 
   unsigned ShAmt = 0;
   if (N0.getOpcode() == ISD::SRL && N0.hasOneUse()) {
-    if (ConstantSDNode *N01 = dyn_cast<ConstantSDNode>(N0.getOperand(1))) {
-      ShAmt = N01->getZExtValue();
+    SDValue SRL = N0;
+    if (auto *ConstShift = dyn_cast<ConstantSDNode>(SRL.getOperand(1))) {
+      ShAmt = ConstShift->getZExtValue();
       unsigned EVTBits = ExtVT.getSizeInBits();
       // Is the shift amount a multiple of size of VT?
       if ((ShAmt & (EVTBits-1)) == 0) {
         N0 = N0.getOperand(0);
         // Is the load width a multiple of size of VT?
         if ((N0.getValueSizeInBits() & (EVTBits-1)) != 0)
           return SDValue();
       }
 
       // At this point, we must have a load or else we can't do the transform.
       if (!isa<LoadSDNode>(N0)) return SDValue();
 
+      auto *LN0 = cast<LoadSDNode>(N0);
+
       // Because a SRL must be assumed to *need* to zero-extend the high bits
       // (as opposed to anyext the high bits), we can't combine the zextload
       // lowering of SRL and an sextload.
-      if (cast<LoadSDNode>(N0)->getExtensionType() == ISD::SEXTLOAD)
+      if (LN0->getExtensionType() == ISD::SEXTLOAD)
         return SDValue();
 
       // If the shift amount is larger than the input type then we're not
       // accessing any of the loaded bytes.  If the load was a zextload/extload
       // then the result of the shift+trunc is zero/undef (handled elsewhere).
-      if (ShAmt >= cast<LoadSDNode>(N0)->getMemoryVT().getSizeInBits())
+      if (ShAmt >= LN0->getMemoryVT().getSizeInBits())
         return SDValue();
+
+      // If the SRL is only used by a masking AND, we may be able to adjust
+      // the ExtVT to make the AND redundant.
+      SDNode *Mask = *(SRL->use_begin());
+      if (Mask->getOpcode() == ISD::AND &&
+          isa<ConstantSDNode>(Mask->getOperand(1))) {
+        const APInt &ShiftMask =
+          cast<ConstantSDNode>(Mask->getOperand(1))->getAPIntValue();
+        if (ShiftMask.isMask()) {
+          EVT MaskedVT = EVT::getIntegerVT(*DAG.getContext(),
+                                           ShiftMask.countTrailingOnes());
+          // Recompute the type.
+          if (TLI.isLoadExtLegal(ExtType, N0.getValueType(), MaskedVT))
+            ExtVT = MaskedVT;
+        }
+      }
     }
   }
 
   // If the load is shifted left (and the result isn't shifted back right),
   // we can fold the truncate through the shift.
   unsigned ShLeftAmt = 0;
   if (ShAmt == 0 && N0.getOpcode() == ISD::SHL && N0.hasOneUse() &&
       ExtVT == VT && TLI.isNarrowingProfitable(N0.getValueType(), VT)) {

llvm/trunk/test/CodeGen/ARM/shift-combine.ll

 define arm_aapcscc i32 @test_lshr_load4_fail(i64* %a) {
 entry:
   %0 = load i64, i64* %a, align 8
   %1 = lshr i64 %0, 8
   %conv = trunc i64 %1 to i32
   ret i32 %conv
 }
 
-; CHECK-LABEL: test_shift8_mask8
+; CHECK-LABEL: test_shift7_mask8
 ; CHECK-BE:         ldr r1, [r0]
 ; CHECK-COMMON:     ldr r1, [r0]
-; CHECK-COMMON:     ubfx r1, r1, #8, #8
+; CHECK-COMMON:     ubfx r1, r1, #7, #8
+; CHECK-COMMON:     str r1, [r0]
+define arm_aapcscc void @test_shift7_mask8(i32* nocapture %p) {
+entry:
+  %0 = load i32, i32* %p, align 4
+  %shl = lshr i32 %0, 7
+  %and = and i32 %shl, 255
+  store i32 %and, i32* %p, align 4
+  ret void
+}
+
+; CHECK-LABEL: test_shift8_mask8
+; CHECK-BE:         ldrb r1, [r0, #2]
+; CHECK-COMMON:     ldrb r1, [r0, #1]
 ; CHECK-COMMON:     str r1, [r0]
 define arm_aapcscc void @test_shift8_mask8(i32* nocapture %p) {
 entry:
   %0 = load i32, i32* %p, align 4
   %shl = lshr i32 %0, 8
   %and = and i32 %shl, 255
   store i32 %and, i32* %p, align 4
   ret void
 }
 
-; CHECK-LABEL: test_shift8_mask16
+; CHECK-LABEL: test_shift8_mask7
 ; CHECK-BE:         ldr r1, [r0]
 ; CHECK-COMMON:     ldr r1, [r0]
-; CHECK-COMMON:     ubfx r1, r1, #8, #16
+; CHECK-COMMON:     ubfx r1, r1, #8, #7
+; CHECK-COMMON:     str r1, [r0]
+define arm_aapcscc void @test_shift8_mask7(i32* nocapture %p) {
+entry:
+  %0 = load i32, i32* %p, align 4
+  %shl = lshr i32 %0, 8
+  %and = and i32 %shl, 127
+  store i32 %and, i32* %p, align 4
+  ret void
+}
+
+; CHECK-LABEL: test_shift9_mask8
+; CHECK-BE:         ldr r1, [r0]
+; CHECK-COMMON:     ldr r1, [r0]
+; CHECK-COMMON:     ubfx r1, r1, #9, #8
+; CHECK-COMMON:     str r1, [r0]
+define arm_aapcscc void @test_shift9_mask8(i32* nocapture %p) {
+entry:
+  %0 = load i32, i32* %p, align 4
+  %shl = lshr i32 %0, 9
+  %and = and i32 %shl, 255
+  store i32 %and, i32* %p, align 4
+  ret void
+}
+
+; CHECK-LABEL: test_shift8_mask16
+; CHECK-ALIGN:      ldr r1, [r0]
+; CHECK-ALIGN:      ubfx r1, r1, #8, #16
+; CHECK-BE:         ldrh r1, [r0, #1]
+; CHECK-ARM:        ldrh r1, [r0, #1]
+; CHECK-THUMB:      ldrh.w r1, [r0, #1]
 ; CHECK-COMMON:     str r1, [r0]
 define arm_aapcscc void @test_shift8_mask16(i32* nocapture %p) {
 entry:
   %0 = load i32, i32* %p, align 4
   %shl = lshr i32 %0, 8
   %and = and i32 %shl, 65535
   store i32 %and, i32* %p, align 4
   ret void
 }
 
+; CHECK-LABEL: test_shift15_mask16
+; CHECK-COMMON:     ldr r1, [r0]
+; CHECK-COMMON:     ubfx r1, r1, #15, #16
+; CHECK-COMMON:     str r1, [r0]
+define arm_aapcscc void @test_shift15_mask16(i32* nocapture %p) {
+entry:
+  %0 = load i32, i32* %p, align 4
+  %shl = lshr i32 %0, 15
+  %and = and i32 %shl, 65535
+  store i32 %and, i32* %p, align 4
+  ret void
+}
+
+; CHECK-LABEL: test_shift16_mask15
+; CHECK-BE:         ldrh r1, [r0]
+; CHECK-COMMON:     ldrh r1, [r0, #2]
+; CHECK-COMMON:     bfc r1, #15, #17
+; CHECK-COMMON:     str r1, [r0]
+define arm_aapcscc void @test_shift16_mask15(i32* nocapture %p) {
+entry:
+  %0 = load i32, i32* %p, align 4
+  %shl = lshr i32 %0, 16
+  %and = and i32 %shl, 32767
+  store i32 %and, i32* %p, align 4
+  ret void
+}
+
+; CHECK-LABEL: test_shift8_mask24
+; CHECK-BE:         ldr r1, [r0]
+; CHECK-COMMON:     ldr r1, [r0]
+; CHECK-ARM:        lsr r1, r1, #8
+; CHECK-THUMB:      lsrs r1, r1, #8
+; CHECK-COMMON:     str r1, [r0]
+define arm_aapcscc void @test_shift8_mask24(i32* nocapture %p) {
+entry:
+  %0 = load i32, i32* %p, align 4
+  %shl = lshr i32 %0, 8
+  %and = and i32 %shl, 16777215
+  store i32 %and, i32* %p, align 4
+  ret void
+}
+
+; CHECK-LABEL: test_shift24_mask16
+; CHECK-BE:         ldrb r1, [r0]
+; CHECK-COMMON:     ldrb r1, [r0, #3]
+; CHECK-COMMON:     str r1, [r0]
+define arm_aapcscc void @test_shift24_mask16(i32* nocapture %p) {
+entry:
+  %0 = load i32, i32* %p, align 4
+  %shl = lshr i32 %0, 24
+  %and = and i32 %shl, 65535
+  store i32 %and, i32* %p, align 4
+  ret void
+}
+
 ; CHECK-LABEL: test_sext_shift8_mask8
 ; CHECK-BE:         ldrb r0, [r0]
 ; CHECK-COMMON:     ldrb r0, [r0, #1]
 ; CHECK-COMMON:     str r0, [r1]
 define arm_aapcscc void @test_sext_shift8_mask8(i16* %p, i32* %q) {
 entry:
   %0 = load i16, i16* %p, align 4
   %1 = sext i16 %0 to i32

 entry:
   %0 = load i16, i16* %p, align 4
   %1 = sext i16 %0 to i32
   %shl = lshr i32 %1, 8
   %and = and i32 %shl, 65535
   store i32 %and, i32* %q, align 4
   ret void
 }
+
+; CHECK-LABEL: trunc_i64_mask_srl
+; CHECK-ARM: ldrh r2, [r1, #4]
+; CHECK-BE: ldrh r2, [r1, #2]
+define i1 @trunc_i64_mask_srl(i32 zeroext %AttrArgNo, i64* %ptr) {
+entry:
+  %bf.load.i = load i64, i64* %ptr, align 8
+  %bf.lshr.i = lshr i64 %bf.load.i, 32
+  %0 = trunc i64 %bf.lshr.i to i32
+  %bf.cast.i = and i32 %0, 65535
+  %cmp.i = icmp ugt i32 %bf.cast.i, %AttrArgNo
+  ret i1 %cmp.i
+}

llvm/trunk/test/CodeGen/PowerPC/trunc-srl-load.ll

+; RUN: llc -mtriple=powerpc64-unknown-unknown %s -o - | FileCheck %s
+
+; CHECK-LABEL: trunc_srl_load
+; CHECK-NOT: lhz 4, 4(0)
+; CHECK: lhz 4, 2(0)
+define dso_local fastcc void @trunc_srl_load(i32 zeroext %AttrArgNo) {
+entry:
+  %bf.load.i = load i64, i64* null, align 8
+  %bf.lshr.i = lshr i64 %bf.load.i, 32
+  %0 = trunc i64 %bf.lshr.i to i32
+  %bf.cast.i = and i32 %0, 65535
+  %cmp.i = icmp ugt i32 %bf.cast.i, %AttrArgNo
+  br i1 %cmp.i, label %exit, label %cond.false
+exit:       ; preds = %entry
+  unreachable
+cond.false:                                       ; preds = %entry
+  unreachable
+}

llvm/trunk/test/CodeGen/X86/h-registers-1.ll

 ; CHECK-NEXT:    .cfi_offset %rbx, -24
 ; CHECK-NEXT:    .cfi_offset %rbp, -16
 ; CHECK-NEXT:    movq %rsi, %rax
 ; CHECK-NEXT:    movq %rdi, %rbx
 ; CHECK-NEXT:    movzbl %bh, %esi
 ; CHECK-NEXT:    movzbl %ah, %eax
 ; CHECK-NEXT:    movq %rax, %r10
 ; CHECK-NEXT:    movzbl %dh, %edx
-; CHECK-NEXT:    movzbl %ch, %eax
-; CHECK-NEXT:    movq %rax, %r11
+; CHECK-NEXT:    movzbl %ch, %ebp
 ; CHECK-NEXT:    movq %r8, %rax
 ; CHECK-NEXT:    movzbl %ah, %ecx
 ; CHECK-NEXT:    movq %r9, %rax
-; CHECK-NEXT:    movzbl %ah, %ebp
-; CHECK-NEXT:    movl {{[0-9]+}}(%rsp), %eax
-; CHECK-NEXT:    movzbl %ah, %eax
-; CHECK-NEXT:    movl {{[0-9]+}}(%rsp), %ebx
-; CHECK-NEXT:    movzbl %bh, %edi
+; CHECK-NEXT:    movzbl %ah, %ebx
+; CHECK-NEXT:    movzbl {{[0-9]+}}(%rsp), %eax
+; CHECK-NEXT:    movzbl {{[0-9]+}}(%rsp), %edi
 ; CHECK-NEXT:    addq %r10, %rsi
-; CHECK-NEXT:    addq %r11, %rdx
+; CHECK-NEXT:    addq %rbp, %rdx
 ; CHECK-NEXT:    addq %rsi, %rdx
-; CHECK-NEXT:    addq %rbp, %rcx
+; CHECK-NEXT:    addq %rbx, %rcx
 ; CHECK-NEXT:    addq %rdi, %rax
 ; CHECK-NEXT:    addq %rcx, %rax
 ; CHECK-NEXT:    addq %rdx, %rax
 ; CHECK-NEXT:    popq %rbx
 ; CHECK-NEXT:    popq %rbp
 ; CHECK-NEXT:    retq
 ;
 ; GNUX32-LABEL: foo:
 ; GNUX32:       # %bb.0:
 ; GNUX32-NEXT:    pushq %rbp
 ; GNUX32-NEXT:    .cfi_def_cfa_offset 16
 ; GNUX32-NEXT:    pushq %rbx
 ; GNUX32-NEXT:    .cfi_def_cfa_offset 24
 ; GNUX32-NEXT:    .cfi_offset %rbx, -24
 ; GNUX32-NEXT:    .cfi_offset %rbp, -16
 ; GNUX32-NEXT:    movq %rsi, %rax
 ; GNUX32-NEXT:    movq %rdi, %rbx
 ; GNUX32-NEXT:    movzbl %bh, %esi
 ; GNUX32-NEXT:    movzbl %ah, %eax
 ; GNUX32-NEXT:    movq %rax, %r10
 ; GNUX32-NEXT:    movzbl %dh, %edx
-; GNUX32-NEXT:    movzbl %ch, %eax
-; GNUX32-NEXT:    movq %rax, %r11
+; GNUX32-NEXT:    movzbl %ch, %ebp
 ; GNUX32-NEXT:    movq %r8, %rax
 ; GNUX32-NEXT:    movzbl %ah, %ecx
 ; GNUX32-NEXT:    movq %r9, %rax
-; GNUX32-NEXT:    movzbl %ah, %ebp
-; GNUX32-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; GNUX32-NEXT:    movzbl %ah, %eax
-; GNUX32-NEXT:    movl {{[0-9]+}}(%esp), %ebx
-; GNUX32-NEXT:    movzbl %bh, %edi
+; GNUX32-NEXT:    movzbl %ah, %ebx
+; GNUX32-NEXT:    movzbl {{[0-9]+}}(%esp), %eax
+; GNUX32-NEXT:    movzbl {{[0-9]+}}(%esp), %edi
 ; GNUX32-NEXT:    addq %r10, %rsi
-; GNUX32-NEXT:    addq %r11, %rdx
+; GNUX32-NEXT:    addq %rbp, %rdx
 ; GNUX32-NEXT:    addq %rsi, %rdx
-; GNUX32-NEXT:    addq %rbp, %rcx
+; GNUX32-NEXT:    addq %rbx, %rcx
 ; GNUX32-NEXT:    addq %rdi, %rax
 ; GNUX32-NEXT:    addq %rcx, %rax
 ; GNUX32-NEXT:    addq %rdx, %rax
 ; GNUX32-NEXT:    popq %rbx
 ; GNUX32-NEXT:    popq %rbp
 ; GNUX32-NEXT:    retq
   %sa = lshr i64 %a, 8
   %A = and i64 %sa, 255