[AArch64][Fix] LdSt optimization generate premature stack-popping

[dnsampaio]

Summary

When moving add and sub to memory operand instructions,
aarch64-ldst-opt would prematurally pop the stack pointer,
before memory instructions that do access the stack using
indirect loads.
e.g.

int foo(int offset){
    int local[4] = {0};
    return local[offset];
}

would generate:

sub     sp, sp, #16            ; Push the stack
mov     x8, sp                 ; Save stack in register
stp     xzr, xzr, [sp], #16    ; Zero initialize stack, and post-increment, making it invalid
------ If an exception goes here, the stack value might be corrupted
ldr     w0, [x8, w0, sxtw #2]  ; Access correct position, but it is not guarded by SP

[efriedma]

I think there's a secondary problem here: if we move an SP adjustment, we might also need to adjust the unwind/debug info. This is particularly nasty on Windows, where messing up the unwind info can actually cause a miscompile.

If you don't want to spend the time to try to figure that out, I'd be happy to just completely forbid optimizing sp adjustments here. It should be rare enough that it doesn't matter much in practice.

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

1844

I'm not sure what you think this mayAlias check is doing?

[dnsampaio]

In D75755#1910057, @efriedma wrote:

I think there's a secondary problem here: if we move an SP adjustment, we might also need to adjust the unwind/debug info. This is particularly nasty on Windows, where messing up the unwind info can actually cause a miscompile.

If you don't want to spend the time to try to figure that out, I'd be happy to just completely forbid optimizing sp adjustments here. It should be rare enough that it doesn't matter much in practice.

Hi @efriedma, thanks for the review.
I would be glad in taking a look into the windows unwind issue, but I don't know anything about it. Could you give me some further details?
Perhaps we can just prevent any SP adjustment for the moment, until I have the windows fix in place.

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

1844

Following a similar code from MachineInstr.mayAlias code, I thought this was just eliminating some trivial cases where the instruction would not access the stack, but now that I see the definition, I'm not so sure.

[dnsampaio]

Do not optimize post-increment SP in windows targets

clang-format-diff

[efriedma]

The problem on Windows is that every instruction in the prologue/epilogue has an associated unwind opcode, represented as a pseudo-instruction. If those pseudo-instructions don't match the actual instructions, the offsets will all be wrong, so the unwind info will be corrupt. In practice, we saw a miscompile involving a call placed immediately after a ret instruction. The check should look something like https://github.com/llvm/llvm-project/blob/6d026c89dc6a3ba405e6767a3aa0bbb6ba27a0c9/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp#L1998 .

DWARF unwind doesn't work like that, of course. But you'll still get very weird results if you set a breakpoint in certain places. I'm pretty sure we have a pseudo-instruction for DWARF stack adjustments? I forget what that looks like.

[dnsampaio]

• Fixed tests
• Avoid optimization when target is Windows with CFI, as we need to correctly update unwind information

[efriedma]

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

1844

I still don't know why you're calling PseudoSourceValue::mayAlias. You don't care whether the access may alias other accesses; you care whether you're accessing the stack.

[dnsampaio]

• Correct mayAlias to isStack

[efriedma]

Despite the name, isStack() doesn't include all accesses to the stack; there's also FixedStack PseudoSourceValues.

That said, I'm not sure it's worth bothering; non-stack PseudoSourceValues are relatively rare, at least on AArch64, so you're not getting much benefit.

[dnsampaio]

• Do not use pseudo-value

[efriedma]

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

1814

"continue"? Should we really be continuing to search past an instruction which modifies sp?

[dnsampaio]

• Do not allow a instruction that may load or store between the LdSt[SP] and the SP update.

[dnsampaio]

• Removed unecessary changes

[dnsampaio]

• Removed incorrect spacing changes

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/IR/DebugLoc.h"
+#include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/DebugCounter.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cassert>
struct AArch64LoadStoreOpt : public MachineFunctionPass {
 
   // Scan the instruction list to find a base register update that can
   // be combined with the current instruction (a load or store) using
   // pre or post indexed addressing with writeback. Scan backwards.
   MachineBasicBlock::iterator
   findMatchingUpdateInsnBackward(MachineBasicBlock::iterator I, unsigned Limit);
 
   // Find an instruction that updates the base register of the ld/st
-  // instruction.
-  bool isMatchingUpdateInsn(MachineInstr &MemMI, MachineInstr &MI,
-                            unsigned BaseReg, int Offset);
+  // instruction. Return update Offset or 0 if not found.
+  int isMatchingUpdateInsn(MachineInstr &MemMI, MachineInstr &MI,
 
   // Merge a pre- or post-index base register update into a ld/st instruction.
   MachineBasicBlock::iterator
   mergeUpdateInsn(MachineBasicBlock::iterator I,
                   MachineBasicBlock::iterator Update, bool IsPreIdx);
 
   // Find and merge zero store instructions.
   bool tryToMergeZeroStInst(MachineBasicBlock::iterator &MBBI);
AArch64LoadStoreOpt::mergeUpdateInsn(MachineBasicBlock::iterator I,
 
   // Erase the old instructions for the block.
   I->eraseFromParent();
   Update->eraseFromParent();
 
   return NextI;
 }
 
-bool AArch64LoadStoreOpt::isMatchingUpdateInsn(MachineInstr &MemMI,
-                                               MachineInstr &MI,
-                                               unsigned BaseReg, int Offset) {
+int AArch64LoadStoreOpt::isMatchingUpdateInsn(MachineInstr &MemMI,
   switch (MI.getOpcode()) {
   default:
     break;
   case AArch64::SUBXri:
   case AArch64::ADDXri:
     // Make sure it's a vanilla immediate operand, not a relocation or
     // anything else we can't handle.
     if (!MI.getOperand(2).isImm())
case AArch64::ADDXri:
     // Scaled offset must fit in the instruction immediate.
     int ScaledOffset = UpdateOffset / Scale;
     if (ScaledOffset > MaxOffset || ScaledOffset < MinOffset)
       break;
 
     // If we have a non-zero Offset, we check that it matches the amount
     // we're adding to the register.
     if (!Offset || Offset == UpdateOffset)
-      return true;
+      return UpdateOffset;
     break;
   }
-  return false;
+  return 0;
 }
 
 MachineBasicBlock::iterator AArch64LoadStoreOpt::findMatchingUpdateInsnForward(
     MachineBasicBlock::iterator I, int UnscaledOffset, unsigned Limit) {
   MachineBasicBlock::iterator E = I->getParent()->end();
   MachineInstr &MemMI = *I;
   MachineBasicBlock::iterator MBBI = I;
 
if (!isTagStore(MemMI) && MemMI.getOpcode() != AArch64::STGPi) {
     }
   }
 
   // Track which register units have been modified and used between the first
   // insn (inclusive) and the second insn.
   ModifiedRegUnits.clear();
   UsedRegUnits.clear();
   ++MBBI;
+
+  // We can't post-increment the stack pointer if any instruction between
+  // the memory access (I) and the increment (MBBI) can access the memory
+  // region defined by [SP, MBBI]. Premature stack
+  const bool BaseRegSP = BaseReg == AArch64::SP;
+
+  if (BaseRegSP) {
+    // FIXME: For now, we always block the optimization over SP in windows
+    // targets as it requires to adjust the unwind/debug info, messing up
+    // the unwind info can actually cause a miscompile.
+    const MCAsmInfo *MAI = I->getMF()->getTarget().getMCAsmInfo();
+    if (MAI->usesWindowsCFI() &&
+        I->getMF()->getFunction().needsUnwindTableEntry())
+      return E;
+  }
+
+  bool MayAccessStack = false;
   for (unsigned Count = 0; MBBI != E && Count < Limit; ++MBBI) {
     MachineInstr &MI = *MBBI;
 
     // Don't count transient instructions towards the search limit since there
     // may be different numbers of them if e.g. debug information is present.
     if (!MI.isTransient())
       ++Count;
 
     // If we found a match, return it.
-    if (isMatchingUpdateInsn(*I, MI, BaseReg, UnscaledOffset))
+    const int Offset = isMatchingUpdateInsn(*I, MI, BaseReg, UnscaledOffset);
+    if (Offset) {
+      // Prevent premature stack-popping
+      if (MayAccessStack && Offset > 0)
+        continue;
       return MBBI;
+    }
 
     // Update the status of what the instruction clobbered and used.
     LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits, UsedRegUnits, TRI);
 
     // Otherwise, if the base register is used or modified, we have no match, so
     // return early.
     if (!ModifiedRegUnits.available(BaseReg) ||
         !UsedRegUnits.available(BaseReg))
       return E;
+
+    // FIXME: For now we conservativly detect any instruction that might access
+    // the stack. We should validate that the access is in the memory region of
+    // two memory operands, (e.g. [SP, MBBI]), but that requires a new may-alias
+    // API. May also need to track if other registers might point to the stack.
+    if (!MayAccessStack && BaseRegSP && MBBI->mayLoadOrStore()) {
+      if (MBBI->memoperands_empty()) {
+        MayAccessStack = true;
+        continue;
+      }
+      MachineFrameInfo *MFI = &(I->getMF()->getFrameInfo());
+      assert(MFI);
+      for (MachineMemOperand *MMO : MBBI->memoperands()) {
+        const PseudoSourceValue *PSVa = MMO->getPseudoValue();
+        MayAccessStack = (!PSVa) || PSVa->mayAlias(MFI);
+        if (MayAccessStack)
+          break;
+      }
+    }
   }
   return E;
 }
 
 MachineBasicBlock::iterator AArch64LoadStoreOpt::findMatchingUpdateInsnBackward(
     MachineBasicBlock::iterator I, unsigned Limit) {
   MachineBasicBlock::iterator B = I->getParent()->begin();
   MachineBasicBlock::iterator E = I->getParent()->end();
do {
     MachineInstr &MI = *MBBI;
 
     // Don't count transient instructions towards the search limit since there
     // may be different numbers of them if e.g. debug information is present.
     if (!MI.isTransient())
       ++Count;
 
     // If we found a match, return it.
-    if (isMatchingUpdateInsn(*I, MI, BaseReg, Offset))
+    if (isMatchingUpdateInsn(*I, MI, BaseReg, Offset) != 0)
       return MBBI;
 
     // Update the status of what the instruction clobbered and used.
     LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits, UsedRegUnits, TRI);
 
     // Otherwise, if the base register is used or modified, we have no match, so
     // return early.
     if (!ModifiedRegUnits.available(BaseReg) ||

llvm/test/CodeGen/AArch64/aarch64-ldst-no-premature-sp-pop.mir

+# RUN: llc -start-before=aarch64-ldst-opt -o - %s | FileCheck %s
+# CHECK-NOT: stp     xzr, xzr, [sp], #16
+# CHECK:     add     sp, sp, #16
+--- |
+  target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64-arm-none-eabi"
+  
+  define hidden i32 @foo(i32 %0) {
+    %2 = alloca [4 x i32], align 4
+    %3 = bitcast [4 x i32]* %2 to i8*
+    call void @llvm.memset.p0i8.i64(i8* nonnull align 4 dereferenceable(16) %3, i8 0, i64 16, i1 false)
+    %4 = sext i32 %0 to i64
+    %5 = getelementptr inbounds [4 x i32], [4 x i32]* %2, i64 0, i64 %4
+    %6 = load i32, i32* %5, align 4
+    ret i32 %6
+  }
+  
+  declare void @llvm.memset.p0i8.i64(i8* nocapture writeonly, i8, i64, i1 immarg) #2
+  declare void @llvm.stackprotector(i8*, i8**) #3
+  
+  !llvm.module.flags = !{!0}
+  !llvm.ident = !{!1}
+  
+  !0 = !{i32 1, !"wchar_size", i32 4}
+  !1 = !{!"clang version 11.0.0 "}
+  !2 = !{!3, !3, i64 0}
+  !3 = !{!"int", !4, i64 0}
+  !4 = !{!"omnipotent char", !5, i64 0}
+  !5 = !{!"Simple C++ TBAA"}
+
+...
+---
+name:            foo
+alignment:       8
+exposesReturnsTwice: false
+legalized:       false
+regBankSelected: false
+selected:        false
+failedISel:      false
+tracksRegLiveness: true
+hasWinCFI:       false
+registers:       []
+liveins:
+  - { reg: '$w0', virtual-reg: '' }
+frameInfo:
+  isFrameAddressTaken: false
+  isReturnAddressTaken: false
+  hasStackMap:     false
+  hasPatchPoint:   false
+  stackSize:       16
+  offsetAdjustment: 0
+  maxAlignment:    8
+  adjustsStack:    false
+  hasCalls:        false
+  stackProtector:  ''
+  maxCallFrameSize: 0
+  cvBytesOfCalleeSavedRegisters: 0
+  hasOpaqueSPAdjustment: false
+  hasVAStart:      false
+  hasMustTailInVarArgFunc: false
+  localFrameSize:  16
+  savePoint:       ''
+  restorePoint:    ''
+fixedStack:      []
+stack:
+  - { id: 0, type: default, offset: -16, size: 16, 
+      alignment: 8, stack-id: default, callee-saved-register: '', callee-saved-restored: true, 
+      local-offset: -16, debug-info-variable: '', debug-info-expression: '', 
+      debug-info-location: '' }
+callSites:       []
+constants:       []
+machineFunctionInfo: {}
+body:             |
+  bb.0 (%ir-block.1):
+    liveins: $w0
+  
+    $sp = frame-setup SUBXri $sp, 16, 0
+    $x8 = ADDXri $sp, 0, 0
+    STRXui $xzr, $sp, 1 :: (store 8 into %ir.3 + 8)
+    STRXui $xzr, $sp, 0 :: (store 8 into %ir.3)
+    renamable $w0 = LDRWroW killed renamable $x8, killed renamable $w0, 1, 1 :: (load 4 from %ir.5, !tbaa !2)
+    $sp = frame-destroy ADDXri $sp, 16, 0
+    RET undef $lr, implicit $w0
+
+...

llvm/test/CodeGen/AArch64/arm64-nvcast.ll

 ; RUN: llc < %s -mtriple=arm64-apple-ios | FileCheck %s
 
-; CHECK-LABEL: _test:
-; CHECK-DAG:  fmov.2d v0, #2.00000000
-; CHECK-DAG: and [[MASK_IDX:x[0-9]+]], x1, #0x3
-; CHECK-DAG:  mov  x9, sp
-; CHECK-DAG:  str  q0, [sp], #16
-; CHECK-DAG:  bfi [[PTR:x[0-9]+]], [[MASK_IDX]], #2, #2
-; CHECK:  ldr s0, {{\[}}[[PTR]]{{\]}}
-; CHECK:  str  s0, [x0]
-
 define void @test(float * %p1, i32 %v1) {
+; CHECK-LABEL: test:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    sub sp, sp, #16 ; =16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    ; kill: def $w1 killed $w1 def $x1
+; CHECK-NEXT:    fmov.2d v0, #2.00000000
+; CHECK-NEXT:    and x8, x1, #0x3
+; CHECK-NEXT:    mov x9, sp
+; CHECK-NEXT:    str q0, [sp]
+; CHECK-NEXT:    bfi x9, x8, #2, #2
+; CHECK-NEXT:    ldr s0, [x9]
+; CHECK-NEXT:    str s0, [x0]
+; CHECK-NEXT:    add sp, sp, #16 ; =16
+; CHECK-NEXT:    ret
 entry:
   %v2 = extractelement <3 x float> <float 0.000000e+00, float 2.000000e+00, float 0.000000e+00>, i32 %v1
   store float %v2, float* %p1, align 4
   ret void
 }
 
-; CHECK-LABEL: _test2
-; CHECK: movi.16b  v0, #63
-; CHECK-DAG: and [[MASK_IDX:x[0-9]+]], x1, #0x3
-; CHECK-DAG: str  q0, [sp], #16
-; CHECK-DAG: mov  x9, sp
-; CHECK-DAG:  bfi [[PTR:x[0-9]+]], [[MASK_IDX]], #2, #2
-; CHECK: ldr s0, {{\[}}[[PTR]]{{\]}}
-; CHECK: str  s0, [x0]
-
 define void @test2(float * %p1, i32 %v1) {
+; CHECK-LABEL: test2:
+; CHECK:       ; %bb.0: ; %entry
+; CHECK-NEXT:    sub sp, sp, #16 ; =16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    ; kill: def $w1 killed $w1 def $x1
+; CHECK-NEXT:    movi.16b v0, #63
+; CHECK-NEXT:    and x8, x1, #0x3
+; CHECK-NEXT:    mov x9, sp
+; CHECK-NEXT:    str q0, [sp]
+; CHECK-NEXT:    bfi x9, x8, #2, #2
+; CHECK-NEXT:    ldr s0, [x9]
+; CHECK-NEXT:    str s0, [x0]
+; CHECK-NEXT:    add sp, sp, #16 ; =16
+; CHECK-NEXT:    ret
 entry:
   %v2 = extractelement <3 x float> <float 0.7470588088035583, float 0.7470588088035583, float 0.7470588088035583>, i32 %v1
   store float %v2, float* %p1, align 4
   ret void
 }
 
 
 %"st1" = type { %"subst1", %"subst1", %"subst1" }

llvm/test/CodeGen/AArch64/arm64-windows-calls.ll

; CHECK: mov x0, xzr
   %1 = load i64, i64* %0, align 4
   ret i64 %1
 }
 
 ; Returns <= 16 bytes should be in X0/X1.
 %struct.S2 = type { i32, i32, i32, i32 }
 define dso_local [2 x i64] @"?f2"() {
 entry:
+; FIXME: Missed optimization, the entire SP push/pop could be removed
 ; CHECK-LABEL: f2
-; CHECK: stp xzr, xzr, [sp], #16
-; CHECK: mov x0, xzr
-; CHECK: mov x1, xzr
+; CHECK:         stp     xzr, xzr, [sp, #-16]!
+; CHECK-NEXT:    mov     x0, xzr
+; CHECK-NEXT:    mov     x1, xzr
+; CHECK-NEXT:    add     sp, sp, #16
 
   %retval = alloca %struct.S2, align 4
   %a = getelementptr inbounds %struct.S2, %struct.S2* %retval, i32 0, i32 0
   store i32 0, i32* %a, align 4
   %b = getelementptr inbounds %struct.S2, %struct.S2* %retval, i32 0, i32 1
   store i32 0, i32* %b, align 4
   %c = getelementptr inbounds %struct.S2, %struct.S2* %retval, i32 0, i32 2
   store i32 0, i32* %c, align 4