diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td --- a/llvm/include/llvm/IR/Intrinsics.td +++ b/llvm/include/llvm/IR/Intrinsics.td @@ -1219,7 +1219,8 @@ []>; def int_coro_async_size_replace : Intrinsic<[], [llvm_ptr_ty, 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], diff --git a/llvm/lib/Transforms/Coroutines/CoroInstr.h b/llvm/lib/Transforms/Coroutines/CoroInstr.h --- a/llvm/lib/Transforms/Coroutines/CoroInstr.h +++ b/llvm/lib/Transforms/Coroutines/CoroInstr.h @@ -523,11 +523,21 @@ /// 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(getArgOperand(StorageArgNoArg)); + return Arg->getZExtValue(); + } + Function *getAsyncContextProjectionFunction() const { return cast( getArgOperand(AsyncContextProjectionArg)->stripPointerCasts()); 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 @@ -759,10 +759,12 @@ // 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(ActiveSuspend); + auto *CalleeContext = + NewF->getArg(ActiveAsyncSuspend->getStorageArgumentIndex()); auto *FramePtrTy = Shape.FrameTy->getPointerTo(); - auto *ProjectionFunc = cast(ActiveSuspend) - ->getAsyncContextProjectionFunction(); + auto *ProjectionFunc = + ActiveAsyncSuspend->getAsyncContextProjectionFunction(); auto DbgLoc = cast(VMap[ActiveSuspend])->getDebugLoc(); // Calling i8* (i8*) @@ -1468,7 +1470,8 @@ 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. @@ -1563,7 +1566,8 @@ // Insert the call to the tail call function and inline it. auto *Fn = Suspend->getMustTailCallFunction(); SmallVector Args(Suspend->args()); - auto FnArgs = ArrayRef(Args).drop_front(3); + auto FnArgs = ArrayRef(Args).drop_front( + CoroSuspendAsyncInst::MustTailCallFuncArg + 1); auto *TailCall = coro::createMustTailCall(Suspend->getDebugLoc(), Fn, FnArgs, Builder); Builder.CreateRetVoid(); diff --git a/llvm/test/Transforms/Coroutines/coro-async-unreachable.ll b/llvm/test/Transforms/Coroutines/coro-async-unreachable.ll --- a/llvm/test/Transforms/Coroutines/coro-async-unreachable.ll +++ b/llvm/test/Transforms/Coroutines/coro-async-unreachable.ll @@ -71,7 +71,7 @@ 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, @@ -97,7 +97,7 @@ 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*) 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 @@ -89,7 +89,7 @@ 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, @@ -204,7 +204,7 @@ 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, @@ -222,7 +222,7 @@ 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, @@ -258,8 +258,8 @@ ; 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 @@ -318,7 +318,7 @@ 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, @@ -364,7 +364,7 @@ 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, @@ -437,11 +437,12 @@ 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 @@ -515,13 +516,13 @@ ; 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*)