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

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;

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

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