Merge pull request #11047 from swiftlang/jepa-stable2

[stable/21.x] Coro: Add a function attribute for resume from call async functlets

Author: AnthonyLatsis

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -989,6 +989,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:

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",