[CoroutineAccessors] Use swiftcorocc if available.

include/swift/ABI/Coro.h

@@ -25,7 +25,7 @@ namespace swift {
 
 enum class CoroAllocatorKind : uint8_t {
   // stacksave/stackrestore
-  Sync = 0,
+  Stack = 0,
   // swift_task_alloc/swift_task_dealloc_through
   Async = 1,
   // malloc/free

include/swift/AST/IRGenOptions.h

@@ -503,6 +503,13 @@ class IRGenOptions {
 
   unsigned UseProfilingMarkerThunks : 1;
 
+  // Whether swiftcorocc should be used for yield_once_2 routines on x86_64.
+  unsigned UseCoroCCX8664 : 1;
+
+  // Whether swiftcorocc should be used for yield_once_2 routines on arm64
+  // variants.
+  unsigned UseCoroCCArm64 : 1;
+
   /// The number of threads for multi-threaded code generation.
   unsigned NumThreads = 0;
 
@@ -606,8 +613,9 @@ class IRGenOptions {
         UseFragileResilientProtocolWitnesses(false), EnableHotColdSplit(false),
         EmitAsyncFramePushPopMetadata(true), EmitTypeMallocForCoroFrame(false),
         EmitYieldOnce2AsYieldOnce(true), AsyncFramePointerAll(false),
-        UseProfilingMarkerThunks(false), DebugInfoForProfiling(false),
-        CmdArgs(), SanitizeCoverage(llvm::SanitizerCoverageOptions()),
+        UseProfilingMarkerThunks(false), UseCoroCCX8664(false),
+        UseCoroCCArm64(false), DebugInfoForProfiling(false), CmdArgs(),
+        SanitizeCoverage(llvm::SanitizerCoverageOptions()),
         TypeInfoFilter(TypeInfoDumpFilter::All),
         PlatformCCallingConvention(llvm::CallingConv::C), UseCASBackend(false),
         CASObjMode(llvm::CASBackendMode::Native) {

include/swift/Option/FrontendOptions.td

@@ -1149,6 +1149,18 @@ def enable_large_loadable_types_reg2mem : Flag<["-"], "enable-large-loadable-typ
 def disable_large_loadable_types_reg2mem : Flag<["-"], "disable-large-loadable-types-reg2mem">,
   HelpText<"Disable large loadable types register to memory pass">;
 
+def enable_x86_64_corocc : Flag<["-"], "enable-x86_64-corocc">,
+  HelpText<"Use swiftcorocc for yield_once_2 routines on x86_64.">;
+
+def disable_x86_64_corocc : Flag<["-"], "disable-x86_64-corocc">,
+  HelpText<"Don't use swiftcorocc for yield_once_2 routines on x86_64.">;
+
+def enable_arm64_corocc : Flag<["-"], "enable-arm64-corocc">,
+  HelpText<"Use swiftcorocc for yield_once_2 routines on arm64 variants.">;
+
+def disable_arm64_corocc : Flag<["-"], "disable-arm64-corocc">,
+  HelpText<"Don't use swiftcorocc for yield_once_2 routines on arm64 variants.">;
+
 let Flags = [FrontendOption, NoDriverOption, HelpHidden, ModuleInterfaceOptionIgnorable] in {
   def enable_pack_metadata_stack_promotion :
     Joined<["-"], "enable-pack-metadata-stack-promotion=">,

lib/Frontend/CompilerInvocation.cpp

@@ -3653,6 +3653,10 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
       Args.hasFlag(OPT_enable_large_loadable_types_reg2mem,
                    OPT_disable_large_loadable_types_reg2mem,
                    Opts.EnableLargeLoadableTypesReg2Mem);
+  Opts.UseCoroCCX8664 = Args.hasFlag(
+      OPT_enable_x86_64_corocc, OPT_disable_x86_64_corocc, Opts.UseCoroCCX8664);
+  Opts.UseCoroCCArm64 = Args.hasFlag(
+      OPT_enable_arm64_corocc, OPT_disable_arm64_corocc, Opts.UseCoroCCArm64);
   Opts.EnableLayoutStringValueWitnesses = Args.hasFlag(OPT_enable_layout_string_value_witnesses,
                                                        OPT_disable_layout_string_value_witnesses,
                                                        Opts.EnableLayoutStringValueWitnesses);

lib/IRGen/GenCall.cpp

@@ -35,6 +35,9 @@
 #include "clang/CodeGen/ModuleBuilder.h"
 #include "clang/Sema/Sema.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/CallingConv.h"
+#include "llvm/IR/Constant.h"
+#include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalPtrAuthInfo.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/Support/Compiler.h"
@@ -220,7 +223,11 @@ IRGenFunction::getDefaultCoroutineAllocatorKind() {
   if (isCoroutine()) {
     return CoroAllocatorKind::Malloc;
   }
-  return CoroAllocatorKind::Sync;
+  if (IGM.SwiftCoroCC != llvm::CallingConv::SwiftCoro) {
+    // If the swiftcorocc isn't available, fall back to malloc.
+    return CoroAllocatorKind::Malloc;
+  }
+  return CoroAllocatorKind::Stack;
 }
 
 llvm::Value *IRGenFunction::getAsyncTask() {
@@ -360,9 +367,10 @@ static llvm::CallingConv::ID getFreestandingConvention(IRGenModule &IGM) {
 
 /// Expand the requirements of the given abstract calling convention
 /// into a "physical" calling convention.
-llvm::CallingConv::ID irgen::expandCallingConv(IRGenModule &IGM,
-                                    SILFunctionTypeRepresentation convention,
-                                    bool isAsync) {
+llvm::CallingConv::ID
+irgen::expandCallingConv(IRGenModule &IGM,
+                         SILFunctionTypeRepresentation convention, bool isAsync,
+                         bool isCalleeAllocatedCoro) {
   switch (convention) {
   case SILFunctionTypeRepresentation::CFunctionPointer:
   case SILFunctionTypeRepresentation::ObjCMethod:
@@ -379,6 +387,8 @@ llvm::CallingConv::ID irgen::expandCallingConv(IRGenModule &IGM,
   case SILFunctionTypeRepresentation::KeyPathAccessorSetter:
   case SILFunctionTypeRepresentation::KeyPathAccessorEquals:
   case SILFunctionTypeRepresentation::KeyPathAccessorHash:
+    if (isCalleeAllocatedCoro && !IGM.getOptions().EmitYieldOnce2AsYieldOnce)
+      return IGM.SwiftCoroCC;
     if (isAsync)
       return IGM.SwiftAsyncCC;
     return getFreestandingConvention(IGM);
@@ -2391,7 +2401,8 @@ Signature SignatureExpansion::getSignature() {
          "C function type without C function info");
 
   auto callingConv =
-      expandCallingConv(IGM, FnType->getRepresentation(), FnType->isAsync());
+      expandCallingConv(IGM, FnType->getRepresentation(), FnType->isAsync(),
+                        FnType->isCalleeAllocatedCoroutine());
 
   Signature result;
   result.Type = llvmType;
@@ -5118,6 +5129,40 @@ void irgen::emitYieldManyCoroutineEntry(
                            allocFn, deallocFn, {});
 }
 
+static llvm::Constant *getCoroAllocWrapperFn(IRGenModule &IGM) {
+  return IGM.getOrCreateHelperFunction(
+      "__swift_coro_alloc_", IGM.Int8PtrTy,
+      {IGM.CoroAllocatorPtrTy, IGM.SizeTy},
+      [](IRGenFunction &IGF) {
+        auto parameters = IGF.collectParameters();
+        auto *allocator = parameters.claimNext();
+        auto *size = parameters.claimNext();
+        auto *nullAllocator = IGF.Builder.CreateCmp(
+            llvm::CmpInst::Predicate::ICMP_EQ, allocator,
+            llvm::ConstantPointerNull::get(
+                cast<llvm::PointerType>(allocator->getType())));
+        auto *poplessReturn = IGF.createBasicBlock("coro.return.popless");
+        auto *normalReturn = IGF.createBasicBlock("coro.return.normal");
+        IGF.Builder.CreateCondBr(nullAllocator, poplessReturn, normalReturn);
+        IGF.Builder.emitBlock(poplessReturn);
+        // Emit the dynamic alloca.
+        auto *alloca =
+            IGF.Builder.IRBuilderBase::CreateAlloca(IGF.IGM.Int8Ty, size);
+        alloca->setAlignment(llvm::Align(MaximumAlignment));
+        auto *ret = IGF.Builder.CreateIntrinsic(
+            alloca->getType(), llvm::Intrinsic::coro_return, {alloca});
+        IGF.Builder.CreateRet(ret);
+        IGF.Builder.emitBlock(normalReturn);
+        auto *call = IGF.Builder.CreateCall(
+            IGF.IGM.getCoroAllocFunctionPointer(), {allocator, size});
+        IGF.Builder.CreateRet(call);
+      },
+      /*setIsNoInline=*/false,
+      /*forPrologue=*/false,
+      /*isPerformanceConstraint=*/false,
+      /*optionalLinkageOverride=*/nullptr, llvm::CallingConv::SwiftCoro);
+}
+
 void irgen::emitYieldOnce2CoroutineEntry(
     IRGenFunction &IGF, LinkEntity coroFunction, CanSILFunctionType fnType,
     NativeCCEntryPointArgumentEmission &emission) {
@@ -5126,7 +5171,11 @@ void irgen::emitYieldOnce2CoroutineEntry(
       IGF.IGM.getAddrOfCoroFunctionPointer(coroFunction));
   llvm::Value *allocator = emission.getCoroutineAllocator();
   IGF.setCoroutineAllocator(allocator);
-  auto allocFn = IGF.IGM.getOpaquePtr(IGF.IGM.getCoroAllocFn());
+  auto isSwiftCoroCCAvailable =
+      IGF.IGM.SwiftCoroCC == llvm::CallingConv::SwiftCoro;
+  auto allocFn = IGF.IGM.getOpaquePtr(isSwiftCoroCCAvailable
+                                          ? getCoroAllocWrapperFn(IGF.IGM)
+                                          : IGF.IGM.getCoroAllocFn());
   auto deallocFn = IGF.IGM.getOpaquePtr(IGF.IGM.getCoroDeallocFn());
   emitRetconCoroutineEntry(
       IGF, fnType, buffer, llvm::Intrinsic::coro_id_retcon_once_dynamic,

lib/IRGen/GenCall.h

@@ -146,9 +146,9 @@ namespace irgen {
   std::pair<llvm::Value *, llvm::Value *>
   getCoroFunctionAndSize(IRGenFunction &IGF, FunctionPointer functionPointer,
                          std::pair<bool, bool> values = {true, true});
-  llvm::CallingConv::ID expandCallingConv(IRGenModule &IGM,
-                                     SILFunctionTypeRepresentation convention,
-                                     bool isAsync);
+  llvm::CallingConv::ID
+  expandCallingConv(IRGenModule &IGM, SILFunctionTypeRepresentation convention,
+                    bool isAsync, bool isCalleeAllocatedCoro);
 
   Signature emitCastOfFunctionPointer(IRGenFunction &IGF, llvm::Value *&fnPtr,
                                       CanSILFunctionType fnType,

lib/IRGen/GenDecl.cpp

@@ -6216,10 +6216,10 @@ IRGenModule::getAddrOfContinuationPrototype(CanSILFunctionType fnType) {
 }
 
 /// Should we be defining the given helper function?
-static llvm::Function *shouldDefineHelper(IRGenModule &IGM,
-                                          llvm::Constant *fn,
-                                          bool setIsNoInline,
-                                          IRLinkage *linkage) {
+static llvm::Function *
+shouldDefineHelper(IRGenModule &IGM, llvm::Constant *fn, bool setIsNoInline,
+                   IRLinkage *linkage,
+                   std::optional<llvm::CallingConv::ID> specialCallingConv) {
   auto *def = dyn_cast<llvm::Function>(fn);
   if (!def) return nullptr;
   if (!def->empty()) return nullptr;
@@ -6231,7 +6231,7 @@ static llvm::Function *shouldDefineHelper(IRGenModule &IGM,
     ApplyIRLinkage(*linkage).to(def);
 
   def->setDoesNotThrow();
-  def->setCallingConv(IGM.DefaultCC);
+  def->setCallingConv(specialCallingConv.value_or(IGM.DefaultCC));
   if (setIsNoInline)
     def->addFnAttr(llvm::Attribute::NoInline);
   return def;
@@ -6245,23 +6245,22 @@ static llvm::Function *shouldDefineHelper(IRGenModule &IGM,
 /// does not collide with any other helper function.  In general, it should
 /// be a Swift-specific C reserved name; that is, it should start with
 //  "__swift".
-llvm::Constant *
-IRGenModule::getOrCreateHelperFunction(StringRef fnName, llvm::Type *resultTy,
-                                       ArrayRef<llvm::Type*> paramTys,
-                        llvm::function_ref<void(IRGenFunction &IGF)> generate,
-                        bool setIsNoInline,
-                        bool forPrologue,
-                        bool isPerformanceConstraint,
-                        IRLinkage *optionalLinkageOverride) {
+llvm::Constant *IRGenModule::getOrCreateHelperFunction(
+    StringRef fnName, llvm::Type *resultTy, ArrayRef<llvm::Type *> paramTys,
+    llvm::function_ref<void(IRGenFunction &IGF)> generate, bool setIsNoInline,
+    bool forPrologue, bool isPerformanceConstraint,
+    IRLinkage *optionalLinkageOverride,
+    std::optional<llvm::CallingConv::ID> specialCallingConv) {
   llvm::FunctionType *fnTy =
     llvm::FunctionType::get(resultTy, paramTys, false);
 
   llvm::Constant *fn =
       cast<llvm::Constant>(
           Module.getOrInsertFunction(fnName, fnTy).getCallee());
 
-  if (llvm::Function *def = shouldDefineHelper(*this, fn, setIsNoInline,
-                                               optionalLinkageOverride)) {
+  if (llvm::Function *def =
+          shouldDefineHelper(*this, fn, setIsNoInline, optionalLinkageOverride,
+                             specialCallingConv)) {
     IRGenFunction IGF(*this, def, isPerformanceConstraint);
     if (DebugInfo && !forPrologue)
       DebugInfo->emitArtificialFunction(IGF, def);

lib/IRGen/GenObjC.cpp

@@ -908,7 +908,8 @@ static llvm::Function *emitObjCPartialApplicationForwarder(IRGenModule &IGM,
                            MANGLE_AS_STRING(OBJC_PARTIAL_APPLY_THUNK_SYM),
                            &IGM.Module);
   fwd->setCallingConv(expandCallingConv(
-      IGM, SILFunctionTypeRepresentation::Thick, false/*isAsync*/));
+      IGM, SILFunctionTypeRepresentation::Thick, false /*isAsync*/,
+      false /*isCalleeAllocatedCoroutine*/));
 
   fwd->setAttributes(attrs);
   // Merge initial attributes with attrs.

lib/IRGen/IRGenModule.cpp

@@ -208,6 +208,24 @@ static void sanityCheckStdlib(IRGenModule &IGM) {
 }
 #endif
 
+static bool getIsCoroCCSupported(const clang::TargetInfo &targetInfo,
+                                 const IRGenOptions &IRGenOpts) {
+  // TODO: CoroutineAccessors: The one caller of this function should just be
+  //                           rewritten as follows (once clang::CC_CoroAsync is
+  //                           defined).
+  //
+  // clangASTContext.getTargetInfo().checkCallingConvention(clang::CC_CoroAsync)
+  if (targetInfo.getTriple().isAArch64() && IRGenOpts.UseCoroCCArm64) {
+    return true;
+  }
+  if (targetInfo.getTriple().getArch() == llvm::Triple::x86_64 &&
+      IRGenOpts.UseCoroCCX8664) {
+    return true;
+  }
+
+  return false;
+}
+
 IRGenModule::IRGenModule(IRGenerator &irgen,
                          std::unique_ptr<llvm::TargetMachine> &&target,
                          SourceFile *SF, StringRef ModuleName,
@@ -602,6 +620,14 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
     AsyncTailCallKind = llvm::CallInst::TCK_Tail;
   }
 
+  bool isCoroCCSupported =
+      getIsCoroCCSupported(clangASTContext.getTargetInfo(), opts);
+  if (isCoroCCSupported) {
+    SwiftCoroCC = llvm::CallingConv::SwiftCoro;
+  } else {
+    SwiftCoroCC = llvm::CallingConv::Swift;
+  }
+
   if (opts.DebugInfoLevel > IRGenDebugInfoLevel::None)
     DebugInfo = IRGenDebugInfo::createIRGenDebugInfo(IRGen.Opts, *CI, *this,
                                                      Module,

lib/IRGen/IRGenModule.h

@@ -876,6 +876,7 @@ class IRGenModule {
   llvm::CallingConv::ID DefaultCC;     /// default calling convention
   llvm::CallingConv::ID SwiftCC;       /// swift calling convention
   llvm::CallingConv::ID SwiftAsyncCC;  /// swift calling convention for async
+  llvm::CallingConv::ID SwiftCoroCC;   /// swift calling convention for callee-allocated coroutines
 
   /// What kind of tail call should be used for async->async calls.
   llvm::CallInst::TailCallKind AsyncTailCallKind;
@@ -1226,14 +1227,13 @@ class IRGenModule {
                                            const PointerAuthEntity &entity,
                                            llvm::Constant *storageAddress);
 
-  llvm::Constant *getOrCreateHelperFunction(StringRef name,
-                                            llvm::Type *resultType,
-                                            ArrayRef<llvm::Type*> paramTypes,
-                        llvm::function_ref<void(IRGenFunction &IGF)> generate,
-                        bool setIsNoInline = false,
-                        bool forPrologue = false,
-                        bool isPerformanceConstraint = false,
-                        IRLinkage *optionalLinkage = nullptr);
+  llvm::Constant *getOrCreateHelperFunction(
+      StringRef name, llvm::Type *resultType, ArrayRef<llvm::Type *> paramTypes,
+      llvm::function_ref<void(IRGenFunction &IGF)> generate,
+      bool setIsNoInline = false, bool forPrologue = false,
+      bool isPerformanceConstraint = false,
+      IRLinkage *optionalLinkage = nullptr,
+      std::optional<llvm::CallingConv::ID> specialCallingConv = std::nullopt);
 
   llvm::Constant *getOrCreateRetainFunction(const TypeInfo &objectTI, SILType t,
                               llvm::Type *llvmType, Atomicity atomicity);

stdlib/public/Concurrency/TaskAlloc.cpp

@@ -87,7 +87,7 @@ CoroAllocator *const swift::_swift_coro_task_allocator = &CoroTaskAllocatorImpl;
 
 CoroAllocator *swift::swift_coro_getGlobalAllocator(CoroAllocatorFlags flags) {
   switch (flags.getKind()) {
-  case CoroAllocatorKind::Sync:
+  case CoroAllocatorKind::Stack:
     return nullptr;
   case CoroAllocatorKind::Async:
     return _swift_coro_task_allocator;
@@ -107,6 +107,8 @@ void *swift::swift_coro_alloc(CoroAllocator *allocator, size_t size) {
 }
 
 void swift::swift_coro_dealloc(CoroAllocator *allocator, void *ptr) {
+  if (!allocator)
+    return;
   // Calls to swift_coro_dealloc are emitted in resume funclets for every
   // live-across dynamic allocation.  Whether such calls immediately deallocate
   // memory depends on the allocator.