[DAGCombine] Remove OneUse restriction when folding (shl (add x, c1), c2)

Author: LiqinWeng

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -10070,17 +10070,33 @@ SDValue DAGCombiner::visitSHL(SDNode *N) {
   // Variant of version done on multiply, except mul by a power of 2 is turned
   // into a shift.
   if ((N0.getOpcode() == ISD::ADD || N0.getOpcode() == ISD::OR) &&
-      N0->hasOneUse() && TLI.isDesirableToCommuteWithShift(N, Level)) {
-    SDValue N01 = N0.getOperand(1);
-    if (SDValue Shl1 =
-            DAG.FoldConstantArithmetic(ISD::SHL, SDLoc(N1), VT, {N01, N1})) {
-      SDValue Shl0 = DAG.getNode(ISD::SHL, SDLoc(N0), VT, N0.getOperand(0), N1);
-      AddToWorklist(Shl0.getNode());
-      SDNodeFlags Flags;
-      // Preserve the disjoint flag for Or.
-      if (N0.getOpcode() == ISD::OR && N0->getFlags().hasDisjoint())
-        Flags.setDisjoint(true);
-      return DAG.getNode(N0.getOpcode(), DL, VT, Shl0, Shl1, Flags);
+      TLI.isDesirableToCommuteWithShift(N, Level)) {
+    // LD/ST will optimize constant Offset extraction, so when AddNode
+    // is used by LD/ST, it can still complete the folding optimization
+    // operation performed above.
+    bool canOptAwlays = false;
+    if (!N0.hasOneUse() && N0.getOpcode() == ISD::ADD) {
+      for (SDNode *Use : N0->uses()) {
+        if (!isa<StoreSDNode>(Use) && !isa<LoadSDNode>(Use) && Use != N) {
+          canOptAwlays = false;
+          break;
+        }
+        canOptAwlays = true;
+      }
+    }
+    if (N0.hasOneUse() || canOptAwlays) {
+      SDValue N01 = N0.getOperand(1);
+      if (SDValue Shl1 =
+              DAG.FoldConstantArithmetic(ISD::SHL, SDLoc(N1), VT, {N01, N1})) {
+        SDValue Shl0 =
+            DAG.getNode(ISD::SHL, SDLoc(N0), VT, N0.getOperand(0), N1);
+        AddToWorklist(Shl0.getNode());
+        SDNodeFlags Flags;
+        // Preserve the disjoint flag for Or.
+        if (N0.getOpcode() == ISD::OR && N0->getFlags().hasDisjoint())
+          Flags.setDisjoint(true);
+        return DAG.getNode(N0.getOpcode(), DL, VT, Shl0, Shl1, Flags);
+      }
     }
   }
 

llvm/test/CodeGen/RISCV/add_shl_constant.ll

@@ -28,11 +28,9 @@ define void @add_shl_addmoreOneUse_in_store(ptr %array1, i32 %a, i32 %b)  {
 ; RV32-LABEL: add_shl_addmoreOneUse_in_store:
 ; RV32:       # %bb.0: # %entry
 ; RV32-NEXT:    addi a3, a1, 5
-; RV32-NEXT:    slli a4, a3, 2
-; RV32-NEXT:    add a4, a0, a4
-; RV32-NEXT:    sw a2, 0(a4)
 ; RV32-NEXT:    slli a1, a1, 2
 ; RV32-NEXT:    add a0, a0, a1
+; RV32-NEXT:    sw a2, 20(a0)
 ; RV32-NEXT:    sw a2, 24(a0)
 ; RV32-NEXT:    sw a3, 140(a0)
 ; RV32-NEXT:    ret

llvm/test/CodeGen/RISCV/riscv-shifted-extend.ll

@@ -2,7 +2,7 @@
 ; RUN: llc -mtriple=riscv64 -verify-machineinstrs < %s \
 ; RUN:   | FileCheck -check-prefix=RV64 %s
 
-define void @test(ptr nocapture noundef writeonly %array1, i32 noundef signext %a, i32 noundef signext %b) {
+define void @test(ptr %array1, i32 signext %a, i32 signext %b) {
 ; RV64-LABEL: test:
 ; RV64:       # %bb.0: # %entry
 ; RV64-NEXT:    addiw a3, a1, 5
@@ -31,7 +31,7 @@ entry:
 }
 
 ; test of jumpping, find add's operand has one more use can simplified
-define void @test1(ptr nocapture noundef %array1, i32 noundef signext %a, i32 noundef signext %b, i32 noundef signext %x) {
+define void @test1(ptr %array1, i32 signext %a, i32 signext %b, i32 signext %x) {
 ; RV64-LABEL: test1:
 ; RV64:       # %bb.0: # %entry
 ; RV64-NEXT:    addiw a4, a1, 5
@@ -66,15 +66,13 @@ entry:
   ret void
 }
 
-define void @test2(ptr nocapture noundef writeonly %array1, i64 noundef %a, i64 noundef %b) local_unnamed_addr #0 {
+define void @test2(ptr %array1, i64 %a, i64 %b) {
 ; RV64-LABEL: test2:
 ; RV64:       # %bb.0: # %entry
 ; RV64-NEXT:    addi a3, a1, 5
-; RV64-NEXT:    slli a4, a3, 3
-; RV64-NEXT:    add a4, a0, a4
-; RV64-NEXT:    sd a2, 0(a4)
 ; RV64-NEXT:    slli a1, a1, 3
 ; RV64-NEXT:    add a0, a1, a0
+; RV64-NEXT:    sd a2, 40(a0)
 ; RV64-NEXT:    sd a2, 48(a0)
 ; RV64-NEXT:    sd a3, 280(a0)
 ; RV64-NEXT:    ret
@@ -91,7 +89,7 @@ entry:
   ret void
 }
 
-define void @test3(ptr nocapture noundef %array1, i64 noundef %a, i64 noundef %b, i64 noundef %x) {
+define void @test3(ptr %array1, i64 %a, i64 %b, i64 %x) {
 ; RV64-LABEL: test3:
 ; RV64:       # %bb.0: # %entry
 ; RV64-NEXT:    addi a4, a1, 5
@@ -100,11 +98,9 @@ define void @test3(ptr nocapture noundef %array1, i64 noundef %a, i64 noundef %b
 ; RV64-NEXT:  # %bb.1: # %entry
 ; RV64-NEXT:    mv a5, a2
 ; RV64-NEXT:  .LBB3_2: # %entry
-; RV64-NEXT:    slli a2, a4, 3
-; RV64-NEXT:    add a2, a0, a2
-; RV64-NEXT:    sd a5, 0(a2)
 ; RV64-NEXT:    slli a1, a1, 3
 ; RV64-NEXT:    add a0, a1, a0
+; RV64-NEXT:    sd a5, 40(a0)
 ; RV64-NEXT:    sd a5, 48(a0)
 ; RV64-NEXT:    sd a4, 280(a0)
 ; RV64-NEXT:    ret