[AArch64][SVE] Restore SP from FP when SVE CSRs and variable sized objects are present

[bsmith]

Summary

Without SVE, after a dynamic stack allocation has modified the SP, it is
presumed that a frame pointer restoration will revert the SP back to
it's correct value prior to any caller stack being restored. However the
SVE frame is restored using the stack pointer directly, as it is located
after the frame pointer. This means that in the presence of a dynamic
stack allocation, any SVE callee state gets corrupted as SP has the
incorrect value when the SVE state is restored.

To address this issue, when variable sized objects and SVE CSRs are
present, treat the stack as having been realigned, hence restoring the
stack pointer from the frame pointerr prior to restoring the SVE state.

[efriedma]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

1961

Do we need to use emitFrameOffset() here?

llvm/test/CodeGen/AArch64/framelayout-sve.mir

634

Can we optimize the stack computation in this case? I can see why you need to set the stack pointer in general, but it doesn't seem necessary here. I guess we can leave that for a followup, though.

[bsmith]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

1961

I'm not sure we should/need to, the allocation of the variable length array that leads to the changing of the stack pointer does not happen during frame lowering, but rather during isel lowering, hence it does not have the usual frame-setup flag etc. As such I'm not sure this, which is just reverting that, should have the frame-destroy flag etc either, however I'm not sure about this.

[bsmith]

• Only restore from base pointer when SVE CSRs are present

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

1961

Should we be doing this when AFI->isStackRealigned() is true? It looks like for that case we recalculate SP from the frame pointer so this'll kick out a redundant instruction. I guess there's no test to show the negative affect of this, it's just something I spotted whilst working on another bug.

[bsmith]

• Rework patch to use same mechanism as stack realignment, rather than using base pointer.

[efriedma]

LGTM with a minor comment.

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

1965

This comment is a bit misleading. For some reading this for the first time, there isn't any presumption that this is for realignment specifically.

Maybe just "If we have stack realignment or variable sized objects on the stack, we must restore the stack pointer using the frame pointer". Or something like that.

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

 }
 
 void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
                                         MachineBasicBlock &MBB) const {
   MachineBasicBlock::iterator MBBI = MBB.getLastNonDebugInstr();
   MachineFrameInfo &MFI = MF.getFrameInfo();
   const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
   const TargetInstrInfo *TII = Subtarget.getInstrInfo();
+  const AArch64RegisterInfo &RegInfo = *Subtarget.getRegisterInfo();
   DebugLoc DL;
   bool NeedsWinCFI = needsWinCFI(MF);
   bool EmitCFI = MF.getInfo<AArch64FunctionInfo>()->needsAsyncDwarfUnwindInfo();
   bool HasWinCFI = false;
   bool IsFunclet = false;
   auto WinCFI = make_scope_exit([&]() { assert(HasWinCFI == MF.hasWinCFI()); });
 
   if (MBB.end() != MBBI) {
void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
   const StackOffset &SVEStackSize = getSVEStackSize(MF);
 
   // If there is a single SP update, insert it before the ret and we're done.
   if (CombineSPBump) {
     assert(!SVEStackSize && "Cannot combine SP bump with SVE");
 
     // When we are about to restore the CSRs, the CFA register is SP again.
     if (EmitCFI && hasFP(MF)) {
-      const AArch64RegisterInfo &RegInfo = *Subtarget.getRegisterInfo();
       unsigned Reg = RegInfo.getDwarfRegNum(AArch64::SP, true);
       unsigned CFIIndex =
           MF.addFrameInst(MCCFIInstruction::cfiDefCfa(nullptr, Reg, NumBytes));
       BuildMI(MBB, LastPopI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
           .addCFIIndex(CFIIndex)
           .setMIFlags(MachineInstr::FrameDestroy);
     }
 
if (int64_t CalleeSavedSize = AFI->getSVECalleeSavedStackSize()) {
     StackOffset CalleeSavedSizeAsOffset =
         StackOffset::getScalable(CalleeSavedSize);
     DeallocateBefore = SVEStackSize - CalleeSavedSizeAsOffset;
     DeallocateAfter = CalleeSavedSizeAsOffset;
   }
 
   // Deallocate the SVE area.
   if (SVEStackSize) {
+    // If we have a base pointer we must restore the stack pointer from it, we
+    // need to explicitly deallocate variable length stack allocations when we
+    // have an SVE stack frame.
+    if (!IsFunclet && RegInfo.hasBasePointer(MF))
+      TII->copyPhysReg(MBB, RestoreBegin, DL, AArch64::SP,
+                       RegInfo.getBaseRegister(), false);
+
     if (AFI->isStackRealigned()) {
       if (int64_t CalleeSavedSize = AFI->getSVECalleeSavedStackSize()) {
         // Set SP to start of SVE callee-save area from which they can
         // be reloaded. The code below will deallocate the stack space
         // space by moving FP -> SP.
         emitFrameOffset(MBB, RestoreBegin, DL, AArch64::SP, AArch64::FP,
                         StackOffset::getScalable(-CalleeSavedSize), TII,
                         MachineInstr::FrameDestroy);
       }
     } else {
       if (AFI->getSVECalleeSavedStackSize()) {

llvm/test/CodeGen/AArch64/framelayout-sve.mir

 # CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa $w29, 32
 # CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $w19, -16
 # CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $w30, -24
 # CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $w29, -32
 # CHECK-NEXT: $sp = frame-setup ADDVL_XXI $sp, -1
 # CHECK-NEXT: $sp = frame-setup SUBXri $sp, 16, 0
 # CHECK-NEXT: $x19 = ADDXri $sp, 0, 0
 # CHECK-NEXT: STRXui $xzr, $x19, 0
+# CHECK-NEXT: $sp = ADDXri $x19, 0, 0
 # CHECK-NEXT: $sp = frame-destroy ADDVL_XXI $sp, 1
 # CHECK-NEXT: $sp = frame-destroy ADDXri $fp, 0, 0
 # CHECK-NEXT: frame-destroy CFI_INSTRUCTION def_cfa $wsp, 32
 # CHECK-NEXT: $x19 = frame-destroy LDRXui $sp, 2
 # CHECK-NEXT: early-clobber $sp, $fp, $lr = frame-destroy LDPXpost $sp, 4
 # CHECK-NEXT: frame-destroy CFI_INSTRUCTION def_cfa_offset 0
 # CHECK-NEXT: frame-destroy CFI_INSTRUCTION restore $w19
 # CHECK-NEXT: frame-destroy CFI_INSTRUCTION restore $w30
 # CHECK-NEXT: frame-destroy CFI_INSTRUCTION restore $w29
 # CHECK-NEXT: RET_ReallyLR

llvm/test/CodeGen/AArch64/sve-alloca.ll

 ; CHECK-NEXT:    and x9, x9, #0xfffffffffffffff0
 ; CHECK-NEXT:    sub x8, x8, x9
 ; CHECK-NEXT:    and x0, x8, #0xffffffffffffffe0
 ; CHECK-NEXT:    mov sp, x0
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    st1d { z1.d }, p0, [x0, #1, mul vl]
 ; CHECK-NEXT:    st1d { z0.d }, p0, [x0]
 ; CHECK-NEXT:    bl bar
+; CHECK-NEXT:    mov sp, x19
 ; CHECK-NEXT:    ldr p15, [sp, #4, mul vl] // 2-byte Folded Reload
 ; CHECK-NEXT:    ldr p14, [sp, #5, mul vl] // 2-byte Folded Reload
 ; CHECK-NEXT:    ldr p13, [sp, #6, mul vl] // 2-byte Folded Reload
 ; CHECK-NEXT:    ldr p12, [sp, #7, mul vl] // 2-byte Folded Reload
 ; CHECK-NEXT:    ldr p11, [sp, #8, mul vl] // 2-byte Folded Reload
 ; CHECK-NEXT:    ldr p10, [sp, #9, mul vl] // 2-byte Folded Reload
 ; CHECK-NEXT:    ldr p9, [sp, #10, mul vl] // 2-byte Folded Reload
 ; CHECK-NEXT:    ldr p8, [sp, #11, mul vl] // 2-byte Folded Reload