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

In some cases, we are too aggressive when folding an add-lsl into an
ldr/str due to an accidental 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...

Author: Dudeldu

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -3232,6 +3232,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();

llvm/test/CodeGen/AArch64/fuse-addr-mode.mir

@@ -0,0 +1,47 @@
+# 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: [[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
+    %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))
+    %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