From 579b7687309c8b86ac18e78605422f4df0bf2bf5 Mon Sep 17 00:00:00 2001 From: Arnold Schwaighofer Date: Fri, 30 Aug 2024 15:35:20 -0700 Subject: [PATCH] Coro: Add a function attribute for resume from call async functlets This change adds a function attribute `async_ret` for any funclet generated from a async function marked with `async_entry` that models a return from an async function. We recognize these continuations by a `__swift_async_resume_project_context` projection function argument to the suspend function. rdar://134460666 --- llvm/lib/Transforms/Coroutines/CoroSplit.cpp | 15 +++++++++++++++ llvm/test/Transforms/Coroutines/coro-async.ll | 5 ++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp index a1055d3d829ce..eb79cf7d7f96f 100644 --- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp +++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp @@ -1002,6 +1002,21 @@ void coro::BaseCloner::create() { // Transfer the original function's attributes. auto FnAttrs = OrigF.getAttributes().getFnAttrs(); NewAttrs = NewAttrs.addFnAttributes(Context, AttrBuilder(Context, FnAttrs)); + + // Mark async funclets that "pop the frame" if the async function is marked + // with the `async_entry` function attribute. Remove the `async_entry` + // attribute from clones. + bool HasAsyncEntryAttribute = OrigF.hasFnAttribute("async_entry"); + auto ProjectionFunctionName = + ActiveAsyncSuspend->getAsyncContextProjectionFunction()->getName(); + bool IsAsyncFramePop = + ProjectionFunctionName == "__swift_async_resume_project_context"; + if (HasAsyncEntryAttribute) { + NewAttrs = NewAttrs.removeFnAttribute(Context, "async_entry"); + if (IsAsyncFramePop) { + NewAttrs = NewAttrs.addFnAttribute(Context, "async_ret"); + } + } break; } case coro::ABI::Retcon: diff --git a/llvm/test/Transforms/Coroutines/coro-async.ll b/llvm/test/Transforms/Coroutines/coro-async.ll index e5d2e6c2c42c6..b91a573aa3347 100644 --- a/llvm/test/Transforms/Coroutines/coro-async.ll +++ b/llvm/test/Transforms/Coroutines/coro-async.ll @@ -58,7 +58,7 @@ entry: } -define swiftcc void @my_async_function(ptr swiftasync %async.ctxt, ptr %task, ptr %actor) presplitcoroutine !dbg !1 { +define swiftcc void @my_async_function(ptr swiftasync %async.ctxt, ptr %task, ptr %actor) presplitcoroutine "async_entry" !dbg !1 { entry: %tmp = alloca { i64, i64 }, align 8 %vector = alloca <4 x double>, align 16 @@ -145,6 +145,7 @@ define void @my_async_function_pa(ptr %ctxt, ptr %task, ptr %actor) { ; CHECK-LABEL: define internal swiftcc void @my_async_functionTQ0_(ptr readonly swiftasync captures(none) %0, ptr %1, ptr readnone captures(none) %2) ; CHECK-O0-LABEL: define internal swiftcc void @my_async_functionTQ0_(ptr swiftasync %0, ptr %1, ptr %2) +; CHECK-SAME: #[[ATTRS:[0-9]+]] ; CHECK-SAME: !dbg ![[SP2:[0-9]+]] { ; CHECK: entryresume.0: ; CHECK: [[CALLER_CONTEXT:%.*]] = load ptr, ptr %0 @@ -514,6 +515,8 @@ declare ptr @hide(ptr) !llvm.dbg.cu = !{!2} !llvm.module.flags = !{!0} +; CHECK: #[[ATTRS]] = {{.*}}async_ret + !0 = !{i32 2, !"Debug Info Version", i32 3} ; CHECK: ![[SP1]] = distinct !DISubprogram(name: "my_async_function", ; CHECK-SAME: linkageName: "my_async_function",