Skip to content

Commit 9fe31d5

Browse files
shahmishalshahmishal
authored
Merge pull request #2553 from aschwaighofer/coro_suspend_async_async_context_index
[coro async] Allow a coro.suspend.async to specify which argument is the context argument
2 parents bf1420a + cc07366 commit 9fe31d5

File tree

6 files changed

+48
-30
lines changed

6 files changed

+48
-30
lines changed

‎llvm/include/llvm/IR/Intrinsics.td

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1212,7 +1212,8 @@ def int_coro_async_resume : Intrinsic<[llvm_ptr_ty],
12121212
[]>;
12131213
def int_coro_async_size_replace : Intrinsic<[], [llvm_ptr_ty, llvm_ptr_ty], []>;
12141214
def int_coro_suspend_async
1215-
: Intrinsic<[llvm_any_ty], [llvm_ptr_ty, llvm_ptr_ty, llvm_vararg_ty], []>;
1215+
: Intrinsic<[llvm_any_ty],
1216+
[llvm_i32_ty, llvm_ptr_ty, llvm_ptr_ty, llvm_vararg_ty], []>;
12161217
def int_coro_prepare_async : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty],
12171218
[IntrNoMem]>;
12181219
def int_coro_begin : Intrinsic<[llvm_ptr_ty], [llvm_token_ty, llvm_ptr_ty],

‎llvm/lib/Transforms/Coroutines/CoroInstr.h

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -523,11 +523,21 @@ inline CoroSaveInst *AnyCoroSuspendInst::getCoroSave() const {
523523

524524
/// This represents the llvm.coro.suspend.async instruction.
525525
class LLVM_LIBRARY_VISIBILITY CoroSuspendAsyncInst : public AnyCoroSuspendInst {
526-
enum { ResumeFunctionArg, AsyncContextProjectionArg, MustTailCallFuncArg };
527-
528526
public:
527+
enum {
528+
StorageArgNoArg,
529+
ResumeFunctionArg,
530+
AsyncContextProjectionArg,
531+
MustTailCallFuncArg
532+
};
533+
529534
void checkWellFormed() const;
530535

536+
unsigned getStorageArgumentIndex() const {
537+
auto *Arg = cast<ConstantInt>(getArgOperand(StorageArgNoArg));
538+
return Arg->getZExtValue();
539+
}
540+
531541
Function *getAsyncContextProjectionFunction() const {
532542
return cast<Function>(
533543
getArgOperand(AsyncContextProjectionArg)->stripPointerCasts());

‎llvm/lib/Transforms/Coroutines/CoroSplit.cpp

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -759,10 +759,12 @@ Value *CoroCloner::deriveNewFramePointer() {
759759
// with the active suspend. The frame is located as a tail to the async
760760
// context header.
761761
case coro::ABI::Async: {
762-
auto *CalleeContext = NewF->getArg(Shape.AsyncLowering.ContextArgNo);
762+
auto *ActiveAsyncSuspend = cast<CoroSuspendAsyncInst>(ActiveSuspend);
763+
auto *CalleeContext =
764+
NewF->getArg(ActiveAsyncSuspend->getStorageArgumentIndex());
763765
auto *FramePtrTy = Shape.FrameTy->getPointerTo();
764-
auto *ProjectionFunc = cast<CoroSuspendAsyncInst>(ActiveSuspend)
765-
->getAsyncContextProjectionFunction();
766+
auto *ProjectionFunc =
767+
ActiveAsyncSuspend->getAsyncContextProjectionFunction();
766768
auto DbgLoc =
767769
cast<CoroSuspendAsyncInst>(VMap[ActiveSuspend])->getDebugLoc();
768770
// Calling i8* (i8*)
@@ -879,12 +881,14 @@ void CoroCloner::create() {
879881
addFramePointerAttrs(NewAttrs, Context, 0,
880882
Shape.FrameSize, Shape.FrameAlign);
881883
break;
882-
case coro::ABI::Async:
883-
if (OrigF.hasParamAttribute(Shape.AsyncLowering.ContextArgNo,
884-
Attribute::SwiftAsync)) {
885-
addAsyncContextAttrs(NewAttrs, Context, Shape.AsyncLowering.ContextArgNo);
884+
case coro::ABI::Async: {
885+
auto *ActiveAsyncSuspend = cast<CoroSuspendAsyncInst>(ActiveSuspend);
886+
auto ContextArgIndex = ActiveAsyncSuspend->getStorageArgumentIndex();
887+
if (OrigF.hasParamAttribute(ContextArgIndex, Attribute::SwiftAsync)) {
888+
addAsyncContextAttrs(NewAttrs, Context, ContextArgIndex);
886889
}
887890
break;
891+
}
888892
case coro::ABI::Retcon:
889893
case coro::ABI::RetconOnce:
890894
// If we have a continuation prototype, just use its attributes,
@@ -1485,7 +1489,8 @@ static void replaceAsyncResumeFunction(CoroSuspendAsyncInst *Suspend,
14851489
auto *Val = Builder.CreateBitOrPointerCast(Continuation, Int8PtrTy);
14861490
ResumeIntrinsic->replaceAllUsesWith(Val);
14871491
ResumeIntrinsic->eraseFromParent();
1488-
Suspend->setOperand(0, UndefValue::get(Int8PtrTy));
1492+
Suspend->setOperand(CoroSuspendAsyncInst::ResumeFunctionArg,
1493+
UndefValue::get(Int8PtrTy));
14891494
}
14901495

14911496
/// Coerce the arguments in \p FnArgs according to \p FnTy in \p CallArgs.
@@ -1580,7 +1585,8 @@ static void splitAsyncCoroutine(Function &F, coro::Shape &Shape,
15801585
// Insert the call to the tail call function and inline it.
15811586
auto *Fn = Suspend->getMustTailCallFunction();
15821587
SmallVector<Value *, 8> Args(Suspend->args());
1583-
auto FnArgs = ArrayRef<Value *>(Args).drop_front(3);
1588+
auto FnArgs = ArrayRef<Value *>(Args).drop_front(
1589+
CoroSuspendAsyncInst::MustTailCallFuncArg + 1);
15841590
auto *TailCall =
15851591
coro::createMustTailCall(Suspend->getDebugLoc(), Fn, FnArgs, Builder);
15861592
Builder.CreateRetVoid();

‎llvm/test/Transforms/Coroutines/coro-async-unreachable.ll

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ entry:
7171
store i8* %async.ctxt, i8** %callee_context.caller_context.addr
7272
%resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
7373
%callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
74-
%res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
74+
%res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
7575
i8* %resume.func_ptr,
7676
i8* %resume_proj_fun,
7777
void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
@@ -97,7 +97,7 @@ entry:
9797
declare i8* @llvm.coro.prepare.async(i8*)
9898
declare token @llvm.coro.id.async(i32, i32, i32, i8*)
9999
declare i8* @llvm.coro.begin(token, i8*)
100-
declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i8*, i8*, ...)
100+
declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i32, i8*, i8*, ...)
101101
declare i8* @llvm.coro.async.context.alloc(i8*, i8*)
102102
declare void @llvm.coro.async.context.dealloc(i8*)
103103
declare swiftcc void @asyncReturn(i8*, %async.task*, %async.actor*)

‎llvm/test/Transforms/Coroutines/coro-async.ll

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ entry:
8989
store i8* %async.ctxt, i8** %callee_context.caller_context.addr
9090
%resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
9191
%callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
92-
%res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
92+
%res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
9393
i8* %resume.func_ptr,
9494
i8* %resume_proj_fun,
9595
void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
@@ -206,7 +206,7 @@ entry:
206206
store i8* %async.ctxt, i8** %callee_context.caller_context.addr
207207
%resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
208208
%callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
209-
%res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
209+
%res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 2,
210210
i8* %resume.func_ptr,
211211
i8* %resume_proj_fun,
212212
void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
@@ -224,7 +224,7 @@ entry:
224224
store i8* %async.ctxt, i8** %callee_context.caller_context.addr.1
225225
%resume_proj_fun.2 = bitcast i8*(i8*)* @resume_context_projection to i8*
226226
%callee.2 = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
227-
%res.2 = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
227+
%res.2 = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
228228
i8* %resume.func_ptr.1,
229229
i8* %resume_proj_fun.2,
230230
void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
@@ -260,8 +260,8 @@ entry:
260260
; CHECK: tail call swiftcc void @asyncSuspend(i8* [[CALLEE_CTXT_RELOAD]]
261261
; CHECK: ret void
262262

263-
; CHECK-LABEL: define internal swiftcc void @my_async_function2.resume.1(i8* nocapture readnone %0, i8* %1, i8* nocapture readonly %2) {
264-
; CHECK: bitcast i8* %2 to i8**
263+
; CHECK-LABEL: define internal swiftcc void @my_async_function2.resume.1(i8* nocapture readonly %0, i8* %1, i8* nocapture readnone %2) {
264+
; CHECK: bitcast i8* %0 to i8**
265265
; CHECK: [[ACTOR_ARG:%.*]] = bitcast i8* %1
266266
; CHECK: tail call swiftcc void @asyncReturn({{.*}}[[ACTOR_ARG]])
267267
; CHECK: ret void
@@ -320,7 +320,7 @@ entry:
320320
store i8* %async.ctxt, i8** %callee_context.caller_context.addr
321321
%resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
322322
%callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
323-
%res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
323+
%res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
324324
i8* %resume.func_ptr,
325325
i8* %resume_proj_fun,
326326
void (i8*, i8*, %async.task*, %async.actor*)* @dont_crash_on_cf_dispatch,
@@ -366,7 +366,7 @@ entry:
366366
store i8* %async.ctxt, i8** %callee_context.caller_context.addr
367367
%resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
368368
%callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
369-
%res = call {i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async(
369+
%res = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
370370
i8* %resume.func_ptr,
371371
i8* %resume_proj_fun,
372372
void (i8*, i8*, %async.task*, %async.actor*)* @dont_crash_on_cf_dispatch,
@@ -439,11 +439,12 @@ entry:
439439
store i8* %async.ctxt, i8** %callee_context.caller_context.addr
440440
%resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
441441
%callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
442-
%res = call {i8*, i8*, i8*, i8*} (i8*, i8*, ...) @llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(
443-
i8* %resume.func_ptr,
444-
i8* %resume_proj_fun,
445-
void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
446-
i8* %callee, i8* %callee_context, %async.task* %task, %async.actor *%actor)
442+
%res = call {i8*, i8*, i8*, i8*} (i32, i8*, i8*, ...)
443+
@llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(i32 0,
444+
i8* %resume.func_ptr,
445+
i8* %resume_proj_fun,
446+
void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
447+
i8* %callee, i8* %callee_context, %async.task* %task, %async.actor *%actor)
447448

448449
call void @llvm.coro.async.context.dealloc(i8* %callee_context)
449450
%continuation_task_arg = extractvalue {i8*, i8*, i8*, i8*} %res, 3
@@ -517,13 +518,13 @@ entry:
517518
; CHECK: store i64* null, i64** [[ALLOCA]]
518519
; CHECK: call void @do_with_swifterror(i64** {{.*}}swifterror{{.*}} [[ALLOCA]])
519520

520-
declare { i8*, i8*, i8*, i8* } @llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(i8*, i8*, ...)
521+
declare { i8*, i8*, i8*, i8* } @llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(i32, i8*, i8*, ...)
521522
declare i8* @llvm.coro.prepare.async(i8*)
522523
declare token @llvm.coro.id.async(i32, i32, i32, i8*)
523524
declare i8* @llvm.coro.begin(token, i8*)
524525
declare i1 @llvm.coro.end.async(i8*, i1, ...)
525526
declare i1 @llvm.coro.end(i8*, i1)
526-
declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i8*, i8*, ...)
527+
declare {i8*, i8*, i8*} @llvm.coro.suspend.async(i32, i8*, i8*, ...)
527528
declare i8* @llvm.coro.async.context.alloc(i8*, i8*)
528529
declare void @llvm.coro.async.context.dealloc(i8*)
529530
declare swiftcc void @asyncReturn(i8*, %async.task*, %async.actor*)

‎llvm/test/Transforms/Coroutines/coro-split-swift-async.ll

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ entry:
6767
%23 = bitcast %T12_Concurrency9MainActorC5_ImplC* %15 to %swift.executor*, !dbg !49
6868
%24 = load %swift.executor*, %swift.executor** %4, align 8, !dbg !49
6969
%25 = load %swift.context*, %swift.context** %5, align 8, !dbg !49
70-
%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
70+
%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
7171
%27 = extractvalue { i8*, i8*, i8* } %26, 0, !dbg !49
7272
%28 = bitcast i8* %27 to %swift.task*, !dbg !49
7373
store %swift.task* %28, %swift.task** %3, align 8, !dbg !49
@@ -136,7 +136,7 @@ entry:
136136
declare extern_weak swiftcc void @swift_task_switch(%swift.task* %0, %swift.executor* %1, %swift.executor* %2) #1
137137

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

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

0 commit comments

Comments
 (0)