Skip to content

[coro async] Allow a coro.suspend.async to specify which argument is the context argument #2553

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Mar 3, 2021
Merged

[coro async] Allow a coro.suspend.async to specify which argument is the context argument #2553

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0a21dec
Allow a coro.suspend.async to specify which argument is the context a…
aschwaighofer Feb 23, 2021
7ecb463
Transfer swiftasync attribute according to suspend.async spec
aschwaighofer Feb 23, 2021
cc07366
Fix coro-split-swift-async.ll test
aschwaighofer Feb 24, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion llvm/include/llvm/IR/Intrinsics.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1211,7 +1211,8 @@ def int_coro_async_resume : Intrinsic<[llvm_ptr_ty],
[],
[]>;
def int_coro_suspend_async
: Intrinsic<[llvm_any_ty], [llvm_ptr_ty, llvm_ptr_ty, llvm_vararg_ty], []>;
: Intrinsic<[llvm_any_ty],
[llvm_i32_ty, llvm_ptr_ty, llvm_ptr_ty, llvm_vararg_ty], []>;
def int_coro_prepare_async : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty],
[IntrNoMem]>;
def int_coro_begin : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty],
Expand Down
14 changes: 12 additions & 2 deletions llvm/lib/Transforms/Coroutines/CoroInstr.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -511,11 +511,21 @@ inline CoroSaveInst *AnyCoroSuspendInst::getCoroSave() const {

/// This represents the llvm.coro.suspend.async instruction.
class LLVM_LIBRARY_VISIBILITY CoroSuspendAsyncInst : public AnyCoroSuspendInst {
enum { ResumeFunctionArg, AsyncContextProjectionArg, MustTailCallFuncArg };

public:
enum {
StorageArgNoArg,
ResumeFunctionArg,
AsyncContextProjectionArg,
MustTailCallFuncArg
};

void checkWellFormed() const;

unsigned getStorageArgumentIndex() const {
auto *Arg = cast<ConstantInt>(getArgOperand(StorageArgNoArg));
return Arg->getZExtValue();
}

Function *getAsyncContextProjectionFunction() const {
return cast<Function>(
getArgOperand(AsyncContextProjectionArg)->stripPointerCasts());
Expand Down
24 changes: 15 additions & 9 deletions llvm/lib/Transforms/Coroutines/CoroSplit.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -759,10 +759,12 @@ Value *CoroCloner::deriveNewFramePointer() {
// with the active suspend. The frame is located as a tail to the async
// context header.
case coro::ABI::Async: {
auto *CalleeContext = NewF->getArg(Shape.AsyncLowering.ContextArgNo);
auto *ActiveAsyncSuspend = cast<CoroSuspendAsyncInst>(ActiveSuspend);
auto *CalleeContext =
NewF->getArg(ActiveAsyncSuspend->getStorageArgumentIndex());
auto *FramePtrTy = Shape.FrameTy->getPointerTo();
auto *ProjectionFunc = cast<CoroSuspendAsyncInst>(ActiveSuspend)
->getAsyncContextProjectionFunction();
auto *ProjectionFunc =
ActiveAsyncSuspend->getAsyncContextProjectionFunction();
auto DbgLoc =
cast<CoroSuspendAsyncInst>(VMap[ActiveSuspend])->getDebugLoc();
// Calling i8* (i8*)
Expand Down Expand Up @@ -880,12 +882,14 @@ void CoroCloner::create() {
addFramePointerAttrs(NewAttrs, Context, 0,
Shape.FrameSize, Shape.FrameAlign);
break;
case coro::ABI::Async:
if (OrigF.hasParamAttribute(Shape.AsyncLowering.ContextArgNo,
Attribute::SwiftAsync)) {
addAsyncContextAttrs(NewAttrs, Context, Shape.AsyncLowering.ContextArgNo);
case coro::ABI::Async: {
auto *ActiveAsyncSuspend = cast<CoroSuspendAsyncInst>(ActiveSuspend);
auto ContextArgIndex = ActiveAsyncSuspend->getStorageArgumentIndex();
if (OrigF.hasParamAttribute(ContextArgIndex, Attribute::SwiftAsync)) {
addAsyncContextAttrs(NewAttrs, Context, ContextArgIndex);
}
break;
}
case coro::ABI::Retcon:
case coro::ABI::RetconOnce:
// If we have a continuation prototype, just use its attributes,
Expand Down Expand Up @@ -1486,7 +1490,8 @@ static void replaceAsyncResumeFunction(CoroSuspendAsyncInst *Suspend,
auto *Val = Builder.CreateBitOrPointerCast(Continuation, Int8PtrTy);
ResumeIntrinsic->replaceAllUsesWith(Val);
ResumeIntrinsic->eraseFromParent();
Suspend->setOperand(0, UndefValue::get(Int8PtrTy));
Suspend->setOperand(CoroSuspendAsyncInst::ResumeFunctionArg,
UndefValue::get(Int8PtrTy));
}

/// Coerce the arguments in \p FnArgs according to \p FnTy in \p CallArgs.
Expand Down Expand Up @@ -1581,7 +1586,8 @@ static void splitAsyncCoroutine(Function &F, coro::Shape &Shape,
// Insert the call to the tail call function and inline it.
auto *Fn = Suspend->getMustTailCallFunction();
SmallVector<Value *, 8> Args(Suspend->args());
auto FnArgs = ArrayRef<Value *>(Args).drop_front(3);
auto FnArgs = ArrayRef<Value *>(Args).drop_front(
CoroSuspendAsyncInst::MustTailCallFuncArg + 1);
auto *TailCall =
coro::createMustTailCall(Suspend->getDebugLoc(), Fn, FnArgs, Builder);
Builder.CreateRetVoid();
Expand Down
4 changes: 2 additions & 2 deletions llvm/test/Transforms/Coroutines/coro-async-unreachable.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ entry:
store i8* %async.ctxt, i8** %callee_context.caller_context.addr
%resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
%callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
%res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
%res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
i8* %resume.func_ptr,
i8* %resume_proj_fun,
void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
Expand All @@ -97,7 +97,7 @@ entry:
declare i8* @llvm.coro.prepare.async(i8*)
declare token @llvm.coro.id.async(i32, i32, i32, i8*)
declare i8* @llvm.coro.begin(token, i8*)
declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i8*, i8*, ...)
declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i32, i8*, i8*, ...)
declare i8* @llvm.coro.async.context.alloc(i8*, i8*)
declare void @llvm.coro.async.context.dealloc(i8*)
declare swiftcc void @asyncReturn(i8*, %async.task*, %async.actor*)
Expand Down
29 changes: 15 additions & 14 deletions llvm/test/Transforms/Coroutines/coro-async.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ entry:
store i8* %async.ctxt, i8** %callee_context.caller_context.addr
%resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
%callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
%res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
%res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
i8* %resume.func_ptr,
i8* %resume_proj_fun,
void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
Expand Down Expand Up @@ -190,7 +190,7 @@ entry:
store i8* %async.ctxt, i8** %callee_context.caller_context.addr
%resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
%callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
%res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
%res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 2,
i8* %resume.func_ptr,
i8* %resume_proj_fun,
void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
Expand All @@ -208,7 +208,7 @@ entry:
store i8* %async.ctxt, i8** %callee_context.caller_context.addr.1
%resume_proj_fun.2 = bitcast i8*(i8*)* @resume_context_projection to i8*
%callee.2 = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
%res.2 = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
%res.2 = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
i8* %resume.func_ptr.1,
i8* %resume_proj_fun.2,
void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
Expand Down Expand Up @@ -244,8 +244,8 @@ entry:
; CHECK: tail call swiftcc void @asyncSuspend(i8* [[CALLEE_CTXT_RELOAD]]
; CHECK: ret void

; CHECK-LABEL: define internal swiftcc void @my_async_function2.resume.1(i8* nocapture readnone %0, i8* %1, i8* nocapture readonly %2) {
; CHECK: bitcast i8* %2 to i8**
; CHECK-LABEL: define internal swiftcc void @my_async_function2.resume.1(i8* nocapture readonly %0, i8* %1, i8* nocapture readnone %2) {
; CHECK: bitcast i8* %0 to i8**
; CHECK: [[ACTOR_ARG:%.*]] = bitcast i8* %1
; CHECK: tail call swiftcc void @asyncReturn({{.*}}[[ACTOR_ARG]])
; CHECK: ret void
Expand Down Expand Up @@ -304,7 +304,7 @@ entry:
store i8* %async.ctxt, i8** %callee_context.caller_context.addr
%resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
%callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
%res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
%res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
i8* %resume.func_ptr,
i8* %resume_proj_fun,
void (i8*, i8*, %async.task*, %async.actor*)* @dont_crash_on_cf_dispatch,
Expand Down Expand Up @@ -350,7 +350,7 @@ entry:
store i8* %async.ctxt, i8** %callee_context.caller_context.addr
%resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
%callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
%res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
%res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
i8* %resume.func_ptr,
i8* %resume_proj_fun,
void (i8*, i8*, %async.task*, %async.actor*)* @dont_crash_on_cf_dispatch,
Expand Down Expand Up @@ -423,11 +423,12 @@ entry:
store i8* %async.ctxt, i8** %callee_context.caller_context.addr
%resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
%callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
%res = call {i8*, i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(
i8* %resume.func_ptr,
i8* %resume_proj_fun,
void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
i8* %callee, i8* %callee_context, %async.task* %task, %async.actor *%actor)
%res = call {i8*, i8*, i8*, i8*} (i32, i8*, i8*, ...)
@llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(i32 0,
i8* %resume.func_ptr,
i8* %resume_proj_fun,
void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
i8* %callee, i8* %callee_context, %async.task* %task, %async.actor *%actor)

call void @llvm.coro.async.context.dealloc(i8* %callee_context)
%continuation_task_arg = extractvalue {i8*, i8*, i8*, i8*} %res, 3
Expand Down Expand Up @@ -501,13 +502,13 @@ entry:
; CHECK: store i64* null, i64** [[ALLOCA]]
; CHECK: call void @do_with_swifterror(i64** {{.*}}swifterror{{.*}} [[ALLOCA]])

declare { i8*, i8*, i8*, i8* } @llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(i8*, i8*, ...)
declare { i8*, i8*, i8*, i8* } @llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(i32, i8*, i8*, ...)
declare i8* @llvm.coro.prepare.async(i8*)
declare token @llvm.coro.id.async(i32, i32, i32, i8*)
declare i8* @llvm.coro.begin(token, i8*)
declare i1 @llvm.coro.end.async(i8*, i1, ...)
declare i1 @llvm.coro.end(i8*, i1)
declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i8*, i8*, ...)
declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i32, i8*, i8*, ...)
declare i8* @llvm.coro.async.context.alloc(i8*, i8*)
declare void @llvm.coro.async.context.dealloc(i8*)
declare swiftcc void @asyncReturn(i8*, %async.task*, %async.actor*)
Expand Down
4 changes: 2 additions & 2 deletions llvm/test/Transforms/Coroutines/coro-split-swift-async.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ entry:
%23 = bitcast %T12_Concurrency9MainActorC5_ImplC* %15 to %swift.executor*, !dbg !49
%24 = load %swift.executor*, %swift.executor** %4, align 8, !dbg !49
%25 = load %swift.context*, %swift.context** %5, align 8, !dbg !49
%26 = call { i8*, i8*, i8* } (i8*, i8*, ...) @llvm.coro.suspend.async(i8* %16, i8* bitcast (i8* (i8*)* @__swift_async_resume_get_context to i8*), i8* bitcast (void (i8*, %swift.executor*, %swift.task*, %swift.executor*, %swift.context*)* @__swift_suspend_point to i8*), i8* %16, %swift.executor* %23, %swift.task* %17, %swift.executor* %24, %swift.context* %25), !dbg !49
%26 = call { i8*, i8*, i8* } (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 2, i8* %16, i8* bitcast (i8* (i8*)* @__swift_async_resume_get_context to i8*), i8* bitcast (void (i8*, %swift.executor*, %swift.task*, %swift.executor*, %swift.context*)* @__swift_suspend_point to i8*), i8* %16, %swift.executor* %23, %swift.task* %17, %swift.executor* %24, %swift.context* %25), !dbg !49
%27 = extractvalue { i8*, i8*, i8* } %26, 0, !dbg !49
%28 = bitcast i8* %27 to %swift.task*, !dbg !49
store %swift.task* %28, %swift.task** %3, align 8, !dbg !49
Expand Down Expand Up @@ -136,7 +136,7 @@ entry:
declare extern_weak swiftcc void @swift_task_switch(%swift.task* %0, %swift.executor* %1, %swift.executor* %2) #1

; Function Attrs: nounwind
declare { i8*, i8*, i8* } @llvm.coro.suspend.async(i8*, i8*, ...) #1
declare { i8*, i8*, i8* } @llvm.coro.suspend.async(i32, i8*, i8*, ...) #1

; Function Attrs: nounwind
declare void @swift_release(%swift.refcounted*) #1
Expand Down