diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp --- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -765,8 +765,8 @@ // context header. case coro::ABI::Async: { auto *ActiveAsyncSuspend = cast(ActiveSuspend); - auto *CalleeContext = - NewF->getArg(ActiveAsyncSuspend->getStorageArgumentIndex()); + auto ContextIdx = ActiveAsyncSuspend->getStorageArgumentIndex() & 0xff; + auto *CalleeContext = NewF->getArg(ContextIdx); auto *FramePtrTy = Shape.FrameTy->getPointerTo(); auto *ProjectionFunc = ActiveAsyncSuspend->getAsyncContextProjectionFunction(); @@ -827,6 +827,13 @@ Attrs = Attrs.addParamAttributes(Context, ParamIndex, ParamAttrs); } +static void addSwiftSelfAttrs(AttributeList &Attrs, LLVMContext &Context, + unsigned ParamIndex) { + AttrBuilder ParamAttrs; + ParamAttrs.addAttribute(Attribute::SwiftSelf); + Attrs = Attrs.addParamAttributes(Context, ParamIndex, ParamAttrs); +} + /// Clone the body of the original function into a resume function of /// some sort. void CoroCloner::create() { @@ -949,8 +956,16 @@ auto *ActiveAsyncSuspend = cast(ActiveSuspend); if (OrigF.hasParamAttribute(Shape.AsyncLowering.ContextArgNo, Attribute::SwiftAsync)) { - auto ContextArgIndex = ActiveAsyncSuspend->getStorageArgumentIndex(); + uint32_t ArgAttributeIndices = + ActiveAsyncSuspend->getStorageArgumentIndex(); + auto ContextArgIndex = ArgAttributeIndices & 0xff; addAsyncContextAttrs(NewAttrs, Context, ContextArgIndex); + + // `swiftasync` must preceed `swiftself` so 0 is not a valid index for + // `swiftself`. + auto SwiftSelfIndex = ArgAttributeIndices >> 8; + if (SwiftSelfIndex) + addSwiftSelfAttrs(NewAttrs, Context, SwiftSelfIndex); } break; } diff --git a/llvm/test/Transforms/Coroutines/coro-async.ll b/llvm/test/Transforms/Coroutines/coro-async.ll --- a/llvm/test/Transforms/Coroutines/coro-async.ll +++ b/llvm/test/Transforms/Coroutines/coro-async.ll @@ -410,7 +410,7 @@ i32 64 ; Initial async context size without space for frame }> -define swiftcc void @polymorphic_suspend_return(i8* %async.ctxt, %async.task* %task, %async.actor* %actor) { +define swiftcc void @polymorphic_suspend_return(i8* swiftasync %async.ctxt, %async.task* %task, %async.actor* %actor) { entry: %tmp = alloca { i64, i64 }, align 8 %proj.1 = getelementptr inbounds { i64, i64 }, { i64, i64 }* %tmp, i64 0, i32 0 @@ -445,7 +445,7 @@ %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*} (i32, i8*, i8*, ...) - @llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(i32 0, + @llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(i32 256, ;; swiftasync at 0 and swiftself at 1 in resume function i8* %resume.func_ptr, i8* %resume_proj_fun, void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply, @@ -464,8 +464,8 @@ unreachable } -; CHECK-LABEL: define swiftcc void @polymorphic_suspend_return(i8* %async.ctxt, %async.task* %task, %async.actor* %actor) -; CHECK-LABEL: define internal swiftcc void @polymorphic_suspend_return.resume.0(i8* {{.*}}%0, i8* {{.*}}%1, i8* {{.*}}%2, i8* {{.*}}%3) +; CHECK-LABEL: define swiftcc void @polymorphic_suspend_return(i8* swiftasync %async.ctxt, %async.task* %task, %async.actor* %actor) +; CHECK-LABEL: define internal swiftcc void @polymorphic_suspend_return.resume.0(i8* {{.*}}swiftasync{{.*}} %0, i8* {{.*}}swiftself{{.*}} %1, i8* {{.*}}%2, i8* {{.*}}%3) ; CHECK: bitcast i8* %3 to %async.task* ; CHECK: }