[AArch64] Lower @llvm.ret.popless in swiftcorocc functions.

On AArch64, swiftcorocc functions are the only functions yet that can
support popless returns.

In the backend, that's done by recognizing the musttail call to
llvm.ret.popless preceding a ret instruction, and asking the target to
adjust that ret to be popless.

Throughout most of the backend, that's not an interesting difference.

In frame lowering, these popless rets now induce several special
behaviors in their (never shrink-wrapped) epilogues, all consequences
of not restoring SP:
- they of course don't do the SP adjustment or restore itself.
- most importantly, they force the epilogue callee-save restores
  to be FP-based rather than SP-based.
- they restore FP/LR last, as we still need the old FP, pointing
  at the frame being destroyed, to do the CSR restoring.
- with ptrauth-returns, they first derive the entry SP from
  FP, into X16, to use as a discriminator for a standalone AUTIB.

rdar://135984630

Author: ahmedbougacha

llvm/include/llvm/CodeGen/GlobalISel/CallLowering.h

@@ -537,6 +537,15 @@ class CallLowering {
     return false;
   }
 
+  /// This hook must be implemented to lower @llvm.ret.popless intrinsics,
+  /// which are required to be musttail, and are effectively annotating a
+  /// return instruction to mark it "popless", i.e., not restoring SP.
+  /// This "adjustment" step runs after lowerReturn, and is only meant to make
+  /// it a little less painful to maintain before we add this to the main hook.
+  virtual bool adjustReturnToPopless(MachineIRBuilder &MIRBuilder) const {
+    report_fatal_error("Popless returns not implemented for this target");
+  }
+
   virtual bool fallBackToDAGISel(const MachineFunction &MF) const {
     return false;
   }

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -4787,6 +4787,10 @@ class TargetLowering : public TargetLoweringBase {
     llvm_unreachable("Not Implemented");
   }
 
+  virtual SDValue adjustReturnPopless(SDValue Chain, SelectionDAG &DAG) const {
+    report_fatal_error("Popless returns not implemented for this target");
+  }
+
   /// Return true if result of the specified node is used by a return node
   /// only. It also compute and return the input chain for the tail call.
   ///

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

@@ -377,7 +377,14 @@ bool IRTranslator::translateRet(const User &U, MachineIRBuilder &MIRBuilder) {
   // The target may mess up with the insertion point, but
   // this is not important as a return is the last instruction
   // of the block anyway.
-  return CLI->lowerReturn(MIRBuilder, Ret, VRegs, FuncInfo, SwiftErrorVReg);
+  bool Success =
+      CLI->lowerReturn(MIRBuilder, Ret, VRegs, FuncInfo, SwiftErrorVReg);
+
+  if (auto *MustTailCI = RI.getParent()->getTerminatingMustTailCall())
+    if (MustTailCI->getIntrinsicID() == Intrinsic::ret_popless)
+      Success &= CLI->adjustReturnToPopless(MIRBuilder);
+
+  return Success;
 }
 
 void IRTranslator::emitBranchForMergedCondition(

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -2198,6 +2198,13 @@ void SelectionDAGBuilder::visitRet(const ReturnInst &I) {
     return;
   }
 
+  // Musttail calls to @llvm.ret.popless are used to annotate the ret as
+  // "popless".  Keep track of it here, and ask the target to do so later.
+  bool IsPoplessReturn = false;
+  if (auto *MustTailCI = I.getParent()->getTerminatingMustTailCall())
+    if (MustTailCI->getIntrinsicID() == Intrinsic::ret_popless)
+      IsPoplessReturn = true;
+
   if (!FuncInfo.CanLowerReturn) {
     unsigned DemoteReg = FuncInfo.DemoteRegister;
     const Function *F = I.getParent()->getParent();
@@ -2336,6 +2343,18 @@ void SelectionDAGBuilder::visitRet(const ReturnInst &I) {
   Chain = DAG.getTargetLoweringInfo().LowerReturn(
       Chain, CallConv, isVarArg, Outs, OutVals, getCurSDLoc(), DAG);
 
+  // If we did find this return instruction to be popless, make it so now.
+  // It's still a normal return in almost all regards, we just need to remember
+  // it's popless, for when we lower the return and emit the epilogue later.
+  // Ideally we'd ask LowerReturn to do that, but the API is enough of a pain
+  // as it is, and all targets would have to learn about that.
+  if (IsPoplessReturn) {
+    SDValue NewChain =
+        DAG.getTargetLoweringInfo().adjustReturnPopless(Chain, DAG);
+    DAG.RemoveDeadNode(Chain.getNode());
+    Chain = NewChain;
+  }
+
   // Verify that the target's LowerReturn behaved as expected.
   assert(Chain.getNode() && Chain.getValueType() == MVT::Other &&
          "LowerReturn didn't return a valid chain!");
@@ -7937,6 +7956,11 @@ void SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I,
     setValue(&I, DAG.getNode(ISD::AND, sdl, PtrVT, Ptr, Mask));
     return;
   }
+
+  case Intrinsic::ret_popless:
+    // We're handling this on the associated ret itself.
+    return;
+
   case Intrinsic::threadlocal_address: {
     setValue(&I, getValue(I.getOperand(0)));
     return;

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -504,6 +504,11 @@ bool AArch64FrameLowering::hasFP(const MachineFunction &MF) const {
       MFI.hasStackMap() || MFI.hasPatchPoint() ||
       RegInfo->hasStackRealignment(MF))
     return true;
+
+  const AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
+  if (AFI->hasPoplessEpilogue())
+    return true;
+
   // With large callframes around we may need to use FP to access the scavenging
   // emergency spillslot.
   //
@@ -1119,6 +1124,12 @@ bool AArch64FrameLowering::canUseAsPrologue(
       return false;
   }
 
+  // If we have some return path that's popless, it needs its own very-special
+  // epilogue, so we can't shrink-wrap it away.
+  // FIXME: this and some of the below checks belong in enableShrinkWrapping.
+  if (AFI->hasPoplessEpilogue())
+    return false;
+
   // Certain stack probing sequences might clobber flags, then we can't use
   // the block as a prologue if the flags register is a live-in.
   if (MF->getInfo<AArch64FunctionInfo>()->hasStackProbing() &&
@@ -1204,6 +1215,12 @@ bool AArch64FrameLowering::shouldCombineCSRLocalStackBump(
 
 bool AArch64FrameLowering::shouldCombineCSRLocalStackBumpInEpilogue(
     MachineBasicBlock &MBB, unsigned StackBumpBytes) const {
+
+  MachineFunction &MF = *MBB.getParent();
+  AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
+  if (AFI->hasPoplessEpilogue())
+    return false;
+
   if (!shouldCombineCSRLocalStackBump(*MBB.getParent(), StackBumpBytes))
     return false;
 
@@ -1560,6 +1577,47 @@ static MachineBasicBlock::iterator convertCalleeSaveRestoreToSPPrePostIncDec(
   return std::prev(MBB.erase(MBBI));
 }
 
+static void fixupCalleeSaveRestoreToFPBased(MachineInstr &MI,
+                                            uint64_t FPSPOffset) {
+  assert(!AArch64InstrInfo::isSEHInstruction(MI));
+
+  unsigned Opc = MI.getOpcode();
+  unsigned Scale;
+  switch (Opc) {
+  case AArch64::STPXi:
+  case AArch64::STRXui:
+  case AArch64::STPDi:
+  case AArch64::STRDui:
+  case AArch64::LDPXi:
+  case AArch64::LDRXui:
+  case AArch64::LDPDi:
+  case AArch64::LDRDui:
+    Scale = 8;
+    break;
+  case AArch64::STPQi:
+  case AArch64::STRQui:
+  case AArch64::LDPQi:
+  case AArch64::LDRQui:
+    Scale = 16;
+    break;
+  default:
+    llvm_unreachable("Unexpected callee-save save/restore opcode!");
+  }
+
+  unsigned OffsetIdx = MI.getNumExplicitOperands() - 1;
+
+  MachineOperand &BaseRegOpnd = MI.getOperand(OffsetIdx - 1);
+  assert(BaseRegOpnd.getReg() == AArch64::SP &&
+         "Unexpected base register in callee-save save/restore instruction!");
+  BaseRegOpnd.setReg(AArch64::FP); // XXX TRI
+
+  // Last operand is immediate offset that needs fixing.
+  MachineOperand &OffsetOpnd = MI.getOperand(OffsetIdx);
+  // All generated opcodes have scaled offsets.
+  assert(FPSPOffset % Scale == 0);
+  OffsetOpnd.setImm(OffsetOpnd.getImm() - FPSPOffset / Scale);
+}
+
 // Fixup callee-save register save/restore instructions to take into account
 // combined SP bump by adding the local stack size to the stack offsets.
 static void fixupCalleeSaveRestoreStackOffset(MachineInstr &MI,
@@ -2298,10 +2356,22 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
   bool EmitCFI = AFI->needsAsyncDwarfUnwindInfo(MF);
   bool HasWinCFI = false;
   bool IsFunclet = false;
+  bool IsSwiftCoroPartialReturn = false;
 
   if (MBB.end() != MBBI) {
     DL = MBBI->getDebugLoc();
     IsFunclet = isFuncletReturnInstr(*MBBI);
+    IsSwiftCoroPartialReturn = MBBI->getOpcode() == AArch64::RET_POPLESS;
+  }
+
+  if (IsSwiftCoroPartialReturn) {
+    // The partial-return intrin/instr requires the swiftcoro cc
+    if (MF.getFunction().getCallingConv() != CallingConv::SwiftCoro)
+      report_fatal_error("llvm.ret.popless requires swiftcorocc");
+    assert(MBBI->getOpcode() == AArch64::RET_POPLESS);
+    BuildMI(MBB, MBBI, DL, TII->get(AArch64::RET_ReallyLR))
+        .setMIFlag(MachineInstr::FrameDestroy);
+    MBB.erase(MBBI);
   }
 
   MachineBasicBlock::iterator EpilogStartI = MBB.end();
@@ -2350,6 +2420,39 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
         if (Info.getReg() != AArch64::LR)
           continue;
         MachineBasicBlock::iterator TI = MBB.getFirstTerminator();
+
+        // When we're doing a popless ret (i.e., that doesn't restore SP), we
+        // can't rely on the exit SP being the same as the entry, but they need
+        // to match for the LR auth to succeed.  Instead, derive the entry SP
+        // from our FP (using a -16 static offset for the size of the frame
+        // record itself), save that into X16, and use that as the discriminator
+        // in an AUTIB.
+        if (IsSwiftCoroPartialReturn) {
+          const auto *TRI = Subtarget.getRegisterInfo();
+
+          MachineBasicBlock::iterator EpilogStartI = MBB.getFirstTerminator();
+          MachineBasicBlock::iterator Begin = MBB.begin();
+          while (EpilogStartI != Begin) {
+            --EpilogStartI;
+            if (!EpilogStartI->getFlag(MachineInstr::FrameDestroy)) {
+              ++EpilogStartI;
+              break;
+            }
+            if (EpilogStartI->readsRegister(AArch64::X16, TRI) ||
+                EpilogStartI->modifiesRegister(AArch64::X16, TRI))
+              report_fatal_error("unable to use x16 for popless ret LR auth");
+          }
+
+          emitFrameOffset(MBB, EpilogStartI, DL, AArch64::X16, AArch64::FP,
+                          StackOffset::getFixed(16), TII,
+                          MachineInstr::FrameDestroy);
+          BuildMI(MBB, TI, DL, TII->get(AArch64::AUTIB), AArch64::LR)
+              .addUse(AArch64::LR)
+              .addUse(AArch64::X16)
+              .setMIFlag(MachineInstr::FrameDestroy);
+          return;
+        }
+
         if (TI != MBB.end() && TI->getOpcode() == AArch64::RET_ReallyLR) {
           // If there is a terminator and it's a RET, we can fold AUTH into it.
           // Be careful to keep the implicitly returned registers.
@@ -2383,6 +2486,7 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
     AFI->setLocalStackSize(NumBytes - PrologueSaveSize);
   if (homogeneousPrologEpilog(MF, &MBB)) {
     assert(!NeedsWinCFI);
+    assert(!IsSwiftCoroPartialReturn);
     auto LastPopI = MBB.getFirstTerminator();
     if (LastPopI != MBB.begin()) {
       auto HomogeneousEpilog = std::prev(LastPopI);
@@ -2404,7 +2508,7 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
   // Assume we can't combine the last pop with the sp restore.
 
   bool CombineAfterCSRBump = false;
-  if (!CombineSPBump && PrologueSaveSize != 0) {
+  if (!CombineSPBump && PrologueSaveSize != 0 && !IsSwiftCoroPartialReturn) {
     MachineBasicBlock::iterator Pop = std::prev(MBB.getFirstTerminator());
     while (Pop->getOpcode() == TargetOpcode::CFI_INSTRUCTION ||
            AArch64InstrInfo::isSEHInstruction(*Pop))
@@ -2440,6 +2544,14 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
         IsSVECalleeSave(LastPopI)) {
       ++LastPopI;
       break;
+    } else if (IsSwiftCoroPartialReturn) {
+      assert(!EmitCFI);
+      assert(hasFP(MF));
+      fixupCalleeSaveRestoreStackOffset(*LastPopI, AFI->getLocalStackSize(),
+                                        NeedsWinCFI, &HasWinCFI);
+      // if FP-based addressing, rewrite CSR restores from SP to FP
+      fixupCalleeSaveRestoreToFPBased(
+          *LastPopI, AFI->getCalleeSaveBaseToFrameRecordOffset());
     } else if (CombineSPBump)
       fixupCalleeSaveRestoreStackOffset(*LastPopI, AFI->getLocalStackSize(),
                                         NeedsWinCFI, &HasWinCFI);
@@ -2459,6 +2571,7 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
   }
 
   if (hasFP(MF) && AFI->hasSwiftAsyncContext()) {
+    assert(!IsSwiftCoroPartialReturn);
     switch (MF.getTarget().Options.SwiftAsyncFramePointer) {
     case SwiftAsyncFramePointerMode::DeploymentBased:
       // Avoid the reload as it is GOT relative, and instead fall back to the
@@ -2492,6 +2605,7 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &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");
+    assert(!IsSwiftCoroPartialReturn);
 
     // When we are about to restore the CSRs, the CFA register is SP again.
     if (EmitCFI && hasFP(MF)) {
@@ -2577,6 +2691,7 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
   }
 
   if (!hasFP(MF)) {
+    assert(!IsSwiftCoroPartialReturn);
     bool RedZone = canUseRedZone(MF);
     // If this was a redzone leaf function, we don't need to restore the
     // stack pointer (but we may need to pop stack args for fastcc).
@@ -2607,6 +2722,9 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
     NumBytes = 0;
   }
 
+  if (IsSwiftCoroPartialReturn)
+    return;
+
   // Restore the original stack pointer.
   // 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
@@ -3449,9 +3567,17 @@ bool AArch64FrameLowering::restoreCalleeSavedRegisters(
   DebugLoc DL;
   SmallVector<RegPairInfo, 8> RegPairs;
   bool NeedsWinCFI = needsWinCFI(MF);
+  bool IsSwiftCoroPartialReturn = false;
 
-  if (MBBI != MBB.end())
+  if (MBBI != MBB.end()) {
     DL = MBBI->getDebugLoc();
+    IsSwiftCoroPartialReturn = MBBI->getOpcode() == AArch64::RET_POPLESS;
+  }
+
+  // The partial-return intrin/instr requires the swiftcoro cc
+  if (IsSwiftCoroPartialReturn &&
+      MF.getFunction().getCallingConv() != CallingConv::SwiftCoro)
+    report_fatal_error("llvm.ret.popless requires swiftcorocc");
 
   computeCalleeSaveRegisterPairs(MF, CSI, TRI, RegPairs, hasFP(MF));
   if (homogeneousPrologEpilog(MF, &MBB)) {
@@ -3464,6 +3590,17 @@ bool AArch64FrameLowering::restoreCalleeSavedRegisters(
     return true;
   }
 
+  // If doing a partial/popless return, CSR restores are from FP, so do it last.
+  if (IsSwiftCoroPartialReturn) {
+    auto IsFPLR = [](const RegPairInfo &c) {
+      return c.Reg1 == AArch64::LR && c.Reg2 == AArch64::FP;
+    };
+    auto FPLRBegin = std::find_if(RegPairs.begin(), RegPairs.end(), IsFPLR);
+    const RegPairInfo FPLRRPI = *FPLRBegin;
+    FPLRBegin = std::remove_if(RegPairs.begin(), RegPairs.end(), IsFPLR);
+    *FPLRBegin = FPLRRPI;
+  }
+
   // For performance reasons restore SVE register in increasing order
   auto IsPPR = [](const RegPairInfo &c) { return c.Type == RegPairInfo::PPR; };
   auto PPRBegin = std::find_if(RegPairs.begin(), RegPairs.end(), IsPPR);
@@ -4796,6 +4933,10 @@ void AArch64FrameLowering::orderFrameObjects(
 
   const AArch64FunctionInfo &AFI = *MF.getInfo<AArch64FunctionInfo>();
   const MachineFrameInfo &MFI = MF.getFrameInfo();
+
+  if (AFI.hasPoplessEpilogue())
+    return;
+
   std::vector<FrameObject> FrameObjects(MFI.getObjectIndexEnd());
   for (auto &Obj : ObjectsToAllocate) {
     FrameObjects[Obj].IsValid = true;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -2560,6 +2560,7 @@ const char *AArch64TargetLowering::getTargetNodeName(unsigned Opcode) const {
     MAKE_CASE(AArch64ISD::AUTH_CALL_RVMARKER)
     MAKE_CASE(AArch64ISD::LOADgot)
     MAKE_CASE(AArch64ISD::RET_GLUE)
+    MAKE_CASE(AArch64ISD::RET_POPLESS)
     MAKE_CASE(AArch64ISD::BRCOND)
     MAKE_CASE(AArch64ISD::CSEL)
     MAKE_CASE(AArch64ISD::CSINV)
@@ -8156,6 +8157,18 @@ bool AArch64TargetLowering::DoesCalleeRestoreStack(CallingConv::ID CallCC,
          CallCC == CallingConv::Tail || CallCC == CallingConv::SwiftTail;
 }
 
+SDValue AArch64TargetLowering::adjustReturnPopless(SDValue RetChain,
+                                                   SelectionDAG &DAG) const {
+  if (RetChain.getOpcode() != AArch64ISD::RET_GLUE)
+    report_fatal_error("Unsupported aarch64 return for popless ret lowering");
+
+  auto *AFI = DAG.getMachineFunction().getInfo<AArch64FunctionInfo>();
+  AFI->setHasPoplessEpilogue();
+
+  return DAG.getNode(AArch64ISD::RET_POPLESS, SDLoc(RetChain),
+                     MVT::Other, RetChain->ops());
+}
+
 // Check if the value is zero-extended from i1 to i8
 static bool checkZExtBool(SDValue Arg, const SelectionDAG &DAG) {
   unsigned SizeInBits = Arg.getValueType().getSizeInBits();

llvm/lib/Target/AArch64/AArch64ISelLowering.h

@@ -91,6 +91,7 @@ enum NodeType : unsigned {
   LOADgot,  // Load from automatically generated descriptor (e.g. Global
             // Offset Table, TLS record).
   RET_GLUE, // Return with a glue operand. Operand 0 is the chain operand.
+  RET_POPLESS, // Same as RET_GLUE, though "popless", = doesn't clean the stack.
   BRCOND,   // Conditional branch instruction; "b.cond".
   CSEL,
   CSINV, // Conditional select invert.
@@ -1094,6 +1095,9 @@ class AArch64TargetLowering : public TargetLowering {
   void saveVarArgRegisters(CCState &CCInfo, SelectionDAG &DAG, const SDLoc &DL,
                            SDValue &Chain) const;
 
+  SDValue adjustReturnPopless(SDValue RetChain,
+                              SelectionDAG &DAG) const override;
+
   bool CanLowerReturn(CallingConv::ID CallConv, MachineFunction &MF,
                       bool isVarArg,
                       const SmallVectorImpl<ISD::OutputArg> &Outs,