[llvm] [AArch64] Allow only LSL to be folded into addressing mode (PR #69235)

[Momchil Velikov via llvm-commits]

https://github.com/momchil-velikov created https://github.com/llvm/llvm-project/pull/69235

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.

>From d33680643201afa2bcacd58aab3977c7ed3594e9 Mon Sep 17 00:00:00 2001
From: Momchil Velikov <momchil.velikov at arm.com>
Date: Mon, 16 Oct 2023 15:46:59 +0100
Subject: [PATCH] [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 e03a94de007c9f5..8f0e272a6fac788 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 000000000000000..b9892fc31bedb5e
--- /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
+}