[Coroutines] Set presplit attribute in Clang and mlir

[ChuanqiXu]

Summary

See https://reviews.llvm.org/D100282 and https://github.com/llvm/llvm-project/issues/49264 for the background.
Simply, coroutine shouldn't be inlined before CoroSplit. And the marker for pre-splited coroutine is created in CoroEarly pass, which ran after AlwaysInliner Pass in O0 pipeline. So that the AlwaysInliner couldn't detect it shouldn't inline a coroutine. So here is the error.

This patch set the presplit attribute in clang and mlir. So the inliner would always detect the attribute before splitting.

[aeubanks]

pass manager changes seem fine, but I'll let somebody else lgtm

llvm/test/Other/new-pm-O0-defaults.ll

34 ↗

(On Diff #394501)

should have -NEXT

[ChuanqiXu]

llvm/test/Other/new-pm-O0-defaults.ll

34 ↗

(On Diff #394501)

In case there is no CHECK-DIS, CHECK-CORO-NEXT doesn't make sense.

[rjmccall]

I don't find it at all weird to say that there's an invariant that certain intrinsics can only appear in functions with a particular attribute, and so the frontend has a responsibility to set the attribute if it's using that intrinsic. We do the same thing with several other intrinsics, such as the constrained FP intrinsics generated for #pragma fenv_access.

[ChuanqiXu]

Address the comments to set the attribute in frontend.

[ChuanqiXu]

@ezhulenev hi, we are discussing if we should set the coroutine attributes in middle-end or front-end. Now in this revision, it would set the coroutine.presplit attribute in frontend, which would break the MLIR test. I have tried to implement it myself today. But I found that it is really not easy to edit MLIR's code. I found that you introduced the Coroutines intrinsic into MLIR. How do you think about set the attribute in the frontend of MLIR? And if you approve this, could you help to add the attributes in MLIR? We could still add the attribute in CoroEarly pass for a short time to make things consistent.

[ezhulenev]

There are two places where in MLIR you can put an attribute to coroutine functions:

1. https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp#L126

This is the point when coroutine functions are created, and you can attach attribute to the func argument

2. https://github.com/llvm/llvm-project/blob/main/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp#L325

Another options is to attach attribute to the op->getParentOfType<FuncOp> when async runtime operations lowered to coro.id intrinsic.

/cc @mehdi_amini to chime in what option is better. I'm not sure though what type of MLIR attribute will be translated to LLVM "coroutine.presplit"="0", UnitAttr? Or it must be IntegerAttr.

I'm on vacation with limjted access to computer until the end of the year, I can take a look at it myself ~late Dec or early Jan.

[ChuanqiXu]

In D115790#3199905, @ezhulenev wrote:

There are two places where in MLIR you can put an attribute to coroutine functions:

1. https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp#L126

This is the point when coroutine functions are created, and you can attach attribute to the func argument

2. https://github.com/llvm/llvm-project/blob/main/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp#L325

Another options is to attach attribute to the op->getParentOfType<FuncOp> when async runtime operations lowered to coro.id intrinsic.

/cc @mehdi_amini to chime in what option is better. I'm not sure though what type of MLIR attribute will be translated to LLVM "coroutine.presplit"="0", UnitAttr? Or it must be IntegerAttr.

I'm on vacation with limjted access to computer until the end of the year, I can take a look at it myself ~late Dec or early Jan.

Hi, thanks for looking into this. In this func, https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp#L126, I tried to write:

func.setAttr("coroutine.presplit", "0");

But it shows that we could only set the attribute for argument and result type instead of the function itself. How could we make it?

[ChuanqiXu]

Set presplit attribute in mlir too

[ChuanqiXu]

@ezhulenev I finally made it! Thanks for your help! @mehdi_amini @ftynse hi, might you help to look if the change in mlir part is good?

[ChuanqiXu]

Clean codes.

[ezhulenev]

mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp

201

nit: builder has an API to simplify attributes construction: https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/Builders.h#L92-L100

[ChuanqiXu]

Address comments.

[ChuanqiXu]

@rjmccall gentle ping~

[rjmccall]

clang/lib/CodeGen/CGCoroutine.cpp

712

The assertion here is not necessary; if it was, we'd need it everywhere.

Please add a comment like "LLVM expects coroutines to have this attribute set coming out of the frontend."

llvm/docs/Coroutines.rst

1179

Please add this paragraph (modified appropriately) to all of the other coro.id.* intrinsics.

llvm/lib/Transforms/Coroutines/CoroEarly.cpp

198

This should turn into the same assertion as above. (Swift will need a change, but that's Swift's problem, and I'll take care of it.)

llvm/lib/Transforms/Coroutines/CoroInternal.h

42

Editorial: don't put "the" before LLVM here

Since this change requires frontend cooperation, and the refactor will *also* require frontend cooperation, should we just do this in one step so that frontends are less disrupted?

[ChuanqiXu]

Address comments.

[rjmccall]

Alright, that's fine.

[ChuanqiXu]

llvm/lib/Transforms/Coroutines/CoroEarly.cpp

198

I don't add the assertion here since I am worrying it may break some builds. Would you be happy to add the assertion when you change swift?

llvm/lib/Transforms/Coroutines/CoroInternal.h

42

I didn't image swift before when I wrote the FIXME. In case we only need to care about clang and mlir, I think it doesn't matter a lot. But it matters if there is other frontends.
I also found that the values "UNPREPARED_FOR_SPLIT"，“PREPARED_FOR_SPLIT” and "ASYNC_RESTART_AFTER_SPLIT" are only meaningful under LegacyPM. @aeubanks , hi what's the status of LegacyPM? Do we need to care about the correctness for the LegacyPM? Or we could give up on it?

[MaskRay]

Early heads-up: we see a failure with llvm/lib/Transforms/Coroutines/CoroEarly.cpp:186 in bool (anonymous namespace)::Lowerer::lowerEarlyIntrinsics(llvm::Function &): F.hasFnAttribute(CORO_PRESPLIT_ATTR) && F.getFnAttribute(CORO_PRESPLIT_ATTR).getValueAsString() == UNPREPARED_FOR_SPLIT && "The frontend uses Swtich-Resumed ABI should emit " "\"coroutine.presplit\" attribute with value \"0\" for the " "coroutine."

A colleague may follow up with the issue or request a revert.

[ChuanqiXu]

In D115790#3221307, @MaskRay wrote:

Early heads-up: we see a failure with llvm/lib/Transforms/Coroutines/CoroEarly.cpp:186 in bool (anonymous namespace)::Lowerer::lowerEarlyIntrinsics(llvm::Function &): F.hasFnAttribute(CORO_PRESPLIT_ATTR) && F.getFnAttribute(CORO_PRESPLIT_ATTR).getValueAsString() == UNPREPARED_FOR_SPLIT && "The frontend uses Swtich-Resumed ABI should emit " "\"coroutine.presplit\" attribute with value \"0\" for the " "coroutine."

A colleague may follow up with the issue or request a revert.

Thanks for reporting this. I have tested this locally with our workloads. I guess there are 3 reasons to hit this:
(1) The frontend (clang or mlir) is updated.
(2) The input IR is generated by an old frontend.
(3) The tested frontend uses switched-resume ABI coroutine intrinsics and the frontend isn't clang nor mlir.

clang/lib/CodeGen/CGCoroutine.cpp

void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
   EmitBlock(RetBB);
   // Emit coro.end before getReturnStmt (and parameter destructors), since
   // resume and destroy parts of the coroutine should not include them.
   llvm::Function *CoroEnd = CGM.getIntrinsic(llvm::Intrinsic::coro_end);
   Builder.CreateCall(CoroEnd, {NullPtr, Builder.getFalse()});
 
   if (Stmt *Ret = S.getReturnStmt())
     EmitStmt(Ret);
+
+  assert(CurFn);
+  CurFn->addFnAttr("coroutine.presplit", "0");
 }
 
 // Emit coroutine intrinsic and patch up arguments of the token type.
 RValue CodeGenFunction::EmitCoroutineIntrinsic(const CallExpr *E,
                                                unsigned int IID) {
   SmallVector<llvm::Value *, 8> Args;
   switch (IID) {
   default:

clang/test/CodeGenCoroutines/coro-always-inline.cpp

 // CHECK: [[CAST1:%[0-9]+]] = bitcast %"struct.std::awaitable"* %ref.tmp{{.*}} to i8*
 // CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 1, i8* [[CAST1]])
 
 // CHECK: [[CAST2:%[0-9]+]] = bitcast %"struct.std::awaitable"* %ref.tmp{{.*}} to i8*
 // CHECK-NEXT: call void @llvm.lifetime.start.p0i8(i64 1, i8* [[CAST2]])
 // CHECK: [[CAST3:%[0-9]+]] = bitcast %"struct.std::awaitable"* %ref.tmp{{.*}} to i8*
 // CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 1, i8* [[CAST3]])
 void foo() { co_return; }
+
+// Check that bar is not inlined even it's marked as always_inline.
+
+// CHECK-LABEL:   define {{.*}} void @_Z3bazv()
+// CHECK:         call void @_Z3barv(
+__attribute__((__always_inline__)) void bar() {
+  co_return;
+}
+void baz() {
+  bar();
+  co_return;
+}

clang/test/CodeGenCoroutines/coro-attributes.cpp

+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -std=c++20 -disable-llvm-passes -emit-llvm %s -o - | FileCheck %s
+#include "Inputs/coroutine.h"
+
+using namespace std;
+
+struct coro {
+  struct promise_type {
+    coro get_return_object();
+    suspend_never initial_suspend();
+    suspend_never final_suspend() noexcept;
+    void return_void();
+    static void unhandled_exception();
+  };
+};
+
+// CHECK: void @_Z3foov() #[[FOO_ATTR_NUM:[0-9]+]]
+// CHECK: attributes #[[FOO_ATTR_NUM]] = { {{.*}} "coroutine.presplit"="0"
+coro foo() {
+  co_await suspend_always{};
+}

llvm/docs/Coroutines.rst

 
 The purpose of this intrinsic is to tie together `coro.id`, `coro.alloc` and
 `coro.begin` belonging to the same coroutine to prevent optimization passes from
 duplicating any of these instructions unless entire body of the coroutine is
 duplicated.
 
 A frontend should emit exactly one `coro.id` intrinsic per coroutine.
 
+The frontend should add attribute `"coroutine.presplit"` with value `"0"` for the coroutine
+containing `coro.id`.
+
 .. _coro.id.async:
 
 'llvm.coro.id.async' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 ::
 
   declare token @llvm.coro.id.async(i32 <context size>, i32 <align>,
                                     i8* <context arg>,

llvm/lib/Transforms/Coroutines/CoroEarly.cpp

if (auto *CB = dyn_cast<CallBase>(&I)) {
       case Intrinsic::coro_noop:
         lowerCoroNoop(cast<IntrinsicInst>(&I));
         break;
       case Intrinsic::coro_id:
         // Mark a function that comes out of the frontend that has a coro.id
         // with a coroutine attribute.
         if (auto *CII = cast<CoroIdInst>(&I)) {
           if (CII->getInfo().isPreSplit()) {
-            F.addFnAttr(CORO_PRESPLIT_ATTR, UNPREPARED_FOR_SPLIT);
+            assert(F.hasFnAttribute(CORO_PRESPLIT_ATTR) &&
+                   F.getFnAttribute(CORO_PRESPLIT_ATTR).getValueAsString() ==
+                       "0" &&
+                   "The frontend uses Swtich-Resumed ABI should emit "
+                   "\"coroutine.presplit\" attribute with value \"0\" for the "
+                   "coroutine.");
             setCannotDuplicate(CII);
             CII->setCoroutineSelf();
             CoroId = cast<CoroIdInst>(&I);
           }
         }
         break;
       case Intrinsic::coro_id_retcon:
       case Intrinsic::coro_id_retcon_once:
       case Intrinsic::coro_id_async:
         F.addFnAttr(CORO_PRESPLIT_ATTR, PREPARED_FOR_SPLIT);
         break;
       case Intrinsic::coro_resume:
         lowerResumeOrDestroy(*CB, CoroSubFnInst::ResumeIndex);
         break;
       case Intrinsic::coro_destroy:
         lowerResumeOrDestroy(*CB, CoroSubFnInst::DestroyIndex);
         break;
       case Intrinsic::coro_promise:

llvm/lib/Transforms/Coroutines/CoroInternal.h

 // twice. First, it lets it go through complete IPO optimization pipeline as a
 // single function. It forces restart of the pipeline by inserting an indirect
 // call to an empty function "coro.devirt.trigger" which is devirtualized by
 // CoroElide pass that triggers a restart of the pipeline by CGPassManager.
 // When CoroSplit pass sees the same coroutine the second time, it splits it up,
 // adds coroutine subfunctions to the SCC to be processed by IPO pipeline.
 // Async lowering similarily triggers a restart of the pipeline after it has
 // split the coroutine.
+//
+// FIXME: Refactor these attributes as LLVM attributes instead of string
+// attributes since these attributes are already used outside the LLVM's
+// coroutine module.
 #define CORO_PRESPLIT_ATTR "coroutine.presplit"
 #define UNPREPARED_FOR_SPLIT "0"
 #define PREPARED_FOR_SPLIT "1"
 #define ASYNC_RESTART_AFTER_SPLIT "2"
 
 #define CORO_DEVIRT_TRIGGER_FN "coro.devirt.trigger"
 
 namespace coro {

llvm/test/Transforms/Coroutines/coro-debug-O2.ll

 ; RUN: opt < %s -passes='function(coro-early),cgscc(coro-split,coro-split),function(sroa)' --reuse-storage-in-coroutine-frame -S | FileCheck %s
 
 ; Checks whether the dbg.declare for `__promise` remains valid under O2.
 
 ; CHECK-LABEL: define internal fastcc void @f.resume({{.*}})
 ; CHECK:       entry.resume:
 ; CHECK:        call void @llvm.dbg.declare(metadata %f.Frame* %FramePtr, metadata ![[PROMISEVAR_RESUME:[0-9]+]], metadata !DIExpression(
 ;
 ; CHECK: ![[PROMISEVAR_RESUME]] = !DILocalVariable(name: "__promise"
 %promise_type = type { i32, i32, double }
 
-define void @f() !dbg !8 {
+define void @f() "coroutine.presplit"="0" !dbg !8  {
 entry:
     %__promise = alloca %promise_type, align 8
     %0 = bitcast %promise_type* %__promise to i8*
     %id = call token @llvm.coro.id(i32 16, i8* %0, i8* null, i8* null)
     %alloc = call i1 @llvm.coro.alloc(token %id)
     br i1 %alloc, label %coro.alloc, label %coro.init
 
 coro.alloc:                                       ; preds = %entry

llvm/test/Transforms/Coroutines/coro-debug-coro-frame.ll

 %promise_type = type { i32, i32, double }
 %struct.big_structure = type { [500 x i8] }
 declare void @produce(%struct.big_structure*)
 declare void @consume(%struct.big_structure*)
 declare void @pi32(i32*)
 declare void @pi64(i64*)
 declare void @pdouble(double*)
 
-define void @f(i32 %a, i32 %b, i64 %c, double %d) !dbg !8 {
+define void @f(i32 %a, i32 %b, i64 %c, double %d) "coroutine.presplit"="0" !dbg !8 {
 entry:
     %__promise = alloca %promise_type, align 8
     %0 = bitcast %promise_type* %__promise to i8*
     %a.alloc = alloca i32, align 4
     %b.alloc = alloca i32, align 4
     %c.alloc = alloca i64, align 4
     %d.alloc = alloca double, align 4
     store i32 %a, i32* %a.alloc
coro.ret:                                         ; preds = %cleanup.cont, %after.coro.free, %final.suspend, %await.suspend, %init.suspend
     %end = call i1 @llvm.coro.end(i8* null, i1 false)
     ret void
 
 unreachable:                                      ; preds = %after.coro.free
     unreachable
 
 }
 
-define void @bar(i32 %a, i64 %c, double %d) !dbg !19 {
+define void @bar(i32 %a, i64 %c, double %d) "coroutine.presplit"="0" !dbg !19 {
 entry:
     %__promise = alloca %promise_type, align 8
     %0 = bitcast %promise_type* %__promise to i8*
     %a.alloc = alloca i32, align 4
     %c.alloc = alloca i64, align 4
     %d.alloc = alloca double, align 4
     store i32 %a, i32* %a.alloc
     store i64 %c, i64* %c.alloc

llvm/test/Transforms/Coroutines/coro-debug-dbg.values-not_used_in_frame.ll

 ; Tests whether resume function would remain dbg.value infomation if corresponding values are not used in the frame.
 ; RUN: opt < %s -passes='function(coro-early),cgscc(coro-split,coro-split)' -S | FileCheck %s
 ;
 ; This file is based on coro-debug-frame-variable.ll.
 ; CHECK:  define internal fastcc void @f.resume(%f.Frame* noalias nonnull align 16 dereferenceable(80) %FramePtr) !dbg ![[RESUME_FN_DBG_NUM:[0-9]+]]
 ; CHECK:       await.ready:
 ; CHECK:         call void @llvm.dbg.value(metadata i32 undef, metadata ![[IVAR_RESUME:[0-9]+]], metadata !DIExpression(
 ; CHECK:         call void @llvm.dbg.value(metadata i32 undef, metadata ![[JVAR_RESUME:[0-9]+]], metadata !DIExpression(
 ;
 ; CHECK: ![[RESUME_FN_DBG_NUM]] = distinct !DISubprogram(name: "foo", linkageName: "_Z3foov"
 ; CHECK: ![[IVAR_RESUME]] = !DILocalVariable(name: "i"
 ; CHECK: ![[JVAR_RESUME]] = !DILocalVariable(name: "j"
 
 source_filename = "../llvm/test/Transforms/Coroutines/coro-debug-dbg.values-O2.ll"
 
-define void @f(i32 %i, i32 %j) !dbg !8 {
+define void @f(i32 %i, i32 %j) "coroutine.presplit"="0" !dbg !8 {
 entry:
   %__promise = alloca i8, align 8
   %x = alloca [10 x i32], align 16
   %id = call token @llvm.coro.id(i32 16, i8* %__promise, i8* null, i8* null)
   %alloc = call i1 @llvm.coro.alloc(token %id)
   br i1 %alloc, label %coro.alloc, label %coro.init
 
 coro.alloc:                                       ; preds = %entry

llvm/test/Transforms/Coroutines/coro-debug-dbg.values.ll

 ; CHECK: ![[IVAR_RESUME]] = !DILocalVariable(name: "i"
 ; CHECK: ![[XVAR_RESUME]] = !DILocalVariable(name: "x"
 ; CHECK: ![[JVAR_RESUME]] = !DILocalVariable(name: "j"
 ; CHECK: ![[SPILL_RESUME]] = !DILocalVariable(name: "produced"
 
 source_filename = "../llvm/test/Transforms/Coroutines/coro-debug-dbg.values-O2.ll"
 declare void @consume(i32)
 
-define void @f(i32 %i, i32 %j) !dbg !8 {
+define void @f(i32 %i, i32 %j) "coroutine.presplit"="0" !dbg !8 {
 entry:
   %__promise = alloca i8, align 8
   %x = alloca [10 x i32], align 16
   %produced = call i32 @value_producer()
   %id = call token @llvm.coro.id(i32 16, i8* %__promise, i8* null, i8* null)
   %alloc = call i1 @llvm.coro.alloc(token %id)
   br i1 %alloc, label %coro.alloc, label %coro.init
 

llvm/test/Transforms/Coroutines/coro-debug-frame-variable.ll

 ; CHECK: ![[JDBGLOC]] = !DILocation(line: 32, column: 7, scope: ![[SCOPE]])
 
 ; CHECK: ![[XVAR_RESUME]] = !DILocalVariable(name: "x"
 ; CHECK: ![[IDBGLOC_RESUME]] = !DILocation(line: 24, column: 7, scope: ![[RESUME_SCOPE:[0-9]+]])
 ; CHECK: ![[RESUME_SCOPE]] = distinct !DILexicalBlock(scope: !8, file: !1, line: 23, column: 12)
 ; CHECK: ![[IVAR_RESUME]] = !DILocalVariable(name: "i"
 ; CHECK: ![[JVAR_RESUME]] = !DILocalVariable(name: "j"
 ; CHECK: ![[JDBGLOC_RESUME]] = !DILocation(line: 32, column: 7, scope: ![[RESUME_SCOPE]])
-define void @f() {
+define void @f() "coroutine.presplit"="0" {
 entry:
   %__promise = alloca i8, align 8
   %i = alloca i32, align 4
   %j = alloca i32, align 4
   %x = alloca [10 x i32], align 16
   %id = call token @llvm.coro.id(i32 16, i8* %__promise, i8* null, i8* null)
   %alloc = call i1 @llvm.coro.alloc(token %id)
   br i1 %alloc, label %coro.alloc, label %coro.init

llvm/test/Transforms/Coroutines/coro-noalias-param.ll

 ; RUN: opt < %s -S -passes=coro-early | FileCheck %s
 %struct.A = type <{ i64, i64, i32, [4 x i8] }>
 
-define void @f(%struct.A* nocapture readonly noalias align 8 %a) {
+define void @f(%struct.A* nocapture readonly noalias align 8 %a) "coroutine.presplit"="0" {
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %size = call i32 @llvm.coro.size.i32()
   %alloc = call i8* @malloc(i32 %size)
   %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
   call void @print(i32 0)
   %s1 = call i8 @llvm.coro.suspend(token none, i1 false)
   switch i8 %s1, label %suspend [i8 0, label %resume 
                                  i8 1, label %cleanup]

llvm/test/Transforms/Coroutines/coro-split-01.ll

 ; Tests that a coroutine is split, inlined into the caller and devirtualized.
 ; RUN: opt < %s -S -enable-coroutines -passes='default<O2>' | FileCheck %s
 
-define i8* @f() {
+define i8* @f() "coroutine.presplit"="0" {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %need.dyn.alloc = call i1 @llvm.coro.alloc(token %id)
   br i1 %need.dyn.alloc, label %dyn.alloc, label %coro.begin
 dyn.alloc:
   %size = call i32 @llvm.coro.size.i32()
   %alloc = call i8* @malloc(i32 %size)
   br label %coro.begin

llvm/test/Transforms/Coroutines/coro-split-recursive.ll

 ; RUN: opt -passes='function(coro-early),cgscc(coro-split)' -S < %s | FileCheck %s
 
 declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*)
 
 declare i8* @llvm.coro.begin(token, i8* writeonly)
 
 declare token @llvm.coro.save(i8*)
 
 declare i8 @llvm.coro.suspend(token, i1)
 
 ; CHECK-LABEL: define void @foo()
 ; CHECK-LABEL: define {{.*}}void @foo.resume(
 ; CHECK: call void @foo()
 ; CHECK-LABEL: define {{.*}}void @foo.destroy(
 
-define void @foo() {
+define void @foo() "coroutine.presplit"="0" {
 entry:
   %__promise = alloca i32, align 8
   %0 = bitcast i32* %__promise to i8*
   %1 = call token @llvm.coro.id(i32 16, i8* %0, i8* null, i8* null)
   %2 = call i8* @llvm.coro.begin(token %1, i8* null)
   br i1 undef, label %if.then154, label %init.suspend
 
 init.suspend:                                     ; preds = %entry

llvm/test/Transforms/Coroutines/ex0.ll

 ; First example from Doc/Coroutines.rst (two block loop)
 ; RUN: opt < %s -enable-coroutines -aa-pipeline=basic-aa -passes='default<O2>' -preserve-alignment-assumptions-during-inlining=false -S | FileCheck %s
 
-define i8* @f(i32 %n) {
+define i8* @f(i32 %n) "coroutine.presplit"="0" {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %size = call i32 @llvm.coro.size.i32()
   %alloc = call i8* @malloc(i32 %size)
   %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
   br label %loop
 
 loop:

llvm/test/Transforms/Coroutines/ex1.ll

 ; First example from Doc/Coroutines.rst (one block loop)
 ; RUN: opt < %s -aa-pipeline=basic-aa -passes='default<O2>' -enable-coroutines -preserve-alignment-assumptions-during-inlining=false -S | FileCheck %s
 
-define i8* @f(i32 %n) {
+define i8* @f(i32 %n) "coroutine.presplit"="0" {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %size = call i32 @llvm.coro.size.i32()
   %alloc = call i8* @malloc(i32 %size)
   %hdl = call noalias i8* @llvm.coro.begin(token %id, i8* %alloc)
   br label %loop
 loop:
   %n.val = phi i32 [ %n, %entry ], [ %inc, %loop ]

llvm/test/Transforms/Coroutines/ex2.ll

 ; Second example from Doc/Coroutines.rst (custom alloc and free functions)
 ; RUN: opt < %s -passes='default<O2>' -enable-coroutines -S | FileCheck %s
 
-define i8* @f(i32 %n) {
+define i8* @f(i32 %n) "coroutine.presplit"="0" {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %need.dyn.alloc = call i1 @llvm.coro.alloc(token %id)
   br i1 %need.dyn.alloc, label %dyn.alloc, label %coro.begin
 dyn.alloc:
   %size = call i32 @llvm.coro.size.i32()
   %alloc = call i8* @CustomAlloc(i32 %size)
   br label %coro.begin

llvm/test/Transforms/Coroutines/ex3.ll

 ; Third example from Doc/Coroutines.rst (two suspend points)
 ; RUN: opt < %s -aa-pipeline=basic-aa -passes='default<O2>' -enable-coroutines -S | FileCheck %s
 
-define i8* @f(i32 %n) {
+define i8* @f(i32 %n) "coroutine.presplit"="0" {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %size = call i32 @llvm.coro.size.i32()
   %need.dyn.alloc = call i1 @llvm.coro.alloc(token %id)
   br i1 %need.dyn.alloc, label %dyn.alloc, label %coro.begin
 dyn.alloc:
   %alloc = call i8* @malloc(i32 %size)
   br label %coro.begin

llvm/test/Transforms/Coroutines/ex4.ll

 ; Fourth example from Doc/Coroutines.rst (coroutine promise)
 ; RUN: opt < %s -passes='default<O2>' -enable-coroutines -S | FileCheck %s
 
-define i8* @f(i32 %n) {
+define i8* @f(i32 %n) "coroutine.presplit"="0" {
 entry:
   %promise = alloca i32
   %pv = bitcast i32* %promise to i8*
   %id = call token @llvm.coro.id(i32 0, i8* %pv, i8* null, i8* null)
   %need.dyn.alloc = call i1 @llvm.coro.alloc(token %id)
   br i1 %need.dyn.alloc, label %dyn.alloc, label %coro.begin
 dyn.alloc:
   %size = call i32 @llvm.coro.size.i32()

llvm/test/Transforms/Coroutines/ex5.ll

 ; Fifth example from Doc/Coroutines.rst (final suspend)
 ; RUN: opt < %s -aa-pipeline=basic-aa -passes='default<O2>' -enable-coroutines -preserve-alignment-assumptions-during-inlining=false -S | FileCheck %s
 
-define i8* @f(i32 %n) {
+define i8* @f(i32 %n) "coroutine.presplit"="0" {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %size = call i32 @llvm.coro.size.i32()
   %alloc = call i8* @malloc(i32 %size)
   %hdl = call noalias i8* @llvm.coro.begin(token %id, i8* %alloc)
   br label %while.cond
 while.cond:
   %n.val = phi i32 [ %n, %entry ], [ %dec, %while.body ]

llvm/test/Transforms/Coroutines/phi-coro-end.ll

 ; Verify that we correctly handle suspend when the coro.end block contains phi
 ; RUN: opt < %s -aa-pipeline=basic-aa -passes='default<O2>' -enable-coroutines -S | FileCheck %s
 
-define i8* @f(i32 %n) {
+define i8* @f(i32 %n) "coroutine.presplit"="0" {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %size = call i32 @llvm.coro.size.i32()
   %alloc = call i8* @malloc(i32 %size)
   %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
   %0 = call i8 @llvm.coro.suspend(token none, i1 false)
   switch i8 %0, label %suspend [i8 0, label %cleanup i8 1, label %cleanup]
 

llvm/test/Transforms/Coroutines/restart-trigger.ll

 ; REQUIRES: asserts
 ; The following tests use the new pass manager, and verify that the coroutine
 ; passes re-run the CGSCC pipeline.
 ; RUN: opt < %s -S -passes='default<O0>' -enable-coroutines -debug-only=coro-split 2>&1 | FileCheck --check-prefix=CHECK-NEWPM %s
 ; RUN: opt < %s -S -passes='default<O1>' -enable-coroutines -debug-only=coro-split 2>&1 | FileCheck --check-prefix=CHECK-NEWPM %s
 
 ; CHECK:      CoroSplit: Processing coroutine 'f' state: 0
 ; CHECK-NEXT: CoroSplit: Processing coroutine 'f' state: 1
 ; CHECK-NEWPM:      CoroSplit: Processing coroutine 'f' state: 0
 ; CHECK-NEWPM-NOT:  CoroSplit: Processing coroutine 'f' state: 1
 
 
-define void @f() {
+define void @f() "coroutine.presplit"="0" {
   %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
   %size = call i32 @llvm.coro.size.i32()
   %alloc = call i8* @malloc(i32 %size)
   %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
   call void @print(i32 0)
   %s1 = call i8 @llvm.coro.suspend(token none, i1 false)
   switch i8 %s1, label %suspend [i8 0, label %resume 
                                  i8 1, label %cleanup]

mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp

if (&block == entryBlock || &block == cleanupBlock ||
       continue;
     Operation *terminator = block.getTerminator();
     if (auto yield = dyn_cast<YieldOp>(terminator)) {
       builder.setInsertionPointToEnd(&block);
       builder.create<BranchOp>(cleanupBlock);
     }
   }
 
+  // The switch-resumed API based coroutine should be marked with
+  // "coroutine.presplit" attribute with value "0" to mark the function as a
+  // coroutine.
+  func->setAttr(
+      "passthrough",
+      builder.getArrayAttr(builder.getArrayAttr({builder.getStringAttr("coroutine.presplit"), 
+                                                 builder.getStringAttr("0")})));
+
   CoroMachinery machinery;
   machinery.func = func;
   machinery.asyncToken = retToken;
   machinery.returnValues = retValues;
   machinery.coroHandle = coroHdlOp.handle();
   machinery.entry = entryBlock;
   machinery.setError = nullptr; // created lazily only if needed
   machinery.cleanup = cleanupBlock;
   machinery.suspend = suspendBlock;

mlir/test/mlir-opt/async.mlir

+// Check if mlir marks the corresponding function with required coroutine attribute.
+//
+// RUN:   mlir-opt %s -async-to-async-runtime                                  \
+// RUN:               -async-runtime-ref-counting                              \
+// RUN:               -async-runtime-ref-counting-opt                          \
+// RUN:               -convert-async-to-llvm                                   \
+// RUN:               -convert-linalg-to-loops                                 \
+// RUN:               -convert-scf-to-std                                      \
+// RUN:               -convert-linalg-to-llvm                                  \
+// RUN:               -convert-memref-to-llvm                                  \
+// RUN:               -convert-arith-to-llvm                                   \
+// RUN:               -convert-std-to-llvm                                     \
+// RUN:               -reconcile-unrealized-casts                              \
+// RUN: | FileCheck %s
+
+// CHECK: llvm.func @async_execute_fn{{.*}}attributes{{.*}}"coroutine.presplit", "0"
+// CHECK: llvm.func @async_execute_fn_0{{.*}}attributes{{.*}}"coroutine.presplit", "0"
+// CHECK: llvm.func @async_execute_fn_1{{.*}}attributes{{.*}}"coroutine.presplit", "0"
+
+func @main() {
+  %i0 = arith.constant 0 : index
+  %i1 = arith.constant 1 : index
+  %i2 = arith.constant 2 : index
+  %i3 = arith.constant 3 : index
+
+  %c0 = arith.constant 0.0 : f32
+  %c1 = arith.constant 1.0 : f32
+  %c2 = arith.constant 2.0 : f32
+  %c3 = arith.constant 3.0 : f32
+  %c4 = arith.constant 4.0 : f32
+
+  %A = memref.alloc() : memref<4xf32>
+  linalg.fill(%c0, %A) : f32, memref<4xf32>
+
+  %U = memref.cast %A :  memref<4xf32> to memref<*xf32>
+  call @print_memref_f32(%U): (memref<*xf32>) -> ()
+
+  memref.store %c1, %A[%i0]: memref<4xf32>
+  call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
+  call @print_memref_f32(%U): (memref<*xf32>) -> ()
+
+  %outer = async.execute {
+    memref.store %c2, %A[%i1]: memref<4xf32>
+    call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
+    call @print_memref_f32(%U): (memref<*xf32>) -> ()
+
+    // No op async region to create a token for testing async dependency.
+    %noop = async.execute {
+      call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
+      async.yield
+    }
+
+    %inner = async.execute [%noop] {
+      memref.store %c3, %A[%i2]: memref<4xf32>
+      call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
+      call @print_memref_f32(%U): (memref<*xf32>) -> ()
+
+      async.yield
+    }
+    async.await %inner : !async.token
+
+    memref.store %c4, %A[%i3]: memref<4xf32>
+    call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
+    call @print_memref_f32(%U): (memref<*xf32>) -> ()
+
+    async.yield
+  }
+  async.await %outer : !async.token
+
+  call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
+  call @print_memref_f32(%U): (memref<*xf32>) -> ()
+
+  memref.dealloc %A : memref<4xf32>
+
+  return
+}
+
+func private @mlirAsyncRuntimePrintCurrentThreadId() -> ()
+
+func private @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }