[AArch64] Support 'swiftcorocc' "popless" calls.

'swiftcorocc' calls are allowed to have "popless" returns, which don't
fully restore the stack, thereby allowing the caller to access some
stack allocations made in the 'swiftcorocc' callee.

Concretely, calls to these functions don't restore SP (but do restore FP).

So the most important characteristic of a 'swiftcorocc' call is that it
forces the caller function to access its stack through FP, like it does
with e.g., variable-size allocas.

Support this on AArch64 by marking the frame as having a popless
call, which we generally honor when we decide whether the frame needs
FP and FP-based addressing, as we do today for variably-sized allocas.

rdar://135984630

Author: ahmedbougacha

llvm/include/llvm/CodeGen/MachineFrameInfo.h

@@ -319,6 +319,10 @@ class MachineFrameInfo {
   /// instructions which manipulate the stack pointer.
   bool HasCopyImplyingStackAdjustment = false;
 
+  /// True if the function contains a call using a calling convention that
+  /// allows it to be "popless", i.e., to not restore SP when the call returns.
+  bool HasPoplessCall = false;
+
   /// True if the function contains a call to the llvm.vastart intrinsic.
   bool HasVAStart = false;
 
@@ -634,6 +638,9 @@ class MachineFrameInfo {
     HasCopyImplyingStackAdjustment = B;
   }
 
+  bool hasPoplessCall() const { return HasPoplessCall; }
+  void setHasPoplessCall(bool B = true) { HasPoplessCall = B; }
+
   /// Returns true if the function calls the llvm.va_start intrinsic.
   bool hasVAStart() const { return HasVAStart; }
   void setHasVAStart(bool B) { HasVAStart = B; }

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -347,7 +347,8 @@ bool AArch64FrameLowering::homogeneousPrologEpilog(
   // Bail on stack adjustment needed on return for simplicity.
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   const TargetRegisterInfo *RegInfo = MF.getSubtarget().getRegisterInfo();
-  if (MFI.hasVarSizedObjects() || RegInfo->hasStackRealignment(MF))
+  if (MFI.hasVarSizedObjects() || RegInfo->hasStackRealignment(MF) ||
+      MFI.hasPoplessCall())
     return false;
   if (Exit && getArgumentStackToRestore(MF, *Exit))
     return false;
@@ -499,6 +500,7 @@ bool AArch64FrameLowering::hasFP(const MachineFunction &MF) const {
   if (MF.getTarget().Options.DisableFramePointerElim(MF))
     return true;
   if (MFI.hasVarSizedObjects() || MFI.isFrameAddressTaken() ||
+      MFI.hasPoplessCall() ||
       MFI.hasStackMap() || MFI.hasPatchPoint() ||
       RegInfo->hasStackRealignment(MF))
     return true;
@@ -1180,6 +1182,9 @@ bool AArch64FrameLowering::shouldCombineCSRLocalStackBump(
   if (MFI.hasVarSizedObjects())
     return false;
 
+  if (MFI.hasPoplessCall())
+    return false;
+
   if (RegInfo->hasStackRealignment(MF))
     return false;
 
@@ -2193,7 +2198,8 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
   StackOffset LocalsSize = SVELocalsSize + StackOffset::getFixed(NumBytes);
   allocateStackSpace(MBB, CalleeSavesBegin, 0, SVECalleeSavesSize, false,
                      nullptr, EmitAsyncCFI && !HasFP, CFAOffset,
-                     MFI.hasVarSizedObjects() || LocalsSize);
+                     MFI.hasVarSizedObjects() || LocalsSize ||
+                         MFI.hasPoplessCall());
   CFAOffset += SVECalleeSavesSize;
 
   if (EmitAsyncCFI)
@@ -2210,7 +2216,8 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
     allocateStackSpace(MBB, CalleeSavesEnd, RealignmentPadding,
                        SVELocalsSize + StackOffset::getFixed(NumBytes),
                        NeedsWinCFI, &HasWinCFI, EmitAsyncCFI && !HasFP,
-                       CFAOffset, MFI.hasVarSizedObjects());
+                       CFAOffset,
+                       MFI.hasVarSizedObjects() || MFI.hasPoplessCall());
   }
 
   // If we need a base pointer, set it up here. It's whatever the value of the
@@ -2531,7 +2538,8 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
     // 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 (AFI->isStackRealigned() || MFI.hasVarSizedObjects() ||
+        MFI.hasPoplessCall()) {
       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
@@ -2603,7 +2611,8 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
   // FIXME: Rather than doing the math here, we should instead just use
   // non-post-indexed loads for the restores if we aren't actually going to
   // be able to save any instructions.
-  if (!IsFunclet && (MFI.hasVarSizedObjects() || AFI->isStackRealigned())) {
+  if (!IsFunclet && (MFI.hasVarSizedObjects() || AFI->isStackRealigned() ||
+                     MFI.hasPoplessCall())) {
     emitFrameOffset(
         MBB, LastPopI, DL, AArch64::SP, AArch64::FP,
         StackOffset::getFixed(-AFI->getCalleeSaveBaseToFrameRecordOffset()),
@@ -2796,7 +2805,7 @@ StackOffset AArch64FrameLowering::resolveFrameOffsetReference(
       bool FPOffsetFits = !ForSimm || FPOffset >= -256;
       PreferFP |= Offset > -FPOffset && !SVEStackSize;
 
-      if (MFI.hasVarSizedObjects()) {
+      if (MFI.hasVarSizedObjects() || MFI.hasPoplessCall()) {
         // If we have variable sized objects, we can use either FP or BP, as the
         // SP offset is unknown. We can use the base pointer if we have one and
         // FP is not preferred. If not, we're stuck with using FP.
@@ -4679,6 +4688,7 @@ StackOffset AArch64FrameLowering::getFrameIndexReferencePreferSP(
 
   // Go to common code if we cannot provide sp + offset.
   if (MFI.hasVarSizedObjects() ||
+      MFI.hasPoplessCall() ||
       MF.getInfo<AArch64FunctionInfo>()->getStackSizeSVE() ||
       MF.getSubtarget().getRegisterInfo()->hasStackRealignment(MF))
     return getFrameIndexReference(MF, FI, FrameReg);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -7205,6 +7205,7 @@ CCAssignFn *AArch64TargetLowering::CCAssignFnForCall(CallingConv::ID CC,
   case CallingConv::CXX_FAST_TLS:
   case CallingConv::Swift:
   case CallingConv::SwiftTail:
+  case CallingConv::SwiftCoro:
   case CallingConv::Tail:
   case CallingConv::GRAAL:
     if (Subtarget->isTargetWindows()) {
@@ -8317,6 +8318,9 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
       CallConv = CallingConv::AArch64_SVE_VectorCall;
   }
 
+  if (CallConv == CallingConv::SwiftCoro)
+    MF.getFrameInfo().setHasPoplessCall();
+
   if (IsTailCall) {
     // Check if it's really possible to do a tail call.
     IsTailCall = isEligibleForTailCallOptimization(CLI);

llvm/lib/Target/AArch64/GISel/AArch64CallLowering.cpp

@@ -1454,6 +1454,9 @@ bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
       return false;
   }
 
+  if (Info.CallConv == CallingConv::SwiftCoro)
+    MF.getFrameInfo().setHasPoplessCall();
+
   if (Info.SwiftErrorVReg) {
     MIB.addDef(AArch64::X21, RegState::Implicit);
     MIRBuilder.buildCopy(Info.SwiftErrorVReg, Register(AArch64::X21));

llvm/test/CodeGen/AArch64/swiftcorocc-call.ll

@@ -0,0 +1,128 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple arm64e-apple-darwin -verify-machineinstrs -o - %s \
+; RUN:   | FileCheck %s
+
+; RUN: llc -mtriple arm64e-apple-darwin -verify-machineinstrs -o - %s \
+; RUN:   -global-isel -global-isel-abort=2 | FileCheck %s
+
+declare i64 @g(ptr, ptr)
+
+define i64 @test_call_to_swiftcoro() #0 {
+; CHECK-LABEL: test_call_to_swiftcoro:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    pacibsp
+; CHECK-NEXT:    stp x26, x25, [sp, #-32]! ; 16-byte Folded Spill
+; CHECK-NEXT:    stp x29, x30, [sp, #16] ; 16-byte Folded Spill
+; CHECK-NEXT:    add x29, sp, #16
+; CHECK-NEXT:    sub sp, sp, #16
+; CHECK-NEXT:    .cfi_def_cfa w29, 16
+; CHECK-NEXT:    .cfi_offset w30, -8
+; CHECK-NEXT:    .cfi_offset w29, -16
+; CHECK-NEXT:    .cfi_offset w25, -24
+; CHECK-NEXT:    .cfi_offset w26, -32
+; CHECK-NEXT:    ; InlineAsm Start
+; CHECK-NEXT:    ; InlineAsm End
+; CHECK-NEXT:    bl _test_call
+; CHECK-NEXT:    sub x0, x29, #24
+; CHECK-NEXT:    sub x1, x29, #32
+; CHECK-NEXT:    bl _g
+; CHECK-NEXT:    sub sp, x29, #16
+; CHECK-NEXT:    ldp x29, x30, [sp, #16] ; 16-byte Folded Reload
+; CHECK-NEXT:    ldp x26, x25, [sp], #32 ; 16-byte Folded Reload
+; CHECK-NEXT:    retab
+  %v1 = alloca i64
+  %v2 = alloca i64
+  call void asm sideeffect "", "~{x25},~{x26}"()
+  %v3 = call swiftcorocc i64 @test_call()
+  %v4 = call i64 @g(ptr %v1, ptr %v2)
+  ret i64 %v4
+}
+
+define i64 @test_call_to_normal() #0 {
+; CHECK-LABEL: test_call_to_normal:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    pacibsp
+; CHECK-NEXT:    sub sp, sp, #48
+; CHECK-NEXT:    stp x26, x25, [sp, #16] ; 16-byte Folded Spill
+; CHECK-NEXT:    stp x29, x30, [sp, #32] ; 16-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 48
+; CHECK-NEXT:    .cfi_offset w30, -8
+; CHECK-NEXT:    .cfi_offset w29, -16
+; CHECK-NEXT:    .cfi_offset w25, -24
+; CHECK-NEXT:    .cfi_offset w26, -32
+; CHECK-NEXT:    ; InlineAsm Start
+; CHECK-NEXT:    ; InlineAsm End
+; CHECK-NEXT:    bl _test_call_normal
+; CHECK-NEXT:    add x0, sp, #8
+; CHECK-NEXT:    mov x1, sp
+; CHECK-NEXT:    bl _g
+; CHECK-NEXT:    ldp x29, x30, [sp, #32] ; 16-byte Folded Reload
+; CHECK-NEXT:    ldp x26, x25, [sp, #16] ; 16-byte Folded Reload
+; CHECK-NEXT:    add sp, sp, #48
+; CHECK-NEXT:    retab
+  %v1 = alloca i64
+  %v2 = alloca i64
+  call void asm sideeffect "", "~{x25},~{x26}"()
+  %v3 = call i64 @test_call_normal()
+  %v4 = call i64 @g(ptr %v1, ptr %v2)
+  ret i64 %v4
+}
+
+define swiftcorocc i64 @test_call() #0 {
+; CHECK-LABEL: test_call:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    pacibsp
+; CHECK-NEXT:    sub sp, sp, #48
+; CHECK-NEXT:    stp x26, x25, [sp, #16] ; 16-byte Folded Spill
+; CHECK-NEXT:    stp x29, x30, [sp, #32] ; 16-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 48
+; CHECK-NEXT:    .cfi_offset w30, -8
+; CHECK-NEXT:    .cfi_offset w29, -16
+; CHECK-NEXT:    .cfi_offset w25, -24
+; CHECK-NEXT:    .cfi_offset w26, -32
+; CHECK-NEXT:    ; InlineAsm Start
+; CHECK-NEXT:    ; InlineAsm End
+; CHECK-NEXT:    add x0, sp, #8
+; CHECK-NEXT:    mov x1, sp
+; CHECK-NEXT:    bl _g
+; CHECK-NEXT:    ldp x29, x30, [sp, #32] ; 16-byte Folded Reload
+; CHECK-NEXT:    ldp x26, x25, [sp, #16] ; 16-byte Folded Reload
+; CHECK-NEXT:    add sp, sp, #48
+; CHECK-NEXT:    retab
+  %v1 = alloca i64
+  %v2 = alloca i64
+  call void asm sideeffect "", "~{x25},~{x26}"()
+  %v3 = call i64 @g(ptr %v1, ptr %v2)
+  ret i64 %v3
+}
+
+define i64 @test_call_normal() #0 {
+; CHECK-LABEL: test_call_normal:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    pacibsp
+; CHECK-NEXT:    sub sp, sp, #48
+; CHECK-NEXT:    stp x26, x25, [sp, #16] ; 16-byte Folded Spill
+; CHECK-NEXT:    stp x29, x30, [sp, #32] ; 16-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 48
+; CHECK-NEXT:    .cfi_offset w30, -8
+; CHECK-NEXT:    .cfi_offset w29, -16
+; CHECK-NEXT:    .cfi_offset w25, -24
+; CHECK-NEXT:    .cfi_offset w26, -32
+; CHECK-NEXT:    ; InlineAsm Start
+; CHECK-NEXT:    ; InlineAsm End
+; CHECK-NEXT:    add x0, sp, #8
+; CHECK-NEXT:    mov x1, sp
+; CHECK-NEXT:    bl _g
+; CHECK-NEXT:    ldp x29, x30, [sp, #32] ; 16-byte Folded Reload
+; CHECK-NEXT:    ldp x26, x25, [sp, #16] ; 16-byte Folded Reload
+; CHECK-NEXT:    add sp, sp, #48
+; CHECK-NEXT:    retab
+  %v1 = alloca i64
+  %v2 = alloca i64
+  call void asm sideeffect "", "~{x25},~{x26}"()
+  %v3 = call i64 @g(ptr %v1, ptr %v2)
+  ret i64 %v3
+}
+
+
+attributes #0 = { "ptrauth-returns" }