Skip to content

Commit 74419d0

Browse files
aschwaighoferaschwaighofer
committed
Lower asynchronous functions to LLVM async coroutine intrinsics
rdar://70097093
1 parent ed5a759 commit 74419d0

File tree

7 files changed

+195
-9
lines changed

7 files changed

+195
-9
lines changed

‎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: 106 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2043,6 +2043,12 @@ class SyncCallEmission final : public CallEmission {
20432043
return origConv.getSILArgumentType(
20442044
index, IGF.IGM.getMaximalTypeExpansionContext());
20452045
}
2046+
2047+
llvm::CallInst *createCall(const FunctionPointer &fn,
2048+
ArrayRef<llvm::Value *> args) override {
2049+
return IGF.Builder.CreateCall(fn, Args);
2050+
}
2051+
20462052
void begin() override { super::begin(); }
20472053
void end() override { super::end(); }
20482054
void setFromCallee() override {
@@ -2237,6 +2243,7 @@ class AsyncCallEmission final : public CallEmission {
22372243
Size contextSize;
22382244
Address context;
22392245
llvm::Value *calleeFunction = nullptr;
2246+
llvm::Value *currentResumeFn = nullptr;
22402247
llvm::Value *thickContext = nullptr;
22412248
Optional<AsyncContextLayout> asyncContextLayout;
22422249

@@ -2337,6 +2344,26 @@ class AsyncCallEmission final : public CallEmission {
23372344
explosion.add(context);
23382345
saveValue(fieldLayout, explosion, isOutlined);
23392346
}
2347+
{ // Return to caller function.
2348+
auto fieldLayout = layout.getResumeParentLayout();
2349+
currentResumeFn = IGF.Builder.CreateIntrinsicCall(
2350+
llvm::Intrinsic::coro_async_resume, {});
2351+
auto fnVal = currentResumeFn;
2352+
// Sign the pointer.
2353+
// TODO: use a distinct schema.
2354+
if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextParent) {
2355+
Address fieldAddr =
2356+
fieldLayout.project(IGF, this->context, /*offsets*/ llvm::None);
2357+
auto authInfo = PointerAuthInfo::emit(
2358+
IGF, schema, fieldAddr.getAddress(), PointerAuthEntity());
2359+
fnVal = emitPointerAuthSign(IGF, fnVal, authInfo);
2360+
}
2361+
fnVal = IGF.Builder.CreateBitCast(fnVal,
2362+
IGF.IGM.TaskContinuationFunctionPtrTy);
2363+
Explosion explosion;
2364+
explosion.add(fnVal);
2365+
saveValue(fieldLayout, explosion, isOutlined);
2366+
}
23402367
{ // caller executor
23412368
Explosion explosion;
23422369
explosion.add(IGF.getAsyncExecutor());
@@ -2396,6 +2423,28 @@ class AsyncCallEmission final : public CallEmission {
23962423
auto address = errorLayout.project(IGF, context, /*offsets*/ llvm::None);
23972424
return address;
23982425
};
2426+
2427+
llvm::CallInst *createCall(const FunctionPointer &fn,
2428+
ArrayRef<llvm::Value *> args) override {
2429+
auto &IGM = IGF.IGM;
2430+
auto &Builder = IGF.Builder;
2431+
// Setup the suspend point.
2432+
SmallVector<llvm::Value *, 8> arguments;
2433+
arguments.push_back(currentResumeFn);
2434+
auto resumeProjFn = IGF.getOrCreateResumePrjFn();
2435+
arguments.push_back(
2436+
Builder.CreateBitOrPointerCast(resumeProjFn, IGM.Int8PtrTy));
2437+
auto dispatchFn = IGF.createAsyncDispatchFn(fn, args);
2438+
arguments.push_back(
2439+
Builder.CreateBitOrPointerCast(dispatchFn, IGM.Int8PtrTy));
2440+
arguments.push_back(
2441+
Builder.CreateBitOrPointerCast(fn.getRawPointer(), IGM.Int8PtrTy));
2442+
for (auto arg: args)
2443+
arguments.push_back(arg);
2444+
auto *id = Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_suspend_async,
2445+
arguments);
2446+
return id;
2447+
}
23992448
};
24002449

24012450
} // end anonymous namespace
@@ -2464,7 +2513,7 @@ llvm::CallInst *CallEmission::emitCallSite() {
24642513
}
24652514

24662515
// TODO: exceptions!
2467-
auto call = IGF.Builder.CreateCall(fn, Args);
2516+
auto call = createCall(fn, Args);
24682517

24692518
// Make coroutines calls opaque to LLVM analysis.
24702519
if (IsCoroutine) {
@@ -3469,6 +3518,29 @@ emitRetconCoroutineEntry(IRGenFunction &IGF, CanSILFunctionType fnType,
34693518
IGF.setCoroutineHandle(hdl);
34703519
}
34713520

3521+
void irgen::emitAsyncFunctionEntry(IRGenFunction &IGF,
3522+
SILFunction *asyncFunction) {
3523+
auto &IGM = IGF.IGM;
3524+
auto size = getAsyncContextLayout(IGM, asyncFunction).getSize();
3525+
auto asyncFuncPointer = IGF.Builder.CreateBitOrPointerCast(
3526+
IGM.getAddrOfAsyncFunctionPointer(asyncFunction, NotForDefinition),
3527+
IGM.Int8PtrTy);
3528+
auto *id = IGF.Builder.CreateIntrinsicCall(
3529+
llvm::Intrinsic::coro_id_async,
3530+
{llvm::ConstantInt::get(IGM.Int32Ty, size.getValue()),
3531+
llvm::ConstantInt::get(IGM.Int32Ty, 16),
3532+
llvm::ConstantInt::get(IGM.Int32Ty, 2), asyncFuncPointer});
3533+
// Call 'llvm.coro.begin', just for consistency with the normal pattern.
3534+
// This serves as a handle that we can pass around to other intrinsics.
3535+
auto hdl = IGF.Builder.CreateIntrinsicCall(
3536+
llvm::Intrinsic::coro_begin,
3537+
{id, llvm::ConstantPointerNull::get(IGM.Int8PtrTy)});
3538+
3539+
// Set the coroutine handle; this also flags that is a coroutine so that
3540+
// e.g. dynamic allocas use the right code generation.
3541+
IGF.setCoroutineHandle(hdl);
3542+
}
3543+
34723544
void irgen::emitYieldOnceCoroutineEntry(
34733545
IRGenFunction &IGF, CanSILFunctionType fnType,
34743546
NativeCCEntryPointArgumentEmission &emission) {
@@ -4484,3 +4556,36 @@ FunctionPointer::getExplosionValue(IRGenFunction &IGF,
44844556
FunctionPointer FunctionPointer::getAsFunction(IRGenFunction &IGF) const {
44854557
return FunctionPointer(KindTy::Function, getPointer(IGF), AuthInfo, Sig);
44864558
}
4559+
4560+
void irgen::emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &asyncLayout,
4561+
CanSILFunctionType fnType) {
4562+
auto contextAddr = asyncLayout.emitCastTo(IGF, IGF.getAsyncContext());
4563+
auto returnToCallerLayout = asyncLayout.getResumeParentLayout();
4564+
auto returnToCallerAddr =
4565+
returnToCallerLayout.project(IGF, contextAddr, llvm::None);
4566+
Explosion fn;
4567+
cast<LoadableTypeInfo>(returnToCallerLayout.getType())
4568+
.loadAsCopy(IGF, returnToCallerAddr, fn);
4569+
llvm::Value *fnVal = fn.claimNext();
4570+
4571+
// TODO: use distinct schema
4572+
if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextParent) {
4573+
Address fieldAddr =
4574+
returnToCallerLayout.project(IGF, contextAddr, /*offsets*/ llvm::None);
4575+
auto authInfo = PointerAuthInfo::emit(IGF, schema, fieldAddr.getAddress(),
4576+
PointerAuthEntity());
4577+
fnVal = emitPointerAuthAuth(IGF, fnVal, authInfo);
4578+
}
4579+
4580+
auto sig = emitCastOfFunctionPointer(IGF, fnVal, fnType);
4581+
FunctionPointer fnPtr(FunctionPointer::KindTy::Function, fnVal,
4582+
PointerAuthInfo(), sig);
4583+
4584+
SmallVector<llvm::Value*, 4> Args;
4585+
// Get the current task, executor, and async context.
4586+
Args.push_back(IGF.getAsyncTask());
4587+
Args.push_back(IGF.getAsyncExecutor());
4588+
Args.push_back(IGF.getAsyncContext());
4589+
auto call = IGF.Builder.CreateCall(fnPtr, Args);
4590+
call->setTailCall();
4591+
}

‎lib/IRGen/GenCall.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -422,6 +422,8 @@ namespace irgen {
422422
AsyncContextLayout layout);
423423
void emitDeallocAsyncContext(IRGenFunction &IGF, Address context, Size size);
424424

425+
void emitAsyncFunctionEntry(IRGenFunction &IGF, SILFunction *asyncFunc);
426+
425427
/// Yield the given values from the current continuation.
426428
///
427429
/// \return an i1 indicating whether the caller wants to unwind this
@@ -435,6 +437,9 @@ namespace irgen {
435437
Executor = 1,
436438
Context = 2,
437439
};
440+
441+
void emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &layout,
442+
CanSILFunctionType fnType);
438443
} // end namespace irgen
439444
} // end namespace swift
440445

‎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+
}

‎lib/IRGen/IRGenFunction.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,10 @@ class IRGenFunction {
132132
virtual llvm::Value *getAsyncExecutor();
133133
virtual llvm::Value *getAsyncContext();
134134

135+
llvm::Function *getOrCreateResumePrjFn();
136+
llvm::Function *createAsyncDispatchFn(const FunctionPointer &fnPtr,
137+
ArrayRef<llvm::Value *> args);
138+
135139
private:
136140
void emitPrologue();
137141
void emitEpilogue();

‎lib/IRGen/IRGenModule.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -591,7 +591,7 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
591591
{RelativeAddressTy, Int32Ty});
592592

593593
AsyncFunctionPointerTy = createStructType(*this, "swift.async_func_pointer",
594-
{RelativeAddressTy, Int32Ty});
594+
{RelativeAddressTy, Int32Ty}, true);
595595
SwiftContextTy = createStructType(*this, "swift.context", {});
596596
SwiftTaskTy = createStructType(*this, "swift.task", {});
597597
SwiftExecutorTy = createStructType(*this, "swift.executor", {});

‎lib/IRGen/IRGenSIL.cpp

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1751,6 +1751,10 @@ static void emitEntryPointArgumentsNativeCC(IRGenSILFunction &IGF,
17511751
break;
17521752
}
17531753

1754+
if (funcTy->isAsync()) {
1755+
emitAsyncFunctionEntry(IGF, IGF.CurSILFn);
1756+
}
1757+
17541758
SILFunctionConventions conv(funcTy, IGF.getSILModule());
17551759

17561760
// The 'self' argument might be in the context position, which is
@@ -3182,7 +3186,7 @@ void IRGenSILFunction::visitUnreachableInst(swift::UnreachableInst *i) {
31823186
Builder.CreateUnreachable();
31833187
}
31843188

3185-
static void emitCoroutineExit(IRGenSILFunction &IGF) {
3189+
static void emitCoroutineOrAsyncExit(IRGenSILFunction &IGF) {
31863190
// The LLVM coroutine representation demands that there be a
31873191
// unique call to llvm.coro.end.
31883192

@@ -3208,12 +3212,13 @@ static void emitCoroutineExit(IRGenSILFunction &IGF) {
32083212

32093213
static void emitReturnInst(IRGenSILFunction &IGF,
32103214
SILType resultTy,
3211-
Explosion &result) {
3215+
Explosion &result,
3216+
CanSILFunctionType fnType) {
32123217
// If we're generating a coroutine, just call coro.end.
3213-
if (IGF.isCoroutine()) {
3218+
if (IGF.isCoroutine() && !IGF.isAsync()) {
32143219
assert(result.empty() &&
32153220
"coroutines do not currently support non-void returns");
3216-
emitCoroutineExit(IGF);
3221+
emitCoroutineOrAsyncExit(IGF);
32173222
return;
32183223
}
32193224

@@ -3249,7 +3254,8 @@ static void emitReturnInst(IRGenSILFunction &IGF,
32493254
cast<LoadableTypeInfo>(fieldLayout.getType())
32503255
.initialize(IGF, result, fieldAddr, /*isOutlined*/ false);
32513256
}
3252-
IGF.Builder.CreateRetVoid();
3257+
emitAsyncReturn(IGF, layout, fnType);
3258+
emitCoroutineOrAsyncExit(IGF);
32533259
} else {
32543260
auto funcLang = IGF.CurSILFn->getLoweredFunctionType()->getLanguage();
32553261
auto swiftCCReturn = funcLang == SILFunctionLanguage::Swift;
@@ -3278,7 +3284,8 @@ void IRGenSILFunction::visitReturnInst(swift::ReturnInst *i) {
32783284
result = std::move(temp);
32793285
}
32803286

3281-
emitReturnInst(*this, i->getOperand()->getType(), result);
3287+
emitReturnInst(*this, i->getOperand()->getType(), result,
3288+
i->getFunction()->getLoweredFunctionType());
32823289
}
32833290

32843291
void IRGenSILFunction::visitThrowInst(swift::ThrowInst *i) {
@@ -3287,6 +3294,14 @@ void IRGenSILFunction::visitThrowInst(swift::ThrowInst *i) {
32873294

32883295
Builder.CreateStore(exn, getCallerErrorResultSlot());
32893296

3297+
// Async functions just return to the continuation.
3298+
if (isAsync()) {
3299+
auto layout = getAsyncContextLayout(*this);
3300+
emitAsyncReturn(*this, layout, i->getFunction()->getLoweredFunctionType());
3301+
emitCoroutineOrAsyncExit(*this);
3302+
return;
3303+
}
3304+
32903305
// Create a normal return, but leaving the return value undefined.
32913306
auto fnTy = CurFn->getType()->getPointerElementType();
32923307
auto retTy = cast<llvm::FunctionType>(fnTy)->getReturnType();
@@ -3300,7 +3315,7 @@ void IRGenSILFunction::visitThrowInst(swift::ThrowInst *i) {
33003315
void IRGenSILFunction::visitUnwindInst(swift::UnwindInst *i) {
33013316
// Just call coro.end; there's no need to distinguish 'unwind'
33023317
// and 'return' at the LLVM level.
3303-
emitCoroutineExit(*this);
3318+
emitCoroutineOrAsyncExit(*this);
33043319
}
33053320

33063321
void IRGenSILFunction::visitYieldInst(swift::YieldInst *i) {

0 commit comments

Comments
 (0)