[AArch64] Decompose constant multiplies used only by memory addresses

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -20433,6 +20433,37 @@ static SDValue performVectorExtCombine(SDNode *N, SelectionDAG &DAG) {
   return SDValue();
 }
 
+static bool isOnlyUsedAsMemoryAddress(SDNode *N) {
+  assert((N->getOpcode() == ISD::ADD || N->getOpcode() == ISD::SUB) &&
+         "Expected add/sub node");
+
+  for (SDUse &Use : N->uses()) {
+    SDNode *User = Use.getUser();
+    switch (User->getOpcode()) {
+    case ISD::LOAD: {
+      auto *Load = cast<LoadSDNode>(User);
+      if (!Load->isUnindexed() || Use.getOperandNo() != 1)
+        return false;
+      break;
+    }
+    case ISD::STORE: {
+      auto *Store = cast<StoreSDNode>(User);
+      if (!Store->isUnindexed() || Use.getOperandNo() != 2)
+        return false;
+      break;
+    }
+    case AArch64ISD::PREFETCH:
+      if (Use.getOperandNo() != 2)
+        return false;
+      break;
+    default:
+      return false;
+    }
+  }
+
+  return true;
+}
+
 static SDValue performMulCombine(SDNode *N, SelectionDAG &DAG,
                                  TargetLowering::DAGCombinerInfo &DCI,
                                  const AArch64Subtarget *Subtarget) {
@@ -20518,9 +20549,12 @@ static SDValue performMulCombine(SDNode *N, SelectionDAG &DAG,
       return SDValue();
     // Conservatively do not lower to shift+add+shift if the mul might be
     // folded into madd or msub.
-    if (N->hasOneUse() && (N->user_begin()->getOpcode() == ISD::ADD ||
-                           N->user_begin()->getOpcode() == ISD::SUB))
-      return SDValue();
+    if (N->hasOneUse()) {
+      SDNode *User = *N->user_begin();
+      if ((User->getOpcode() == ISD::ADD || User->getOpcode() == ISD::SUB) &&
+          !isOnlyUsedAsMemoryAddress(User))
+        return SDValue();
+    }
   }
   // Use ShiftedConstValue instead of ConstValue to support both shift+add/sub
   // and shift+add+shift.

llvm/test/CodeGen/AArch64/mul-const-addressing-mode.ll

@@ -0,0 +1,137 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc < %s -mtriple=aarch64 | FileCheck %s
+
+define i32 @ldrw_3x(ptr %p, i64 %x) {
+; CHECK-LABEL: ldrw_3x:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x1, x1, lsl #1
+; CHECK-NEXT:    ldr w0, [x0, x8, lsl #2]
+; CHECK-NEXT:    ret
+  %off = mul i64 %x, 12
+  %addr = getelementptr inbounds i8, ptr %p, i64 %off
+  %v = load i32, ptr %addr, align 4
+  ret i32 %v
+}
+
+define i32 @ldrw_5x(ptr %p, i64 %x) {
+; CHECK-LABEL: ldrw_5x:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x1, x1, lsl #2
+; CHECK-NEXT:    ldr w0, [x0, x8, lsl #2]
+; CHECK-NEXT:    ret
+  %off = mul i64 %x, 20
+  %addr = getelementptr inbounds i8, ptr %p, i64 %off
+  %v = load i32, ptr %addr, align 4
+  ret i32 %v
+}
+
+define i32 @ldrw_6x(ptr %p, i64 %x) {
+; CHECK-LABEL: ldrw_6x:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x1, x1, lsl #1
+; CHECK-NEXT:    lsl x8, x8, #3
+; CHECK-NEXT:    ldr w0, [x0, x8]
+; CHECK-NEXT:    ret
+  %off = mul i64 %x, 24
+  %addr = getelementptr inbounds i8, ptr %p, i64 %off
+  %v = load i32, ptr %addr, align 4
+  ret i32 %v
+}
+
+define void @strw_3x(ptr %p, i64 %x, i32 %v) {
+; CHECK-LABEL: strw_3x:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x1, x1, lsl #1
+; CHECK-NEXT:    str w2, [x0, x8, lsl #2]
+; CHECK-NEXT:    ret
+  %off = mul i64 %x, 12
+  %addr = getelementptr inbounds i8, ptr %p, i64 %off
+  store i32 %v, ptr %addr, align 4
+  ret void
+}
+
+define i64 @ldrx_3x(ptr %p, i64 %x) {
+; CHECK-LABEL: ldrx_3x:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x1, x1, lsl #1
+; CHECK-NEXT:    ldr x0, [x0, x8, lsl #3]
+; CHECK-NEXT:    ret
+  %off = mul i64 %x, 24
+  %addr = getelementptr inbounds i8, ptr %p, i64 %off
+  %v = load i64, ptr %addr, align 8
+  ret i64 %v
+}
+
+define i64 @ldrx_5x(ptr %p, i64 %x) {
+; CHECK-LABEL: ldrx_5x:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x1, x1, lsl #2
+; CHECK-NEXT:    ldr x0, [x0, x8, lsl #3]
+; CHECK-NEXT:    ret
+  %off = mul i64 %x, 40
+  %addr = getelementptr inbounds i8, ptr %p, i64 %off
+  %v = load i64, ptr %addr, align 8
+  ret i64 %v
+}
+
+define i64 @ldrx_6x(ptr %p, i64 %x) {
+; CHECK-LABEL: ldrx_6x:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x1, x1, lsl #1
+; CHECK-NEXT:    lsl x8, x8, #4
+; CHECK-NEXT:    ldr x0, [x0, x8]
+; CHECK-NEXT:    ret
+  %off = mul i64 %x, 48
+  %addr = getelementptr inbounds i8, ptr %p, i64 %off
+  %v = load i64, ptr %addr, align 8
+  ret i64 %v
+}
+
+define void @strx_3x(ptr %p, i64 %x, i64 %v) {
+; CHECK-LABEL: strx_3x:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x1, x1, lsl #1
+; CHECK-NEXT:    str x2, [x0, x8, lsl #3]
+; CHECK-NEXT:    ret
+  %off = mul i64 %x, 24
+  %addr = getelementptr inbounds i8, ptr %p, i64 %off
+  store i64 %v, ptr %addr, align 8
+  ret void
+}
+
+define void @store_add_as_value_and_addr(i64 %x, i64 %y) {
+; CHECK-LABEL: store_add_as_value_and_addr:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #24 // =0x18
+; CHECK-NEXT:    mul x8, x0, x8
+; CHECK-NEXT:    add x9, x8, x1
+; CHECK-NEXT:    str x9, [x8, x1]
+; CHECK-NEXT:    ret
+  %mul = mul i64 %x, 24
+  %add = add i64 %mul, %y
+  %addr = inttoptr i64 %add to ptr
+  store i64 %add, ptr %addr, align 8
+  ret void
+}
+
+define i64 @arithmetic_mul_add(i64 %x, i64 %y) {
+; CHECK-LABEL: arithmetic_mul_add:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #24 // =0x18
+; CHECK-NEXT:    madd x0, x0, x8, x1
+; CHECK-NEXT:    ret
+  %mul = mul i64 %x, 24
+  %add = add i64 %mul, %y
+  ret i64 %add
+}
+
+define i64 @arithmetic_mul_sub(i64 %x, i64 %y) {
+; CHECK-LABEL: arithmetic_mul_sub:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #24 // =0x18
+; CHECK-NEXT:    msub x0, x0, x8, x1
+; CHECK-NEXT:    ret
+  %mul = mul i64 %x, 24
+  %sub = sub i64 %y, %mul
+  ret i64 %sub
+}