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

[wwei]

Summary

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 while it will use getImm() directly.

[efriedma]

llvm/test/CodeGen/AArch64/arm64-ldst-opt-match.ll

20 ↗

(On Diff #292785)

This testcase seems extremely fragile; can we use an MIR testcase instead? (There are a bunch of in-tree examples of using -run-pass=aarch64-ldst-opt.)

[efriedma]

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

1190

Also, maybe fix this comment to clarify what the problem is. The store should be an opcode that can use an immediate operand. The problem is that the "immediate" might actually be the address of a global variable.

[wwei]

Update test case by using MIR

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

1190

fixed.

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

do {
     // may be different numbers of them if e.g. debug information is present.
     if (!MI.isTransient())
       ++Count;
 
     // If the load instruction reads directly from the address to which the
     // 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;
       return true;
     }
 
     if (MI.isCall())
       return false;

llvm/test/CodeGen/AArch64/ldst-opt-non-imm-offset.mir

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