[DAGCombine] Fold SRA of a load into a narrower sign-extending load

[bjope]

Summary

An sra is basically sign-extending a narrower value. Fold away the
shift by doing a sextload of a narrower value, when it is legal to
reduce the load width accordingly.

[bjope]

Fixup comment that had been edited by mistake.

[spatel]

Nice transform! Might be worth including an IR proof in the description to make it clearer:
https://alive2.llvm.org/ce/z/yYuasA

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12173

The various shift amount variables are difficult to follow. This is really a question for the existing code - can we make it clearer (either through renaming or code comments) how "ShAmt", "ShiftAmt" and "ShLeftAmt" are different?

If we can improve that, is it possible to make a small helper/lambda, so we can share the bailout conditions for SRA/SRL instead of duplicating code?

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

73–76

This comment is not accurate. We are replicating (splatting) the sign bit of the loaded i16 across 32-bits, so there's still a shift.
In IR, instcombine would transform this into:

define i32 @sra_of_sextload_no_fold(i16* %p) {
  %load = load i16, i16* %p, align 2
  %1 = ashr i16 %load, 15
  %shift = sext i16 %1 to i32
  ret i32 %shift
}

[bjope]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12173

Yes. I thought the old names were kind of confusing. So I added the SRA separately. I'll make a separate patch, trying to improve the existing code as well.

One reason I did not share things with SRL (I started off trying to do it that way) is that the solution for SRL seemed to be divided into two parts. I believe the second part (that includes the hasOneUse guard) can be triggered also as being nestled inside an AND (N pointing at the AND and N0 being the SRL). Although, we can probably do some code sharing for the first part regardless of that, at least now when I understand that the unexplained "N0 = SDValue(N, 0)" only should be done only for SRL and not for SRA.

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

73–76

Right, the comment is supposed to say "so we can't fold away the shift".

[bjope]

Rebased upon https://reviews.llvm.org/D117104

The new parent revision contains some refactoring of the old SRL folds, as well as updated code comments.

[bjope]

Rebased.

[RKSimon]

Does this need to be rebased after rG46cacdbb21c2?

Looks a lot nicer now - @spatel any more comments?

[spatel]

In D116930#3265911, @RKSimon wrote:

Does this need to be rebased after rG46cacdbb21c2?

Looks a lot nicer now - @spatel any more comments?

LGTM. This part is relatively straightforward now - assuming we got the preceding steps right (might want to give that a little time to bake in the wild).

[bjope]

@spatel / @RKSimon : Thanks for all the feedback!
A rebase should be trivial (it's already rebased on the patch that updated the code common with SRL). I'll wait until tomorrow before landing this part.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (N1C && !N1C->isOpaque())
     if (SDValue NewSRA = visitShiftByConstant(N))
       return NewSRA;
 
   // Try to transform this shift into a multiply-high if
   // it matches the appropriate pattern detected in combineShiftToMULH.
   if (SDValue MULH = combineShiftToMULH(N, DAG, TLI))
     return MULH;
 
+  // Attempt to convert a sra of a load into a narrower sign-extending load.
+  if (SDValue NarrowLoad = reduceLoadWidth(N))
+    return NarrowLoad;
+
   return SDValue();
 }
 
 SDValue DAGCombiner::visitSRL(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
   if (SDValue V = DAG.simplifyShift(N0, N1))
     return V;
SDValue DAGCombiner::reduceLoadWidth(SDNode *N) {
   // to indicate that the narrowed load should be left-shifted ShAmt bits to get
   // the result.
   bool HasShiftedOffset = false;
   // Special case: SIGN_EXTEND_INREG is basically truncating to ExtVT then
   // extended to VT.
   if (Opc == ISD::SIGN_EXTEND_INREG) {
     ExtType = ISD::SEXTLOAD;
     ExtVT = cast<VTSDNode>(N->getOperand(1))->getVT();
-  } else if (Opc == ISD::SRL) {
-    // Another special-case: SRL is basically zero-extending a narrower value,
-    // or it may be shifting a higher subword, half or byte into the lowest
-    // bits.
+  } else if (Opc == ISD::SRL || Opc == ISD::SRA) {
+    // Another special-case: SRL/SRA is basically zero/sign-extending a narrower
+    // value, or it may be shifting a higher subword, half or byte into the
+    // lowest bits.
 
     // Only handle shift with constant shift amount, and the shiftee must be a
     // load.
     auto *LN = dyn_cast<LoadSDNode>(N0);
     auto *N1C = dyn_cast<ConstantSDNode>(N->getOperand(1));
     if (!N1C || !LN)
       return SDValue();
     // If the shift amount is larger than the memory type then we're not
     // accessing any of the loaded bytes.
     ShAmt = N1C->getZExtValue();
     uint64_t MemoryWidth = LN->getMemoryVT().getScalarSizeInBits();
     if (MemoryWidth <= ShAmt)
       return SDValue();
-    // Attempt to fold away the SRL by using ZEXTLOAD.
-    ExtType = ISD::ZEXTLOAD;
+    // Attempt to fold away the SRL by using ZEXTLOAD and SRA by using SEXTLOAD.
+    ExtType = Opc == ISD::SRL ? ISD::ZEXTLOAD : ISD::SEXTLOAD;
     ExtVT = EVT::getIntegerVT(*DAG.getContext(), MemoryWidth - ShAmt);
     // If original load is a SEXTLOAD then we can't simply replace it by a
     // ZEXTLOAD (we could potentially replace it by a more narrow SEXTLOAD
-    // followed by a ZEXT, but that is not handled at the moment).
-    if (LN->getExtensionType() == ISD::SEXTLOAD)
+    // followed by a ZEXT, but that is not handled at the moment). Similarly if
+    // the original load is a ZEXTLOAD and we want to use a SEXTLOAD.
+    if ((LN->getExtensionType() == ISD::SEXTLOAD ||
+         LN->getExtensionType() == ISD::ZEXTLOAD) &&
+        LN->getExtensionType() != ExtType)
       return SDValue();
   } else if (Opc == ISD::AND) {
     // An AND with a constant mask is the same as a truncate + zero-extend.
     auto AndC = dyn_cast<ConstantSDNode>(N->getOperand(1));
     if (!AndC)
       return SDValue();
 
     const APInt &Mask = AndC->getAPIntValue();

llvm/test/CodeGen/PowerPC/pr13891.ll

 ; RUN: llc -verify-machineinstrs < %s | FileCheck %s
 target datalayout = "E-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v128:128:128-n32:64"
 target triple = "powerpc64-unknown-linux-gnu"
 
 %struct.foo = type { i8, i8 }
 
 define void @_Z5check3foos(%struct.foo* nocapture byval(%struct.foo) %f, i16 signext %i) noinline {
 ; CHECK-LABEL: _Z5check3foos:
 ; CHECK: sth 3, {{[0-9]+}}(1)
-; CHECK: lha {{[0-9]+}}, {{[0-9]+}}(1)
+; CHECK: lbz {{[0-9]+}}, {{[0-9]+}}(1)
 entry:
   %0 = bitcast %struct.foo* %f to i16*
   %1 = load i16, i16* %0, align 2
   %bf.val.sext = ashr i16 %1, 8
   %cmp = icmp eq i16 %bf.val.sext, %i
   br i1 %cmp, label %if.end, label %if.then
 
 if.then:                                          ; preds = %entry

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

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=x86_64-unknown-unknown | FileCheck %s --check-prefix=CHECK
 
-; FIXME: fold (sra (load i32), 16)) -> (sextload i16)
+; fold (sra (load i32), 16)) -> (sextload i16)
 define i32 @sra_half(i32* %p) {
 ; CHECK-LABEL: sra_half:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movl (%rdi), %eax
-; CHECK-NEXT:    sarl $16, %eax
+; CHECK-NEXT:    movswl 2(%rdi), %eax
 ; CHECK-NEXT:    retq
   %load = load i32, i32* %p
   %shift = ashr i32 %load, 16
   ret i32 %shift
 }
 
 ; Vector version not folded.
 define <4 x i32> @sra_half_vec(<4 x i32>* %p) {
 ; CHECK-LABEL: sra_half_vec:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movdqa (%rdi), %xmm0
 ; CHECK-NEXT:    psrad $16, %xmm0
 ; CHECK-NEXT:    retq
   %load = load <4 x i32>, <4 x i32>* %p
   %shift = ashr <4 x i32> %load, <i32 16, i32 16, i32 16, i32 16>
   ret <4 x i32> %shift
 }
 
-; FIXME: fold (sra (load i64), 48)) -> (sextload i16)
+; fold (sra (load i64), 48)) -> (sextload i16)
 define i64 @sra_large_shift(i64* %r) {
 ; CHECK-LABEL: sra_large_shift:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movq (%rdi), %rax
-; CHECK-NEXT:    sarq $48, %rax
+; CHECK-NEXT:    movswq 6(%rdi), %rax
 ; CHECK-NEXT:    retq
   %t0 = load i64, i64* %r
   %conv = ashr i64 %t0, 48
   ret i64 %conv
 }
 
 ; Negative test, no fold expected.
 define i32 @sra_small_shift(i32* %p) {

 ; CHECK-NEXT:    movzbl 1(%rdi), %eax
 ; CHECK-NEXT:    retq
   %load = load i16, i16* %p
   %zext = zext i16 %load to i32
   %shift = ashr i32 %zext, 8
   ret i32 %shift
 }
 
-; FIXME: fold (sra (sextload i16 to i32), 8) -> (sextload i8)
+; fold (sra (sextload i16 to i32), 8) -> (sextload i8)
 define i32 @sra_of_sextload(i16* %p) {
 ; CHECK-LABEL: sra_of_sextload:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movswl (%rdi), %eax
-; CHECK-NEXT:    sarl $8, %eax
+; CHECK-NEXT:    movsbl 1(%rdi), %eax
 ; CHECK-NEXT:    retq
   %load = load i16, i16* %p
   %sext = sext i16 %load to i32
   %shift = ashr i32 %sext, 8
   ret i32 %shift
 }
 
 ; Negative test. If the shift amount is larger than the memory type then
 ; we're not accessing any of the loaded bytes (only the extended bits). So the
 ; shift is expected to remain.
 define i32 @sra_of_sextload_no_fold(i16* %p) {
 ; CHECK-LABEL: sra_of_sextload_no_fold:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    movswl (%rdi), %eax
 ; CHECK-NEXT:    sarl $16, %eax
 ; CHECK-NEXT:    retq
   %load = load i16, i16* %p
   %sext = sext i16 %load to i32
   %shift = ashr i32 %sext, 16
   ret i32 %shift
 }
 
-; FIXME: Fold even if SRA has multiple uses.
+; Fold even if SRA has multiple uses.
 define i32 @sra_to_sextload_multiple_sra_uses(i32* %p) {
 ; CHECK-LABEL: sra_to_sextload_multiple_sra_uses:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movl (%rdi), %ecx
-; CHECK-NEXT:    sarl $16, %ecx
+; CHECK-NEXT:    movswl 2(%rdi), %ecx
 ; CHECK-NEXT:    movl %ecx, %eax
 ; CHECK-NEXT:    xorl $6, %eax
 ; CHECK-NEXT:    orl %ecx, %eax
 ; CHECK-NEXT:    retq
   %load = load i32, i32* %p
   %shift = ashr i32 %load, 16
   %use1 = xor i32 %shift, 6
   %use2 = or i32 %shift, %use1
   ret i32 %use2
 }