[llvm] [AArch64] Fix invalid address-mode folding (PR #142167)

[via llvm-commits]

https://github.com/Dudeldu updated https://github.com/llvm/llvm-project/pull/142167

>From 09070d0122accd334fbdf7997bbf878a8854789e Mon Sep 17 00:00:00 2001
From: Florian Drescher <florian.drescher at tum.de>
Date: Fri, 30 May 2025 17:10:45 +0200
Subject: [PATCH 1/2] [AArch64] Fix invalid address-mode folding

In some cases we are too aggressive when folding an add-lsl into an
ldr/str due to an accidential truncation of the 64-bit scale to 32 bit.
In cases where we shift by more than 31 bits (which is valid for 64-bit registers)
we just drop the shift...
---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp |  2 +
 llvm/test/CodeGen/AArch64/fuse-addr-mode.mir | 39 ++++++++++++++++++++
 2 files changed, 41 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/fuse-addr-mode.mir

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index d1e0d37e33e4e..9ac3727aad1f1 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -3229,6 +3229,8 @@ bool AArch64InstrInfo::canFoldIntoAddrMode(const MachineInstr &MemI,
           ExtAddrMode::Formula Form = ExtAddrMode::Formula::Basic) -> bool {
     if (MemI.getOperand(2).getImm() != 0)
       return false;
+    if ((unsigned)Scale != Scale)
+      return false;
     if (!isLegalAddressingMode(NumBytes, /* Offset */ 0, Scale))
       return false;
     AM.BaseReg = AddrI.getOperand(1).getReg();
diff --git a/llvm/test/CodeGen/AArch64/fuse-addr-mode.mir b/llvm/test/CodeGen/AArch64/fuse-addr-mode.mir
new file mode 100644
index 0000000000000..69bdbc9809ed2
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/fuse-addr-mode.mir
@@ -0,0 +1,39 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
+# RUN: llc -mtriple=aarch64 -o - %s -run-pass machine-sink | FileCheck %s
+
+# we want to fuse an addition with lsl into an ldr but we have to be careful with
+# the shift distance: we can only represent specific shift distances: e.g: 3
+# but nothing large like 32
+
+--- |
+  define dso_local i64 @fuse_shift_add_into_addr_mode()  {
+  entry:
+    ret i64 0
+  }
+
+---
+name:            fuse_shift_add_into_addr_mode
+body:             |
+  bb.1.entry:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: fuse_shift_add_into_addr_mode
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64common = COPY $x0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
+    ; CHECK-NEXT: [[LDRXroX:%[0-9]+]]:gpr64 = LDRXroX [[COPY]], [[COPY1]], 0, 1 :: (load (s64))
+    ; CHECK-NEXT: [[ADDXrs:%[0-9]+]]:gpr64common = ADDXrs [[COPY]], [[COPY1]], 5
+    ; CHECK-NEXT: [[LDRXui:%[0-9]+]]:gpr64 = LDRXui [[ADDXrs]], 0 :: (load (s64))
+    ; CHECK-NEXT: [[ADDXrs1:%[0-9]+]]:gpr64common = ADDXrs [[COPY]], [[COPY1]], 32
+    ; CHECK-NEXT: [[LDRXui1:%[0-9]+]]:gpr64 = LDRXui [[ADDXrs1]], 0 :: (load (s64))
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %0:gpr64 = COPY $x0
+    %1:gpr64 = COPY $x1
+    %2:gpr64common = ADDXrs %0, %1, 3
+    %3:gpr64 = LDRXui %2, 0 :: (load (s64))
+    %4:gpr64common = ADDXrs %0, %1, 5
+    %5:gpr64 = LDRXui %4, 0 :: (load (s64))
+    %6:gpr64common = ADDXrs %0, %1, 32
+    %7:gpr64 = LDRXui %6, 0 :: (load (s64))
+    RET_ReallyLR implicit $x0

>From e32bfcb1a0b4cfd10d91bc65a6ef0e7b095e282c Mon Sep 17 00:00:00 2001
From: Florian Drescher <florian.drescher at tum.de>
Date: Mon, 2 Jun 2025 08:38:26 +0200
Subject: [PATCH 2/2] Add more testcases

---
 llvm/test/CodeGen/AArch64/fuse-addr-mode.mir | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/llvm/test/CodeGen/AArch64/fuse-addr-mode.mir b/llvm/test/CodeGen/AArch64/fuse-addr-mode.mir
index 69bdbc9809ed2..5d218d0a2f24a 100644
--- a/llvm/test/CodeGen/AArch64/fuse-addr-mode.mir
+++ b/llvm/test/CodeGen/AArch64/fuse-addr-mode.mir
@@ -27,6 +27,10 @@ body:             |
     ; CHECK-NEXT: [[LDRXui:%[0-9]+]]:gpr64 = LDRXui [[ADDXrs]], 0 :: (load (s64))
     ; CHECK-NEXT: [[ADDXrs1:%[0-9]+]]:gpr64common = ADDXrs [[COPY]], [[COPY1]], 32
     ; CHECK-NEXT: [[LDRXui1:%[0-9]+]]:gpr64 = LDRXui [[ADDXrs1]], 0 :: (load (s64))
+    ; CHECK-NEXT: [[ADDXrs2:%[0-9]+]]:gpr64common = ADDXrs [[COPY]], [[COPY1]], 35
+    ; CHECK-NEXT: [[LDRXui2:%[0-9]+]]:gpr64 = LDRXui [[ADDXrs1]], 0 :: (load (s64))
+    ; CHECK-NEXT: [[ADDXrs3:%[0-9]+]]:gpr64common = ADDXrs [[COPY]], [[COPY1]], 63
+    ; CHECK-NEXT: [[LDRXui3:%[0-9]+]]:gpr64 = LDRXui [[ADDXrs1]], 0 :: (load (s64))
     ; CHECK-NEXT: RET_ReallyLR implicit $x0
     %0:gpr64 = COPY $x0
     %1:gpr64 = COPY $x1
@@ -36,4 +40,8 @@ body:             |
     %5:gpr64 = LDRXui %4, 0 :: (load (s64))
     %6:gpr64common = ADDXrs %0, %1, 32
     %7:gpr64 = LDRXui %6, 0 :: (load (s64))
+    %8:gpr64common = ADDXrs %0, %1, 35
+    %9:gpr64 = LDRXui %6, 0 :: (load (s64))
+    %10:gpr64common = ADDXrs %0, %1, 63
+    %11:gpr64 = LDRXui %6, 0 :: (load (s64))
     RET_ReallyLR implicit $x0