From d33680643201afa2bcacd58aab3977c7ed3594e9 Mon Sep 17 00:00:00 2001
From: Momchil Velikov <momchil.velikov@arm.com>
Date: Mon, 16 Oct 2023 15:46:59 +0100
Subject: [PATCH 1/2] [AArch64] Allow only LSL to be folded into addressing
 mode

There was an error in decoding shift type, which permitted shift types
other than LSL to be (incorrectly) folded into the addressing mode of
a load/store instruction.
---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp    |  5 ++++-
 .../GlobalISel/sink-and-fold-illegal-shift.ll   | 17 +++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AArch64/GlobalISel/sink-and-fold-illegal-shift.ll

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index e03a94de007c9..8f0e272a6fac7 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -2978,7 +2978,10 @@ bool AArch64InstrInfo::canFoldIntoAddrMode(const MachineInstr &MemI,
 
     // Don't fold the add if the result would be slower, unless optimising for
     // size.
-    int64_t Shift = AddrI.getOperand(3).getImm();
+    unsigned Shift = static_cast<unsigned>(AddrI.getOperand(3).getImm());
+    if (AArch64_AM::getShiftType(Shift) != AArch64_AM::ShiftExtendType::LSL)
+      return false;
+    Shift = AArch64_AM::getShiftValue(Shift);
     if (!OptSize) {
       if ((Shift != 2 && Shift != 3) || !Subtarget.hasAddrLSLFast())
         return false;
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/sink-and-fold-illegal-shift.ll b/llvm/test/CodeGen/AArch64/GlobalISel/sink-and-fold-illegal-shift.ll
new file mode 100644
index 0000000000000..b9892fc31bedb
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/sink-and-fold-illegal-shift.ll
@@ -0,0 +1,17 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 3
+; RUN: llc -global-isel --aarch64-enable-sink-fold=true < %s | FileCheck %s
+
+target triple = "aarch64-linux"
+
+; Test a non-LSL shift cannot be folded into the addressing mode.
+define void @f(ptr %p, i64 %i) optsize {
+; CHECK-LABEL: f:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x0, x1, asr #32
+; CHECK-NEXT:    strb wzr, [x8]
+; CHECK-NEXT:    ret
+	%d = ashr i64 %i, 32
+	%a = getelementptr i8, ptr %p, i64 %d
+	store i8 0, ptr %a
+	ret void
+}

From b59bf1e5af3d5dd97b8a81f3a2424785c5edfe13 Mon Sep 17 00:00:00 2001
From: Momchil Velikov <momchil.velikov@arm.com>
Date: Tue, 17 Oct 2023 10:33:26 +0100
Subject: [PATCH 2/2] Add a MIR test

---
 .../AArch64/sink-and-fold-illegal-shift.mir   | 95 +++++++++++++++++++
 1 file changed, 95 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/sink-and-fold-illegal-shift.mir

diff --git a/llvm/test/CodeGen/AArch64/sink-and-fold-illegal-shift.mir b/llvm/test/CodeGen/AArch64/sink-and-fold-illegal-shift.mir
new file mode 100644
index 0000000000000..d2f6a3ab1aeeb
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/sink-and-fold-illegal-shift.mir
@@ -0,0 +1,95 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 3
+# RUN: llc --run-pass=machine-sink --aarch64-enable-sink-fold=true %s -o - | FileCheck %s
+--- |
+  source_filename = "../llvm/test/CodeGen/AArch64/GlobalISel/sink-and-fold-illegal-shift.ll"
+  target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64-linux"
+
+  define void @f(ptr %p, i64 %i) #0 {
+    %d = ashr i64 %i, 32
+    %a = getelementptr i8, ptr %p, i64 %d
+    store i8 0, ptr %a, align 1
+    ret void
+  }
+
+  attributes #0 = { optsize }
+
+...
+---
+name:            f
+alignment:       4
+exposesReturnsTwice: false
+legalized:       true
+regBankSelected: true
+selected:        true
+failedISel:      false
+tracksRegLiveness: true
+hasWinCFI:       false
+callsEHReturn:   false
+callsUnwindInit: false
+hasEHCatchret:   false
+hasEHScopes:     false
+hasEHFunclets:   false
+isOutlined:      false
+debugInstrRef:   false
+failsVerification: false
+tracksDebugUserValues: false
+registers:
+  - { id: 0, class: gpr64, preferred-register: '' }
+  - { id: 1, class: gpr64, preferred-register: '' }
+  - { id: 2, class: gpr, preferred-register: '' }
+  - { id: 3, class: gpr, preferred-register: '' }
+  - { id: 4, class: gpr64common, preferred-register: '' }
+  - { id: 5, class: _, preferred-register: '' }
+  - { id: 6, class: gpr, preferred-register: '' }
+  - { id: 7, class: gpr64, preferred-register: '' }
+liveins:
+  - { reg: '$x0', virtual-reg: '' }
+  - { reg: '$x1', virtual-reg: '' }
+frameInfo:
+  isFrameAddressTaken: false
+  isReturnAddressTaken: false
+  hasStackMap:     false
+  hasPatchPoint:   false
+  stackSize:       0
+  offsetAdjustment: 0
+  maxAlignment:    1
+  adjustsStack:    false
+  hasCalls:        false
+  stackProtector:  ''
+  functionContext: ''
+  maxCallFrameSize: 0
+  cvBytesOfCalleeSavedRegisters: 0
+  hasOpaqueSPAdjustment: false
+  hasVAStart:      false
+  hasMustTailInVarArgFunc: false
+  hasTailCall:     false
+  localFrameSize:  0
+  savePoint:       ''
+  restorePoint:    ''
+fixedStack:      []
+stack:           []
+entry_values:    []
+callSites:       []
+debugValueSubstitutions: []
+constants:       []
+machineFunctionInfo: {}
+body:             |
+  bb.1 (%ir-block.0):
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: f
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY $x1
+    ; CHECK-NEXT: [[ADDXrs:%[0-9]+]]:gpr64common = ADDXrs [[COPY]], [[COPY1]], 160
+    ; CHECK-NEXT: STRBBui $wzr, [[ADDXrs]], 0 :: (store (s8) into %ir.a)
+    ; CHECK-NEXT: RET_ReallyLR
+    %0:gpr64 = COPY $x0
+    %1:gpr64 = COPY $x1
+    %4:gpr64common = ADDXrs %0, %1, 160
+    STRBBui $wzr, %4, 0 :: (store (s8) into %ir.a)
+    RET_ReallyLR
+
+...