[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

1840

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

1840

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

1840

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

1812

"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>
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].
+  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;
+  }
+
   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))
       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 we are optimizing SP, do not allow instructions that may load or store
+    // in between the load and the optimized value update.
     if (!ModifiedRegUnits.available(BaseReg) ||
-        !UsedRegUnits.available(BaseReg))
+        !UsedRegUnits.available(BaseReg) ||
+        (BaseRegSP && MBBI->mayLoadOrStore()))
       return E;
   }
   return E;
 }
 
 MachineBasicBlock::iterator AArch64LoadStoreOpt::findMatchingUpdateInsnBackward(
     MachineBasicBlock::iterator I, unsigned Limit) {
   MachineBasicBlock::iterator B = I->getParent()->begin();
   MachineBasicBlock::iterator E = I->getParent()->end();
   MachineInstr &MemMI = *I;
   MachineBasicBlock::iterator MBBI = I;
 
   Register BaseReg = getLdStBaseOp(MemMI).getReg();
   int Offset = getLdStOffsetOp(MemMI).getImm();
 
   // If the load/store is the first instruction in the block, there's obviously
   // not any matching update. Ditto if the memory offset isn't zero.
   if (MBBI == B || Offset != 0)
     return E;
   // If the base register overlaps a destination register, we can't
   // merge the update.
   if (!isTagStore(MemMI)) {
     bool IsPairedInsn = isPairedLdSt(MemMI);
     for (unsigned i = 0, e = IsPairedInsn ? 2 : 1; i != e; ++i) {
       Register DestReg = getLdStRegOp(MemMI, i).getReg();
       if (DestReg == BaseReg || TRI->isSubRegister(BaseReg, DestReg))
         return E;
     }

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