[AArch64] Allow only LSL to be folded into addressing mode

[momchil-velikov]

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.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

@llvm/pr-subscribers-llvm-globalisel

Author: Momchil Velikov (momchil-velikov)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/69235.diff

2 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64InstrInfo.cpp (+4-1)
• (added) llvm/test/CodeGen/AArch64/GlobalISel/sink-and-fold-illegal-shift.ll (+17)

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
+}

[davemgreen]

It might be a good idea to add this as a mir test, as GISel could change in the future to be more like SDAG.

Otherwise LGTM

[momchil-velikov]

It might be a good idea to add this as a mir test, as GISel could change in the future to be more like SDAG.

Ah, indeed! Done.