[mlir:Async] Use LLVM coro operations to lower from async.coro

[ezhulenev]

Summary

Instead of using function calls to LLVM intrinsics, use intr operations defined in the LLVM dialect.

Leftovers from: https://reviews.llvm.org/D94923

[ezhulenev]

Rebased

[ezhulenev]

Rebase

[ezhulenev]

Rebase

mlir/include/mlir/Dialect/Async/IR/Async.h

 
 namespace mlir {
 namespace async {
 
 namespace detail {
 struct ValueTypeStorage;
 } // namespace detail
 
+//===----------------------------------------------------------------------===//
+// Async dialect types.
+//===----------------------------------------------------------------------===//
+
 /// The token type to represent asynchronous operation completion.
 class TokenType : public Type::TypeBase<TokenType, Type, TypeStorage> {
 public:
   using Base::Base;
 };
 
 /// The value type to represent values returned from asynchronous operations.
 class ValueType
     : public Type::TypeBase<ValueType, Type, detail::ValueTypeStorage> {
 public:
   using Base::Base;
 
   /// Get or create an async ValueType with the provided value type.
   static ValueType get(Type valueType);
 
   Type getValueType();
 };
 
 /// The group type to represent async tokens or values grouped together.
 class GroupType : public Type::TypeBase<GroupType, Type, TypeStorage> {
 public:
   using Base::Base;
 };
 
-// -------------------------------------------------------------------------- //
+//===----------------------------------------------------------------------===//
+// LLVM coroutines types.
+//===----------------------------------------------------------------------===//
+
+/// The type identifying a switched-resume coroutine.
+class CoroIdType : public Type::TypeBase<CoroIdType, Type, TypeStorage> {
+public:
+  using Base::Base;
+};
+
+/// The coroutine handle type which is a pointer to the coroutine frame.
+class CoroHandleType
+    : public Type::TypeBase<CoroHandleType, Type, TypeStorage> {
+public:
+  using Base::Base;
+};
+
+/// The coroutine saved state type.
+class CoroStateType : public Type::TypeBase<CoroStateType, Type, TypeStorage> {
+public:
+  using Base::Base;
+};
+
+//===----------------------------------------------------------------------===//
 // Helper functions of Async dialect transformations.
-// -------------------------------------------------------------------------- //
+//===----------------------------------------------------------------------===//
 
 /// Returns true if the type is reference counted. All async dialect types are
 /// reference counted at runtime.
 inline bool isRefCounted(Type type) {
   return type.isa<TokenType, ValueType, GroupType>();
 }
 
 } // namespace async
 } // namespace mlir
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Async/IR/AsyncOps.h.inc"
 
 #include "mlir/Dialect/Async/IR/AsyncOpsDialect.h.inc"
 
 #endif // MLIR_DIALECT_ASYNC_IR_ASYNC_H

mlir/include/mlir/Dialect/Async/IR/AsyncBase.td

 
 def Async_AnyValueOrTokenType : AnyTypeOf<[Async_AnyValueType,
                                            Async_TokenType]>;
 
 def Async_AnyAsyncType : AnyTypeOf<[Async_AnyValueType,
                                     Async_TokenType,
                                     Async_GroupType]>;
 
+//===----------------------------------------------------------------------===//
+// Types for lowering to LLVM + Async Runtime via the LLVM Coroutines.
+//===----------------------------------------------------------------------===//
+
+// LLVM coroutines intrinsics use `token` and `i8*` types to represent coroutine
+// identifiers and handles. To define type-safe Async Runtime operations and
+// build a properly typed intermediate IR during the Async to LLVM lowering we
+// define a separate types for values that can be produced by LLVM intrinsics.
+
+def Async_CoroIdType : DialectType<AsyncDialect,
+    CPred<"$_self.isa<::mlir::async::CoroIdType>()">, "coro.id type">,
+    BuildableType<"$_builder.getType<::mlir::async::CoroIdType>()"> {
+  let description = [{
+    `async.coro.id` is a type identifying a switched-resume coroutine.
+  }];
+}
+
+def Async_CoroHandleType : DialectType<AsyncDialect,
+    CPred<"$_self.isa<::mlir::async::CoroHandleType>()">, "coro.handle type">,
+    BuildableType<"$_builder.getType<::mlir::async::CoroHandleType>()"> {
+  let description = [{
+    `async.coro.handle` is a handle to the coroutine (pointer to the coroutine
+    frame) that can be passed around to resume or destroy the coroutine.
+  }];
+}
+
+def Async_CoroStateType : DialectType<AsyncDialect,
+    CPred<"$_self.isa<::mlir::async::CoroStateType>()">, "coro.state type">,
+    BuildableType<"$_builder.getType<::mlir::async::CoroStateType>()"> {
+  let description = [{
+    `async.coro.state` is a saved coroutine state that should be passed to the
+    coroutine suspension operation.
+  }];
+}
+
 #endif // ASYNC_BASE_TD

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

let description = [{
 
     ```mlir
     %0 = async.create_group
     ...
     async.await_all %0
     ```
   }];
 
-  let arguments = (ins );
   let results = (outs Async_GroupType:$result);
 
   let assemblyFormat = "attr-dict";
 }
 
 def Async_AddToGroupOp : Async_Op<"add_to_group", []> {
   let summary = "adds and async token or value to the group";
   let description = [{
def Async_AwaitAllOp : Async_Op<"await_all", []> {
 
   let arguments = (ins Async_GroupType:$operand);
   let results = (outs);
 
   let assemblyFormat = "$operand attr-dict";
 }
 
 //===----------------------------------------------------------------------===//
-// Async Dialect Automatic Reference Counting Operations.
+// Async Dialect LLVM Coroutines Operations.
 //===----------------------------------------------------------------------===//
 
+// Async to LLVM dialect lowering converts async tasks (regions inside async
+// execute operations) to LLVM coroutines [1], and relies on switched-resume
+// lowering [2] to produce an asynchronous executable.
+//
+// We define LLVM coro intrinsics in the async dialect to facilitate progressive
+// lowering with verifiable and type-safe IR during the multi-step lowering
+// pipeline. First we convert from high level async operations (e.g. execute) to
+// the explicit calls to coro intrinsics and runtime API, and then finalize
+// lowering to LLVM with a simple dialect conversion pass.
+//
+// [1] https://llvm.org/docs/Coroutines.html
+// [2] https://llvm.org/docs/Coroutines.html#switched-resume-lowering
+
+def Async_CoroIdOp : Async_Op<"coro.id"> {
+  let summary = "returns a switched-resume coroutine identifier";
+  let description = [{
+    The `async.coro.id` returns a switched-resume coroutine identifier.
+  }];
+
+  let results = (outs Async_CoroIdType:$id);
+  let assemblyFormat = "attr-dict";
+}
+
+def Async_CoroBeginOp : Async_Op<"coro.begin"> {
+  let summary = "returns a handle to the coroutine";
+  let description = [{
+    The `async.coro.begin` allocates a coroutine frame and returns a handle to
+    the coroutine.
+  }];
+
+  let arguments = (ins Async_CoroIdType:$id);
+  let results = (outs Async_CoroHandleType:$handle);
+  let assemblyFormat = "$id attr-dict";
+}
+
+def Async_CoroFreeOp : Async_Op<"coro.free"> {
+  let summary = "deallocates the coroutine frame";
+  let description = [{
+    The `async.coro.free` deallocates the coroutine frame created by the
+    async.coro.begin operation.
+  }];
+
+  let arguments = (ins Async_CoroIdType:$id,
+                       Async_CoroHandleType:$handle);
+  let assemblyFormat = "$id `,` $handle attr-dict";
+}
+
+def Async_CoroEndOp : Async_Op<"coro.end"> {
+  let summary = "marks the end of the coroutine in the suspend block";
+  let description = [{
+    The `async.coro.end` marks the point where a coroutine needs to return
+    control back to the caller if it is not an initial invocation of the
+    coroutine. It the start part of the coroutine is is no-op.
+  }];
+
+  let arguments = (ins Async_CoroHandleType:$handle);
+  let assemblyFormat = "$handle attr-dict";
+}
+
+def Async_CoroSaveOp : Async_Op<"coro.save"> {
+  let summary = "saves the coroutine state";
+  let description = [{
+    The `async.coro.saves` saves the coroutine state.
+  }];
+
+  let arguments = (ins Async_CoroHandleType:$handle);
+  let results = (outs Async_CoroStateType:$state);
+  let assemblyFormat = "$handle attr-dict";
+}
+
+def Async_CoroSuspendOp : Async_Op<"coro.suspend", [Terminator]> {
+  let summary = "suspends the coroutine";
+  let description = [{
+    The `async.coro.suspend` suspends the coroutine and transfers control to the
+    `suspend` successor. If suspended coroutine later resumed it will transfer
+    control to the `resume` successor. If it is destroyed it will transfer
+    control to the the `cleanup` successor.
+
+    In switched-resume lowering coroutine can be already in resumed state when
+    suspend operation is called, in this case control will be transferred to the
+    `resume` successor skipping the `suspend` successor.
+  }];
+
+  let arguments = (ins Async_CoroStateType:$state);
+  let successors = (successor AnySuccessor:$suspendDest,
+                              AnySuccessor:$resumeDest,
+                              AnySuccessor:$cleanupDest);
+  let assemblyFormat =
+    "$state `,` $suspendDest `,` $resumeDest  `,` $cleanupDest attr-dict";
+}
+
+//===----------------------------------------------------------------------===//
+// Async Dialect Runtime Operations.
+//===----------------------------------------------------------------------===//
+
+// The following operations are intermediate async dialect operations to help
+// lowering from high level async operation like `async.execute` to the Async
+// Runtime API defined in the `ExecutionEngine/AsyncRuntime.h`.
+
+def Async_RuntimeCreateOp : Async_Op<"runtime.create"> {
+  let summary = "creates an async runtime value (token, value or group)";
+  let description = [{
+    The `async.runtime.create` operation creates an async dialect value
+    (token, value or group). Tokens and values are created in non-ready state.
+    Groups are created in empty state.
+  }];
+
+  let results = (outs Async_AnyAsyncType:$result);
+  let assemblyFormat = "attr-dict `:` type($result)";
+}
+
+def Async_RuntimeSetAvailableOp : Async_Op<"runtime.set_available"> {
+  let summary = "switches token or value available state";
+  let description = [{
+    The `async.runtime.set_available` operation switches async token or value
+    state to available.
+  }];
+
+  let arguments = (ins Async_AnyValueOrTokenType:$operand);
+  let assemblyFormat = "$operand attr-dict `:` type($operand)";
+}
+
+def Async_RuntimeAwaitOp : Async_Op<"runtime.await"> {
+  let summary = "blocks the caller thread until the operand becomes available";
+  let description = [{
+    The `async.runtime.await` operation blocks the caller thread until the
+    operand becomes available.
+  }];
+
+  let arguments = (ins Async_AnyAsyncType:$operand);
+  let assemblyFormat = "$operand attr-dict `:` type($operand)";
+}
+
+def Async_RuntimeResumeOp : Async_Op<"runtime.resume"> {
+  let summary = "resumes the coroutine on a thread managed by the runtime";
+  let description = [{
+    The `async.runtime.resume` operation resumes the coroutine on a thread
+    managed by the runtime.
+  }];
+
+  let arguments = (ins Async_CoroHandleType:$handle);
+  let assemblyFormat = "$handle attr-dict";
+}
+
+def Async_RuntimeAwaitAndResumeOp : Async_Op<"runtime.await_and_resume"> {
+  let summary = "awaits the async operand and resumes the coroutine";
+  let description = [{
+    The `async.runtime.await_and_resume` operation awaits for the operand to
+    become available and resumes the coroutine on a thread managed by the
+    runtime.
+  }];
+
+  let arguments = (ins Async_AnyAsyncType:$operand,
+                       Async_CoroHandleType:$handle);
+  let assemblyFormat = "$operand `,` $handle attr-dict `:` type($operand)";
+}
+
+def Async_RuntimeStoreOp : Async_Op<"runtime.store",
+      [TypesMatchWith<"type of 'value' matches element type of 'storage'",
+                     "storage", "value",
+                     "$_self.cast<ValueType>().getValueType()">]> {
+  let summary = "stores the value into the runtime async.value";
+  let description = [{
+    The `async.runtime.store` operation stores the value into the runtime
+    async.value storage.
+  }];
+
+  let arguments = (ins AnyType:$value,
+                       Async_AnyValueType:$storage);
+  let assemblyFormat = "$value `,` $storage attr-dict `:` type($storage)";
+}
+
+def Async_RuntimeLoadOp : Async_Op<"runtime.load",
+      [TypesMatchWith<"type of 'value' matches element type of 'storage'",
+                     "storage", "result",
+                     "$_self.cast<ValueType>().getValueType()">]> {
+  let summary = "loads the value from the runtime async.value";
+  let description = [{
+    The `async.runtime.load` operation loads the value from the runtime
+    async.value storage.
+  }];
+
+  let arguments = (ins Async_AnyValueType:$storage);
+  let results = (outs AnyType:$result);
+  let assemblyFormat = "$storage attr-dict `:` type($storage)";
+}
+
+def Async_RuntimeAddToGroupOp : Async_Op<"runtime.add_to_group", []> {
+  let summary = "adds and async token or value to the group";
+  let description = [{
+    The `async.runtime.add_to_group` adds an async token or value to the async
+    group. Returns the rank of the added element in the group.
+  }];
+
+  let arguments = (ins Async_AnyValueOrTokenType:$operand,
+                       Async_GroupType:$group);
+  let results = (outs Index:$rank);
+
+  let assemblyFormat = "$operand `,` $group attr-dict `:` type($operand)";
+}
+
 // All async values (values, tokens, groups) are reference counted at runtime
 // and automatically destructed when reference count drops to 0.
 //
 // All values are semantically created with a reference count of +1 and it is
 // the responsibility of the last async value user to drop reference count.
 //
 // Async values created when:
 //   1. Operation returns async result (e.g. the result of an `async.execute`).
 //   2. Async value passed in as a block argument.
 //
 // It is the responsibility of the async value user to extend the lifetime by
 // adding a +1 reference, if the reference counted value captured by the
 // asynchronously executed region (`async.execute` operation), and drop it after
 // the last nested use.
 //
 // Reference counting operations can be added to the IR using automatic
 // reference count pass, that relies on liveness analysis to find the last uses
 // of all reference counted values and automatically inserts
 // `drop_ref` operations.
 //
 // See `AsyncRefCountingPass` documentation for the implementation details.
 
-def Async_AddRefOp : Async_Op<"add_ref"> {
+def Async_RuntimeAddRefOp : Async_Op<"runtime.add_ref"> {
   let summary = "adds a reference to async value";
   let description = [{
-    The `async.add_ref` operation adds a reference(s) to async value (token,
-    value or group).
+    The `async.runtime.add_ref` operation adds a reference(s) to async value
+    (token, value or group).
   }];
 
   let arguments = (ins Async_AnyAsyncType:$operand,
                        Confined<I32Attr, [IntPositive]>:$count);
-  let results = (outs );
 
   let assemblyFormat = [{
     $operand attr-dict `:` type($operand)
   }];
 }
 
-def Async_DropRefOp : Async_Op<"drop_ref"> {
+def Async_RuntimeDropRefOp : Async_Op<"runtime.drop_ref"> {
   let summary = "drops a reference to async value";
   let description = [{
-    The `async.drop_ref` operation drops a reference(s) to async value (token,
-    value or group).
+    The `async.runtime.drop_ref` operation drops a reference(s) to async value
+    (token, value or group).
   }];
 
   let arguments = (ins Async_AnyAsyncType:$operand,
                        Confined<I32Attr, [IntPositive]>:$count);
-  let results = (outs );
 
   let assemblyFormat = [{
     $operand attr-dict `:` type($operand)
   }];
 }
 
 #endif // ASYNC_OPS

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

 /// async.value type, we use opaque pointers (!llvm.ptr<i8>) instead. After
 /// lowering all async data types become opaque pointers at runtime.
 struct AsyncAPI {
   // All async types are lowered to opaque i8* LLVM pointers at runtime.
   static LLVM::LLVMPointerType opaquePointerType(MLIRContext *ctx) {
     return LLVM::LLVMPointerType::get(IntegerType::get(ctx, 8));
   }
 
+  static LLVM::LLVMTokenType tokenType(MLIRContext *ctx) {
+    return LLVM::LLVMTokenType::get(ctx);
+  }
+
   static FunctionType addOrDropRefFunctionType(MLIRContext *ctx) {
     auto ref = opaquePointerType(ctx);
     auto count = IntegerType::get(ctx, 32);
     return FunctionType::get(ctx, {ref, count}, {});
   }
 
   static FunctionType createTokenFunctionType(MLIRContext *ctx) {
     return FunctionType::get(ctx, {}, {TokenType::get(ctx)});
addFuncDecl(kAwaitTokenAndExecute,
               AsyncAPI::awaitTokenAndExecuteFunctionType(ctx));
   addFuncDecl(kAwaitValueAndExecute,
               AsyncAPI::awaitValueAndExecuteFunctionType(ctx));
   addFuncDecl(kAwaitAllAndExecute,
               AsyncAPI::awaitAllAndExecuteFunctionType(ctx));
 }
 
 //===----------------------------------------------------------------------===//
-// LLVM coroutines intrinsics declarations.
+// Add malloc/free declarations to the module.
 //===----------------------------------------------------------------------===//
 
-static constexpr const char *kCoroId = "llvm.coro.id";
-static constexpr const char *kCoroSizeI64 = "llvm.coro.size.i64";
-static constexpr const char *kCoroBegin = "llvm.coro.begin";
-static constexpr const char *kCoroSave = "llvm.coro.save";
-static constexpr const char *kCoroSuspend = "llvm.coro.suspend";
-static constexpr const char *kCoroEnd = "llvm.coro.end";
-static constexpr const char *kCoroFree = "llvm.coro.free";
-static constexpr const char *kCoroResume = "llvm.coro.resume";
+static constexpr const char *kMalloc = "malloc";
+static constexpr const char *kFree = "free";
 
 static void addLLVMFuncDecl(ModuleOp module, ImplicitLocOpBuilder &builder,
                             StringRef name, Type ret, ArrayRef<Type> params) {
   if (module.lookupSymbol(name))
     return;
   Type type = LLVM::LLVMFunctionType::get(ret, params);
   builder.create<LLVM::LLVMFuncOp>(name, type);
 }
 
-/// Adds coroutine intrinsics declarations to the module.
-static void addCoroutineIntrinsicsDeclarations(ModuleOp module) {
-  using namespace mlir::LLVM;
-
-  MLIRContext *ctx = module.getContext();
-  ImplicitLocOpBuilder builder(module.getLoc(),
-                               module.getBody()->getTerminator());
-
-  auto token = LLVMTokenType::get(ctx);
-  auto voidTy = LLVMVoidType::get(ctx);
-
-  auto i8 = IntegerType::get(ctx, 8);
-  auto i1 = IntegerType::get(ctx, 1);
-  auto i32 = IntegerType::get(ctx, 32);
-  auto i64 = IntegerType::get(ctx, 64);
-  auto i8Ptr = LLVMPointerType::get(i8);
-
-  addLLVMFuncDecl(module, builder, kCoroId, token, {i32, i8Ptr, i8Ptr, i8Ptr});
-  addLLVMFuncDecl(module, builder, kCoroSizeI64, i64, {});
-  addLLVMFuncDecl(module, builder, kCoroBegin, i8Ptr, {token, i8Ptr});
-  addLLVMFuncDecl(module, builder, kCoroSave, token, {i8Ptr});
-  addLLVMFuncDecl(module, builder, kCoroSuspend, i8, {token, i1});
-  addLLVMFuncDecl(module, builder, kCoroEnd, i1, {i8Ptr, i1});
-  addLLVMFuncDecl(module, builder, kCoroFree, i8Ptr, {token, i8Ptr});
-  addLLVMFuncDecl(module, builder, kCoroResume, voidTy, {i8Ptr});
-}
-
-//===----------------------------------------------------------------------===//
-// Add malloc/free declarations to the module.
-//===----------------------------------------------------------------------===//
-
-static constexpr const char *kMalloc = "malloc";
-static constexpr const char *kFree = "free";
-
 /// Adds malloc/free declarations to the module.
 static void addCRuntimeDeclarations(ModuleOp module) {
   using namespace mlir::LLVM;
 
   MLIRContext *ctx = module.getContext();
   ImplicitLocOpBuilder builder(module.getLoc(),
                                module.getBody()->getTerminator());
 
static void addResumeFunction(ModuleOp module) {
 
   auto resumeOp = moduleBuilder.create<LLVM::LLVMFuncOp>(
       loc, kResume, LLVM::LLVMFunctionType::get(voidTy, {i8Ptr}));
   resumeOp.setPrivate();
 
   auto *block = resumeOp.addEntryBlock();
   auto blockBuilder = ImplicitLocOpBuilder::atBlockEnd(loc, block);
 
-  blockBuilder.create<LLVM::CallOp>(TypeRange(),
-                                    blockBuilder.getSymbolRefAttr(kCoroResume),
-                                    resumeOp.getArgument(0));
-
+  blockBuilder.create<LLVM::CoroResumeOp>(resumeOp.getArgument(0));
   blockBuilder.create<LLVM::ReturnOp>(ValueRange());
 }
 
 //===----------------------------------------------------------------------===//
 // async.execute op outlining to the coroutine functions.
 //===----------------------------------------------------------------------===//
 
 /// Function targeted for coroutine transformation has two additional blocks at
 /// the end: coroutine cleanup and coroutine suspension.
 ///
 /// async.await op lowering additionaly creates a resume block for each
 /// operation to enable non-blocking waiting via coroutine suspension.
 namespace {
 struct CoroMachinery {
   // Async execute region returns a completion token, and an async value for
   // each yielded value.
   //
   //   %token, %result = async.execute -> !async.value<T> {
   //     %0 = constant ... : T
   //     async.yield %0 : T
   //   }
   Value asyncToken; // token representing completion of the async region
   llvm::SmallVector<Value, 4> returnValues; // returned async values
 
-  Value coroHandle;
-  Block *cleanup;
-  Block *suspend;
+  Value coroHandle; // coroutine handle (!async.coro.handle value)
+  Block *cleanup;   // coroutine cleanup block
+  Block *suspend;   // coroutine suspension block
 };
 } // namespace
 
-/// Builds an coroutine template compatible with LLVM coroutines lowering.
+/// Builds an coroutine template compatible with LLVM coroutines switched-resume
+/// lowering using `async.runtime.*` and `async.coro.*` operations.
+///
+/// See LLVM coroutines documentation: https://llvm.org/docs/Coroutines.html
 ///
 ///  - `entry` block sets up the coroutine.
 ///  - `cleanup` block cleans up the coroutine state.
 ///  - `suspend block after the @llvm.coro.end() defines what value will be
 ///    returned to the initial caller of a coroutine. Everything before the
 ///    @llvm.coro.end() will be executed at every suspension point.
 ///
 /// Coroutine structure (only the important bits):
 ///
 ///   func @async_execute_fn(<function-arguments>)
 ///        -> (!async.token, !async.value<T>)
 ///   {
-///     ^entryBlock(<function-arguments>):
+///     ^entry(<function-arguments>):
 ///       %token = <async token> : !async.token    // create async runtime token
 ///       %value = <async value> : !async.value<T> // create async value
-///       %hdl = llvm.call @llvm.coro.id(...)      // create a coroutine handle
+///       %id = async.coro.id                      // create a coroutine id
+///       %hdl = async.coro.begin %id              // create a coroutine handle
 ///       br ^cleanup
 ///
 ///     ^cleanup:
-///       llvm.call @llvm.coro.free(...)  // delete coroutine state
+///       async.coro.free %hdl // delete the coroutine state
 ///       br ^suspend
 ///
 ///     ^suspend:
-///       llvm.call @llvm.coro.end(...)  // marks the end of a coroutine
+///       async.coro.end %hdl // marks the end of a coroutine
 ///       return %token, %value : !async.token, !async.value<T>
 ///   }
 ///
 /// The actual code for the async.execute operation body region will be inserted
 /// before the entry block terminator.
 ///
 ///
 static CoroMachinery setupCoroMachinery(FuncOp func) {
   assert(func.getBody().empty() && "Function must have empty body");
 
   MLIRContext *ctx = func.getContext();
-
-  auto token = LLVM::LLVMTokenType::get(ctx);
-  auto i1 = IntegerType::get(ctx, 1);
-  auto i32 = IntegerType::get(ctx, 32);
-  auto i64 = IntegerType::get(ctx, 64);
-  auto i8Ptr = LLVM::LLVMPointerType::get(IntegerType::get(ctx, 8));
-
   Block *entryBlock = func.addEntryBlock();
-  Location loc = func.getBody().getLoc();
 
-  auto builder = ImplicitLocOpBuilder::atBlockBegin(loc, entryBlock);
+  auto builder = ImplicitLocOpBuilder::atBlockBegin(func->getLoc(), entryBlock);
 
   // ------------------------------------------------------------------------ //
-  // Allocate async tokens/values that we will return from a ramp function.
+  // Allocate async token/values that we will return from a ramp function.
   // ------------------------------------------------------------------------ //
-  auto createToken = builder.create<CallOp>(kCreateToken, TokenType::get(ctx));
-
-  // Async value operands and results must be convertible to LLVM types. This is
-  // verified before the function outlining.
-  LLVMTypeConverter converter(ctx);
-
-  // Returns the size requirements for the async value storage.
-  // http://nondot.org/sabre/LLVMNotes/SizeOf-OffsetOf-VariableSizedStructs.txt
-  auto sizeOf = [&](ValueType valueType) -> Value {
-    auto storedType = converter.convertType(valueType.getValueType());
-    auto storagePtrType = LLVM::LLVMPointerType::get(storedType);
-
-    // %Size = getelementptr %T* null, int 1
-    // %SizeI = ptrtoint %T* %Size to i32
-    auto nullPtr = builder.create<LLVM::NullOp>(loc, storagePtrType);
-    auto one = builder.create<LLVM::ConstantOp>(loc, i32,
-                                                builder.getI32IntegerAttr(1));
-    auto gep = builder.create<LLVM::GEPOp>(loc, storagePtrType, nullPtr,
-                                           one.getResult());
-    return builder.create<LLVM::PtrToIntOp>(loc, i32, gep);
-  };
+  auto retToken = builder.create<RuntimeCreateOp>(TokenType::get(ctx)).result();
 
-  // We use the `async.value` type as a return type although it does not match
-  // the `kCreateValue` function signature, because it will be later lowered to
-  // the runtime type (opaque i8* pointer).
-  llvm::SmallVector<CallOp, 4> createValues;
-  for (auto resultType : func.getCallableResults().drop_front(1))
-    createValues.emplace_back(builder.create<CallOp>(
-        loc, kCreateValue, resultType, sizeOf(resultType.cast<ValueType>())));
-
-  auto createdValues = llvm::map_range(
-      createValues, [](CallOp call) { return call.getResult(0); });
-  llvm::SmallVector<Value, 4> returnValues(createdValues.begin(),
-                                           createdValues.end());
+  llvm::SmallVector<Value, 4> retValues;
+  for (auto resType : func.getCallableResults().drop_front())
+    retValues.emplace_back(builder.create<RuntimeCreateOp>(resType).result());
 
   // ------------------------------------------------------------------------ //
-  // Initialize coroutine: allocate frame, get coroutine handle.
+  // Initialize coroutine: get coroutine id and coroutine handle.
   // ------------------------------------------------------------------------ //
-
-  // Constants for initializing coroutine frame.
-  auto constZero =
-      builder.create<LLVM::ConstantOp>(i32, builder.getI32IntegerAttr(0));
-  auto constFalse =
-      builder.create<LLVM::ConstantOp>(i1, builder.getBoolAttr(false));
-  auto nullPtr = builder.create<LLVM::NullOp>(i8Ptr);
-
-  // Get coroutine id: @llvm.coro.id
-  auto coroId = builder.create<LLVM::CallOp>(
-      token, builder.getSymbolRefAttr(kCoroId),
-      ValueRange({constZero, nullPtr, nullPtr, nullPtr}));
-
-  // Get coroutine frame size: @llvm.coro.size.i64
-  auto coroSize = builder.create<LLVM::CallOp>(
-      i64, builder.getSymbolRefAttr(kCoroSizeI64), ValueRange());
-
-  // Allocate memory for coroutine frame.
-  auto coroAlloc =
-      builder.create<LLVM::CallOp>(i8Ptr, builder.getSymbolRefAttr(kMalloc),
-                                   ValueRange(coroSize.getResult(0)));
-
-  // Begin a coroutine: @llvm.coro.begin
-  auto coroHdl = builder.create<LLVM::CallOp>(
-      i8Ptr, builder.getSymbolRefAttr(kCoroBegin),
-      ValueRange({coroId.getResult(0), coroAlloc.getResult(0)}));
+  auto coroIdOp = builder.create<CoroIdOp>(CoroIdType::get(ctx));
+  auto coroHdlOp =
+      builder.create<CoroBeginOp>(CoroHandleType::get(ctx), coroIdOp.id());
 
   Block *cleanupBlock = func.addBlock();
   Block *suspendBlock = func.addBlock();
 
   // ------------------------------------------------------------------------ //
   // Coroutine cleanup block: deallocate coroutine frame, free the memory.
   // ------------------------------------------------------------------------ //
   builder.setInsertionPointToStart(cleanupBlock);
+  builder.create<CoroFreeOp>(coroIdOp.id(), coroHdlOp.handle());
 
-  // Get a pointer to the coroutine frame memory: @llvm.coro.free.
-  auto coroMem = builder.create<LLVM::CallOp>(
-      i8Ptr, builder.getSymbolRefAttr(kCoroFree),
-      ValueRange({coroId.getResult(0), coroHdl.getResult(0)}));
-
-  // Free the memory.
-  builder.create<LLVM::CallOp>(TypeRange(), builder.getSymbolRefAttr(kFree),
-                               ValueRange(coroMem.getResult(0)));
   // Branch into the suspend block.
   builder.create<BranchOp>(suspendBlock);
 
   // ------------------------------------------------------------------------ //
   // Coroutine suspend block: mark the end of a coroutine and return allocated
   // async token.
   // ------------------------------------------------------------------------ //
   builder.setInsertionPointToStart(suspendBlock);
 
-  // Mark the end of a coroutine: @llvm.coro.end.
-  builder.create<LLVM::CallOp>(i1, builder.getSymbolRefAttr(kCoroEnd),
-                               ValueRange({coroHdl.getResult(0), constFalse}));
+  // Mark the end of a coroutine: async.coro.end
+  builder.create<CoroEndOp>(coroHdlOp.handle());
 
   // Return created `async.token` and `async.values` from the suspend block.
   // This will be the return value of a coroutine ramp function.
-  SmallVector<Value, 4> ret{createToken.getResult(0)};
-  ret.insert(ret.end(), returnValues.begin(), returnValues.end());
-  builder.create<ReturnOp>(loc, ret);
+  SmallVector<Value, 4> ret{retToken};
+  ret.insert(ret.end(), retValues.begin(), retValues.end());
+  builder.create<ReturnOp>(ret);
 
   // Branch from the entry block to the cleanup block to create a valid CFG.
   builder.setInsertionPointToEnd(entryBlock);
-
   builder.create<BranchOp>(cleanupBlock);
 
   // `async.await` op lowering will create resume blocks for async
   // continuations, and will conditionally branch to cleanup or suspend blocks.
 
   CoroMachinery machinery;
-  machinery.asyncToken = createToken.getResult(0);
-  machinery.returnValues = returnValues;
-  machinery.coroHandle = coroHdl.getResult(0);
+  machinery.asyncToken = retToken;
+  machinery.returnValues = retValues;
+  machinery.coroHandle = coroHdlOp.handle();
   machinery.cleanup = cleanupBlock;
   machinery.suspend = suspendBlock;
   return machinery;
 }
 
-/// Add a LLVM coroutine suspension point to the end of suspended block, to
-/// resume execution in resume block. The caller is responsible for creating the
-/// two suspended/resume blocks with the desired ops contained in each block.
-/// This function merely provides the required control flow logic.
-///
-/// `coroState` must be a value returned from the call to @llvm.coro.save(...)
-/// intrinsic (saved coroutine state).
-///
-/// Before:
-///
-///   ^bb0:
-///     "opBefore"(...)
-///     "op"(...)
-///   ^cleanup: ...
-///   ^suspend: ...
-///   ^resume:
-///     "op"(...)
-///
-/// After:
-///
-///   ^bb0:
-///     "opBefore"(...)
-///     %suspend = llmv.call @llvm.coro.suspend(...)
-///     switch %suspend [-1: ^suspend, 0: ^resume, 1: ^cleanup]
-///   ^resume:
-///     "op"(...)
-///   ^cleanup: ...
-///   ^suspend: ...
-///
-static void addSuspensionPoint(CoroMachinery coro, Value coroState,
-                               Operation *op, Block *suspended, Block *resume,
-                               OpBuilder &builder) {
-  Location loc = op->getLoc();
-  MLIRContext *ctx = op->getContext();
-  auto i1 = IntegerType::get(ctx, 1);
-  auto i8 = IntegerType::get(ctx, 8);
-
-  // Add a coroutine suspension in place of original `op` in the split block.
-  OpBuilder::InsertionGuard guard(builder);
-  builder.setInsertionPointToEnd(suspended);
-
-  auto constFalse =
-      builder.create<LLVM::ConstantOp>(loc, i1, builder.getBoolAttr(false));
-
-  // Suspend a coroutine: @llvm.coro.suspend
-  auto coroSuspend = builder.create<LLVM::CallOp>(
-      loc, i8, builder.getSymbolRefAttr(kCoroSuspend),
-      ValueRange({coroState, constFalse}));
-
-  // After a suspension point decide if we should branch into resume, cleanup
-  // or suspend block of the coroutine (see @llvm.coro.suspend return code
-  // documentation).
-  auto constZero =
-      builder.create<LLVM::ConstantOp>(loc, i8, builder.getI8IntegerAttr(0));
-  auto constNegOne =
-      builder.create<LLVM::ConstantOp>(loc, i8, builder.getI8IntegerAttr(-1));
-
-  Block *resumeOrCleanup = builder.createBlock(resume);
-
-  // Suspend the coroutine ...?
-  builder.setInsertionPointToEnd(suspended);
-  auto isNegOne = builder.create<LLVM::ICmpOp>(
-      loc, LLVM::ICmpPredicate::eq, coroSuspend.getResult(0), constNegOne);
-  builder.create<LLVM::CondBrOp>(loc, isNegOne, /*trueDest=*/coro.suspend,
-                                 /*falseDest=*/resumeOrCleanup);
-
-  // ... or resume or cleanup the coroutine?
-  builder.setInsertionPointToStart(resumeOrCleanup);
-  auto isZero = builder.create<LLVM::ICmpOp>(
-      loc, LLVM::ICmpPredicate::eq, coroSuspend.getResult(0), constZero);
-  builder.create<LLVM::CondBrOp>(loc, isZero, /*trueDest=*/resume,
-                                 /*falseDest=*/coro.cleanup);
-}
-
 /// Outline the body region attached to the `async.execute` op into a standalone
 /// function.
 ///
 /// Note that this is not reversible transformation.
 static std::pair<FuncOp, CoroMachinery>
 outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
   ModuleOp module = execute->getParentOfType<ModuleOp>();
 
outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
   // TODO: Derive outlined function name from the parent FuncOp (support
   // multiple nested async.execute operations).
   FuncOp func = FuncOp::create(loc, kAsyncFnPrefix, funcType, funcAttrs);
   symbolTable.insert(func, Block::iterator(module.getBody()->getTerminator()));
 
   SymbolTable::setSymbolVisibility(func, SymbolTable::Visibility::Private);
 
   // Prepare a function for coroutine lowering by adding entry/cleanup/suspend
-  // blocks, adding llvm.coro instrinsics and setting up control flow.
+  // blocks, adding async.coro operations and setting up control flow.
   CoroMachinery coro = setupCoroMachinery(func);
 
   // Suspend async function at the end of an entry block, and resume it using
-  // Async execute API (execution will be resumed in a thread managed by the
-  // async runtime).
+  // Async resume operation (execution will be resumed in a thread managed by
+  // the async runtime).
   Block *entryBlock = &func.getBlocks().front();
   auto builder = ImplicitLocOpBuilder::atBlockTerminator(loc, entryBlock);
 
-  // A pointer to coroutine resume intrinsic wrapper.
-  auto resumeFnTy = AsyncAPI::resumeFunctionType(ctx);
-  auto resumePtr = builder.create<LLVM::AddressOfOp>(
-      LLVM::LLVMPointerType::get(resumeFnTy), kResume);
-
-  // Save the coroutine state: @llvm.coro.save
-  auto coroSave = builder.create<LLVM::CallOp>(
-      LLVM::LLVMTokenType::get(ctx), builder.getSymbolRefAttr(kCoroSave),
-      ValueRange({coro.coroHandle}));
+  // Save the coroutine state: async.coro.save
+  auto coroSaveOp =
+      builder.create<CoroSaveOp>(CoroStateType::get(ctx), coro.coroHandle);
 
-  // Call async runtime API to execute a coroutine in the managed thread.
-  SmallVector<Value, 2> executeArgs = {coro.coroHandle, resumePtr.res()};
-  builder.create<CallOp>(TypeRange(), kExecute, executeArgs);
+  // Pass coroutine to the runtime to be resumed on a runtime managed thread.
+  builder.create<RuntimeResumeOp>(coro.coroHandle);
 
-  // Split the entry block before the terminator.
+  // Split the entry block before the terminator (branch to suspend block).
   auto *terminatorOp = entryBlock->getTerminator();
   Block *suspended = terminatorOp->getBlock();
   Block *resume = suspended->splitBlock(terminatorOp);
-  addSuspensionPoint(coro, coroSave.getResult(0), terminatorOp, suspended,
-                     resume, builder);
+
+  // Add async.coro.suspend as a suspended block terminator.
+  builder.setInsertionPointToEnd(suspended);
+  builder.create<CoroSuspendOp>(coroSaveOp.state(), coro.suspend, resume,
+                                coro.cleanup);
 
   size_t numDependencies = execute.dependencies().size();
   size_t numOperands = execute.operands().size();
 
   // Await on all dependencies before starting to execute the body region.
   builder.setInsertionPointToStart(resume);
   for (size_t i = 0; i < numDependencies; ++i)
     builder.create<AwaitOp>(func.getArgument(i));
outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
   return {func, coro};
 }
 
 //===----------------------------------------------------------------------===//
 // Convert Async dialect types to LLVM types.
 //===----------------------------------------------------------------------===//
 
 namespace {
-
 /// AsyncRuntimeTypeConverter only converts types from the Async dialect to
 /// their runtime type (opaque pointers) and does not convert any other types.
 class AsyncRuntimeTypeConverter : public TypeConverter {
 public:
   AsyncRuntimeTypeConverter() {
     addConversion([](Type type) { return type; });
     addConversion(convertAsyncTypes);
   }
 
   static Optional<Type> convertAsyncTypes(Type type) {
     if (type.isa<TokenType, GroupType, ValueType>())
       return AsyncAPI::opaquePointerType(type.getContext());
+
+    if (type.isa<CoroIdType, CoroStateType>())
+      return AsyncAPI::tokenType(type.getContext());
+    if (type.isa<CoroHandleType>())
+      return AsyncAPI::opaquePointerType(type.getContext());
+
     return llvm::None;
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// Convert return operations that return async values from async regions.
+// Convert async.coro.id to @llvm.coro.id intrinsic.
 //===----------------------------------------------------------------------===//
 
 namespace {
-class ReturnOpOpConversion : public ConversionPattern {
+class CoroIdOpConversion : public OpConversionPattern<CoroIdOp> {
 public:
-  explicit ReturnOpOpConversion(TypeConverter &converter, MLIRContext *ctx)
-      : ConversionPattern(ReturnOp::getOperationName(), 1, converter, ctx) {}
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+  matchAndRewrite(CoroIdOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<ReturnOp>(op, operands);
+    auto token = AsyncAPI::tokenType(op->getContext());
+    auto i8Ptr = AsyncAPI::opaquePointerType(op->getContext());
+    auto loc = op->getLoc();
+
+    // Constants for initializing coroutine frame.
+    auto constZero = rewriter.create<LLVM::ConstantOp>(
+        loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(0));
+    auto nullPtr = rewriter.create<LLVM::NullOp>(loc, i8Ptr);
+
+    // Get coroutine id: @llvm.coro.id.
+    rewriter.replaceOpWithNewOp<LLVM::CoroIdOp>(
+        op, token, ValueRange({constZero, nullPtr, nullPtr, nullPtr}));
+
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Convert async.coro.begin to @llvm.coro.begin intrinsic.
+//===----------------------------------------------------------------------===//
+
+namespace {
+class CoroBeginOpConversion : public OpConversionPattern<CoroBeginOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(CoroBeginOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto i8Ptr = AsyncAPI::opaquePointerType(op->getContext());
+    auto loc = op->getLoc();
+
+    // Get coroutine frame size: @llvm.coro.size.i64.
+    auto coroSize =
+        rewriter.create<LLVM::CoroSizeOp>(loc, rewriter.getI64Type());
+
+    // Allocate memory for the coroutine frame.
+    auto coroAlloc = rewriter.create<LLVM::CallOp>(
+        loc, i8Ptr, rewriter.getSymbolRefAttr(kMalloc),
+        ValueRange(coroSize.getResult()));
+
+    // Begin a coroutine: @llvm.coro.begin.
+    auto coroId = CoroBeginOpAdaptor(operands).id();
+    rewriter.replaceOpWithNewOp<LLVM::CoroBeginOp>(
+        op, i8Ptr, ValueRange({coroId, coroAlloc.getResult(0)}));
+
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Convert async.coro.free to @llvm.coro.free intrinsic.
+//===----------------------------------------------------------------------===//
+
+namespace {
+class CoroFreeOpConversion : public OpConversionPattern<CoroFreeOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(CoroFreeOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto i8Ptr = AsyncAPI::opaquePointerType(op->getContext());
+    auto loc = op->getLoc();
+
+    // Get a pointer to the coroutine frame memory: @llvm.coro.free.
+    auto coroMem = rewriter.create<LLVM::CoroFreeOp>(loc, i8Ptr, operands);
+
+    // Free the memory.
+    rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, TypeRange(),
+                                              rewriter.getSymbolRefAttr(kFree),
+                                              ValueRange(coroMem.getResult()));
+
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Convert async.coro.end to @llvm.coro.end intrinsic.
+//===----------------------------------------------------------------------===//
+
+namespace {
+class CoroEndOpConversion : public OpConversionPattern<CoroEndOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(CoroEndOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    // We are not in the block that is part of the unwind sequence.
+    auto constFalse = rewriter.create<LLVM::ConstantOp>(
+        op->getLoc(), rewriter.getI1Type(), rewriter.getBoolAttr(false));
+
+    // Mark the end of a coroutine: @llvm.coro.end.
+    auto coroHdl = CoroEndOpAdaptor(operands).handle();
+    rewriter.create<LLVM::CoroEndOp>(op->getLoc(), rewriter.getI1Type(),
+                                     ValueRange({coroHdl, constFalse}));
+    rewriter.eraseOp(op);
+
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Convert async.coro.save to @llvm.coro.save intrinsic.
+//===----------------------------------------------------------------------===//
+
+namespace {
+class CoroSaveOpConversion : public OpConversionPattern<CoroSaveOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(CoroSaveOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Save the coroutine state: @llvm.coro.save
+    rewriter.replaceOpWithNewOp<LLVM::CoroSaveOp>(
+        op, AsyncAPI::tokenType(op->getContext()), operands);
+
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Convert async.coro.suspend to @llvm.coro.suspend intrinsic.
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// Convert async.coro.suspend to the @llvm.coro.suspend intrinsic call, and
+/// branch to the appropriate block based on the return code.
+///
+/// Before:
+///
+///   ^suspended:
+///     "opBefore"(...)
+///     async.coro.suspend %state, ^suspend, ^resume, ^cleanup
+///   ^resume:
+///     "op"(...)
+///   ^cleanup: ...
+///   ^suspend: ...
+///
+/// After:
+///
+///   ^suspended:
+///     "opBefore"(...)
+///     %suspend = llmv.intr.coro.suspend ...
+///     switch %suspend [-1: ^suspend, 0: ^resume, 1: ^cleanup]
+///   ^resume:
+///     "op"(...)
+///   ^cleanup: ...
+///   ^suspend: ...
+///
+class CoroSuspendOpConversion : public OpConversionPattern<CoroSuspendOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(CoroSuspendOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto i8 = rewriter.getIntegerType(8);
+    auto i32 = rewriter.getI32Type();
+    auto loc = op->getLoc();
+
+    // This is not a final suspension point.
+    auto constFalse = rewriter.create<LLVM::ConstantOp>(
+        loc, rewriter.getI1Type(), rewriter.getBoolAttr(false));
+
+    // Suspend a coroutine: @llvm.coro.suspend
+    auto coroState = CoroSuspendOpAdaptor(operands).state();
+    auto coroSuspend = rewriter.create<LLVM::CoroSuspendOp>(
+        loc, i8, ValueRange({coroState, constFalse}));
+
+    // Cast return code to i32.
+
+    // After a suspension point decide if we should branch into resume, cleanup
+    // or suspend block of the coroutine (see @llvm.coro.suspend return code
+    // documentation).
+    llvm::SmallVector<int32_t, 2> caseValues = {0, 1};
+    llvm::SmallVector<Block *, 2> caseDest = {op.resumeDest(),
+                                              op.cleanupDest()};
+    rewriter.replaceOpWithNewOp<LLVM::SwitchOp>(
+        op, rewriter.create<LLVM::SExtOp>(loc, i32, coroSuspend.getResult()),
+        /*defaultDestination=*/op.suspendDest(),
+        /*defaultOperands=*/ValueRange(),
+        /*caseValues=*/caseValues,
+        /*caseDestinations=*/caseDest,
+        /*caseOperands=*/ArrayRef<ValueRange>(),
+        /*branchWeights=*/ArrayRef<int32_t>());
+
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Convert async.runtime.create to the corresponding runtime API call.
+//
+// To allocate storage for the async values we use getelementptr trick:
+// http://nondot.org/sabre/LLVMNotes/SizeOf-OffsetOf-VariableSizedStructs.txt
+//===----------------------------------------------------------------------===//
+
+namespace {
+class RuntimeCreateOpLowering : public OpConversionPattern<RuntimeCreateOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(RuntimeCreateOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    TypeConverter *converter = getTypeConverter();
+    Type resultType = op->getResultTypes()[0];
+
+    // Tokens and Groups lowered to function calls without arguments.
+    if (resultType.isa<TokenType>() || resultType.isa<GroupType>()) {
+      rewriter.replaceOpWithNewOp<CallOp>(
+          op, resultType.isa<TokenType>() ? kCreateToken : kCreateGroup,
+          converter->convertType(resultType));
+      return success();
+    }
+
+    // To create a value we need to compute the storage requirement.
+    if (auto value = resultType.dyn_cast<ValueType>()) {
+      // Returns the size requirements for the async value storage.
+      auto sizeOf = [&](ValueType valueType) -> Value {
+        auto loc = op->getLoc();
+        auto i32 = rewriter.getI32Type();
+
+        auto storedType = converter->convertType(valueType.getValueType());
+        auto storagePtrType = LLVM::LLVMPointerType::get(storedType);
+
+        // %Size = getelementptr %T* null, int 1
+        // %SizeI = ptrtoint %T* %Size to i32
+        auto nullPtr = rewriter.create<LLVM::NullOp>(loc, storagePtrType);
+        auto one = rewriter.create<LLVM::ConstantOp>(
+            loc, i32, rewriter.getI32IntegerAttr(1));
+        auto gep = rewriter.create<LLVM::GEPOp>(loc, storagePtrType, nullPtr,
+                                                one.getResult());
+        return rewriter.create<LLVM::PtrToIntOp>(loc, i32, gep);
+      };
+
+      rewriter.replaceOpWithNewOp<CallOp>(op, kCreateValue, resultType,
+                                          sizeOf(value));
+
+      return success();
+    }
+
+    return rewriter.notifyMatchFailure(op, "unsupported async type");
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Convert async.runtime.set_available to the corresponding runtime API call.
+//===----------------------------------------------------------------------===//
+
+namespace {
+class RuntimeSetAvailableOpLowering
+    : public OpConversionPattern<RuntimeSetAvailableOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(RuntimeSetAvailableOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type operandType = op.operand().getType();
+
+    if (operandType.isa<TokenType>() || operandType.isa<ValueType>()) {
+      rewriter.create<CallOp>(op->getLoc(),
+                              operandType.isa<TokenType>() ? kEmplaceToken
+                                                           : kEmplaceValue,
+                              TypeRange(), operands);
+      rewriter.eraseOp(op);
+      return success();
+    }
+
+    return rewriter.notifyMatchFailure(op, "unsupported async type");
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Convert async.runtime.await to the corresponding runtime API call.
+//===----------------------------------------------------------------------===//
+
+namespace {
+class RuntimeAwaitOpLowering : public OpConversionPattern<RuntimeAwaitOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(RuntimeAwaitOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type operandType = op.operand().getType();
+
+    StringRef apiFuncName;
+    if (operandType.isa<TokenType>())
+      apiFuncName = kAwaitToken;
+    else if (operandType.isa<ValueType>())
+      apiFuncName = kAwaitValue;
+    else if (operandType.isa<GroupType>())
+      apiFuncName = kAwaitGroup;
+    else
+      return rewriter.notifyMatchFailure(op, "unsupported async type");
+
+    rewriter.create<CallOp>(op->getLoc(), apiFuncName, TypeRange(), operands);
+    rewriter.eraseOp(op);
+
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Convert async.runtime.await_and_resume to the corresponding runtime API call.
+//===----------------------------------------------------------------------===//
+
+namespace {
+class RuntimeAwaitAndResumeOpLowering
+    : public OpConversionPattern<RuntimeAwaitAndResumeOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(RuntimeAwaitAndResumeOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type operandType = op.operand().getType();
+
+    StringRef apiFuncName;
+    if (operandType.isa<TokenType>())
+      apiFuncName = kAwaitTokenAndExecute;
+    else if (operandType.isa<ValueType>())
+      apiFuncName = kAwaitValueAndExecute;
+    else if (operandType.isa<GroupType>())
+      apiFuncName = kAwaitAllAndExecute;
+    else
+      return rewriter.notifyMatchFailure(op, "unsupported async type");
+
+    Value operand = RuntimeAwaitAndResumeOpAdaptor(operands).operand();
+    Value handle = RuntimeAwaitAndResumeOpAdaptor(operands).handle();
+
+    // A pointer to coroutine resume intrinsic wrapper.
+    auto resumeFnTy = AsyncAPI::resumeFunctionType(op->getContext());
+    auto resumePtr = rewriter.create<LLVM::AddressOfOp>(
+        op->getLoc(), LLVM::LLVMPointerType::get(resumeFnTy), kResume);
+
+    rewriter.create<CallOp>(op->getLoc(), apiFuncName, TypeRange(),
+                            ValueRange({operand, handle, resumePtr.res()}));
+    rewriter.eraseOp(op);
+
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// Async reference counting ops lowering (`async.add_ref` and `async.drop_ref`
-// to the corresponding API calls).
+// Convert async.runtime.resume to the corresponding runtime API call.
 //===----------------------------------------------------------------------===//
 
 namespace {
+class RuntimeResumeOpLowering : public OpConversionPattern<RuntimeResumeOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(RuntimeResumeOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    // A pointer to coroutine resume intrinsic wrapper.
+    auto resumeFnTy = AsyncAPI::resumeFunctionType(op->getContext());
+    auto resumePtr = rewriter.create<LLVM::AddressOfOp>(
+        op->getLoc(), LLVM::LLVMPointerType::get(resumeFnTy), kResume);
+
+    // Call async runtime API to execute a coroutine in the managed thread.
+    auto coroHdl = RuntimeResumeOpAdaptor(operands).handle();
+    rewriter.replaceOpWithNewOp<CallOp>(op, TypeRange(), kExecute,
+                                        ValueRange({coroHdl, resumePtr.res()}));
+
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Convert async.runtime.store to the corresponding runtime API call.
+//===----------------------------------------------------------------------===//
+
+namespace {
+class RuntimeStoreOpLowering : public OpConversionPattern<RuntimeStoreOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(RuntimeStoreOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op->getLoc();
+
+    // Get a pointer to the async value storage from the runtime.
+    auto i8Ptr = AsyncAPI::opaquePointerType(rewriter.getContext());
+    auto storage = RuntimeStoreOpAdaptor(operands).storage();
+    auto storagePtr = rewriter.create<CallOp>(loc, kGetValueStorage,
+                                              TypeRange(i8Ptr), storage);
+
+    // Cast from i8* to the LLVM pointer type.
+    auto valueType = op.value().getType();
+    auto llvmValueType = getTypeConverter()->convertType(valueType);
+    auto castedStoragePtr = rewriter.create<LLVM::BitcastOp>(
+        loc, LLVM::LLVMPointerType::get(llvmValueType),
+        storagePtr.getResult(0));
+
+    // Store the yielded value into the async value storage.
+    auto value = RuntimeStoreOpAdaptor(operands).value();
+    rewriter.create<LLVM::StoreOp>(loc, value, castedStoragePtr.getResult());
+
+    // Erase the original runtime store operation.
+    rewriter.eraseOp(op);
+
+    return success();
+  }
+};
+} // namespace
 
+//===----------------------------------------------------------------------===//
+// Convert async.runtime.load to the corresponding runtime API call.
+//===----------------------------------------------------------------------===//
+
+namespace {
+class RuntimeLoadOpLowering : public OpConversionPattern<RuntimeLoadOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(RuntimeLoadOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op->getLoc();
+
+    // Get a pointer to the async value storage from the runtime.
+    auto i8Ptr = AsyncAPI::opaquePointerType(rewriter.getContext());
+    auto storage = RuntimeLoadOpAdaptor(operands).storage();
+    auto storagePtr = rewriter.create<CallOp>(loc, kGetValueStorage,
+                                              TypeRange(i8Ptr), storage);
+
+    // Cast from i8* to the LLVM pointer type.
+    auto valueType = op.result().getType();
+    auto llvmValueType = getTypeConverter()->convertType(valueType);
+    auto castedStoragePtr = rewriter.create<LLVM::BitcastOp>(
+        loc, LLVM::LLVMPointerType::get(llvmValueType),
+        storagePtr.getResult(0));
+
+    // Load from the casted pointer.
+    rewriter.replaceOpWithNewOp<LLVM::LoadOp>(op, castedStoragePtr.getResult());
+
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Convert async.runtime.add_to_group to the corresponding runtime API call.
+//===----------------------------------------------------------------------===//
+
+namespace {
+class RuntimeAddToGroupOpLowering
+    : public OpConversionPattern<RuntimeAddToGroupOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(RuntimeAddToGroupOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Currently we can only add tokens to the group.
+    if (!op.operand().getType().isa<TokenType>())
+      return rewriter.notifyMatchFailure(op, "only token type is supported");
+
+    // Replace with a runtime API function call.
+    rewriter.replaceOpWithNewOp<CallOp>(op, kAddTokenToGroup,
+                                        rewriter.getI64Type(), operands);
+
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Async reference counting ops lowering (`async.runtime.add_ref` and
+// `async.runtime.drop_ref` to the corresponding API calls).
+//===----------------------------------------------------------------------===//
+
+namespace {
 template <typename RefCountingOp>
-class RefCountingOpLowering : public ConversionPattern {
+class RefCountingOpLowering : public OpConversionPattern<RefCountingOp> {
 public:
   explicit RefCountingOpLowering(TypeConverter &converter, MLIRContext *ctx,
                                  StringRef apiFunctionName)
-      : ConversionPattern(RefCountingOp::getOperationName(), 1, converter, ctx),
+      : OpConversionPattern<RefCountingOp>(converter, ctx),
         apiFunctionName(apiFunctionName) {}
 
   LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+  matchAndRewrite(RefCountingOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    RefCountingOp refCountingOp = cast<RefCountingOp>(op);
-
-    auto count = rewriter.create<ConstantOp>(
-        op->getLoc(), rewriter.getI32Type(),
-        rewriter.getI32IntegerAttr(refCountingOp.count()));
+    auto count =
+        rewriter.create<ConstantOp>(op->getLoc(), rewriter.getI32Type(),
+                                    rewriter.getI32IntegerAttr(op.count()));
 
+    auto operand = typename RefCountingOp::Adaptor(operands).operand();
     rewriter.replaceOpWithNewOp<CallOp>(op, TypeRange(), apiFunctionName,
-                                        ValueRange({operands[0], count}));
+                                        ValueRange({operand, count}));
 
     return success();
   }
 
 private:
   StringRef apiFunctionName;
 };
 
-/// async.drop_ref op lowering to mlirAsyncRuntimeDropRef function call.
-class AddRefOpLowering : public RefCountingOpLowering<AddRefOp> {
+class RuntimeAddRefOpLowering : public RefCountingOpLowering<RuntimeAddRefOp> {
 public:
-  explicit AddRefOpLowering(TypeConverter &converter, MLIRContext *ctx)
+  explicit RuntimeAddRefOpLowering(TypeConverter &converter, MLIRContext *ctx)
       : RefCountingOpLowering(converter, ctx, kAddRef) {}
 };
 
-/// async.create_group op lowering to mlirAsyncRuntimeCreateGroup function call.
-class DropRefOpLowering : public RefCountingOpLowering<DropRefOp> {
+class RuntimeDropRefOpLowering
+    : public RefCountingOpLowering<RuntimeDropRefOp> {
 public:
-  explicit DropRefOpLowering(TypeConverter &converter, MLIRContext *ctx)
+  explicit RuntimeDropRefOpLowering(TypeConverter &converter, MLIRContext *ctx)
       : RefCountingOpLowering(converter, ctx, kDropRef) {}
 };
-
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// async.create_group op lowering to mlirAsyncRuntimeCreateGroup function call.
+// Convert return operations that return async values from async regions.
 //===----------------------------------------------------------------------===//
 
 namespace {
-class CreateGroupOpLowering : public ConversionPattern {
+class ReturnOpOpConversion : public OpConversionPattern<ReturnOp> {
 public:
-  explicit CreateGroupOpLowering(TypeConverter &converter, MLIRContext *ctx)
-      : ConversionPattern(CreateGroupOp::getOperationName(), 1, converter,
-                          ctx) {}
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+  matchAndRewrite(ReturnOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    auto retTy = GroupType::get(op->getContext());
-    rewriter.replaceOpWithNewOp<CallOp>(op, kCreateGroup, retTy);
+    rewriter.replaceOpWithNewOp<ReturnOp>(op, operands);
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// async.add_to_group op lowering to runtime function call.
+// Convert async.create_group operation to async.runtime.create
 //===----------------------------------------------------------------------===//
 
 namespace {
-class AddToGroupOpLowering : public ConversionPattern {
+class CreateGroupOpLowering : public OpConversionPattern<CreateGroupOp> {
 public:
-  explicit AddToGroupOpLowering(TypeConverter &converter, MLIRContext *ctx)
-      : ConversionPattern(AddToGroupOp::getOperationName(), 1, converter, ctx) {
-  }
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+  matchAndRewrite(CreateGroupOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    // Currently we can only add tokens to the group.
-    auto addToGroup = cast<AddToGroupOp>(op);
-    if (!addToGroup.operand().getType().isa<TokenType>())
-      return failure();
-
-    auto i64 = IntegerType::get(op->getContext(), 64);
-    rewriter.replaceOpWithNewOp<CallOp>(op, kAddTokenToGroup, i64, operands);
+    rewriter.replaceOpWithNewOp<RuntimeCreateOp>(
+        op, GroupType::get(op->getContext()));
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// async.await and async.await_all op lowerings to the corresponding async
-// runtime function calls.
+// Convert async.add_to_group operation to async.runtime.add_to_group.
 //===----------------------------------------------------------------------===//
 
 namespace {
+class AddToGroupOpLowering : public OpConversionPattern<AddToGroupOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
 
+  LogicalResult
+  matchAndRewrite(AddToGroupOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<RuntimeAddToGroupOp>(
+        op, rewriter.getIndexType(), operands);
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Convert async.await and async.await_all operations to the async.runtime.await
+// or async.runtime.await_and_resume operations.
+//===----------------------------------------------------------------------===//
+
+namespace {
 template <typename AwaitType, typename AwaitableType>
-class AwaitOpLoweringBase : public ConversionPattern {
-protected:
-  explicit AwaitOpLoweringBase(
-      TypeConverter &converter, MLIRContext *ctx,
-      const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions,
-      StringRef blockingAwaitFuncName, StringRef coroAwaitFuncName)
-      : ConversionPattern(AwaitType::getOperationName(), 1, converter, ctx),
-        outlinedFunctions(outlinedFunctions),
-        blockingAwaitFuncName(blockingAwaitFuncName),
-        coroAwaitFuncName(coroAwaitFuncName) {}
+class AwaitOpLoweringBase : public OpConversionPattern<AwaitType> {
+  using AwaitAdaptor = typename AwaitType::Adaptor;
 
 public:
+  AwaitOpLoweringBase(
+      MLIRContext *ctx,
+      const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions)
+      : OpConversionPattern<AwaitType>(ctx),
+        outlinedFunctions(outlinedFunctions) {}
+
   LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+  matchAndRewrite(AwaitType op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     // We can only await on one the `AwaitableType` (for `await` it can be
     // a `token` or a `value`, for `await_all` it must be a `group`).
-    auto await = cast<AwaitType>(op);
-    if (!await.operand().getType().template isa<AwaitableType>())
-      return failure();
+    if (!op.operand().getType().template isa<AwaitableType>())
+      return rewriter.notifyMatchFailure(op, "unsupported awaitable type");
 
     // Check if await operation is inside the outlined coroutine function.
-    auto func = await->template getParentOfType<FuncOp>();
+    auto func = op->template getParentOfType<FuncOp>();
     auto outlined = outlinedFunctions.find(func);
     const bool isInCoroutine = outlined != outlinedFunctions.end();
 
     Location loc = op->getLoc();
+    Value operand = AwaitAdaptor(operands).operand();
 
-    // Inside regular function we convert await operation to the blocking
-    // async API await function call.
+    // Inside regular functions we use the blocking wait operation to wait for
+    // the async object (token, value or group) to become available.
     if (!isInCoroutine)
-      rewriter.create<CallOp>(loc, TypeRange(), blockingAwaitFuncName,
-                              ValueRange(operands[0]));
+      rewriter.create<RuntimeAwaitOp>(loc, operand);
 
     // Inside the coroutine we convert await operation into coroutine suspension
     // point, and resume execution asynchronously.
     if (isInCoroutine) {
       const CoroMachinery &coro = outlined->getSecond();
+      Block *suspended = op->getBlock();
 
       ImplicitLocOpBuilder builder(loc, op, rewriter.getListener());
       MLIRContext *ctx = op->getContext();
 
-      // A pointer to coroutine resume intrinsic wrapper.
-      auto resumeFnTy = AsyncAPI::resumeFunctionType(ctx);
-      auto resumePtr = builder.create<LLVM::AddressOfOp>(
-          LLVM::LLVMPointerType::get(resumeFnTy), kResume);
-
-      // Save the coroutine state: @llvm.coro.save
-      auto coroSave = builder.create<LLVM::CallOp>(
-          LLVM::LLVMTokenType::get(ctx), builder.getSymbolRefAttr(kCoroSave),
-          ValueRange(coro.coroHandle));
-
-      // Call async runtime API to resume a coroutine in the managed thread when
-      // the async await argument becomes ready.
-      SmallVector<Value, 3> awaitAndExecuteArgs = {operands[0], coro.coroHandle,
-                                                   resumePtr.res()};
-      builder.create<CallOp>(TypeRange(), coroAwaitFuncName,
-                             awaitAndExecuteArgs);
-
-      Block *suspended = op->getBlock();
+      // Save the coroutine state and resume on a runtime managed thread when
+      // the operand becomes available.
+      auto coroSaveOp =
+          builder.create<CoroSaveOp>(CoroStateType::get(ctx), coro.coroHandle);
+      builder.create<RuntimeAwaitAndResumeOp>(operand, coro.coroHandle);
 
       // Split the entry block before the await operation.
       Block *resume = rewriter.splitBlock(suspended, Block::iterator(op));
-      addSuspensionPoint(coro, coroSave.getResult(0), op, suspended, resume,
-                         builder);
+
+      // Add async.coro.suspend as a suspended block terminator.
+      builder.setInsertionPointToEnd(suspended);
+      builder.create<CoroSuspendOp>(coroSaveOp.state(), coro.suspend, resume,
+                                    coro.cleanup);
 
       // Make sure that replacement value will be constructed in resume block.
       rewriter.setInsertionPointToStart(resume);
     }
 
-    // Replace or erase the await operation with the new value.
-    if (Value replaceWith = getReplacementValue(op, operands[0], rewriter))
+    // Erase or replace the await operation with the new value.
+    if (Value replaceWith = getReplacementValue(op, operand, rewriter))
       rewriter.replaceOp(op, replaceWith);
     else
       rewriter.eraseOp(op);
 
     return success();
   }
 
-  virtual Value getReplacementValue(Operation *op, Value operand,
+  virtual Value getReplacementValue(AwaitType op, Value operand,
                                     ConversionPatternRewriter &rewriter) const {
     return Value();
   }
 
 private:
   const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions;
-  StringRef blockingAwaitFuncName;
-  StringRef coroAwaitFuncName;
 };
 
 /// Lowering for `async.await` with a token operand.
 class AwaitTokenOpLowering : public AwaitOpLoweringBase<AwaitOp, TokenType> {
   using Base = AwaitOpLoweringBase<AwaitOp, TokenType>;
 
 public:
-  explicit AwaitTokenOpLowering(
-      TypeConverter &converter, MLIRContext *ctx,
-      const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions)
-      : Base(converter, ctx, outlinedFunctions, kAwaitToken,
-             kAwaitTokenAndExecute) {}
+  using Base::Base;
 };
 
 /// Lowering for `async.await` with a value operand.
 class AwaitValueOpLowering : public AwaitOpLoweringBase<AwaitOp, ValueType> {
   using Base = AwaitOpLoweringBase<AwaitOp, ValueType>;
 
 public:
-  explicit AwaitValueOpLowering(
-      TypeConverter &converter, MLIRContext *ctx,
-      const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions)
-      : Base(converter, ctx, outlinedFunctions, kAwaitValue,
-             kAwaitValueAndExecute) {}
+  using Base::Base;
 
   Value
-  getReplacementValue(Operation *op, Value operand,
+  getReplacementValue(AwaitOp op, Value operand,
                       ConversionPatternRewriter &rewriter) const override {
-    Location loc = op->getLoc();
-    auto i8Ptr = AsyncAPI::opaquePointerType(rewriter.getContext());
-
-    // Get the underlying value type from the `async.value`.
-    auto await = cast<AwaitOp>(op);
-    auto valueType = await.operand().getType().cast<ValueType>().getValueType();
-
-    // Get a pointer to an async value storage from the runtime.
-    auto storage = rewriter.create<CallOp>(loc, kGetValueStorage,
-                                           TypeRange(i8Ptr), operand);
-
-    // Cast from i8* to the pointer pointer to LLVM type.
-    auto llvmValueType = getTypeConverter()->convertType(valueType);
-    auto castedStorage = rewriter.create<LLVM::BitcastOp>(
-        loc, LLVM::LLVMPointerType::get(llvmValueType), storage.getResult(0));
-
     // Load from the async value storage.
-    return rewriter.create<LLVM::LoadOp>(loc, castedStorage.getResult());
+    auto valueType = operand.getType().cast<ValueType>().getValueType();
+    return rewriter.create<RuntimeLoadOp>(op->getLoc(), valueType, operand);
   }
 };
 
 /// Lowering for `async.await_all` operation.
 class AwaitAllOpLowering : public AwaitOpLoweringBase<AwaitAllOp, GroupType> {
   using Base = AwaitOpLoweringBase<AwaitAllOp, GroupType>;
 
 public:
-  explicit AwaitAllOpLowering(
-      TypeConverter &converter, MLIRContext *ctx,
-      const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions)
-      : Base(converter, ctx, outlinedFunctions, kAwaitGroup,
-             kAwaitAllAndExecute) {}
+  using Base::Base;
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// async.yield op lowerings to the corresponding async runtime function calls.
+// Convert async.yield operation to async.runtime operations.
 //===----------------------------------------------------------------------===//
 
-class YieldOpLowering : public ConversionPattern {
+class YieldOpLowering : public OpConversionPattern<async::YieldOp> {
 public:
-  explicit YieldOpLowering(
-      TypeConverter &converter, MLIRContext *ctx,
+  YieldOpLowering(
+      MLIRContext *ctx,
       const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions)
-      : ConversionPattern(async::YieldOp::getOperationName(), 1, converter,
-                          ctx),
+      : OpConversionPattern<async::YieldOp>(ctx),
         outlinedFunctions(outlinedFunctions) {}
 
   LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+  matchAndRewrite(async::YieldOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     // Check if yield operation is inside the outlined coroutine function.
     auto func = op->template getParentOfType<FuncOp>();
     auto outlined = outlinedFunctions.find(func);
     if (outlined == outlinedFunctions.end())
-      return op->emitOpError(
-          "async.yield is not inside the outlined coroutine function");
+      return rewriter.notifyMatchFailure(
+          op, "operation is not inside the outlined async.execute function");
 
     Location loc = op->getLoc();
     const CoroMachinery &coro = outlined->getSecond();
 
-    // Store yielded values into the async values storage and emplace them.
-    auto i8Ptr = AsyncAPI::opaquePointerType(rewriter.getContext());
-
+    // Store yielded values into the async values storage and switch async
+    // values state to available.
     for (auto tuple : llvm::zip(operands, coro.returnValues)) {
-      // Store `yieldValue` into the `asyncValue` storage.
       Value yieldValue = std::get<0>(tuple);
       Value asyncValue = std::get<1>(tuple);
-
-      // Get an opaque i8* pointer to an async value storage from the runtime.
-      auto storage = rewriter.create<CallOp>(loc, kGetValueStorage,
-                                             TypeRange(i8Ptr), asyncValue);
-
-      // Cast storage pointer to the yielded value type.
-      auto castedStorage = rewriter.create<LLVM::BitcastOp>(
-          loc, LLVM::LLVMPointerType::get(yieldValue.getType()),
-          storage.getResult(0));
-
-      // Store the yielded value into the async value storage.
-      rewriter.create<LLVM::StoreOp>(loc, yieldValue,
-                                     castedStorage.getResult());
-
-      // Emplace the `async.value` to mark it ready.
-      rewriter.create<CallOp>(loc, kEmplaceValue, TypeRange(), asyncValue);
+      rewriter.create<RuntimeStoreOp>(loc, yieldValue, asyncValue);
+      rewriter.create<RuntimeSetAvailableOp>(loc, asyncValue);
     }
 
-    // Emplace the completion token to mark it ready.
-    rewriter.create<CallOp>(loc, kEmplaceToken, TypeRange(), coro.asyncToken);
-
-    // Original operation was replaced by the function call(s).
-    rewriter.eraseOp(op);
+    // Switch the coroutine completion token to available state.
+    rewriter.replaceOpWithNewOp<RuntimeSetAvailableOp>(op, coro.asyncToken);
 
     return success();
   }
 
 private:
   const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions;
 };
 
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct ConvertAsyncToLLVMPass
     : public ConvertAsyncToLLVMBase<ConvertAsyncToLLVMPass> {
   void runOnOperation() override;
 };
+} // namespace
 
 void ConvertAsyncToLLVMPass::runOnOperation() {
   ModuleOp module = getOperation();
   SymbolTable symbolTable(module);
 
   MLIRContext *ctx = &getContext();
 
   // Outline all `async.execute` body regions into async functions (coroutines).
void ConvertAsyncToLLVMPass::runOnOperation() {
   LLVM_DEBUG({
     llvm::dbgs() << "Outlined " << outlinedFunctions.size()
                  << " async functions\n";
   });
 
   // Add declarations for all functions required by the coroutines lowering.
   addResumeFunction(module);
   addAsyncRuntimeApiDeclarations(module);
-  addCoroutineIntrinsicsDeclarations(module);
   addCRuntimeDeclarations(module);
 
+  // ------------------------------------------------------------------------ //
+  // Lower async operations to async.runtime operations.
+  // ------------------------------------------------------------------------ //
+  OwningRewritePatternList asyncPatterns;
+
+  // Async lowering does not use type converter because it must preserve all
+  // types for async.runtime operations.
+  asyncPatterns.insert<CreateGroupOpLowering, AddToGroupOpLowering>(ctx);
+  asyncPatterns.insert<AwaitTokenOpLowering, AwaitValueOpLowering,
+                       AwaitAllOpLowering, YieldOpLowering>(ctx,
+                                                            outlinedFunctions);
+
+  // All high level async operations must be lowered to the runtime operations.
+  ConversionTarget runtimeTarget(*ctx);
+  runtimeTarget.addLegalDialect<AsyncDialect>();
+  runtimeTarget.addIllegalOp<CreateGroupOp, AddToGroupOp>();
+  runtimeTarget.addIllegalOp<ExecuteOp, AwaitOp, AwaitAllOp, async::YieldOp>();
+
+  if (failed(applyPartialConversion(module, runtimeTarget,
+                                    std::move(asyncPatterns)))) {
+    signalPassFailure();
+    return;
+  }
+
+  // ------------------------------------------------------------------------ //
+  // Lower async.runtime and async.coro operations to Async Runtime API and
+  // LLVM coroutine intrinsics.
+  // ------------------------------------------------------------------------ //
+
   // Convert async dialect types and operations to LLVM dialect.
   AsyncRuntimeTypeConverter converter;
   OwningRewritePatternList patterns;
 
   // Convert async types in function signatures and function calls.
   populateFuncOpTypeConversionPattern(patterns, ctx, converter);
   populateCallOpTypeConversionPattern(patterns, ctx, converter);
 
   // Convert return operations inside async.execute regions.
   patterns.insert<ReturnOpOpConversion>(converter, ctx);
 
-  // Lower async operations to async runtime API calls.
-  patterns.insert<AddRefOpLowering, DropRefOpLowering>(converter, ctx);
-  patterns.insert<CreateGroupOpLowering, AddToGroupOpLowering>(converter, ctx);
-
-  // Use LLVM type converter to automatically convert between the async value
-  // payload type and LLVM type when loading/storing from/to the async
-  // value storage which is an opaque i8* pointer using LLVM load/store ops.
-  patterns
-      .insert<AwaitTokenOpLowering, AwaitValueOpLowering, AwaitAllOpLowering>(
-          llvmConverter, ctx, outlinedFunctions);
-  patterns.insert<YieldOpLowering>(llvmConverter, ctx, outlinedFunctions);
+  // Lower async.runtime operations to the async runtime API calls.
+  patterns.insert<RuntimeSetAvailableOpLowering, RuntimeAwaitOpLowering,
+                  RuntimeAwaitAndResumeOpLowering, RuntimeResumeOpLowering,
+                  RuntimeAddToGroupOpLowering, RuntimeAddRefOpLowering,
+                  RuntimeDropRefOpLowering>(converter, ctx);
+
+  // Lower async.runtime operations that rely on LLVM type converter to convert
+  // from async value payload type to the LLVM type.
+  patterns.insert<RuntimeCreateOpLowering, RuntimeStoreOpLowering,
+                  RuntimeLoadOpLowering>(llvmConverter, ctx);
+
+  // Lower async coroutine operations to LLVM coroutine intrinsics.
+  patterns.insert<CoroIdOpConversion, CoroBeginOpConversion,
+                  CoroFreeOpConversion, CoroEndOpConversion,
+                  CoroSaveOpConversion, CoroSuspendOpConversion>(converter,
+                                                                 ctx);
 
   ConversionTarget target(*ctx);
   target.addLegalOp<ConstantOp>();
   target.addLegalDialect<LLVM::LLVMDialect>();
 
-  // All operations from Async dialect must be lowered to the runtime API calls.
+  // All operations from Async dialect must be lowered to the runtime API and
+  // LLVM intrinsics calls.
   target.addIllegalDialect<AsyncDialect>();
 
   // Add dynamic legality constraints to apply conversions defined above.
   target.addDynamicallyLegalOp<FuncOp>(
       [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); });
   target.addDynamicallyLegalOp<ReturnOp>(
       [&](ReturnOp op) { return converter.isLegal(op.getOperandTypes()); });
   target.addDynamicallyLegalOp<CallOp>([&](CallOp op) {
     return converter.isSignatureLegal(op.getCalleeType());
   });
 
   if (failed(applyPartialConversion(module, target, std::move(patterns))))
     signalPassFailure();
 }
-} // namespace
+
+//===----------------------------------------------------------------------===//
+// Patterns for structural type conversions for the Async dialect operations.
+//===----------------------------------------------------------------------===//
 
 namespace {
 class ConvertExecuteOpTypes : public OpConversionPattern<ExecuteOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ExecuteOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {

mlir/lib/Dialect/Async/IR/Async.cpp

 using namespace mlir;
 using namespace mlir::async;
 
 void AsyncDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc"
       >();
-  addTypes<TokenType>();
-  addTypes<ValueType>();
-  addTypes<GroupType>();
+  addTypes<TokenType, ValueType, GroupType>();           // async types
+  addTypes<CoroIdType, CoroHandleType, CoroStateType>(); // coro types
 }
 
 /// Parse a type registered to this dialect.
 Type AsyncDialect::parseType(DialectAsmParser &parser) const {
   StringRef keyword;
   if (parser.parseKeyword(&keyword))
     return Type();
 
   if (keyword == "token")
     return TokenType::get(getContext());
 
+  if (keyword == "group")
+    return GroupType::get(getContext());
+
   if (keyword == "value") {
     Type ty;
     if (parser.parseLess() || parser.parseType(ty) || parser.parseGreater()) {
       parser.emitError(parser.getNameLoc(), "failed to parse async value type");
       return Type();
     }
     return ValueType::get(ty);
   }
 
+  if (keyword == "coro.id")
+    return CoroIdType::get(getContext());
+
+  if (keyword == "coro.handle")
+    return CoroHandleType::get(getContext());
+
+  if (keyword == "coro.state")
+    return CoroStateType::get(getContext());
+
   parser.emitError(parser.getNameLoc(), "unknown async type: ") << keyword;
   return Type();
 }
 
 /// Print a type registered to this dialect.
 void AsyncDialect::printType(Type type, DialectAsmPrinter &os) const {
   TypeSwitch<Type>(type)
       .Case<TokenType>([&](TokenType) { os << "token"; })
       .Case<ValueType>([&](ValueType valueTy) {
         os << "value<";
         os.printType(valueTy.getValueType());
         os << '>';
       })
       .Case<GroupType>([&](GroupType) { os << "group"; })
+      .Case<CoroIdType>([&](CoroIdType) { os << "coro.id"; })
+      .Case<CoroHandleType>([&](CoroHandleType) { os << "coro.handle"; })
+      .Case<CoroStateType>([&](CoroStateType) { os << "coro.state"; })
       .Default([](Type) { llvm_unreachable("unexpected 'async' type kind"); });
 }
 
 //===----------------------------------------------------------------------===//
 /// ValueType
 //===----------------------------------------------------------------------===//
 
 namespace mlir {

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

LogicalResult AsyncRefCountingPass::addAutomaticRefCounting(Value value) {
   else
     builder.setInsertionPointToStart(value.getParentBlock());
 
   Location loc = value.getLoc();
   auto i32 = IntegerType::get(ctx, 32);
 
   // Drop the reference count immediately if the value has no uses.
   if (value.getUses().empty()) {
-    builder.create<DropRefOp>(loc, value, IntegerAttr::get(i32, 1));
+    builder.create<RuntimeDropRefOp>(loc, value, IntegerAttr::get(i32, 1));
     return success();
   }
 
   // Use liveness analysis to find the placement of `drop_ref`operation.
   auto liveness = getAnalysis<Liveness>();
 
   // We analyse only the blocks of the region that defines the `value`, and do
   // not check nested blocks attached to operations.
for (Operation *lastUser : lastUsers) {
 
     // We can't currently handle other types of terminators.
     if (lastUser->hasTrait<OpTrait::IsTerminator>())
       return lastUser->emitError() << "async reference counting can't handle "
                                       "terminators that are not ReturnLike";
 
     // Add a drop_ref immediately after the last user.
     builder.setInsertionPointAfter(lastUser);
-    builder.create<DropRefOp>(loc, value, IntegerAttr::get(i32, 1));
+    builder.create<RuntimeDropRefOp>(loc, value, IntegerAttr::get(i32, 1));
   }
 
   // ------------------------------------------------------------------------ //
   // Find blocks where the `value` is in `liveOut` set, however it is not in
   // the `liveIn` set of all successors. If the `value` is not in the successor
   // `liveIn` set, we add a `drop_ref` to the beginning of it.
   // ------------------------------------------------------------------------ //
 
for (Block *successor : block.getSuccessors()) {
         dropRefSuccessors.insert(successor);
     }
   }
 
   // Drop reference in all successor blocks that do not have the `value` in
   // their `liveIn` set.
   for (Block *dropRefSuccessor : dropRefSuccessors) {
     builder.setInsertionPointToStart(dropRefSuccessor);
-    builder.create<DropRefOp>(loc, value, IntegerAttr::get(i32, 1));
+    builder.create<RuntimeDropRefOp>(loc, value, IntegerAttr::get(i32, 1));
   }
 
   // ------------------------------------------------------------------------ //
   // Find all `async.execute` operation that take `value` as an operand
   // (dependency token or async value), or capture implicitly by the nested
   // region. Each `async.execute` operation will require `add_ref` operation
   // to keep all captured values alive until it will finish its execution.
   // ------------------------------------------------------------------------ //
for (Operation *user : value.getUsers()) {
     trackAsyncExecute(user);
   }
 
   // Process all `async.execute` operations capturing `value`.
   for (ExecuteOp execute : executeOperations) {
     // Add a reference before the execute operation to keep the reference
     // counted alive before the async region completes execution.
     builder.setInsertionPoint(execute.getOperation());
-    builder.create<AddRefOp>(loc, value, IntegerAttr::get(i32, 1));
+    builder.create<RuntimeAddRefOp>(loc, value, IntegerAttr::get(i32, 1));
 
     // Drop the reference inside the async region before completion.
     OpBuilder executeBuilder = OpBuilder::atBlockTerminator(execute.getBody());
-    executeBuilder.create<DropRefOp>(loc, value, IntegerAttr::get(i32, 1));
+    executeBuilder.create<RuntimeDropRefOp>(loc, value,
+                                            IntegerAttr::get(i32, 1));
   }
 
   return success();
 }
 
 void AsyncRefCountingPass::runOnFunction() {
   FuncOp func = getFunction();
 
   // Check that we do not have explicit `add_ref` or `drop_ref` in the IR
   // because otherwise automatic reference counting will produce incorrect
   // results.
   WalkResult refCountingWalk = func.walk([&](Operation *op) -> WalkResult {
-    if (isa<AddRefOp, DropRefOp>(op))
+    if (isa<RuntimeAddRefOp, RuntimeDropRefOp>(op))
       return op->emitError() << "explicit reference counting is not supported";
     return WalkResult::advance();
   });
 
   if (refCountingWalk.wasInterrupted())
     signalPassFailure();
 
   // Add reference counting to block arguments.

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

AsyncRefCountingOptimizationPass::optimizeReferenceCounting(Value value) {
   // operations in the nested regions that do use the `value`.
 
   struct UserInfo {
     Operation *operation;
     bool hasExecuteUser;
   };
 
   struct BlockUsersInfo {
-    llvm::SmallVector<AddRefOp, 4> addRefs;
-    llvm::SmallVector<DropRefOp, 4> dropRefs;
+    llvm::SmallVector<RuntimeAddRefOp, 4> addRefs;
+    llvm::SmallVector<RuntimeDropRefOp, 4> dropRefs;
     llvm::SmallVector<UserInfo, 4> users;
   };
 
   llvm::DenseMap<Block *, BlockUsersInfo> blockUsers;
 
   auto updateBlockUsersInfo = [&](UserInfo user) {
     BlockUsersInfo &info = blockUsers[user.operation->getBlock()];
     info.users.push_back(user);
 
-    if (auto addRef = dyn_cast<AddRefOp>(user.operation))
+    if (auto addRef = dyn_cast<RuntimeAddRefOp>(user.operation))
       info.addRefs.push_back(addRef);
-    if (auto dropRef = dyn_cast<DropRefOp>(user.operation))
+    if (auto dropRef = dyn_cast<RuntimeDropRefOp>(user.operation))
       info.dropRefs.push_back(dropRef);
   };
 
   for (Operation *user : value.getUsers()) {
     bool isAsyncUser = isa<ExecuteOp>(user);
 
     while (user->getParentRegion() != definingRegion) {
       updateBlockUsersInfo({user, isAsyncUser});
for (auto &kv : blockUsers) {
     BlockUsersInfo &info = preprocessBlockUsersInfo(kv.second);
 
     // Find all cancellable pairs first and erase them later to keep all
     // pointers in the `info` valid until the end.
     //
     // Mapping from `dropRef.getOperation()` to `addRef.getOperation()`.
     llvm::SmallDenseMap<Operation *, Operation *> cancellable;
 
-    for (AddRefOp addRef : info.addRefs) {
-      for (DropRefOp dropRef : info.dropRefs) {
+    for (RuntimeAddRefOp addRef : info.addRefs) {
+      for (RuntimeDropRefOp dropRef : info.dropRefs) {
         // `drop_ref` operation after the `add_ref` with matching count.
         if (dropRef.count() != addRef.count() ||
             dropRef->isBeforeInBlock(addRef.getOperation()))
           continue;
 
         // `drop_ref` was already marked for removal.
         if (cancellable.find(dropRef.getOperation()) != cancellable.end())
           continue;

mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir

+// RUN: mlir-opt %s -convert-async-to-llvm | FileCheck %s
+
+// CHECK-LABEL: @coro_id
+func @coro_id() {
+  // CHECK: %0 = llvm.mlir.constant(0 : i32) : i32
+  // CHECK: %1 = llvm.mlir.null : !llvm.ptr<i8>
+  // CHECK: %2 = llvm.intr.coro.id %0, %1, %1, %1 : !llvm.token
+  %0 = async.coro.id
+  return
+}
+
+// CHECK-LABEL: @coro_begin
+func @coro_begin() {
+  // CHECK: %[[ID:.*]] = llvm.intr.coro.id
+  %0 = async.coro.id
+  // CHECK: %[[SIZE:.*]] = llvm.intr.coro.size : i64
+  // CHECK: %[[ALLOC:.*]] = llvm.call @malloc(%[[SIZE]])
+  // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin %[[ID]], %[[ALLOC]]
+  %1 = async.coro.begin %0
+  return
+}
+
+// CHECK-LABEL: @coro_free
+func @coro_free() {
+  // CHECK: %[[ID:.*]] = llvm.intr.coro.id
+  %0 = async.coro.id
+  // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
+  %1 = async.coro.begin %0
+  // CHECK: %[[MEM:.*]] = llvm.intr.coro.free %[[ID]], %[[HDL]]
+  // CHECK: llvm.call @free(%[[MEM]])
+  async.coro.free %0, %1
+  return
+}
+
+// CHECK-LABEL: @coro_end
+func @coro_end() {
+  %0 = async.coro.id
+  // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
+  %1 = async.coro.begin %0
+  // CHECK: %[[FALSE:.*]] = llvm.mlir.constant(false) : i1
+  // CHECK: llvm.intr.coro.end %[[HDL]], %[[FALSE]]
+  async.coro.end %1
+  return
+}
+
+// CHECK-LABEL: @coro_save
+func @coro_save() {
+  %0 = async.coro.id
+  // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
+  %1 = async.coro.begin %0
+  // CHECK: llvm.intr.coro.save %[[HDL]]
+  %2 = async.coro.save %1
+  return
+}
+
+// CHECK-LABEL: @coro_suspend
+func @coro_suspend() {
+  %0 = async.coro.id
+  // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
+  %1 = async.coro.begin %0
+  // CHECK: %[[STATE:.*]] = llvm.intr.coro.save %[[HDL]]
+  %2 = async.coro.save %1
+
+  // CHECK: %[[FINAL:.*]] = llvm.mlir.constant(false) : i1
+  // CHECK: %[[RET:.*]] = llvm.intr.coro.suspend %[[STATE]], %[[FINAL]]
+  // CHECK: %[[SEXT:.*]] = llvm.sext %[[RET]] : i8 to i32
+  // CHECK: llvm.switch %[[SEXT]], ^[[SUSPEND:[b0-9]+]]
+  // CHECK-NEXT: 0: ^[[RESUME:[b0-9]+]]
+  // CHECK-NEXT: 1: ^[[CLEANUP:[b0-9]+]]
+  async.coro.suspend %2, ^suspend, ^resume, ^cleanup
+^resume:
+  // CHECK: ^[[RESUME]]
+  // CHECK:   return {coro.resume}
+  return { coro.resume }
+^cleanup:
+  // CHECK: ^[[CLEANUP]]
+  // CHECK:   return {coro.cleanup}
+  return { coro.cleanup }
+^suspend:
+  // CHECK: ^[[SUSPEND]]
+  // CHECK:   return {coro.suspend}
+  return { coro.suspend }
+}

mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm.mlir

+// RUN: mlir-opt %s -convert-async-to-llvm | FileCheck %s
+
+// CHECK-LABEL: @create_token
+func @create_token() {
+  // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
+  %0 = async.runtime.create : !async.token
+  return
+}
+
+// CHECK-LABEL: @create_value
+func @create_value() {
+  // CHECK: %[[NULL:.*]] = llvm.mlir.null : !llvm.ptr<f32>
+  // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[OFFSET:.*]] = llvm.getelementptr %[[NULL]][%[[ONE]]]
+  // CHECK: %[[SIZE:.*]] = llvm.ptrtoint %[[OFFSET]]
+  // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue(%[[SIZE]])
+  %0 = async.runtime.create : !async.value<f32>
+  return
+}
+
+// CHECK-LABEL: @create_group
+func @create_group() {
+  // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup
+  %0 = async.runtime.create : !async.group
+  return
+}
+
+// CHECK-LABEL: @set_token_available
+func @set_token_available() {
+  // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
+  %0 = async.runtime.create : !async.token
+  // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]])
+  async.runtime.set_available %0 : !async.token
+  return
+}
+
+// CHECK-LABEL: @set_value_available
+func @set_value_available() {
+  // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
+  %0 = async.runtime.create : !async.value<f32>
+  // CHECK: call @mlirAsyncRuntimeEmplaceValue(%[[VALUE]])
+  async.runtime.set_available %0 : !async.value<f32>
+  return
+}
+
+// CHECK-LABEL: @await_token
+func @await_token() {
+  // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
+  %0 = async.runtime.create : !async.token
+  // CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]])
+  async.runtime.await %0 : !async.token
+  return
+}
+
+// CHECK-LABEL: @await_value
+func @await_value() {
+  // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
+  %0 = async.runtime.create : !async.value<f32>
+  // CHECK: call @mlirAsyncRuntimeAwaitValue(%[[VALUE]])
+  async.runtime.await %0 : !async.value<f32>
+  return
+}
+
+// CHECK-LABEL: @await_group
+func @await_group() {
+  // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup
+  %0 = async.runtime.create : !async.group
+  // CHECK: call @mlirAsyncRuntimeAwaitAllInGroup(%[[GROUP]])
+  async.runtime.await %0 : !async.group
+  return
+}
+
+// CHECK-LABEL: @await_and_resume_token
+func @await_and_resume_token() {
+  %0 = async.coro.id
+  // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
+  %1 = async.coro.begin %0
+  // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
+  %2 = async.runtime.create : !async.token
+  // CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume
+  // CHECK: call @mlirAsyncRuntimeAwaitTokenAndExecute
+  // CHECK-SAME: (%[[TOKEN]], %[[HDL]], %[[RESUME]])
+  async.runtime.await_and_resume %2, %1 : !async.token
+  return
+}
+
+// CHECK-LABEL: @await_and_resume_value
+func @await_and_resume_value() {
+  %0 = async.coro.id
+  // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
+  %1 = async.coro.begin %0
+  // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
+  %2 = async.runtime.create : !async.value<f32>
+  // CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume
+  // CHECK: call @mlirAsyncRuntimeAwaitValueAndExecute
+  // CHECK-SAME: (%[[VALUE]], %[[HDL]], %[[RESUME]])
+  async.runtime.await_and_resume %2, %1 : !async.value<f32>
+  return
+}
+
+// CHECK-LABEL: @await_and_resume_group
+func @await_and_resume_group() {
+  %0 = async.coro.id
+  // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
+  %1 = async.coro.begin %0
+  // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateGroup
+  %2 = async.runtime.create : !async.group
+  // CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume
+  // CHECK: call @mlirAsyncRuntimeAwaitAllInGroupAndExecute
+  // CHECK-SAME: (%[[TOKEN]], %[[HDL]], %[[RESUME]])
+  async.runtime.await_and_resume %2, %1 : !async.group
+  return
+}
+
+// CHECK-LABEL: @resume
+func @resume() {
+  %0 = async.coro.id
+  // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
+  %1 = async.coro.begin %0
+  // CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume
+  // CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]], %[[RESUME]])
+  async.runtime.resume %1
+  return
+}
+
+// CHECK-LABEL: @store
+func @store() {
+  // CHECK: %[[CST:.*]] = constant 1.0
+  %0 = constant 1.0 : f32
+  // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
+  %1 = async.runtime.create : !async.value<f32>
+  // CHECK: %[[P0:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[VALUE]])
+  // CHECK: %[[P1:.*]] = llvm.bitcast %[[P0]] : !llvm.ptr<i8> to !llvm.ptr<f32>
+  // CHECK: llvm.store %[[CST]], %[[P1]]
+  async.runtime.store %0, %1 : !async.value<f32>
+  return
+}
+
+// CHECK-LABEL: @load
+func @load() -> f32 {
+  // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
+  %0 = async.runtime.create : !async.value<f32>
+  // CHECK: %[[P0:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[VALUE]])
+  // CHECK: %[[P1:.*]] = llvm.bitcast %[[P0]] : !llvm.ptr<i8> to !llvm.ptr<f32>
+  // CHECK: %[[VALUE:.*]] = llvm.load %[[P1]]
+  %1 = async.runtime.load %0 : !async.value<f32>
+  // CHECK: return %[[VALUE]] : f32
+  return %1 : f32
+}
+
+// CHECK-LABEL: @add_token_to_group
+func @add_token_to_group() {
+  // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
+  %0 = async.runtime.create : !async.token
+  // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup
+  %1 = async.runtime.create : !async.group
+  // CHECK: call @mlirAsyncRuntimeAddTokenToGroup(%[[TOKEN]], %[[GROUP]])
+  async.runtime.add_to_group %0, %1 : !async.token
+  return
+}

mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir

 // RUN: mlir-opt %s -split-input-file -convert-async-to-llvm | FileCheck %s
 
 // CHECK-LABEL: reference_counting
 func @reference_counting(%arg0: !async.token) {
   // CHECK: %[[C2:.*]] = constant 2 : i32
   // CHECK: call @mlirAsyncRuntimeAddRef(%arg0, %[[C2]])
-  async.add_ref %arg0 {count = 2 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 2 : i32} : !async.token
 
   // CHECK: %[[C1:.*]] = constant 1 : i32
   // CHECK: call @mlirAsyncRuntimeDropRef(%arg0, %[[C1]])
-  async.drop_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
 
   return
 }
 
 // -----
 
 // CHECK-LABEL: execute_no_async_args
 func @execute_no_async_args(%arg0: f32, %arg1: memref<1xf32>) {

 }
 
 // Function outlined from the async.execute operation.
 // CHECK-LABEL: func private @async_execute_fn(%arg0: f32, %arg1: memref<1xf32>)
 // CHECK-SAME: -> !llvm.ptr<i8>
 
 // Create token for return op, and mark a function as a coroutine.
 // CHECK: %[[RET:.*]] = call @mlirAsyncRuntimeCreateToken()
-// CHECK: %[[HDL:.*]] = llvm.call @llvm.coro.begin
+// CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
 
 // Pass a suspended coroutine to the async runtime.
+// CHECK: %[[STATE:.*]] = llvm.intr.coro.save
 // CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume
-// CHECK: %[[STATE:.*]] = llvm.call @llvm.coro.save
 // CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]], %[[RESUME]])
-// CHECK: %[[SUSPENDED:.*]] = llvm.call @llvm.coro.suspend(%[[STATE]]
+// CHECK: %[[SUSPENDED:.*]] = llvm.intr.coro.suspend %[[STATE]]
 
 // Decide the next block based on the code returned from suspend.
-// CHECK: %[[ZERO:.*]] = llvm.mlir.constant(0 : i8)
-// CHECK: %[[NONE:.*]] = llvm.mlir.constant(-1 : i8)
-// CHECK: %[[IS_NONE:.*]] = llvm.icmp "eq" %[[SUSPENDED]], %[[NONE]]
-// CHECK: llvm.cond_br %[[IS_NONE]], ^[[SUSPEND:.*]], ^[[RESUME_OR_CLEANUP:.*]]
-
-// Decide if branch to resume or cleanup block.
-// CHECK: ^[[RESUME_OR_CLEANUP]]:
-// CHECK: %[[IS_ZERO:.*]] = llvm.icmp "eq" %[[SUSPENDED]], %[[ZERO]]
-// CHECK: llvm.cond_br %[[IS_ZERO]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]
+// CHECK: %[[SEXT:.*]] = llvm.sext %[[SUSPENDED]] : i8 to i32
+// CHECK: llvm.switch %[[SEXT]], ^[[SUSPEND:[b0-9]+]]
+// CHECK-NEXT: 0: ^[[RESUME:[b0-9]+]]
+// CHECK-NEXT: 1: ^[[CLEANUP:[b0-9]+]]
 
 // Resume coroutine after suspension.
 // CHECK: ^[[RESUME]]:
 // CHECK: store %arg0, %arg1[%c0] : memref<1xf32>
 // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET]])
 
 // Delete coroutine.
 // CHECK: ^[[CLEANUP]]:
-// CHECK: %[[MEM:.*]] = llvm.call @llvm.coro.free
+// CHECK: %[[MEM:.*]] = llvm.intr.coro.free
 // CHECK: llvm.call @free(%[[MEM]])
 
 // Suspend coroutine, and also a return statement for ramp function.
 // CHECK: ^[[SUSPEND]]:
-// CHECK: llvm.call @llvm.coro.end
+// CHECK: llvm.intr.coro.end
 // CHECK: return %[[RET]]
 
 // -----
 
 // CHECK-LABEL: nested_async_execute
 func @nested_async_execute(%arg0: f32, %arg1: f32, %arg2: memref<1xf32>) {
   // CHECK: %[[TOKEN:.*]] = call @async_execute_fn_0(%arg0, %arg2, %arg1)
   %token0 = async.execute {
func @nested_async_execute(%arg0: f32, %arg1: f32, %arg2: memref<1xf32>) {
   async.await %token0 : !async.token
   return
 }
 
 // Function outlined from the inner async.execute operation.
 // CHECK-LABEL: func private @async_execute_fn(%arg0: f32, %arg1: memref<1xf32>, %arg2: index)
 // CHECK-SAME: -> !llvm.ptr<i8>
 // CHECK: %[[RET_0:.*]] = call @mlirAsyncRuntimeCreateToken()
-// CHECK: %[[HDL_0:.*]] = llvm.call @llvm.coro.begin
+// CHECK: %[[HDL_0:.*]] = llvm.intr.coro.begin
 // CHECK: call @mlirAsyncRuntimeExecute
-// CHECK: llvm.call @llvm.coro.suspend
+// CHECK: llvm.intr.coro.suspend
 // CHECK: store %arg0, %arg1[%arg2] : memref<1xf32>
 // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_0]])
 
 // Function outlined from the outer async.execute operation.
 // CHECK-LABEL: func private @async_execute_fn_0(%arg0: f32, %arg1: memref<1xf32>, %arg2: f32)
 // CHECK-SAME: -> !llvm.ptr<i8>
 // CHECK: %[[RET_1:.*]] = call @mlirAsyncRuntimeCreateToken()
-// CHECK: %[[HDL_1:.*]] = llvm.call @llvm.coro.begin
+// CHECK: %[[HDL_1:.*]] = llvm.intr.coro.begin
 
 // Suspend coroutine in the beginning.
 // CHECK: call @mlirAsyncRuntimeExecute
-// CHECK: llvm.call @llvm.coro.suspend
+// CHECK: llvm.intr.coro.suspend
 
 // Suspend coroutine second time waiting for the completion of inner execute op.
 // CHECK: %[[TOKEN_1:.*]] = call @async_execute_fn
-// CHECK: llvm.call @llvm.coro.save
+// CHECK: llvm.intr.coro.save
 // CHECK: call @mlirAsyncRuntimeAwaitTokenAndExecute(%[[TOKEN_1]], %[[HDL_1]]
-// CHECK: llvm.call @llvm.coro.suspend
+// CHECK: llvm.intr.coro.suspend
 
 // Emplace result token after second resumption.
 // CHECK: store %arg2, %arg1[%c0] : memref<1xf32>
 // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_1]])
 
 // -----
 
 // CHECK-LABEL: async_execute_token_dependency
func @async_execute_token_dependency(%arg0: f32, %arg1: memref<1xf32>) {
   }
   return
 }
 
 // Function outlined from the first async.execute operation.
 // CHECK-LABEL: func private @async_execute_fn(%arg0: f32, %arg1: memref<1xf32>)
 // CHECK-SAME: -> !llvm.ptr<i8>
 // CHECK: %[[RET_0:.*]] = call @mlirAsyncRuntimeCreateToken()
-// CHECK: %[[HDL_0:.*]] = llvm.call @llvm.coro.begin
+// CHECK: %[[HDL_0:.*]] = llvm.intr.coro.begin
 // CHECK: call @mlirAsyncRuntimeExecute
-// CHECK: llvm.call @llvm.coro.suspend
+// CHECK: llvm.intr.coro.suspend
 // CHECK: store %arg0, %arg1[%c0] : memref<1xf32>
 // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_0]])
 
 // Function outlined from the second async.execute operation with dependency.
 // CHECK-LABEL: func private @async_execute_fn_0(%arg0: !llvm.ptr<i8>, %arg1: f32, %arg2: memref<1xf32>)
 // CHECK-SAME: -> !llvm.ptr<i8>
 // CHECK: %[[RET_1:.*]] = call @mlirAsyncRuntimeCreateToken()
-// CHECK: %[[HDL_1:.*]] = llvm.call @llvm.coro.begin
+// CHECK: %[[HDL_1:.*]] = llvm.intr.coro.begin
 
 // Suspend coroutine in the beginning.
 // CHECK: call @mlirAsyncRuntimeExecute(%[[HDL_1]],
-// CHECK: llvm.call @llvm.coro.suspend
+// CHECK: llvm.intr.coro.suspend
 
 // Suspend coroutine second time waiting for the completion of token dependency.
-// CHECK: llvm.call @llvm.coro.save
+// CHECK: llvm.intr.coro.save
 // CHECK: call @mlirAsyncRuntimeAwaitTokenAndExecute(%arg0, %[[HDL_1]],
-// CHECK: llvm.call @llvm.coro.suspend
+// CHECK: llvm.intr.coro.suspend
 
 // Emplace result token after second resumption.
 // CHECK: store %arg1, %arg2[%c0] : memref<1xf32>
 // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_1]])
 
 // -----
 
 // CHECK-LABEL: async_group_await_all
func @async_group_await_all(%arg0: f32, %arg1: memref<1xf32>) {
   async.await_all %0
 
   return
 }
 
 // Function outlined from the async.execute operation.
 // CHECK: func private @async_execute_fn_0(%arg0: !llvm.ptr<i8>)
 // CHECK: %[[RET_1:.*]] = call @mlirAsyncRuntimeCreateToken()
-// CHECK: %[[HDL_1:.*]] = llvm.call @llvm.coro.begin
+// CHECK: %[[HDL_1:.*]] = llvm.intr.coro.begin
 
 // Suspend coroutine in the beginning.
 // CHECK: call @mlirAsyncRuntimeExecute(%[[HDL_1]],
-// CHECK: llvm.call @llvm.coro.suspend
+// CHECK: llvm.intr.coro.suspend
 
 // Suspend coroutine second time waiting for the group.
-// CHECK: llvm.call @llvm.coro.save
+// CHECK: llvm.intr.coro.save
 // CHECK: call @mlirAsyncRuntimeAwaitAllInGroupAndExecute(%arg0, %[[HDL_1]],
-// CHECK: llvm.call @llvm.coro.suspend
+// CHECK: llvm.intr.coro.suspend
 
 // Emplace result token.
 // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_1]])
 
 // -----
 
 // CHECK-LABEL: execute_and_return_f32
 func @execute_and_return_f32() -> f32 {
// CHECK: %[[RET:.*]]:2 = call @async_execute_fn
 
   return %0 : f32
 }
 
 // Function outlined from the async.execute operation.
 // CHECK-LABEL: func private @async_execute_fn()
 // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken()
 // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
-// CHECK: %[[HDL:.*]] = llvm.call @llvm.coro.begin
+// CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
 
 // Suspend coroutine in the beginning.
 // CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]],
-// CHECK: llvm.call @llvm.coro.suspend
+// CHECK: llvm.intr.coro.suspend
 
 // Emplace result value.
 // CHECK: %[[CST:.*]] = constant 1.230000e+02 : f32
 // CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[VALUE]])
 // CHECK: %[[ST_F32:.*]] = llvm.bitcast %[[STORAGE]]
 // CHECK: llvm.store %[[CST]], %[[ST_F32]] : !llvm.ptr<f32>
 // CHECK: call @mlirAsyncRuntimeEmplaceValue(%[[VALUE]])
 

 }
 
 // Function outlined from the first async.execute operation.
 // CHECK-LABEL: func private @async_execute_fn()
 
 // Function outlined from the second async.execute operation.
 // CHECK-LABEL: func private @async_execute_fn_0(%arg0: !llvm.ptr<i8>)
 // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken()
-// CHECK: %[[HDL:.*]] = llvm.call @llvm.coro.begin
+// CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
 
 // Suspend coroutine in the beginning.
 // CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]],
-// CHECK: llvm.call @llvm.coro.suspend
+// CHECK: llvm.intr.coro.suspend
 
 // Suspend coroutine second time waiting for the async operand.
-// CHECK: llvm.call @llvm.coro.save
+// CHECK: llvm.intr.coro.save
 // CHECK: call @mlirAsyncRuntimeAwaitValueAndExecute(%arg0, %[[HDL]],
-// CHECK: llvm.call @llvm.coro.suspend
+// CHECK: llvm.intr.coro.suspend
 
 // Get the operand value storage, cast to f32 and add the value.
 // CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%arg0)
 // CHECK: %[[ST_F32:.*]] = llvm.bitcast %[[STORAGE]]
 // CHECK: %[[LOADED:.*]] = llvm.load %[[ST_F32]] :  !llvm.ptr<f32>
 // CHECK: addf %[[LOADED]], %[[LOADED]] : f32
 
 // Emplace result token.
 // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]])
 
 

mlir/test/Dialect/Async/async-ref-counting-optimization.mlir

 // RUN: mlir-opt %s -async-ref-counting-optimization | FileCheck %s
 
 // CHECK-LABEL: @cancellable_operations_0
 func @cancellable_operations_0(%arg0: !async.token) {
-  // CHECK-NOT: async.add_ref
-  // CHECK-NOT: async.drop_ref
-  async.add_ref %arg0 {count = 1 : i32} : !async.token
-  async.drop_ref %arg0 {count = 1 : i32} : !async.token
+  // CHECK-NOT: async.runtime.add_ref
+  // CHECK-NOT: async.runtime.drop_ref
+  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @cancellable_operations_1
 func @cancellable_operations_1(%arg0: !async.token) {
-  // CHECK-NOT: async.add_ref
+  // CHECK-NOT: async.runtime.add_ref
   // CHECK: async.execute
-  async.add_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
   async.execute [%arg0] {
-    // CHECK: async.drop_ref
-    async.drop_ref %arg0 {count = 1 : i32} : !async.token
+    // CHECK: async.runtime.drop_ref
+    async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
     // CHECK-NEXT: async.yield
     async.yield
   }
-  // CHECK-NOT: async.drop_ref
-  async.drop_ref %arg0 {count = 1 : i32} : !async.token
+  // CHECK-NOT: async.runtime.drop_ref
+  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @cancellable_operations_2
 func @cancellable_operations_2(%arg0: !async.token) {
   // CHECK: async.await
   // CHECK-NEXT: async.await
   // CHECK-NEXT: async.await
   // CHECK-NEXT: return
-  async.add_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
   async.await %arg0 : !async.token
-  async.drop_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
   async.await %arg0 : !async.token
-  async.add_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
   async.await %arg0 : !async.token
-  async.drop_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
   return
 }
 
 // CHECK-LABEL: @cancellable_operations_3
 func @cancellable_operations_3(%arg0: !async.token) {
   // CHECK-NOT: add_ref
-  async.add_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
   %token = async.execute {
     async.await %arg0 : !async.token
-    // CHECK: async.drop_ref
-    async.drop_ref %arg0 {count = 1 : i32} : !async.token
+    // CHECK: async.runtime.drop_ref
+    async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
     async.yield
   }
-  // CHECK-NOT: async.drop_ref
-  async.drop_ref %arg0 {count = 1 : i32} : !async.token
+  // CHECK-NOT: async.runtime.drop_ref
+  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
   // CHECK: async.await
   async.await %arg0 : !async.token
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @not_cancellable_operations_0
 func @not_cancellable_operations_0(%arg0: !async.token, %arg1: i1) {
   // It is unsafe to cancel `add_ref` / `drop_ref` pair because it is possible
   // that the body of the `async.execute` operation will run before the await
   // operation in the function body, and will destroy the `%arg0` token.
   // CHECK: add_ref
-  async.add_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
   %token = async.execute {
     // CHECK: async.await
     async.await %arg0 : !async.token
-    // CHECK: async.drop_ref
-    async.drop_ref %arg0 {count = 1 : i32} : !async.token
+    // CHECK: async.runtime.drop_ref
+    async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
     // CHECK: async.yield
     async.yield
   }
   // CHECK: async.await
   async.await %arg0 : !async.token
   // CHECK: drop_ref
-  async.drop_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @not_cancellable_operations_1
 func @not_cancellable_operations_1(%arg0: !async.token, %arg1: i1) {
   // Same reason as above, although `async.execute` is inside the nested
   // region or "regular" operation.
   //
   // NOTE: This test is not correct w.r.t. reference counting, and at runtime
   // would leak %arg0 value if %arg1 is false. IR like this will not be
   // constructed by automatic reference counting pass, because it would
-  // place `async.add_ref` right before the `async.execute` inside `scf.if`.
+  // place `async.runtime.add_ref` right before the `async.execute`
+  // inside `scf.if`.
   
-  // CHECK: async.add_ref
-  async.add_ref %arg0 {count = 1 : i32} : !async.token
+  // CHECK: async.runtime.add_ref
+  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
   scf.if %arg1 {
     %token = async.execute {
       async.await %arg0 : !async.token
-      // CHECK: async.drop_ref
-      async.drop_ref %arg0 {count = 1 : i32} : !async.token
+      // CHECK: async.runtime.drop_ref
+      async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
       async.yield
     }
   }
   // CHECK: async.await
   async.await %arg0 : !async.token
-  // CHECK: async.drop_ref
-  async.drop_ref %arg0 {count = 1 : i32} : !async.token
+  // CHECK: async.runtime.drop_ref
+  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
   // CHECK: return
   return
 }

mlir/test/Dialect/Async/async-ref-counting.mlir

 // RUN: mlir-opt %s -async-ref-counting | FileCheck %s
 
 // CHECK-LABEL: @cond
 func private @cond() -> i1
 
 // CHECK-LABEL: @token_arg_no_uses
 func @token_arg_no_uses(%arg0: !async.token) {
-  // CHECK: async.drop_ref %arg0 {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %arg0 {count = 1 : i32}
   return
 }
 
 // CHECK-LABEL: @token_arg_conditional_await
 func @token_arg_conditional_await(%arg0: !async.token, %arg1: i1) {
   cond_br %arg1, ^bb1, ^bb2
 ^bb1:
-  // CHECK: async.drop_ref %arg0 {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %arg0 {count = 1 : i32}
   return
 ^bb2:
   // CHECK: async.await %arg0
-  // CHECK: async.drop_ref %arg0 {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %arg0 {count = 1 : i32}
   async.await %arg0 : !async.token
   return
 }
 
 // CHECK-LABEL: @token_no_uses
 func @token_no_uses() {
   // CHECK: %[[TOKEN:.*]] = async.execute
-  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
   %token = async.execute {
     async.yield
   }
   return
 }
 
 // CHECK-LABEL: @token_return
 func @token_return() -> !async.token {
   // CHECK: %[[TOKEN:.*]] = async.execute
   %token = async.execute {
     async.yield
   }
   // CHECK: return %[[TOKEN]]
   return %token : !async.token
 }
 
 // CHECK-LABEL: @token_await
 func @token_await() {
   // CHECK: %[[TOKEN:.*]] = async.execute
   %token = async.execute {
     async.yield
   }
   // CHECK: async.await %[[TOKEN]]
   async.await %token : !async.token
-  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @token_await_and_return
 func @token_await_and_return() -> !async.token {
   // CHECK: %[[TOKEN:.*]] = async.execute
   %token = async.execute {
     async.yield
   }
   // CHECK: async.await %[[TOKEN]]
-  // CHECK-NOT: async.drop_ref
+  // CHECK-NOT: async.runtime.drop_ref
   async.await %token : !async.token
   // CHECK: return %[[TOKEN]]
   return %token : !async.token
 }
 
 // CHECK-LABEL: @token_await_inside_scf_if
 func @token_await_inside_scf_if(%arg0: i1) {
   // CHECK: %[[TOKEN:.*]] = async.execute
   %token = async.execute {
     async.yield
   }
   // CHECK: scf.if %arg0 {
   scf.if %arg0 {
     // CHECK: async.await %[[TOKEN]]
     async.await %token : !async.token
   }
   // CHECK: }
-  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @token_conditional_await
 func @token_conditional_await(%arg0: i1) {
   // CHECK: %[[TOKEN:.*]] = async.execute
   %token = async.execute {
     async.yield
   }
   cond_br %arg0, ^bb1, ^bb2
 ^bb1:
-  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
   return
 ^bb2:
   // CHECK: async.await %[[TOKEN]]
-  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
   async.await %token : !async.token
   return
 }
 
 // CHECK-LABEL: @token_await_in_the_loop
 func @token_await_in_the_loop() {
   // CHECK: %[[TOKEN:.*]] = async.execute
   %token = async.execute {
     async.yield
   }
   br ^bb1
 ^bb1:
   // CHECK: async.await %[[TOKEN]]
   async.await %token : !async.token
   %0 = call @cond(): () -> (i1)
   cond_br %0, ^bb1, ^bb2
 ^bb2:
-  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
   return
 }
 
 // CHECK-LABEL: @token_defined_in_the_loop
 func @token_defined_in_the_loop() {
   br ^bb1
 ^bb1:
   // CHECK: %[[TOKEN:.*]] = async.execute
   %token = async.execute {
     async.yield
   }
   // CHECK: async.await %[[TOKEN]]
-  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
   async.await %token : !async.token
   %0 = call @cond(): () -> (i1)
   cond_br %0, ^bb1, ^bb2
 ^bb2:
   return
 }
 
 // CHECK-LABEL: @token_capture
 func @token_capture() {
   // CHECK: %[[TOKEN:.*]] = async.execute
   %token = async.execute {
     async.yield
   }
 
-  // CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32}
   // CHECK: %[[TOKEN_0:.*]] = async.execute
   %token_0 = async.execute {
-    // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+    // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
     // CHECK-NEXT: async.yield
     async.await %token : !async.token
     async.yield
   }
-  // CHECK: async.drop_ref %[[TOKEN_0]] {count = 1 : i32}
-  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN_0]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @token_nested_capture
 func @token_nested_capture() {
   // CHECK: %[[TOKEN:.*]] = async.execute
   %token = async.execute {
     async.yield
   }
 
-  // CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32}
   // CHECK: %[[TOKEN_0:.*]] = async.execute
   %token_0 = async.execute {
-    // CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
+    // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32}
     // CHECK: %[[TOKEN_1:.*]] = async.execute
     %token_1 = async.execute {
-      // CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
+      // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32}
       // CHECK: %[[TOKEN_2:.*]] = async.execute
       %token_2 = async.execute {
         // CHECK: async.await %[[TOKEN]]
-        // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+        // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
         async.await %token : !async.token
         async.yield
       }
-      // CHECK: async.drop_ref %[[TOKEN_2]] {count = 1 : i32}
-      // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+      // CHECK: async.runtime.drop_ref %[[TOKEN_2]] {count = 1 : i32}
+      // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
       async.yield
     }
-    // CHECK: async.drop_ref %[[TOKEN_1]] {count = 1 : i32}
-    // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+    // CHECK: async.runtime.drop_ref %[[TOKEN_1]] {count = 1 : i32}
+    // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
     async.yield
   }
-  // CHECK: async.drop_ref %[[TOKEN_0]] {count = 1 : i32}
-  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN_0]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @token_dependency
 func @token_dependency() {
   // CHECK: %[[TOKEN:.*]] = async.execute
   %token = async.execute {
     async.yield
   }
 
-  // CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32}
   // CHECK: %[[TOKEN_0:.*]] = async.execute
   %token_0 = async.execute[%token] {
-    // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+    // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
     // CHECK-NEXT: async.yield
     async.yield
   }
 
   // CHECK: async.await %[[TOKEN]]
-  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
   async.await %token : !async.token
   // CHECK: async.await %[[TOKEN_0]]
-  // CHECK: async.drop_ref %[[TOKEN_0]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN_0]] {count = 1 : i32}
   async.await %token_0 : !async.token
 
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @value_operand
 func @value_operand() -> f32 {
   // CHECK: %[[TOKEN:.*]], %[[RESULTS:.*]] = async.execute
   %token, %results = async.execute -> !async.value<f32> {
     %0 = constant 0.0 : f32
     async.yield %0 : f32
   }
 
-  // CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
-  // CHECK: async.add_ref %[[RESULTS]] {count = 1 : i32}
+  // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.add_ref %[[RESULTS]] {count = 1 : i32}
   // CHECK: %[[TOKEN_0:.*]] = async.execute
   %token_0 = async.execute[%token](%results as %arg0 : !async.value<f32>)  {
-    // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
-    // CHECK: async.drop_ref %[[RESULTS]] {count = 1 : i32}
+    // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+    // CHECK: async.runtime.drop_ref %[[RESULTS]] {count = 1 : i32}
     // CHECK: async.yield
     async.yield
   }
 
   // CHECK: async.await %[[TOKEN]]
-  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
   async.await %token : !async.token
 
   // CHECK: async.await %[[TOKEN_0]]
-  // CHECK: async.drop_ref %[[TOKEN_0]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN_0]] {count = 1 : i32}
   async.await %token_0 : !async.token
 
   // CHECK: async.await %[[RESULTS]]
-  // CHECK: async.drop_ref %[[RESULTS]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[RESULTS]] {count = 1 : i32}
   %0 = async.await %results : !async.value<f32>
 
   // CHECK: return
   return %0 : f32
 }

mlir/test/Dialect/Async/coro.mlir

+// RUN: mlir-opt %s | FileCheck %s --dump-input=always
+
+// CHECK-LABEL: @coro_id
+func @coro_id() -> !async.coro.id {
+  // CHECK: %0 = async.coro.id
+  // CHECK: return %0 : !async.coro.id
+  %0 = async.coro.id
+  return %0 : !async.coro.id
+}
+
+// CHECK-LABEL: @coro_handle
+func @coro_handle(%arg0: !async.coro.id) -> !async.coro.handle {
+  // CHECK: %0 = async.coro.begin %arg0
+  // CHECK: return %0 : !async.coro.handle
+  %0 = async.coro.begin %arg0
+  return %0 : !async.coro.handle
+}
+
+// CHECK-LABEL: @coro_free
+func @coro_free(%arg0: !async.coro.id, %arg1: !async.coro.handle) {
+  // CHECK: async.coro.free %arg0, %arg1
+  async.coro.free %arg0, %arg1
+  return
+}
+
+// CHECK-LABEL: @coro_end
+func @coro_end(%arg0: !async.coro.handle) {
+  // CHECK: async.coro.end %arg0
+  async.coro.end %arg0
+  return
+}
+
+// CHECK-LABEL: @coro_save
+func @coro_save(%arg0: !async.coro.handle) -> !async.coro.state {
+  // CHECK: %0 = async.coro.save %arg0
+  %0 = async.coro.save %arg0
+  // CHECK: return %0 : !async.coro.state
+  return %0 : !async.coro.state
+}
+
+// CHECK-LABEL: @coro_suspend
+func @coro_suspend(%arg0: !async.coro.state) {
+  // CHECK:      async.coro.suspend %arg0
+  // CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]
+  async.coro.suspend %arg0, ^suspend, ^resume, ^cleanup
+^resume:
+  // CHECK: ^[[RESUME]]
+  // CHECK: return {coro.resume}
+  return { coro.resume }
+^cleanup:
+  // CHECK: ^[[CLEANUP]]
+  // CHECK: return {coro.cleanup}
+  return { coro.cleanup }
+^suspend:
+  // CHECK: ^[[SUSPEND]]
+  // CHECK: return {coro.suspend}
+  return { coro.suspend }
+}

mlir/test/Dialect/Async/ops.mlir

func @create_group_and_await_all(%arg0: !async.token, %arg1: !async.value<f32>) -> index {
   // CHECK: async.add_to_group %arg1
   %1 = async.add_to_group %arg0, %0 : !async.token
   %2 = async.add_to_group %arg1, %0 : !async.value<f32>
   async.await_all %0
 
   %3 = addi %1, %2 : index
   return %3 : index
 }
-
-// CHECK-LABEL: @add_ref
-func @add_ref(%arg0: !async.token) {
-  // CHECK: async.add_ref %arg0 {count = 1 : i32}
-  async.add_ref %arg0 {count = 1 : i32} : !async.token
-  return
-}
-
-// CHECK-LABEL: @drop_ref
-func @drop_ref(%arg0: !async.token) {
-  // CHECK: async.drop_ref %arg0 {count = 1 : i32}
-  async.drop_ref %arg0 {count = 1 : i32} : !async.token
-  return
-}

mlir/test/Dialect/Async/runtime.mlir

+// RUN: mlir-opt %s | FileCheck %s --dump-input=always
+
+// CHECK-LABEL: @create_token
+func @create_token() -> !async.token {
+  // CHECK: %0 = async.runtime.create : !async.token
+  %0 = async.runtime.create : !async.token
+  // CHECK: return %0 : !async.token
+  return %0 : !async.token
+}
+
+// CHECK-LABEL: @create_value
+func @create_value() -> !async.value<f32> {
+  // CHECK: %0 = async.runtime.create : !async.value<f32>
+  %0 = async.runtime.create : !async.value<f32>
+  // CHECK: return %0 : !async.value<f32>
+  return %0 : !async.value<f32>
+}
+
+// CHECK-LABEL: @create_group
+func @create_group() -> !async.group {
+  // CHECK: %0 = async.runtime.create : !async.group
+  %0 = async.runtime.create : !async.group
+  // CHECK: return %0 : !async.group
+  return %0 : !async.group
+}
+
+// CHECK-LABEL: @set_token_available
+func @set_token_available(%arg0: !async.token) {
+  // CHECK: async.runtime.set_available %arg0 : !async.token
+  async.runtime.set_available %arg0 : !async.token
+  return
+}
+
+// CHECK-LABEL: @set_value_available
+func @set_value_available(%arg0: !async.value<f32>) {
+  // CHECK: async.runtime.set_available %arg0 : !async.value<f32>
+  async.runtime.set_available %arg0 : !async.value<f32>
+  return
+}
+
+// CHECK-LABEL: @await_token
+func @await_token(%arg0: !async.token) {
+  // CHECK: async.runtime.await %arg0 : !async.token
+  async.runtime.await %arg0 : !async.token
+  return
+}
+
+// CHECK-LABEL: @await_value
+func @await_value(%arg0: !async.value<f32>) {
+  // CHECK: async.runtime.await %arg0 : !async.value<f32>
+  async.runtime.await %arg0 : !async.value<f32>
+  return
+}
+
+// CHECK-LABEL: @await_group
+func @await_group(%arg0: !async.group) {
+  // CHECK: async.runtime.await %arg0 : !async.group
+  async.runtime.await %arg0 : !async.group
+  return
+}
+
+// CHECK-LABEL: @await_and_resume_token
+func @await_and_resume_token(%arg0: !async.token,
+                             %arg1: !async.coro.handle) {
+  // CHECK: async.runtime.await_and_resume %arg0, %arg1 : !async.token
+  async.runtime.await_and_resume %arg0, %arg1 : !async.token
+  return
+}
+
+// CHECK-LABEL: @await_and_resume_value
+func @await_and_resume_value(%arg0: !async.value<f32>,
+                             %arg1: !async.coro.handle) {
+  // CHECK: async.runtime.await_and_resume %arg0, %arg1 : !async.value<f32>
+  async.runtime.await_and_resume %arg0, %arg1 : !async.value<f32>
+  return
+}
+
+// CHECK-LABEL: @await_and_resume_group
+func @await_and_resume_group(%arg0: !async.group,
+                             %arg1: !async.coro.handle) {
+  // CHECK: async.runtime.await_and_resume %arg0, %arg1 : !async.group
+  async.runtime.await_and_resume %arg0, %arg1 : !async.group
+  return
+}
+
+// CHECK-LABEL: @resume
+func @resume(%arg0: !async.coro.handle) {
+  // CHECK: async.runtime.resume %arg0
+  async.runtime.resume %arg0
+  return
+}
+
+// CHECK-LABEL: @store
+func @store(%arg0: f32, %arg1: !async.value<f32>) {
+  // CHECK: async.runtime.store %arg0, %arg1 : !async.value<f32>
+  async.runtime.store %arg0, %arg1 : !async.value<f32>
+  return
+}
+
+// CHECK-LABEL: @load
+func @load(%arg0: !async.value<f32>) -> f32 {
+  // CHECK: %0 = async.runtime.load %arg0 : !async.value<f32>
+  // CHECK: return %0 : f32
+  %0 = async.runtime.load %arg0 : !async.value<f32>
+  return %0 : f32
+}
+
+// CHECK-LABEL: @add_to_group
+func @add_to_group(%arg0: !async.token, %arg1: !async.value<f32>,
+                   %arg2: !async.group) {
+  // CHECK: async.runtime.add_to_group %arg0, %arg2 : !async.token
+  async.runtime.add_to_group %arg0, %arg2 : !async.token
+  // CHECK: async.runtime.add_to_group %arg1, %arg2 : !async.value<f32>
+  async.runtime.add_to_group %arg1, %arg2 : !async.value<f32>
+  return
+}
+
+// CHECK-LABEL: @add_ref
+func @add_ref(%arg0: !async.token) {
+  // CHECK: async.runtime.add_ref %arg0 {count = 1 : i32}
+  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
+  return
+}
+
+// CHECK-LABEL: @drop_ref
+func @drop_ref(%arg0: !async.token) {
+  // CHECK: async.runtime.drop_ref %arg0 {count = 1 : i32}
+  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
+  return
+}