[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

632

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

1958

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

if (int64_t CalleeSavedSize = AFI->getSVECalleeSavedStackSize()) {
     StackOffset CalleeSavedSizeAsOffset =
         StackOffset::getScalable(CalleeSavedSize);
     DeallocateBefore = SVEStackSize - CalleeSavedSizeAsOffset;
     DeallocateAfter = CalleeSavedSizeAsOffset;
   }
 
   // Deallocate the SVE area.
   if (SVEStackSize) {
-    if (AFI->isStackRealigned()) {
+    // If we have stack realignment or variable sized objects on the stack,
+    // restore the stack pointer from the frame pointer prior to SVE CSR
+    // restoration.
+    if (AFI->isStackRealigned() || MFI.hasVarSizedObjects()) {
       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 = 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:    addvl sp, x29, #-18
 ; 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

 ; CHECK-NEXT:    ldr z15, [sp, #10, mul vl] // 16-byte Folded Reload
 ; CHECK-NEXT:    ldr z14, [sp, #11, mul vl] // 16-byte Folded Reload
 ; CHECK-NEXT:    ldr z13, [sp, #12, mul vl] // 16-byte Folded Reload
 ; CHECK-NEXT:    ldr z12, [sp, #13, mul vl] // 16-byte Folded Reload
 ; CHECK-NEXT:    ldr z11, [sp, #14, mul vl] // 16-byte Folded Reload
 ; CHECK-NEXT:    ldr z10, [sp, #15, mul vl] // 16-byte Folded Reload
 ; CHECK-NEXT:    ldr z9, [sp, #16, mul vl] // 16-byte Folded Reload
 ; CHECK-NEXT:    ldr z8, [sp, #17, mul vl] // 16-byte Folded Reload
-; CHECK-NEXT:    addvl sp, sp, #18
 ; CHECK-NEXT:    mov sp, x29
 ; CHECK-NEXT:    ldp x28, x19, [sp, #16] // 16-byte Folded Reload
 ; CHECK-NEXT:    ldp x29, x30, [sp], #32 // 16-byte Folded Reload
 ; CHECK-NEXT:  .LBB0_2: // %if.end
 ; CHECK-NEXT:    .cfi_def_cfa wsp, 0
 ; CHECK-NEXT:    .cfi_same_value b8
 ; CHECK-NEXT:    .cfi_same_value b9
 ; CHECK-NEXT:    .cfi_same_value b10