[DAGCombiner] Adjust some checks in DAGCombiner::reduceLoadWidth

[bjope]

Summary

In code review for D117104 two slightly weird checks were found
in DAGCombiner::reduceLoadWidth. They were typically checking
if BitsA was a mulitple of BitsB by looking at (BitsA & (BitsB - 1)),
but such a comparison actually only make sense if BitsB is a power
of two.

The checks were related to the code that attempted to shrink a load
based on the fact that the loaded value would be right shifted.

Afaict the legality of the value types is checked later (typically in
isLegalNarrowLdSt), so the existing checks were both overly
conservative as well as being wrong whenever ExtVTBits wasn't a
power of two. The latter was a situation triggered by a number of
lit tests so we could not just assert on ExtVTBIts being a power of
two).

When attempting to simply remove the checks I found some problems,
that seems to have been guarded by the checks (maybe just out of
luck). A typical example would be a pattern like this:

t1 = load i96* ptr
t2 = srl t1, 64
t3 = truncate t2 to i64

When DAGCombine is visiting the truncate reduceLoadWidth is called
attempting to narrow the load to 64 bits (ExtVT := MVT::i64). Then
the SRL is detected and we set ShAmt to 64.

In the past we've bailed out due to i96 not being a multiple of 64.
If we simply remove that check then we would end up replacing the
load with a new load that would read 64 bits but with a base pointer
adjusted by 64 bits. So we would read 32 bits the wasn't accessed by
the original load.
This patch will instead utilize the fact that the logical left shift
can be folded away by using a zextload. Thus, the pattern above will
now be combined into

t3 = load i32* ptr+offset, zext to i64

Another case is shown in the new X86/combine-srl-load.ll test case:

t1 = load i32* ptr
t2 = srl i32 t1, 8
t3 = truncate t2 to i16

In the past we bailed out due to the shift count (8) not being a
multiple of 16. Now the narrowing kicks in and we get

t3 = load i16* ptr+offset

[spatel]

llvm/test/CodeGen/PowerPC/pr39478.ll

6 ↗

(On Diff #400319)

This is a generic missed fold. I added more tests and a proposed fix:
D117508

[bjope]

Rebased (now including the BSWAP fold from D117508 to avoid regression in a PPC test case).

Also added a new X86 test as a regression test for new fold when the shift
count isn't a multiple of the narrowed load width.

[bjope]

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

7 ↗

(On Diff #400773)

I've assumed that DAGCombiner::isLegalNarrowLdSt would have bailed out if this would result in a unaligned memory access for the target (it does use TLI.allowsMemoryAccess).

Had perhaps been nice with a test case for some target where this would matter. Any suggestions?

[bjope]

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

7 ↗

(On Diff #400773)

Well, I've verified it by using riscv64 as a target. Not sure if I need to add a test case for that.

[spatel]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12219

happend -> happen

12225–12226

Why set these variables if we are exiting the function? Can we specify in the comment which function does the narrowing of the load for this pattern?

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

4 ↗

(On Diff #400773)

It would better to pre-commit this test with the baseline CHECKs. I don't think there's a specific test file for x86 used for verifying this type of fold, but "shift-folding.ll" has a test that looks similar, so you could add it to that file.

[bjope]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12225–12226

Oh, that return statement looks like a mistake. I started out by simply returning here, but then I realized that it should be OK to use ZEXTLOAD here, as long as that assert never triggers about replacing sextload by zextload. The idea was to remove the return statement again, not including it in the patch.

The example with (i64 (truncate (i96 (srl (load x), 64)))) would however be optimzed also without this special case. But then it would start off by first rewriting srl+load into a larger zextload, which later is combined with the truncate to form a smaller zextload. By removing the return here we would trigger the full fold already based on the truncate.

[bjope]

Rebased given parent patch with test cases for pre-commit (D117588).

[spatel]

LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

 //===- DAGCombiner.cpp - Implement a DAG node combiner --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
if (Opc == ISD::SRL || N0.getOpcode() == ISD::SRL) {
     auto *SRL1C = dyn_cast<ConstantSDNode>(SRL.getOperand(1));
     if (!SRL1C || !LN)
       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).
     ShAmt = SRL1C->getZExtValue();
-    if (ShAmt >= LN->getMemoryVT().getSizeInBits())
+    uint64_t MemoryWidth = LN->getMemoryVT().getSizeInBits();
+    if (ShAmt >= MemoryWidth)
       return SDValue();
 
     // 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 (LN->getExtensionType() == ISD::SEXTLOAD)
       return SDValue();
 
-    unsigned ExtVTBits = ExtVT.getScalarSizeInBits();
-    // Is the shift amount a multiple of size of ExtVT?
-    if ((ShAmt & (ExtVTBits - 1)) != 0)
-      return SDValue();
-    // Is the load width a multiple of size of ExtVT?
-    if ((SRL.getScalarValueSizeInBits() & (ExtVTBits - 1)) != 0)
+    // Avoid reading outside the memory accessed by the original load (could
+    // happened if we only adjust the load base pointer by ShAmt). Instead we
+    // try to narrow the load even further. The typical scenario here is:
+    //   (i64 (truncate (i96 (srl (load x), 64)))) ->
+    //     (i64 (truncate (i96 (zextload (load i32 + offset) from i32))))
+    if (ExtVT.getScalarSizeInBits() > MemoryWidth - ShAmt) {
+      // Don't replace sextload by zextload.
+      if (ExtType == ISD::SEXTLOAD)
         return SDValue();
+      // Narrow the load.
+      ExtType = ISD::ZEXTLOAD;
+      ExtVT = EVT::getIntegerVT(*DAG.getContext(), MemoryWidth - ShAmt);
+    }
 
     // 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 (SRL.hasOneUse() && Mask->getOpcode() == ISD::AND &&
         isa<ConstantSDNode>(Mask->getOperand(1))) {
       const APInt& ShiftMask = Mask->getConstantOperandAPInt(1);
       if (ShiftMask.isMask()) {
         EVT MaskedVT = EVT::getIntegerVT(*DAG.getContext(),
                                          ShiftMask.countTrailingOnes());
         // If the mask is smaller, recompute the type.
-        if ((ExtVTBits > MaskedVT.getScalarSizeInBits()) &&
+        if ((ExtVT.getScalarSizeInBits() > MaskedVT.getScalarSizeInBits()) &&
             TLI.isLoadExtLegal(ExtType, SRL.getValueType(), MaskedVT))
           ExtVT = MaskedVT;
       }
     }
 
     N0 = SRL.getOperand(0);
   }
 

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

 define arm_aapcscc i32 @test_lshr_load64_4_unaligned(i64* %a) {
 ; CHECK-ARM-LABEL: test_lshr_load64_4_unaligned:
 ; CHECK-ARM:       @ %bb.0: @ %entry
 ; CHECK-ARM-NEXT:    ldr r0, [r0, #2]
 ; CHECK-ARM-NEXT:    bx lr
 ;
 ; CHECK-BE-LABEL: test_lshr_load64_4_unaligned:
 ; CHECK-BE:       @ %bb.0: @ %entry
-; CHECK-BE-NEXT:    ldr r1, [r0]
-; CHECK-BE-NEXT:    ldrh r0, [r0, #4]
-; CHECK-BE-NEXT:    orr r0, r0, r1, lsl #16
+; CHECK-BE-NEXT:    ldr r0, [r0, #2]
 ; CHECK-BE-NEXT:    bx lr
 ;
 ; CHECK-THUMB-LABEL: test_lshr_load64_4_unaligned:
 ; CHECK-THUMB:       @ %bb.0: @ %entry
 ; CHECK-THUMB-NEXT:    ldr.w r0, [r0, #2]
 ; CHECK-THUMB-NEXT:    bx lr
 ;
 ; CHECK-ALIGN-LABEL: test_lshr_load64_4_unaligned:

 define arm_aapcscc i32 @test_lshr_load64_1_lsb(i64* %a) {
 ; CHECK-ARM-LABEL: test_lshr_load64_1_lsb:
 ; CHECK-ARM:       @ %bb.0: @ %entry
 ; CHECK-ARM-NEXT:    ldr r0, [r0, #3]
 ; CHECK-ARM-NEXT:    bx lr
 ;
 ; CHECK-BE-LABEL: test_lshr_load64_1_lsb:
 ; CHECK-BE:       @ %bb.0: @ %entry
-; CHECK-BE-NEXT:    ldr r1, [r0]
-; CHECK-BE-NEXT:    ldrb r0, [r0, #4]
-; CHECK-BE-NEXT:    orr r0, r0, r1, lsl #8
+; CHECK-BE-NEXT:    ldr r0, [r0, #1]
 ; CHECK-BE-NEXT:    bx lr
 ;
 ; CHECK-THUMB-LABEL: test_lshr_load64_1_lsb:
 ; CHECK-THUMB:       @ %bb.0: @ %entry
 ; CHECK-THUMB-NEXT:    ldr.w r0, [r0, #3]
 ; CHECK-THUMB-NEXT:    bx lr
 ;
 ; CHECK-ALIGN-LABEL: test_lshr_load64_1_lsb:
entry:
   %1 = lshr i64 %0, 48
   %conv = trunc i64 %1 to i32
   ret i32 %conv
 }
 
 define arm_aapcscc i32 @test_lshr_load4_fail(i64* %a) {
 ; CHECK-ARM-LABEL: test_lshr_load4_fail:
 ; CHECK-ARM:       @ %bb.0: @ %entry
-; CHECK-ARM-NEXT:    ldrd r0, r1, [r0]
-; CHECK-ARM-NEXT:    lsr r0, r0, #8
-; CHECK-ARM-NEXT:    orr r0, r0, r1, lsl #24
+; CHECK-ARM-NEXT:    ldr r0, [r0, #1]
 ; CHECK-ARM-NEXT:    bx lr
 ;
 ; CHECK-BE-LABEL: test_lshr_load4_fail:
 ; CHECK-BE:       @ %bb.0: @ %entry
-; CHECK-BE-NEXT:    ldrd r0, r1, [r0]
-; CHECK-BE-NEXT:    lsr r1, r1, #8
-; CHECK-BE-NEXT:    orr r0, r1, r0, lsl #24
+; CHECK-BE-NEXT:    ldr r0, [r0, #3]
 ; CHECK-BE-NEXT:    bx lr
 ;
 ; CHECK-THUMB-LABEL: test_lshr_load4_fail:
 ; CHECK-THUMB:       @ %bb.0: @ %entry
-; CHECK-THUMB-NEXT:    ldrd r0, r1, [r0]
-; CHECK-THUMB-NEXT:    lsrs r0, r0, #8
-; CHECK-THUMB-NEXT:    orr.w r0, r0, r1, lsl #24
+; CHECK-THUMB-NEXT:    ldr.w r0, [r0, #1]
 ; CHECK-THUMB-NEXT:    bx lr
 ;
 ; CHECK-ALIGN-LABEL: test_lshr_load4_fail:
 ; CHECK-ALIGN:       @ %bb.0: @ %entry
 ; CHECK-ALIGN-NEXT:    ldrd r0, r1, [r0]
 ; CHECK-ALIGN-NEXT:    lsrs r0, r0, #8
 ; CHECK-ALIGN-NEXT:    orr.w r0, r0, r1, lsl #24
 ; CHECK-ALIGN-NEXT:    bx lr

llvm/test/CodeGen/X86/shift-folding.ll

; CHECK-NEXT:    retl
   ret i32 %xor
 }
 
 ; Should be possible to adjust the pointer and narrow the load to 16 bits.
 define i16 @srl_load_narrowing1(i32* %arg) {
 ; CHECK-LABEL: srl_load_narrowing1:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; CHECK-NEXT:    movl (%eax), %eax
-; CHECK-NEXT:    shrl $8, %eax
-; CHECK-NEXT:    # kill: def $ax killed $ax killed $eax
+; CHECK-NEXT:    movzwl 1(%eax), %eax
 ; CHECK-NEXT:    retl
   %tmp1 = load i32, i32* %arg, align 1
   %tmp2 = lshr i32 %tmp1, 8
   %tmp3 = trunc i32 %tmp2 to i16
   ret i16 %tmp3
 }
 
 define i16 @srl_load_narrowing2(i32* %arg) {