Skip to content

Commit f0fc662

Browse files
aschwaighoferaschwaighofer
authored
Merge pull request swiftlang#34731 from aschwaighofer/irgen_async_coroutine_splitting
Add llvm coroutine splitting
2 parents c7f4f07 + e51575b commit f0fc662

File tree

48 files changed

+877
-267
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

48 files changed

+877
-267
lines changed

‎include/swift/Runtime/Concurrency.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,10 @@ SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
145145
NearestTaskDeadline
146146
swift_task_getNearestDeadline(AsyncTask *task);
147147

148+
// TODO: Remove this hack.
149+
SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
150+
void swift_task_run(AsyncTask *taskToRun);
151+
148152
}
149153

150154
#endif

‎lib/IRGen/Address.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@ class StackAddress {
112112

113113
/// In a normal function, the result of llvm.stacksave or null.
114114
/// In a coroutine, the result of llvm.coro.alloca.alloc.
115+
/// In an async function, the result of the taskAlloc call.
115116
llvm::Value *ExtraInfo;
116117

117118
public:

‎lib/IRGen/CallEmission.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ namespace irgen {
3030
class Explosion;
3131
class LoadableTypeInfo;
3232
struct WitnessMetadata;
33+
class FunctionPointer;
3334

3435
/// A plan for emitting a series of calls.
3536
class CallEmission {
@@ -70,6 +71,9 @@ class CallEmission {
7071
virtual FunctionPointer getCalleeFunctionPointer() = 0;
7172
llvm::CallInst *emitCallSite();
7273

74+
virtual llvm::CallInst *createCall(const FunctionPointer &fn,
75+
ArrayRef<llvm::Value *> args) = 0;
76+
7377
CallEmission(IRGenFunction &IGF, llvm::Value *selfValue, Callee &&callee)
7478
: IGF(IGF), selfValue(selfValue), CurCallee(std::move(callee)) {}
7579

‎lib/IRGen/GenCall.cpp

Lines changed: 115 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -330,10 +330,6 @@ AsyncContextLayout::AsyncContextLayout(
330330
#endif
331331
}
332332

333-
static Size getAsyncContextSize(AsyncContextLayout layout) {
334-
return layout.getSize();
335-
}
336-
337333
static Alignment getAsyncContextAlignment(IRGenModule &IGM) {
338334
return IGM.getPointerAlignment();
339335
}
@@ -2043,6 +2039,12 @@ class SyncCallEmission final : public CallEmission {
20432039
return origConv.getSILArgumentType(
20442040
index, IGF.IGM.getMaximalTypeExpansionContext());
20452041
}
2042+
2043+
llvm::CallInst *createCall(const FunctionPointer &fn,
2044+
ArrayRef<llvm::Value *> args) override {
2045+
return IGF.Builder.CreateCall(fn, Args);
2046+
}
2047+
20462048
void begin() override { super::begin(); }
20472049
void end() override { super::end(); }
20482050
void setFromCallee() override {
@@ -2234,9 +2236,9 @@ class AsyncCallEmission final : public CallEmission {
22342236
using super = CallEmission;
22352237

22362238
Address contextBuffer;
2237-
Size contextSize;
22382239
Address context;
22392240
llvm::Value *calleeFunction = nullptr;
2241+
llvm::Value *currentResumeFn = nullptr;
22402242
llvm::Value *thickContext = nullptr;
22412243
Optional<AsyncContextLayout> asyncContextLayout;
22422244

@@ -2278,8 +2280,7 @@ class AsyncCallEmission final : public CallEmission {
22782280
CurCallee.getFunctionPointer(), thickContext);
22792281
auto *dynamicContextSize =
22802282
IGF.Builder.CreateZExt(dynamicContextSize32, IGF.IGM.SizeTy);
2281-
std::tie(contextBuffer, contextSize) = emitAllocAsyncContext(
2282-
IGF, layout, dynamicContextSize, getAsyncContextSize(layout));
2283+
contextBuffer = emitAllocAsyncContext(IGF, dynamicContextSize);
22832284
context = layout.emitCastTo(IGF, contextBuffer.getAddress());
22842285
if (layout.canHaveError()) {
22852286
auto fieldLayout = layout.getErrorLayout();
@@ -2292,7 +2293,7 @@ class AsyncCallEmission final : public CallEmission {
22922293
void end() override {
22932294
assert(contextBuffer.isValid());
22942295
assert(context.isValid());
2295-
emitDeallocAsyncContext(IGF, contextBuffer, contextSize);
2296+
emitDeallocAsyncContext(IGF, contextBuffer);
22962297
super::end();
22972298
}
22982299
void setFromCallee() override {
@@ -2337,6 +2338,26 @@ class AsyncCallEmission final : public CallEmission {
23372338
explosion.add(context);
23382339
saveValue(fieldLayout, explosion, isOutlined);
23392340
}
2341+
{ // Return to caller function.
2342+
auto fieldLayout = layout.getResumeParentLayout();
2343+
currentResumeFn = IGF.Builder.CreateIntrinsicCall(
2344+
llvm::Intrinsic::coro_async_resume, {});
2345+
auto fnVal = currentResumeFn;
2346+
// Sign the pointer.
2347+
// TODO: use a distinct schema.
2348+
if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextParent) {
2349+
Address fieldAddr =
2350+
fieldLayout.project(IGF, this->context, /*offsets*/ llvm::None);
2351+
auto authInfo = PointerAuthInfo::emit(
2352+
IGF, schema, fieldAddr.getAddress(), PointerAuthEntity());
2353+
fnVal = emitPointerAuthSign(IGF, fnVal, authInfo);
2354+
}
2355+
fnVal = IGF.Builder.CreateBitCast(fnVal,
2356+
IGF.IGM.TaskContinuationFunctionPtrTy);
2357+
Explosion explosion;
2358+
explosion.add(fnVal);
2359+
saveValue(fieldLayout, explosion, isOutlined);
2360+
}
23402361
{ // caller executor
23412362
Explosion explosion;
23422363
explosion.add(IGF.getAsyncExecutor());
@@ -2396,6 +2417,28 @@ class AsyncCallEmission final : public CallEmission {
23962417
auto address = errorLayout.project(IGF, context, /*offsets*/ llvm::None);
23972418
return address;
23982419
};
2420+
2421+
llvm::CallInst *createCall(const FunctionPointer &fn,
2422+
ArrayRef<llvm::Value *> args) override {
2423+
auto &IGM = IGF.IGM;
2424+
auto &Builder = IGF.Builder;
2425+
// Setup the suspend point.
2426+
SmallVector<llvm::Value *, 8> arguments;
2427+
arguments.push_back(currentResumeFn);
2428+
auto resumeProjFn = IGF.getOrCreateResumePrjFn();
2429+
arguments.push_back(
2430+
Builder.CreateBitOrPointerCast(resumeProjFn, IGM.Int8PtrTy));
2431+
auto dispatchFn = IGF.createAsyncDispatchFn(fn, args);
2432+
arguments.push_back(
2433+
Builder.CreateBitOrPointerCast(dispatchFn, IGM.Int8PtrTy));
2434+
arguments.push_back(
2435+
Builder.CreateBitOrPointerCast(fn.getRawPointer(), IGM.Int8PtrTy));
2436+
for (auto arg: args)
2437+
arguments.push_back(arg);
2438+
auto *id = Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_suspend_async,
2439+
arguments);
2440+
return id;
2441+
}
23992442
};
24002443

24012444
} // end anonymous namespace
@@ -2464,7 +2507,7 @@ llvm::CallInst *CallEmission::emitCallSite() {
24642507
}
24652508

24662509
// TODO: exceptions!
2467-
auto call = IGF.Builder.CreateCall(fn, Args);
2510+
auto call = createCall(fn, Args);
24682511

24692512
// Make coroutines calls opaque to LLVM analysis.
24702513
if (IsCoroutine) {
@@ -3469,6 +3512,28 @@ emitRetconCoroutineEntry(IRGenFunction &IGF, CanSILFunctionType fnType,
34693512
IGF.setCoroutineHandle(hdl);
34703513
}
34713514

3515+
void irgen::emitAsyncFunctionEntry(IRGenFunction &IGF,
3516+
SILFunction *asyncFunction) {
3517+
auto &IGM = IGF.IGM;
3518+
auto size = getAsyncContextLayout(IGM, asyncFunction).getSize();
3519+
auto asyncFuncPointer = IGF.Builder.CreateBitOrPointerCast(
3520+
IGM.getAddrOfAsyncFunctionPointer(asyncFunction), IGM.Int8PtrTy);
3521+
auto *id = IGF.Builder.CreateIntrinsicCall(
3522+
llvm::Intrinsic::coro_id_async,
3523+
{llvm::ConstantInt::get(IGM.Int32Ty, size.getValue()),
3524+
llvm::ConstantInt::get(IGM.Int32Ty, 16),
3525+
llvm::ConstantInt::get(IGM.Int32Ty, 2), asyncFuncPointer});
3526+
// Call 'llvm.coro.begin', just for consistency with the normal pattern.
3527+
// This serves as a handle that we can pass around to other intrinsics.
3528+
auto hdl = IGF.Builder.CreateIntrinsicCall(
3529+
llvm::Intrinsic::coro_begin,
3530+
{id, llvm::ConstantPointerNull::get(IGM.Int8PtrTy)});
3531+
3532+
// Set the coroutine handle; this also flags that is a coroutine so that
3533+
// e.g. dynamic allocas use the right code generation.
3534+
IGF.setCoroutineHandle(hdl);
3535+
}
3536+
34723537
void irgen::emitYieldOnceCoroutineEntry(
34733538
IRGenFunction &IGF, CanSILFunctionType fnType,
34743539
NativeCCEntryPointArgumentEmission &emission) {
@@ -3518,28 +3583,6 @@ void irgen::emitDeallocYieldManyCoroutineBuffer(IRGenFunction &IGF,
35183583
IGF.Builder.CreateLifetimeEnd(buffer, bufferSize);
35193584
}
35203585

3521-
Address irgen::emitTaskAlloc(IRGenFunction &IGF, llvm::Value *size,
3522-
Alignment alignment) {
3523-
auto *call = IGF.Builder.CreateCall(IGF.IGM.getTaskAllocFn(),
3524-
{IGF.getAsyncTask(), size});
3525-
call->setDoesNotThrow();
3526-
call->setCallingConv(IGF.IGM.SwiftCC);
3527-
call->addAttribute(llvm::AttributeList::FunctionIndex,
3528-
llvm::Attribute::ReadNone);
3529-
auto address = Address(call, alignment);
3530-
return address;
3531-
}
3532-
3533-
void irgen::emitTaskDealloc(IRGenFunction &IGF, Address address,
3534-
llvm::Value *size) {
3535-
auto *call = IGF.Builder.CreateCall(
3536-
IGF.IGM.getTaskDeallocFn(), {IGF.getAsyncTask(), address.getAddress()});
3537-
call->setDoesNotThrow();
3538-
call->setCallingConv(IGF.IGM.SwiftCC);
3539-
call->addAttribute(llvm::AttributeList::FunctionIndex,
3540-
llvm::Attribute::ReadNone);
3541-
}
3542-
35433586
void irgen::emitTaskCancel(IRGenFunction &IGF, llvm::Value *task) {
35443587
if (task->getType() != IGF.IGM.SwiftTaskPtrTy) {
35453588
task = IGF.Builder.CreateBitCast(task, IGF.IGM.SwiftTaskPtrTy);
@@ -3591,28 +3634,17 @@ llvm::Value *irgen::emitTaskCreate(
35913634
return result;
35923635
}
35933636

3594-
std::pair<Address, Size> irgen::emitAllocAsyncContext(IRGenFunction &IGF,
3595-
AsyncContextLayout layout,
3596-
llvm::Value *sizeValue,
3597-
Size sizeLowerBound) {
3637+
Address irgen::emitAllocAsyncContext(IRGenFunction &IGF,
3638+
llvm::Value *sizeValue) {
35983639
auto alignment = getAsyncContextAlignment(IGF.IGM);
3599-
auto address = emitTaskAlloc(IGF, sizeValue, alignment);
3600-
IGF.Builder.CreateLifetimeStart(address, sizeLowerBound);
3601-
return {address, sizeLowerBound};
3602-
}
3603-
3604-
std::pair<Address, Size>
3605-
irgen::emitAllocAsyncContext(IRGenFunction &IGF, AsyncContextLayout layout) {
3606-
auto size = getAsyncContextSize(layout);
3607-
auto *sizeValue = llvm::ConstantInt::get(IGF.IGM.SizeTy, size.getValue());
3608-
return emitAllocAsyncContext(IGF, layout, sizeValue, size);
3640+
auto address = IGF.emitTaskAlloc(sizeValue, alignment);
3641+
IGF.Builder.CreateLifetimeStart(address, Size(-1) /*dynamic size*/);
3642+
return address;
36093643
}
36103644

3611-
void irgen::emitDeallocAsyncContext(IRGenFunction &IGF, Address context,
3612-
Size size) {
3613-
auto *sizeValue = llvm::ConstantInt::get(IGF.IGM.SizeTy, size.getValue());
3614-
emitTaskDealloc(IGF, context, sizeValue);
3615-
IGF.Builder.CreateLifetimeEnd(context, size);
3645+
void irgen::emitDeallocAsyncContext(IRGenFunction &IGF, Address context) {
3646+
IGF.emitTaskDealloc(context);
3647+
IGF.Builder.CreateLifetimeEnd(context, Size(-1) /*dynamic size*/);
36163648
}
36173649

36183650
llvm::Value *irgen::emitYield(IRGenFunction &IGF,
@@ -4484,3 +4516,36 @@ FunctionPointer::getExplosionValue(IRGenFunction &IGF,
44844516
FunctionPointer FunctionPointer::getAsFunction(IRGenFunction &IGF) const {
44854517
return FunctionPointer(KindTy::Function, getPointer(IGF), AuthInfo, Sig);
44864518
}
4519+
4520+
void irgen::emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &asyncLayout,
4521+
CanSILFunctionType fnType) {
4522+
auto contextAddr = asyncLayout.emitCastTo(IGF, IGF.getAsyncContext());
4523+
auto returnToCallerLayout = asyncLayout.getResumeParentLayout();
4524+
auto returnToCallerAddr =
4525+
returnToCallerLayout.project(IGF, contextAddr, llvm::None);
4526+
Explosion fn;
4527+
cast<LoadableTypeInfo>(returnToCallerLayout.getType())
4528+
.loadAsCopy(IGF, returnToCallerAddr, fn);
4529+
llvm::Value *fnVal = fn.claimNext();
4530+
4531+
// TODO: use distinct schema
4532+
if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextParent) {
4533+
Address fieldAddr =
4534+
returnToCallerLayout.project(IGF, contextAddr, /*offsets*/ llvm::None);
4535+
auto authInfo = PointerAuthInfo::emit(IGF, schema, fieldAddr.getAddress(),
4536+
PointerAuthEntity());
4537+
fnVal = emitPointerAuthAuth(IGF, fnVal, authInfo);
4538+
}
4539+
4540+
auto sig = emitCastOfFunctionPointer(IGF, fnVal, fnType);
4541+
FunctionPointer fnPtr(FunctionPointer::KindTy::Function, fnVal,
4542+
PointerAuthInfo(), sig);
4543+
4544+
SmallVector<llvm::Value*, 4> Args;
4545+
// Get the current task, executor, and async context.
4546+
Args.push_back(IGF.getAsyncTask());
4547+
Args.push_back(IGF.getAsyncExecutor());
4548+
Args.push_back(IGF.getAsyncContext());
4549+
auto call = IGF.Builder.CreateCall(fnPtr, Args);
4550+
call->setTailCall();
4551+
}

‎lib/IRGen/GenCall.h

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -397,9 +397,6 @@ namespace irgen {
397397
CanSILFunctionType coroutineType,
398398
NativeCCEntryPointArgumentEmission &emission);
399399

400-
Address emitTaskAlloc(IRGenFunction &IGF, llvm::Value *size,
401-
Alignment alignment);
402-
void emitTaskDealloc(IRGenFunction &IGF, Address address, llvm::Value *size);
403400
void emitTaskCancel(IRGenFunction &IGF, llvm::Value *task);
404401

405402
/// Emit a class to swift_task_create[_f] with the given flags, parent task,
@@ -408,19 +405,11 @@ namespace irgen {
408405
IRGenFunction &IGF, llvm::Value *flags, llvm::Value *parentTask,
409406
llvm::Value *taskFunction, llvm::Value *localContextInfo);
410407

411-
/// Allocate task local storage for the specified layout but using the
412-
/// provided dynamic size. Allowing the size to be specified dynamically is
413-
/// necessary for applies of thick functions the sizes of whose async contexts
414-
/// are dependent on the underlying, already partially applied, called
415-
/// function. The provided sizeLowerBound will be used to track the lifetime
416-
/// of the allocation that is known statically.
417-
std::pair<Address, Size> emitAllocAsyncContext(IRGenFunction &IGF,
418-
AsyncContextLayout layout,
419-
llvm::Value *sizeValue,
420-
Size sizeLowerBound);
421-
std::pair<Address, Size> emitAllocAsyncContext(IRGenFunction &IGF,
422-
AsyncContextLayout layout);
423-
void emitDeallocAsyncContext(IRGenFunction &IGF, Address context, Size size);
408+
/// Allocate task local storage for the provided dynamic size.
409+
Address emitAllocAsyncContext(IRGenFunction &IGF, llvm::Value *sizeValue);
410+
void emitDeallocAsyncContext(IRGenFunction &IGF, Address context);
411+
412+
void emitAsyncFunctionEntry(IRGenFunction &IGF, SILFunction *asyncFunc);
424413

425414
/// Yield the given values from the current continuation.
426415
///
@@ -435,6 +424,9 @@ namespace irgen {
435424
Executor = 1,
436425
Context = 2,
437426
};
427+
428+
void emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &layout,
429+
CanSILFunctionType fnType);
438430
} // end namespace irgen
439431
} // end namespace swift
440432

‎lib/IRGen/GenFunc.cpp

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2372,3 +2372,56 @@ void irgen::emitBlockHeader(IRGenFunction &IGF,
23722372
IGF.Builder.CreateStore(descriptorVal,
23732373
IGF.Builder.CreateStructGEP(headerAddr, 4, layout));
23742374
}
2375+
2376+
llvm::Function *IRGenFunction::getOrCreateResumePrjFn() {
2377+
auto name = "__swift_async_resume_project_context";
2378+
return cast<llvm::Function>(IGM.getOrCreateHelperFunction(
2379+
name, IGM.Int8PtrTy, {IGM.Int8PtrTy},
2380+
[&](IRGenFunction &IGF) {
2381+
auto it = IGF.CurFn->arg_begin();
2382+
auto &Builder = IGF.Builder;
2383+
auto addr = Builder.CreateBitOrPointerCast(&(*it), IGF.IGM.Int8PtrPtrTy);
2384+
Address callerContextAddr(addr, IGF.IGM.getPointerAlignment());
2385+
auto callerContext = Builder.CreateLoad(callerContextAddr);
2386+
Builder.CreateRet(callerContext);
2387+
},
2388+
false /*isNoInline*/));
2389+
}
2390+
2391+
llvm::Function *
2392+
IRGenFunction::createAsyncDispatchFn(const FunctionPointer &fnPtr,
2393+
ArrayRef<llvm::Value *> args) {
2394+
SmallVector<llvm::Type*, 8> argTys;
2395+
argTys.push_back(IGM.Int8PtrTy); // Function pointer to be called.
2396+
for (auto arg : args) {
2397+
auto *ty = arg->getType();
2398+
argTys.push_back(ty);
2399+
}
2400+
auto calleeFnPtrType = fnPtr.getRawPointer()->getType();
2401+
auto *dispatchFnTy =
2402+
llvm::FunctionType::get(IGM.VoidTy, argTys, false /*vaargs*/);
2403+
llvm::SmallString<40> name;
2404+
llvm::raw_svector_ostream(name) << "__swift_suspend_dispatch_" << args.size();
2405+
llvm::Function *dispatch =
2406+
llvm::Function::Create(dispatchFnTy, llvm::Function::InternalLinkage,
2407+
llvm::StringRef(name), &IGM.Module);
2408+
dispatch->setCallingConv(IGM.DefaultCC);
2409+
dispatch->setDoesNotThrow();
2410+
IRGenFunction dispatchIGF(IGM, dispatch);
2411+
if (IGM.DebugInfo)
2412+
IGM.DebugInfo->emitArtificialFunction(dispatchIGF, dispatch);
2413+
auto &Builder = dispatchIGF.Builder;
2414+
auto it = dispatchIGF.CurFn->arg_begin(), end = dispatchIGF.CurFn->arg_end();
2415+
llvm::Value *ptrArg = &*(it++);
2416+
SmallVector<llvm::Value *, 8> callArgs;
2417+
for (; it != end; ++it) {
2418+
callArgs.push_back(&*it);
2419+
}
2420+
ptrArg = Builder.CreateBitOrPointerCast(ptrArg, calleeFnPtrType);
2421+
auto callee = FunctionPointer(fnPtr.getKind(), ptrArg, fnPtr.getAuthInfo(),
2422+
fnPtr.getSignature());
2423+
auto call = Builder.CreateCall(callee, callArgs);
2424+
call->setTailCall();
2425+
Builder.CreateRetVoid();
2426+
return dispatch;
2427+
}

0 commit comments

Comments
 (0)