[PATCH] D159198: [DAG] Fold (shl (ext (add x, c1)), c2) -> (add (shl (ext x), c2), c1 << c2)

[Simon Pilgrim via Phabricator via llvm-commits]

RKSimon created this revision.
RKSimon added reviewers: dmgreen, goldstein.w.n, craig.topper.
Herald added subscribers: luismarques, pengfei, s.egerton, PkmX, simoncook, hiraditya, kristof.beyls, arichardson.
Herald added a project: All.
RKSimon requested review of this revision.
Herald added a subscriber: wangpc.
Herald added a project: LLVM.

Assuming the ADD is nsw/nuz then it may be sign/zero-extended to merge with a SHL op in a similar fold to the existing `(shl (add x, c1), c2) -> (add (shl x, c2), c1 << c2)` fold.

This is most useful for helping to expose address math for X86, but has also touched several aarch64 test cases as well - I think they are benign but would like some confirmation.

@craig.topper RISCV uses isDesirableToCommuteWithShift to prevent creating shifted adds with larger offsets - should this be extended to peek through sext/zext nodes as well? If so I'll need to add suitable test coverage. IIRC there's implicit extension on RISCV targets that I'm not very familiar with.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D159198

Files:
  llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
  llvm/test/CodeGen/AArch64/arm64-shifted-sext.ll
  llvm/test/CodeGen/AArch64/arm64-trunc-store.ll
  llvm/test/CodeGen/X86/addr-mode-matcher-2.ll


Index: llvm/test/CodeGen/X86/addr-mode-matcher-2.ll
===================================================================
--- llvm/test/CodeGen/X86/addr-mode-matcher-2.ll
+++ llvm/test/CodeGen/X86/addr-mode-matcher-2.ll
@@ -52,8 +52,8 @@
 ; X64-NEXT:    .p2align 4, 0x90
 ; X64-NEXT:  .LBB0_2: # =>This Inner Loop Header: Depth=1
 ; X64-NEXT:    cltq
-; X64-NEXT:    leaq 4(,%rax,4), %rax
-; X64-NEXT:    leaq (%rax,%rax,4), %rdi
+; X64-NEXT:    shlq $2, %rax
+; X64-NEXT:    leaq 20(%rax,%rax,4), %rdi
 ; X64-NEXT:    callq bar at PLT
 ; X64-NEXT:    jmp .LBB0_2
   br i1 %0, label %9, label %3
Index: llvm/test/CodeGen/AArch64/arm64-trunc-store.ll
===================================================================
--- llvm/test/CodeGen/AArch64/arm64-trunc-store.ll
+++ llvm/test/CodeGen/AArch64/arm64-trunc-store.ll
@@ -20,10 +20,10 @@
 ; CHECK-LABEL: fct32:
 ; CHECK:       // %bb.0: // %bb
 ; CHECK-NEXT:    adrp x8, :got:zptr32
-; CHECK-NEXT:    sub w9, w0, #1
 ; CHECK-NEXT:    ldr x8, [x8, :got_lo12:zptr32]
 ; CHECK-NEXT:    ldr x8, [x8]
-; CHECK-NEXT:    str w1, [x8, w9, sxtw #2]
+; CHECK-NEXT:    add x8, x8, w0, sxtw #2
+; CHECK-NEXT:    stur w1, [x8, #-4]
 ; CHECK-NEXT:    ret
 bb:
   %.pre37 = load ptr, ptr @zptr32, align 8
@@ -39,10 +39,10 @@
 ; CHECK-LABEL: fct16:
 ; CHECK:       // %bb.0: // %bb
 ; CHECK-NEXT:    adrp x8, :got:zptr16
-; CHECK-NEXT:    sub w9, w0, #1
 ; CHECK-NEXT:    ldr x8, [x8, :got_lo12:zptr16]
 ; CHECK-NEXT:    ldr x8, [x8]
-; CHECK-NEXT:    strh w1, [x8, w9, sxtw #1]
+; CHECK-NEXT:    add x8, x8, w0, sxtw #1
+; CHECK-NEXT:    sturh w1, [x8, #-2]
 ; CHECK-NEXT:    ret
 bb:
   %.pre37 = load ptr, ptr @zptr16, align 8
Index: llvm/test/CodeGen/AArch64/arm64-shifted-sext.ll
===================================================================
--- llvm/test/CodeGen/AArch64/arm64-shifted-sext.ll
+++ llvm/test/CodeGen/AArch64/arm64-shifted-sext.ll
@@ -275,8 +275,9 @@
 define i64 @extendedLeftShiftintToint64By4(i32 %a) nounwind readnone ssp {
 ; CHECK-LABEL: extendedLeftShiftintToint64By4:
 ; CHECK:       ; %bb.0: ; %entry
-; CHECK-NEXT:    add w8, w0, #1
-; CHECK-NEXT:    sbfiz x0, x8, #4, #32
+; CHECK-NEXT:    ; kill: def $w0 killed $w0 def $x0
+; CHECK-NEXT:    sbfiz x8, x0, #4, #32
+; CHECK-NEXT:    add x0, x8, #16
 ; CHECK-NEXT:    ret
 entry:
   %inc = add nsw i32 %a, 1
Index: llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
===================================================================
--- llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -9927,6 +9927,30 @@
     }
   }
 
+  // fold (shl (sext (add_nsw x, c1)), c2) -> (add (shl (sext x), c2), c1 << c2)
+  // fold (shl (zext (add_nuw x, c1)), c2) -> (add (shl (zext x), c2), c1 << c2)
+  if ((N0.getOpcode() == ISD::SIGN_EXTEND ||
+       N0.getOpcode() == ISD::ZERO_EXTEND) &&
+      N0.getOperand(0).getOpcode() == ISD::ADD && N0->hasOneUse() &&
+      N0.getOperand(0)->hasOneUse() &&
+      TLI.isDesirableToCommuteWithShift(N, Level)) {
+    SDValue Add = N0.getOperand(0);
+    bool IsSext = N0.getOpcode() == ISD::SIGN_EXTEND;
+    if ((IsSext && Add->getFlags().hasNoSignedWrap()) ||
+        (!IsSext && Add->getFlags().hasNoUnsignedWrap())) {
+      SDLoc DL(N0);
+      if (SDValue ExtC = DAG.FoldConstantArithmetic(N0.getOpcode(), DL, VT,
+                                                    {Add.getOperand(1)})) {
+        if (SDValue ShlC =
+                DAG.FoldConstantArithmetic(ISD::SHL, DL, VT, {ExtC, N1})) {
+          SDValue ExtX = DAG.getNode(N0.getOpcode(), DL, VT, Add.getOperand(0));
+          SDValue ShlX = DAG.getNode(ISD::SHL, DL, VT, ExtX, N1);
+          return DAG.getNode(ISD::ADD, DL, VT, ShlX, ShlC);
+        }
+      }
+    }
+  }
+
   // fold (shl (mul x, c1), c2) -> (mul x, c1 << c2)
   if (N0.getOpcode() == ISD::MUL && N0->hasOneUse()) {
     SDValue N01 = N0.getOperand(1);


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D159198.554690.patch
Type: text/x-patch
Size: 3885 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20230830/4c4f269b/attachment.bin>