[AArch64] Fix ldst optimization of non-immediate store offset

When matching store instruction for ldst opt, we should make sure store instr is in 'reg+imm' form as load instr,
otherwise, it will have assertion in isLdOffsetInRangeOfSt since it will use getImm() directly.

Reviewed By: efriedma

Differential Revision: https://reviews.llvm.org/D87905

diff --git a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
index 8d90a9a..ea2e848 100644 (file)
--- a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
+++ b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
@@ -1186,8 +1186,10 @@ bool AArch64LoadStoreOpt::findMatchingStore(
     // store instruction writes and the stored value is not modified, we can
     // promote the load. Since we do not handle stores with pre-/post-index,
     // it's unnecessary to check if BaseReg is modified by the store itself.
+    // Also we can't handle stores without an immediate offset operand,
+    // while the operand might be the address for a global variable.
     if (MI.mayStore() && isMatchingStore(LoadMI, MI) &&
-        BaseReg == getLdStBaseOp(MI).getReg() &&
+        BaseReg == getLdStBaseOp(MI).getReg() && getLdStOffsetOp(MI).isImm() &&
         isLdOffsetInRangeOfSt(LoadMI, MI, TII) &&
         ModifiedRegUnits.available(getLdStRegOp(MI).getReg())) {
       StoreI = MBBI;

diff --git a/llvm/test/CodeGen/AArch64/ldst-opt-non-imm-offset.mir b/llvm/test/CodeGen/AArch64/ldst-opt-non-imm-offset.mir
new file mode 100644 (file)
index 0000000..86dff69
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/ldst-opt-non-imm-offset.mir
@@ -0,0 +1,27 @@
+# RUN: llc -mtriple=aarch64 -run-pass=aarch64-ldst-opt %s -verify-machineinstrs -o - | FileCheck %s
+--- |
+  @g = common dso_local global i32 0, align 4
+
+  define i32 @test() {
+  entry:
+    store i32 0, i32* @g, align 4
+    %0 = load i32, i32* undef, align 4
+    ret i32 %0
+  }
+
+...
+---
+# Don't crash when there's no immediate operand for store.
+# CHECK-LABEL: name: test
+# CHECK: STRWui $wzr
+# CHECK: LDRWui
+name:            test
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    renamable $x8 = ADRP target-flags(aarch64-page) @g
+    STRWui $wzr, killed renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @g :: (store 4 into @g)
+    renamable $w0 = LDRWui undef renamable $x8, 0 :: (load 4 from `i32* undef`)
+    RET_ReallyLR implicit $w0
+
+...