Skip to content

Commit f8b60bb

Browse files
nate-chandlernate-chandler
authored
Merge pull request #79685 from nate-chandler/general-coro/20250227/1
[CoroutineAccessors] Use swiftcorocc if available.
2 parents 7f8de34 + d1f1b4c commit f8b60bb

File tree

12 files changed

+173
-40
lines changed

12 files changed

+173
-40
lines changed

‎include/swift/ABI/Coro.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ namespace swift {
2525

2626
enum class CoroAllocatorKind : uint8_t {
2727
// stacksave/stackrestore
28-
Sync = 0,
28+
Stack = 0,
2929
// swift_task_alloc/swift_task_dealloc_through
3030
Async = 1,
3131
// malloc/free

‎include/swift/AST/IRGenOptions.h

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -507,6 +507,13 @@ class IRGenOptions {
507507

508508
unsigned UseProfilingMarkerThunks : 1;
509509

510+
// Whether swiftcorocc should be used for yield_once_2 routines on x86_64.
511+
unsigned UseCoroCCX8664 : 1;
512+
513+
// Whether swiftcorocc should be used for yield_once_2 routines on arm64
514+
// variants.
515+
unsigned UseCoroCCArm64 : 1;
516+
510517
/// The number of threads for multi-threaded code generation.
511518
unsigned NumThreads = 0;
512519

@@ -610,8 +617,9 @@ class IRGenOptions {
610617
UseFragileResilientProtocolWitnesses(false), EnableHotColdSplit(false),
611618
EmitAsyncFramePushPopMetadata(true), EmitTypeMallocForCoroFrame(false),
612619
EmitYieldOnce2AsYieldOnce(true), AsyncFramePointerAll(false),
613-
UseProfilingMarkerThunks(false), DebugInfoForProfiling(false),
614-
CmdArgs(), SanitizeCoverage(llvm::SanitizerCoverageOptions()),
620+
UseProfilingMarkerThunks(false), UseCoroCCX8664(false),
621+
UseCoroCCArm64(false), DebugInfoForProfiling(false), CmdArgs(),
622+
SanitizeCoverage(llvm::SanitizerCoverageOptions()),
615623
TypeInfoFilter(TypeInfoDumpFilter::All),
616624
PlatformCCallingConvention(llvm::CallingConv::C), UseCASBackend(false),
617625
CASObjMode(llvm::CASBackendMode::Native) {

‎include/swift/Option/FrontendOptions.td

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1149,6 +1149,18 @@ def enable_large_loadable_types_reg2mem : Flag<["-"], "enable-large-loadable-typ
11491149
def disable_large_loadable_types_reg2mem : Flag<["-"], "disable-large-loadable-types-reg2mem">,
11501150
HelpText<"Disable large loadable types register to memory pass">;
11511151

1152+
def enable_x86_64_corocc : Flag<["-"], "enable-x86_64-corocc">,
1153+
HelpText<"Use swiftcorocc for yield_once_2 routines on x86_64.">;
1154+
1155+
def disable_x86_64_corocc : Flag<["-"], "disable-x86_64-corocc">,
1156+
HelpText<"Don't use swiftcorocc for yield_once_2 routines on x86_64.">;
1157+
1158+
def enable_arm64_corocc : Flag<["-"], "enable-arm64-corocc">,
1159+
HelpText<"Use swiftcorocc for yield_once_2 routines on arm64 variants.">;
1160+
1161+
def disable_arm64_corocc : Flag<["-"], "disable-arm64-corocc">,
1162+
HelpText<"Don't use swiftcorocc for yield_once_2 routines on arm64 variants.">;
1163+
11521164
let Flags = [FrontendOption, NoDriverOption, HelpHidden, ModuleInterfaceOptionIgnorable] in {
11531165
def enable_pack_metadata_stack_promotion :
11541166
Joined<["-"], "enable-pack-metadata-stack-promotion=">,

‎lib/Frontend/CompilerInvocation.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3653,6 +3653,10 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
36533653
Args.hasFlag(OPT_enable_large_loadable_types_reg2mem,
36543654
OPT_disable_large_loadable_types_reg2mem,
36553655
Opts.EnableLargeLoadableTypesReg2Mem);
3656+
Opts.UseCoroCCX8664 = Args.hasFlag(
3657+
OPT_enable_x86_64_corocc, OPT_disable_x86_64_corocc, Opts.UseCoroCCX8664);
3658+
Opts.UseCoroCCArm64 = Args.hasFlag(
3659+
OPT_enable_arm64_corocc, OPT_disable_arm64_corocc, Opts.UseCoroCCArm64);
36563660
Opts.EnableLayoutStringValueWitnesses = Args.hasFlag(OPT_enable_layout_string_value_witnesses,
36573661
OPT_disable_layout_string_value_witnesses,
36583662
Opts.EnableLayoutStringValueWitnesses);

‎lib/IRGen/GenCall.cpp

Lines changed: 55 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,9 @@
3535
#include "clang/CodeGen/ModuleBuilder.h"
3636
#include "clang/Sema/Sema.h"
3737
#include "llvm/Analysis/ValueTracking.h"
38+
#include "llvm/IR/CallingConv.h"
39+
#include "llvm/IR/Constant.h"
40+
#include "llvm/IR/Function.h"
3841
#include "llvm/IR/GlobalPtrAuthInfo.h"
3942
#include "llvm/IR/GlobalValue.h"
4043
#include "llvm/Support/Compiler.h"
@@ -220,7 +223,11 @@ IRGenFunction::getDefaultCoroutineAllocatorKind() {
220223
if (isCoroutine()) {
221224
return CoroAllocatorKind::Malloc;
222225
}
223-
return CoroAllocatorKind::Sync;
226+
if (IGM.SwiftCoroCC != llvm::CallingConv::SwiftCoro) {
227+
// If the swiftcorocc isn't available, fall back to malloc.
228+
return CoroAllocatorKind::Malloc;
229+
}
230+
return CoroAllocatorKind::Stack;
224231
}
225232

226233
llvm::Value *IRGenFunction::getAsyncTask() {
@@ -360,9 +367,10 @@ static llvm::CallingConv::ID getFreestandingConvention(IRGenModule &IGM) {
360367

361368
/// Expand the requirements of the given abstract calling convention
362369
/// into a "physical" calling convention.
363-
llvm::CallingConv::ID irgen::expandCallingConv(IRGenModule &IGM,
364-
SILFunctionTypeRepresentation convention,
365-
bool isAsync) {
370+
llvm::CallingConv::ID
371+
irgen::expandCallingConv(IRGenModule &IGM,
372+
SILFunctionTypeRepresentation convention, bool isAsync,
373+
bool isCalleeAllocatedCoro) {
366374
switch (convention) {
367375
case SILFunctionTypeRepresentation::CFunctionPointer:
368376
case SILFunctionTypeRepresentation::ObjCMethod:
@@ -379,6 +387,8 @@ llvm::CallingConv::ID irgen::expandCallingConv(IRGenModule &IGM,
379387
case SILFunctionTypeRepresentation::KeyPathAccessorSetter:
380388
case SILFunctionTypeRepresentation::KeyPathAccessorEquals:
381389
case SILFunctionTypeRepresentation::KeyPathAccessorHash:
390+
if (isCalleeAllocatedCoro && !IGM.getOptions().EmitYieldOnce2AsYieldOnce)
391+
return IGM.SwiftCoroCC;
382392
if (isAsync)
383393
return IGM.SwiftAsyncCC;
384394
return getFreestandingConvention(IGM);
@@ -2391,7 +2401,8 @@ Signature SignatureExpansion::getSignature() {
23912401
"C function type without C function info");
23922402

23932403
auto callingConv =
2394-
expandCallingConv(IGM, FnType->getRepresentation(), FnType->isAsync());
2404+
expandCallingConv(IGM, FnType->getRepresentation(), FnType->isAsync(),
2405+
FnType->isCalleeAllocatedCoroutine());
23952406

23962407
Signature result;
23972408
result.Type = llvmType;
@@ -5118,6 +5129,40 @@ void irgen::emitYieldManyCoroutineEntry(
51185129
allocFn, deallocFn, {});
51195130
}
51205131

5132+
static llvm::Constant *getCoroAllocWrapperFn(IRGenModule &IGM) {
5133+
return IGM.getOrCreateHelperFunction(
5134+
"__swift_coro_alloc_", IGM.Int8PtrTy,
5135+
{IGM.CoroAllocatorPtrTy, IGM.SizeTy},
5136+
[](IRGenFunction &IGF) {
5137+
auto parameters = IGF.collectParameters();
5138+
auto *allocator = parameters.claimNext();
5139+
auto *size = parameters.claimNext();
5140+
auto *nullAllocator = IGF.Builder.CreateCmp(
5141+
llvm::CmpInst::Predicate::ICMP_EQ, allocator,
5142+
llvm::ConstantPointerNull::get(
5143+
cast<llvm::PointerType>(allocator->getType())));
5144+
auto *poplessReturn = IGF.createBasicBlock("coro.return.popless");
5145+
auto *normalReturn = IGF.createBasicBlock("coro.return.normal");
5146+
IGF.Builder.CreateCondBr(nullAllocator, poplessReturn, normalReturn);
5147+
IGF.Builder.emitBlock(poplessReturn);
5148+
// Emit the dynamic alloca.
5149+
auto *alloca =
5150+
IGF.Builder.IRBuilderBase::CreateAlloca(IGF.IGM.Int8Ty, size);
5151+
alloca->setAlignment(llvm::Align(MaximumAlignment));
5152+
auto *ret = IGF.Builder.CreateIntrinsic(
5153+
alloca->getType(), llvm::Intrinsic::coro_return, {alloca});
5154+
IGF.Builder.CreateRet(ret);
5155+
IGF.Builder.emitBlock(normalReturn);
5156+
auto *call = IGF.Builder.CreateCall(
5157+
IGF.IGM.getCoroAllocFunctionPointer(), {allocator, size});
5158+
IGF.Builder.CreateRet(call);
5159+
},
5160+
/*setIsNoInline=*/false,
5161+
/*forPrologue=*/false,
5162+
/*isPerformanceConstraint=*/false,
5163+
/*optionalLinkageOverride=*/nullptr, llvm::CallingConv::SwiftCoro);
5164+
}
5165+
51215166
void irgen::emitYieldOnce2CoroutineEntry(
51225167
IRGenFunction &IGF, LinkEntity coroFunction, CanSILFunctionType fnType,
51235168
NativeCCEntryPointArgumentEmission &emission) {
@@ -5126,7 +5171,11 @@ void irgen::emitYieldOnce2CoroutineEntry(
51265171
IGF.IGM.getAddrOfCoroFunctionPointer(coroFunction));
51275172
llvm::Value *allocator = emission.getCoroutineAllocator();
51285173
IGF.setCoroutineAllocator(allocator);
5129-
auto allocFn = IGF.IGM.getOpaquePtr(IGF.IGM.getCoroAllocFn());
5174+
auto isSwiftCoroCCAvailable =
5175+
IGF.IGM.SwiftCoroCC == llvm::CallingConv::SwiftCoro;
5176+
auto allocFn = IGF.IGM.getOpaquePtr(isSwiftCoroCCAvailable
5177+
? getCoroAllocWrapperFn(IGF.IGM)
5178+
: IGF.IGM.getCoroAllocFn());
51305179
auto deallocFn = IGF.IGM.getOpaquePtr(IGF.IGM.getCoroDeallocFn());
51315180
emitRetconCoroutineEntry(
51325181
IGF, fnType, buffer, llvm::Intrinsic::coro_id_retcon_once_dynamic,

‎lib/IRGen/GenCall.h

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -146,9 +146,9 @@ namespace irgen {
146146
std::pair<llvm::Value *, llvm::Value *>
147147
getCoroFunctionAndSize(IRGenFunction &IGF, FunctionPointer functionPointer,
148148
std::pair<bool, bool> values = {true, true});
149-
llvm::CallingConv::ID expandCallingConv(IRGenModule &IGM,
150-
SILFunctionTypeRepresentation convention,
151-
bool isAsync);
149+
llvm::CallingConv::ID
150+
expandCallingConv(IRGenModule &IGM, SILFunctionTypeRepresentation convention,
151+
bool isAsync, bool isCalleeAllocatedCoro);
152152

153153
Signature emitCastOfFunctionPointer(IRGenFunction &IGF, llvm::Value *&fnPtr,
154154
CanSILFunctionType fnType,

‎lib/IRGen/GenDecl.cpp

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -6216,10 +6216,10 @@ IRGenModule::getAddrOfContinuationPrototype(CanSILFunctionType fnType) {
62166216
}
62176217

62186218
/// Should we be defining the given helper function?
6219-
static llvm::Function *shouldDefineHelper(IRGenModule &IGM,
6220-
llvm::Constant *fn,
6221-
bool setIsNoInline,
6222-
IRLinkage *linkage) {
6219+
static llvm::Function *
6220+
shouldDefineHelper(IRGenModule &IGM, llvm::Constant *fn, bool setIsNoInline,
6221+
IRLinkage *linkage,
6222+
std::optional<llvm::CallingConv::ID> specialCallingConv) {
62236223
auto *def = dyn_cast<llvm::Function>(fn);
62246224
if (!def) return nullptr;
62256225
if (!def->empty()) return nullptr;
@@ -6231,7 +6231,7 @@ static llvm::Function *shouldDefineHelper(IRGenModule &IGM,
62316231
ApplyIRLinkage(*linkage).to(def);
62326232

62336233
def->setDoesNotThrow();
6234-
def->setCallingConv(IGM.DefaultCC);
6234+
def->setCallingConv(specialCallingConv.value_or(IGM.DefaultCC));
62356235
if (setIsNoInline)
62366236
def->addFnAttr(llvm::Attribute::NoInline);
62376237
return def;
@@ -6245,23 +6245,22 @@ static llvm::Function *shouldDefineHelper(IRGenModule &IGM,
62456245
/// does not collide with any other helper function. In general, it should
62466246
/// be a Swift-specific C reserved name; that is, it should start with
62476247
// "__swift".
6248-
llvm::Constant *
6249-
IRGenModule::getOrCreateHelperFunction(StringRef fnName, llvm::Type *resultTy,
6250-
ArrayRef<llvm::Type*> paramTys,
6251-
llvm::function_ref<void(IRGenFunction &IGF)> generate,
6252-
bool setIsNoInline,
6253-
bool forPrologue,
6254-
bool isPerformanceConstraint,
6255-
IRLinkage *optionalLinkageOverride) {
6248+
llvm::Constant *IRGenModule::getOrCreateHelperFunction(
6249+
StringRef fnName, llvm::Type *resultTy, ArrayRef<llvm::Type *> paramTys,
6250+
llvm::function_ref<void(IRGenFunction &IGF)> generate, bool setIsNoInline,
6251+
bool forPrologue, bool isPerformanceConstraint,
6252+
IRLinkage *optionalLinkageOverride,
6253+
std::optional<llvm::CallingConv::ID> specialCallingConv) {
62566254
llvm::FunctionType *fnTy =
62576255
llvm::FunctionType::get(resultTy, paramTys, false);
62586256

62596257
llvm::Constant *fn =
62606258
cast<llvm::Constant>(
62616259
Module.getOrInsertFunction(fnName, fnTy).getCallee());
62626260

6263-
if (llvm::Function *def = shouldDefineHelper(*this, fn, setIsNoInline,
6264-
optionalLinkageOverride)) {
6261+
if (llvm::Function *def =
6262+
shouldDefineHelper(*this, fn, setIsNoInline, optionalLinkageOverride,
6263+
specialCallingConv)) {
62656264
IRGenFunction IGF(*this, def, isPerformanceConstraint);
62666265
if (DebugInfo && !forPrologue)
62676266
DebugInfo->emitArtificialFunction(IGF, def);

‎lib/IRGen/GenObjC.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -908,7 +908,8 @@ static llvm::Function *emitObjCPartialApplicationForwarder(IRGenModule &IGM,
908908
MANGLE_AS_STRING(OBJC_PARTIAL_APPLY_THUNK_SYM),
909909
&IGM.Module);
910910
fwd->setCallingConv(expandCallingConv(
911-
IGM, SILFunctionTypeRepresentation::Thick, false/*isAsync*/));
911+
IGM, SILFunctionTypeRepresentation::Thick, false /*isAsync*/,
912+
false /*isCalleeAllocatedCoroutine*/));
912913

913914
fwd->setAttributes(attrs);
914915
// Merge initial attributes with attrs.

‎lib/IRGen/IRGenModule.cpp

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,24 @@ static void sanityCheckStdlib(IRGenModule &IGM) {
208208
}
209209
#endif
210210

211+
static bool getIsCoroCCSupported(const clang::TargetInfo &targetInfo,
212+
const IRGenOptions &IRGenOpts) {
213+
// TODO: CoroutineAccessors: The one caller of this function should just be
214+
// rewritten as follows (once clang::CC_CoroAsync is
215+
// defined).
216+
//
217+
// clangASTContext.getTargetInfo().checkCallingConvention(clang::CC_CoroAsync)
218+
if (targetInfo.getTriple().isAArch64() && IRGenOpts.UseCoroCCArm64) {
219+
return true;
220+
}
221+
if (targetInfo.getTriple().getArch() == llvm::Triple::x86_64 &&
222+
IRGenOpts.UseCoroCCX8664) {
223+
return true;
224+
}
225+
226+
return false;
227+
}
228+
211229
IRGenModule::IRGenModule(IRGenerator &irgen,
212230
std::unique_ptr<llvm::TargetMachine> &&target,
213231
SourceFile *SF, StringRef ModuleName,
@@ -602,6 +620,14 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
602620
AsyncTailCallKind = llvm::CallInst::TCK_Tail;
603621
}
604622

623+
bool isCoroCCSupported =
624+
getIsCoroCCSupported(clangASTContext.getTargetInfo(), opts);
625+
if (isCoroCCSupported) {
626+
SwiftCoroCC = llvm::CallingConv::SwiftCoro;
627+
} else {
628+
SwiftCoroCC = llvm::CallingConv::Swift;
629+
}
630+
605631
if (opts.DebugInfoLevel > IRGenDebugInfoLevel::None)
606632
DebugInfo = IRGenDebugInfo::createIRGenDebugInfo(IRGen.Opts, *CI, *this,
607633
Module,

‎lib/IRGen/IRGenModule.h

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -876,6 +876,7 @@ class IRGenModule {
876876
llvm::CallingConv::ID DefaultCC; /// default calling convention
877877
llvm::CallingConv::ID SwiftCC; /// swift calling convention
878878
llvm::CallingConv::ID SwiftAsyncCC; /// swift calling convention for async
879+
llvm::CallingConv::ID SwiftCoroCC; /// swift calling convention for callee-allocated coroutines
879880

880881
/// What kind of tail call should be used for async->async calls.
881882
llvm::CallInst::TailCallKind AsyncTailCallKind;
@@ -1226,14 +1227,13 @@ class IRGenModule {
12261227
const PointerAuthEntity &entity,
12271228
llvm::Constant *storageAddress);
12281229

1229-
llvm::Constant *getOrCreateHelperFunction(StringRef name,
1230-
llvm::Type *resultType,
1231-
ArrayRef<llvm::Type*> paramTypes,
1232-
llvm::function_ref<void(IRGenFunction &IGF)> generate,
1233-
bool setIsNoInline = false,
1234-
bool forPrologue = false,
1235-
bool isPerformanceConstraint = false,
1236-
IRLinkage *optionalLinkage = nullptr);
1230+
llvm::Constant *getOrCreateHelperFunction(
1231+
StringRef name, llvm::Type *resultType, ArrayRef<llvm::Type *> paramTypes,
1232+
llvm::function_ref<void(IRGenFunction &IGF)> generate,
1233+
bool setIsNoInline = false, bool forPrologue = false,
1234+
bool isPerformanceConstraint = false,
1235+
IRLinkage *optionalLinkage = nullptr,
1236+
std::optional<llvm::CallingConv::ID> specialCallingConv = std::nullopt);
12371237

12381238
llvm::Constant *getOrCreateRetainFunction(const TypeInfo &objectTI, SILType t,
12391239
llvm::Type *llvmType, Atomicity atomicity);

‎stdlib/public/Concurrency/TaskAlloc.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ CoroAllocator *const swift::_swift_coro_task_allocator = &CoroTaskAllocatorImpl;
8787

8888
CoroAllocator *swift::swift_coro_getGlobalAllocator(CoroAllocatorFlags flags) {
8989
switch (flags.getKind()) {
90-
case CoroAllocatorKind::Sync:
90+
case CoroAllocatorKind::Stack:
9191
return nullptr;
9292
case CoroAllocatorKind::Async:
9393
return _swift_coro_task_allocator;
@@ -107,6 +107,8 @@ void *swift::swift_coro_alloc(CoroAllocator *allocator, size_t size) {
107107
}
108108

109109
void swift::swift_coro_dealloc(CoroAllocator *allocator, void *ptr) {
110+
if (!allocator)
111+
return;
110112
// Calls to swift_coro_dealloc are emitted in resume funclets for every
111113
// live-across dynamic allocation. Whether such calls immediately deallocate
112114
// memory depends on the allocator.

0 commit comments

Comments
 (0)