[llvm] a42bea2 - [coro async] Allow a coro.suspend.async to specify which argument is the context argument
Arnold Schwaighofer via llvm-commits
llvm-commits at lists.llvm.org
Wed Mar 3 08:28:42 PST 2021
Author: Arnold Schwaighofer
Date: 2021-03-03T08:27:37-08:00
New Revision: a42bea211a3c99edfc56ab27874924d8374a11c6
URL: https://github.com/llvm/llvm-project/commit/a42bea211a3c99edfc56ab27874924d8374a11c6
DIFF: https://github.com/llvm/llvm-project/commit/a42bea211a3c99edfc56ab27874924d8374a11c6.diff
LOG: [coro async] Allow a coro.suspend.async to specify which argument is the context argument
Before we used the same argument as the entry point. The resume partial
function might want to use a different ABI for its context argument
Differential Revision: https://reviews.llvm.org/D97333
Added:
Modified:
llvm/include/llvm/IR/Intrinsics.td
llvm/lib/Transforms/Coroutines/CoroInstr.h
llvm/lib/Transforms/Coroutines/CoroSplit.cpp
llvm/test/Transforms/Coroutines/coro-async-unreachable.ll
llvm/test/Transforms/Coroutines/coro-async.ll
Removed:
################################################################################
diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
index d1044df24558..385fdfa34b1b 100644
--- a/llvm/include/llvm/IR/Intrinsics.td
+++ b/llvm/include/llvm/IR/Intrinsics.td
@@ -1219,7 +1219,8 @@ def int_coro_async_resume : Intrinsic<[llvm_ptr_ty],
[]>;
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
index 7f6e95f48b9c..5ed800d67fe9 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInstr.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInstr.h
@@ -523,11 +523,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());
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 52c06e87c9ff..0b8c495669a4 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -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*)
@@ -1468,7 +1470,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.
@@ -1563,7 +1566,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();
diff --git a/llvm/test/Transforms/Coroutines/coro-async-unreachable.ll b/llvm/test/Transforms/Coroutines/coro-async-unreachable.ll
index 0ca3a6ce162a..492bc56e535f 100644
--- a/llvm/test/Transforms/Coroutines/coro-async-unreachable.ll
+++ b/llvm/test/Transforms/Coroutines/coro-async-unreachable.ll
@@ -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,
@@ -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*)
diff --git a/llvm/test/Transforms/Coroutines/coro-async.ll b/llvm/test/Transforms/Coroutines/coro-async.ll
index fb00a47fdd07..cdefced06883 100644
--- a/llvm/test/Transforms/Coroutines/coro-async.ll
+++ b/llvm/test/Transforms/Coroutines/coro-async.ll
@@ -89,7 +89,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,
@@ -204,7 +204,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,
@@ -222,7 +222,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,
@@ -258,8 +258,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
@@ -318,7 +318,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,
@@ -364,7 +364,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,
@@ -437,11 +437,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
@@ -515,13 +516,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*)
More information about the llvm-commits
mailing list