[mlir] Async: lowering async.value to LLVM

[ezhulenev]

Summary

1. Add new methods to Async runtime API to support yielding async values
2. Add lowering from async.yield with value payload to the new runtime API calls

async.value lowering requires that payload type is convertible to LLVM and supported by llvm.mlir.cast (DialectCast) operation.

[csigg]

mlir/include/mlir/ExecutionEngine/AsyncRuntime.h

75

Typo.

77

Could you document this parameter?

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

133

Could you add brief comments what AsyncTokens/Values/Groups are?

[ezhulenev]

Add documentation

[ezhulenev]

Fix style warnings

[ezhulenev]

Rebase

[ezhulenev]

Add missing targets to CMakeLists

[imaihal]

Hi, @ezhulenev
I'm testing MLIR on System z (Triple = "s390x-ibm-linux", CPU = "z14"), but I realized installation test (cmake --build . --target check-mlir) failed after this patch.
Endianness of System z is big-endian. Do you think this patch depends on the endianness? Do you have any suggestions to fix this?

******************** TEST 'MLIR :: mlir-cpu-runner/async-value.mlir' FAILED ********************
Script:
--
: 'RUN: at line 1';     /home/imaihal/work/llvm-project/build/bin/mlir-opt /home/imaihal/docker/imaihal-ubuntu/work/llvm-project/mlir/te
st/mlir-cpu-runner/async-value.mlir -async-ref-counting                                                     -convert-async-to-llvm      
                                            -convert-vector-to-llvm                                                 -convert-std-to-llvm
                                      | mlir-cpu-runner                                                           -e main -entry-point-r
esult=void -O0                                    -shared-libs=/home/imaihal/work/llvm-project/build/lib/libmlir_c_runner_utils.so      
 -shared-libs=/home/imaihal/work/llvm-project/build/lib/libmlir_runner_utils.so         -shared-libs=/home/imaihal/work/llvm-project/bui
ld/lib/libmlir_async_runtime.so    | /home/imaihal/work/llvm-project/build/bin/FileCheck /home/imaihal/docker/imaihal-ubuntu/work/llvm-p
roject/mlir/test/mlir-cpu-runner/async-value.mlir --dump-input=always
--
Exit Code: 2

Command Output (stderr):
--
LLVM ERROR: out of memory
Allocation failed
PLEASE submit a bug report to https://bugs.llvm.org/ and include the crash backtrace.
Stack dump:
0.      Program arguments: mlir-cpu-runner -e main -entry-point-result=void -O0 -shared-libs=/home/imaihal/work/llvm-project/build/lib/l
ibmlir_c_runner_utils.so -shared-libs=/home/imaihal/work/llvm-project/build/lib/libmlir_runner_utils.so -shared-libs=/home/imaihal/work/
llvm-project/build/lib/libmlir_async_runtime.so
 #0 0x000002aa38d54f18 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/home/imaihal/docker/imaihal-ubuntu/work/llvm-project/build/
bin/mlir-cpu-runner+0x2d4f18)
 #1 0x000002aa38d52dc6 llvm::sys::RunSignalHandlers() (/home/imaihal/docker/imaihal-ubuntu/work/llvm-project/build/bin/mlir-cpu-runner+0
x2d2dc6)
 #2 0x000002aa38d52f82 SignalHandler(int) (/home/imaihal/docker/imaihal-ubuntu/work/llvm-project/build/bin/mlir-cpu-runner+0x2d2f82)
 #3 0x000002aa3d689b2e 
 #4 0x000003ff96fbdef4 raise (/lib/s390x-linux-gnu/libc.so.6+0x3def4)
 #5 0x000003ff96fbf37a abort (/lib/s390x-linux-gnu/libc.so.6+0x3f37a)
 #6 0x000002aa38d6a22e llvm::report_bad_alloc_error(char const*, bool) (/home/imaihal/docker/imaihal-ubuntu/work/llvm-project/build/bin/
mlir-cpu-runner+0x2ea22e)
 #7 0x000002aa38d6a26c (/home/imaihal/docker/imaihal-ubuntu/work/llvm-project/build/bin/mlir-cpu-runner+0x2ea26c)
 #8 0x000003ff973a957a operator new(unsigned long) (/usr/lib/s390x-linux-gnu/libstdc++.so.6+0xa957a)
 #9 0x000003ff96d82a80 mlirAsyncRuntimeCreateValue (/home/imaihal/work/llvm-project/build/lib/libmlir_async_runtime.so+0x2a80)
#10 0x000003ff97749200 
#11 0x000003ff97749460 
#12 0x000002aa38e666be compileAndExecute((anonymous namespace)::Options&, mlir::ModuleOp, llvm::StringRef, (anonymous namespace)::CompileAndExecuteConfig, void**) (/home/imaihal/docker/imaihal-ubuntu/work/llvm-project/build/bin/mlir-cpu-runner+0x3e66be)
#13 0x000002aa38e66c0c compileAndExecuteVoidFunction((anonymous namespace)::Options&, mlir::ModuleOp, llvm::StringRef, (anonymous namespace)::CompileAndExecuteConfig) (/home/imaihal/docker/imaihal-ubuntu/work/llvm-project/build/bin/mlir-cpu-runner+0x3e6c0c)
#14 0x000002aa38e6b3a0 mlir::JitRunnerMain(int, char**, mlir::JitRunnerConfig) (/home/imaihal/docker/imaihal-ubuntu/work/llvm-project/build/bin/mlir-cpu-runner+0x3eb3a0)
#15 0x000002aa38cd4588 main (/home/imaihal/docker/imaihal-ubuntu/work/llvm-project/build/bin/mlir-cpu-runner+0x254588)
#16 0x000003ff96fa3aca __libc_start_main (/lib/s390x-linux-gnu/libc.so.6+0x23aca)
#17 0x000002aa38d40efc _start (/home/imaihal/docker/imaihal-ubuntu/work/llvm-project/build/bin/mlir-cpu-runner+0x2c0efc)
#18 0x0000000000000000 
FileCheck error: '<stdin>' is empty.
FileCheck command line:  /home/imaihal/work/llvm-project/build/bin/FileCheck /home/imaihal/docker/imaihal-ubuntu/work/llvm-project/mlir/test/mlir-cpu-runner/async-value.mlir --dump-input=always

I simplified the async.value.mlir as follows, but I still got the error.

func @main() {

  // ------------------------------------------------------------------------ //
  // Blocking async.await outside of the async.execute.
  // ------------------------------------------------------------------------ //
  %token, %result = async.execute -> !async.value<f32> {
    %0 = constant 123.456 : f32
    async.yield %0 : f32
  }
  %1 = async.await %result : !async.value<f32>

  // CHECK: 123.456
  vector.print %1 : f32

  return
}

func private @print_memref_f32(memref<*xf32>)
  attributes { llvm.emit_c_interface }

[ezhulenev]

Can the problem be in calculation of the value storage size? (https://github.com/llvm/llvm-project/blob/d480f968ad8b56d3ee4a6b6df5532d485b0ad01e/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp#L518-L542) I'd check what the value of size passed in here: https://github.com/llvm/llvm-project/blob/main/mlir/lib/ExecutionEngine/AsyncRuntime.cpp#L205-L208

[imaihal]

In D93592#2717318, @ezhulenev wrote:

Can the problem be in calculation of the value storage size?

@ezhulenev Thanks for replying. It seems the size has problem, but, I see strange behavior.
When the size is printed out in mlirAsyncRuntimeCreateValue(), it seems correct (the size was 4), but the execution failed. When I wrote it directory in the argument like AsyncValue *value = new AsyncValue(getDefaultAsyncRuntime(), 4);, it worked without the error.

Also, when I wrote the size directory here as follows, it worked.

AsyncValue(AsyncRuntime *runtime, int32_t size)
      : RefCounted(runtime, /*count=*/2), storage(4) {}

[imaihal]

In D93592#2715960, @imaihal wrote:

Hi, @ezhulenev
I'm testing MLIR on System z (Triple = "s390x-ibm-linux", CPU = "z14"), but I realized installation test (cmake --build . --target check-mlir) failed after this patch.
Endianness of System z is big-endian. Do you think this patch depends on the endianness? Do you have any suggestions to fix this?

Update:
I found an workaround to avoid this error. In gcc(gcc (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0), I needed to use -O0 instead of -O3 in this file. When using clang, I didn't got this error. This might be a bug in gcc.

mlir/include/mlir/ExecutionEngine/AsyncRuntime.h

 //===----------------------------------------------------------------------===//
 
 // Runtime implementation of `async.token` data type.
 typedef struct AsyncToken AsyncToken;
 
 // Runtime implementation of `async.group` data type.
 typedef struct AsyncGroup AsyncGroup;
 
+// Runtime implementation of `async.value` data type.
+typedef struct AsyncValue AsyncValue;
+
+// Async value payload stored in a memory owned by the async.value.
+using ValueStorage = void *;
+
 // Async runtime uses LLVM coroutines to represent asynchronous tasks. Task
 // function is a coroutine handle and a resume function that continue coroutine
 // execution from a suspension point.
 using CoroHandle = void *;           // coroutine handle
 using CoroResume = void (*)(void *); // coroutine resume function
 
 // Async runtime uses reference counting to manage the lifetime of async values
 // (values of async types like tokens, values and groups).
 using RefCountedObjPtr = void *;
 
 // Adds references to reference counted runtime object.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
     mlirAsyncRuntimeAddRef(RefCountedObjPtr, int32_t);
 
 // Drops references from reference counted runtime object.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
     mlirAsyncRuntimeDropRef(RefCountedObjPtr, int32_t);
 
 // Create a new `async.token` in not-ready state.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT AsyncToken *mlirAsyncRuntimeCreateToken();
 
+// Create a new `async.valie` in not-ready state.
+extern "C" MLIR_ASYNCRUNTIME_EXPORT AsyncValue *
+    mlirAsyncRuntimeCreateValue(int32_t);
+
 // Create a new `async.group` in empty state.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT AsyncGroup *mlirAsyncRuntimeCreateGroup();
 
 extern "C" MLIR_ASYNCRUNTIME_EXPORT int64_t
 mlirAsyncRuntimeAddTokenToGroup(AsyncToken *, AsyncGroup *);
 
 // Switches `async.token` to ready state and runs all awaiters.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
 mlirAsyncRuntimeEmplaceToken(AsyncToken *);
 
+// Switches `async.value` to ready state and runs all awaiters.
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeEmplaceValue(AsyncValue *);
+
 // Blocks the caller thread until the token becomes ready.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
 mlirAsyncRuntimeAwaitToken(AsyncToken *);
 
+// Blocks the caller thread until the value becomes ready.
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeAwaitValue(AsyncValue *);
+
 // Blocks the caller thread until the elements in the group become ready.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
 mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *);
 
+// Returns a pointer to the storage owned by the async value.
+extern "C" MLIR_ASYNCRUNTIME_EXPORT ValueStorage
+mlirAsyncRuntimeGetValueStorage(AsyncValue *);
+
 // Executes the task (coro handle + resume function) in one of the threads
 // managed by the runtime.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void mlirAsyncRuntimeExecute(CoroHandle,
                                                                  CoroResume);
 
 // Executes the task (coro handle + resume function) in one of the threads
 // managed by the runtime after the token becomes ready.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
 mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *, CoroHandle, CoroResume);
 
 // Executes the task (coro handle + resume function) in one of the threads
+// managed by the runtime after the value becomes ready.
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeAwaitValueAndExecute(AsyncValue *, CoroHandle, CoroResume);
+
+// Executes the task (coro handle + resume function) in one of the threads
 // managed by the runtime after the all members of the group become ready.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
 mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *, CoroHandle, CoroResume);
 
 //===----------------------------------------------------------------------===//
 // Small async runtime support library for testing.
 //===----------------------------------------------------------------------===//
 
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void mlirAsyncRuntimePrintCurrentThreadId();
 
 } // namespace runtime
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

 //===- AsyncToLLVM.cpp - Convert Async to LLVM dialect --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 
 #include "../PassDetail.h"
+#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/StandardOps/Transforms/FuncConversions.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/FormatVariadic.h"
 
 #define DEBUG_TYPE "convert-async-to-llvm"
 
 using namespace mlir;
 using namespace mlir::async;
 
 // Prefix for functions outlined from `async.execute` op regions.
 static constexpr const char kAsyncFnPrefix[] = "async_execute_fn";
 
 //===----------------------------------------------------------------------===//
 // Async Runtime C API declaration.
 //===----------------------------------------------------------------------===//
 
 static constexpr const char *kAddRef = "mlirAsyncRuntimeAddRef";
 static constexpr const char *kDropRef = "mlirAsyncRuntimeDropRef";
 static constexpr const char *kCreateToken = "mlirAsyncRuntimeCreateToken";
+static constexpr const char *kCreateValue = "mlirAsyncRuntimeCreateValue";
 static constexpr const char *kCreateGroup = "mlirAsyncRuntimeCreateGroup";
 static constexpr const char *kEmplaceToken = "mlirAsyncRuntimeEmplaceToken";
+static constexpr const char *kEmplaceValue = "mlirAsyncRuntimeEmplaceValue";
 static constexpr const char *kAwaitToken = "mlirAsyncRuntimeAwaitToken";
+static constexpr const char *kAwaitValue = "mlirAsyncRuntimeAwaitValue";
 static constexpr const char *kAwaitGroup = "mlirAsyncRuntimeAwaitAllInGroup";
 static constexpr const char *kExecute = "mlirAsyncRuntimeExecute";
+static constexpr const char *kGetValueStorage =
+    "mlirAsyncRuntimeGetValueStorage";
 static constexpr const char *kAddTokenToGroup =
     "mlirAsyncRuntimeAddTokenToGroup";
-static constexpr const char *kAwaitAndExecute =
+static constexpr const char *kAwaitTokenAndExecute =
     "mlirAsyncRuntimeAwaitTokenAndExecute";
+static constexpr const char *kAwaitValueAndExecute =
+    "mlirAsyncRuntimeAwaitValueAndExecute";
 static constexpr const char *kAwaitAllAndExecute =
     "mlirAsyncRuntimeAwaitAllInGroupAndExecute";
 
 namespace {
-// Async Runtime API function types.
+/// Async Runtime API function types.
+///
+/// Because we can't create API function signature for type parametrized
+/// async.value type, we use opaque pointers (!llvm.ptr<i8>) instead. After
+/// lowering all async data types become opaque pointers at runtime.
 struct AsyncAPI {
   static FunctionType addOrDropRefFunctionType(MLIRContext *ctx) {
     auto ref = LLVM::LLVMType::getInt8PtrTy(ctx);
     auto count = IntegerType::get(32, ctx);
     return FunctionType::get({ref, count}, {}, ctx);
   }
 
   static FunctionType createTokenFunctionType(MLIRContext *ctx) {
     return FunctionType::get({}, {TokenType::get(ctx)}, ctx);
   }
 
+  static FunctionType createValueFunctionType(MLIRContext *ctx) {
+    auto i32 = IntegerType::get(32, ctx);
+    auto value = LLVM::LLVMType::getInt8PtrTy(ctx);
+    return FunctionType::get({i32}, {value}, ctx);
+  }
+
   static FunctionType createGroupFunctionType(MLIRContext *ctx) {
     return FunctionType::get({}, {GroupType::get(ctx)}, ctx);
   }
 
+  static FunctionType getValueStorageFunctionType(MLIRContext *ctx) {
+    auto value = LLVM::LLVMType::getInt8PtrTy(ctx);
+    auto storage = LLVM::LLVMType::getInt8PtrTy(ctx);
+    return FunctionType::get({value}, {storage}, ctx);
+  }
+
   static FunctionType emplaceTokenFunctionType(MLIRContext *ctx) {
     return FunctionType::get({TokenType::get(ctx)}, {}, ctx);
   }
 
+  static FunctionType emplaceValueFunctionType(MLIRContext *ctx) {
+    auto value = LLVM::LLVMType::getInt8PtrTy(ctx);
+    return FunctionType::get({value}, {}, ctx);
+  }
+
   static FunctionType awaitTokenFunctionType(MLIRContext *ctx) {
     return FunctionType::get({TokenType::get(ctx)}, {}, ctx);
   }
 
+  static FunctionType awaitValueFunctionType(MLIRContext *ctx) {
+    auto value = LLVM::LLVMType::getInt8PtrTy(ctx);
+    return FunctionType::get({value}, {}, ctx);
+  }
+
   static FunctionType awaitGroupFunctionType(MLIRContext *ctx) {
     return FunctionType::get({GroupType::get(ctx)}, {}, ctx);
   }
 
   static FunctionType executeFunctionType(MLIRContext *ctx) {
     auto hdl = LLVM::LLVMType::getInt8PtrTy(ctx);
     auto resume = resumeFunctionType(ctx).getPointerTo();
     return FunctionType::get({hdl, resume}, {}, ctx);
   }
 
   static FunctionType addTokenToGroupFunctionType(MLIRContext *ctx) {
     auto i64 = IntegerType::get(64, ctx);
     return FunctionType::get({TokenType::get(ctx), GroupType::get(ctx)}, {i64},
                              ctx);
   }
 
-  static FunctionType awaitAndExecuteFunctionType(MLIRContext *ctx) {
+  static FunctionType awaitTokenAndExecuteFunctionType(MLIRContext *ctx) {
     auto hdl = LLVM::LLVMType::getInt8PtrTy(ctx);
     auto resume = resumeFunctionType(ctx).getPointerTo();
     return FunctionType::get({TokenType::get(ctx), hdl, resume}, {}, ctx);
   }
 
+  static FunctionType awaitValueAndExecuteFunctionType(MLIRContext *ctx) {
+    auto value = LLVM::LLVMType::getInt8PtrTy(ctx);
+    auto hdl = LLVM::LLVMType::getInt8PtrTy(ctx);
+    auto resume = resumeFunctionType(ctx).getPointerTo();
+    return FunctionType::get({value, hdl, resume}, {}, ctx);
+  }
+
   static FunctionType awaitAllAndExecuteFunctionType(MLIRContext *ctx) {
     auto hdl = LLVM::LLVMType::getInt8PtrTy(ctx);
     auto resume = resumeFunctionType(ctx).getPointerTo();
     return FunctionType::get({GroupType::get(ctx), hdl, resume}, {}, ctx);
   }
 
   // Auxiliary coroutine resume intrinsic wrapper.
   static LLVM::LLVMType resumeFunctionType(MLIRContext *ctx) {
     auto voidTy = LLVM::LLVMType::getVoidTy(ctx);
     auto i8Ptr = LLVM::LLVMType::getInt8PtrTy(ctx);
     return LLVM::LLVMType::getFunctionTy(voidTy, {i8Ptr}, false);
   }
 };
 } // namespace
 
-// Adds Async Runtime C API declarations to the module.
+/// Adds Async Runtime C API declarations to the module.
 static void addAsyncRuntimeApiDeclarations(ModuleOp module) {
   auto builder = OpBuilder::atBlockTerminator(module.getBody());
 
   auto addFuncDecl = [&](StringRef name, FunctionType type) {
     if (module.lookupSymbol(name))
       return;
     builder.create<FuncOp>(module.getLoc(), name, type).setPrivate();
   };
 
   MLIRContext *ctx = module.getContext();
   addFuncDecl(kAddRef, AsyncAPI::addOrDropRefFunctionType(ctx));
   addFuncDecl(kDropRef, AsyncAPI::addOrDropRefFunctionType(ctx));
   addFuncDecl(kCreateToken, AsyncAPI::createTokenFunctionType(ctx));
+  addFuncDecl(kCreateValue, AsyncAPI::createValueFunctionType(ctx));
   addFuncDecl(kCreateGroup, AsyncAPI::createGroupFunctionType(ctx));
   addFuncDecl(kEmplaceToken, AsyncAPI::emplaceTokenFunctionType(ctx));
+  addFuncDecl(kEmplaceValue, AsyncAPI::emplaceValueFunctionType(ctx));
   addFuncDecl(kAwaitToken, AsyncAPI::awaitTokenFunctionType(ctx));
+  addFuncDecl(kAwaitValue, AsyncAPI::awaitValueFunctionType(ctx));
   addFuncDecl(kAwaitGroup, AsyncAPI::awaitGroupFunctionType(ctx));
   addFuncDecl(kExecute, AsyncAPI::executeFunctionType(ctx));
+  addFuncDecl(kGetValueStorage, AsyncAPI::getValueStorageFunctionType(ctx));
   addFuncDecl(kAddTokenToGroup, AsyncAPI::addTokenToGroupFunctionType(ctx));
-  addFuncDecl(kAwaitAndExecute, AsyncAPI::awaitAndExecuteFunctionType(ctx));
+  addFuncDecl(kAwaitTokenAndExecute,
+              AsyncAPI::awaitTokenAndExecuteFunctionType(ctx));
+  addFuncDecl(kAwaitValueAndExecute,
+              AsyncAPI::awaitValueAndExecuteFunctionType(ctx));
   addFuncDecl(kAwaitAllAndExecute,
               AsyncAPI::awaitAllAndExecuteFunctionType(ctx));
 }
 
 //===----------------------------------------------------------------------===//
 // LLVM coroutines intrinsics declarations.
 //===----------------------------------------------------------------------===//
 

 }
 
 //===----------------------------------------------------------------------===//
 // Coroutine resume function wrapper.
 //===----------------------------------------------------------------------===//
 
 static constexpr const char *kResume = "__resume";
 
-// A function that takes a coroutine handle and calls a `llvm.coro.resume`
-// intrinsics. We need this function to be able to pass it to the async
-// runtime execute API.
+/// A function that takes a coroutine handle and calls a `llvm.coro.resume`
+/// intrinsics. We need this function to be able to pass it to the async
+/// runtime execute API.
 static void addResumeFunction(ModuleOp module) {
   MLIRContext *ctx = module.getContext();
 
   OpBuilder moduleBuilder(module.getBody()->getTerminator());
   Location loc = module.getLoc();
 
   if (module.lookupSymbol(kResume))
     return;
static void addResumeFunction(ModuleOp module) {
 
   blockBuilder.create<LLVM::ReturnOp>(loc, 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.
+/// 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 {
-  Value asyncToken;
+  // 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;
 };
 } // namespace
 
-// Builds an coroutine template compatible with LLVM coroutines lowering.
-//
-//  - `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 {
-//     ^entryBlock(<function-arguments>):
-//       %token = <async token> : !async.token // create async runtime token
-//       %hdl = llvm.call @llvm.coro.id(...)   // create a coroutine handle
-//       br ^cleanup
-//
-//     ^cleanup:
-//       llvm.call @llvm.coro.free(...)        // delete coroutine state
-//       br ^suspend
-//
-//     ^suspend:
-//       llvm.call @llvm.coro.end(...)         // marks the end of a coroutine
-//       return %token : !async.token
-//   }
-//
-// The actual code for the async.execute operation body region will be inserted
-// before the entry block terminator.
-//
-//
+/// Builds an coroutine template compatible with LLVM coroutines lowering.
+///
+///  - `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>):
+///       %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
+///       br ^cleanup
+///
+///     ^cleanup:
+///       llvm.call @llvm.coro.free(...)  // delete coroutine state
+///       br ^suspend
+///
+///     ^suspend:
+///       llvm.call @llvm.coro.end(...)  // 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 = LLVM::LLVMType::getInt1Ty(ctx);
   auto i32 = LLVM::LLVMType::getInt32Ty(ctx);
   auto i64 = LLVM::LLVMType::getInt64Ty(ctx);
   auto i8Ptr = LLVM::LLVMType::getInt8PtrTy(ctx);
 
   Block *entryBlock = func.addEntryBlock();
   Location loc = func.getBody().getLoc();
 
   OpBuilder builder = OpBuilder::atBlockBegin(entryBlock);
 
   // ------------------------------------------------------------------------ //
   // Allocate async tokens/values that we will return from a ramp function.
   // ------------------------------------------------------------------------ //
   auto createToken =
       builder.create<CallOp>(loc, 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 = storedType.cast<LLVM::LLVMType>().getPointerTo();
+
+    // %Size = getelementptr %T* null, int 1
+    // %SizeI = ptrtoint %T* %Size to i32
+    auto nullPtr = builder.create<LLVM::NullOp>(loc, storagePtrType);
+    auto i32Type = LLVM::LLVMType::getInt32Ty(ctx);
+    auto one = builder.create<LLVM::ConstantOp>(loc, i32Type,
+                                                builder.getI32IntegerAttr(1));
+    auto gep = builder.create<LLVM::GEPOp>(loc, storagePtrType, nullPtr,
+                                           one.getResult());
+    auto size = builder.create<LLVM::PtrToIntOp>(loc, i32Type, gep);
+
+    // Cast to std type because runtime API defined using std types.
+    return builder.create<LLVM::DialectCastOp>(loc, builder.getI32Type(),
+                                               size.getResult());
+  };
+
+  // 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());
+
   // ------------------------------------------------------------------------ //
   // Initialize coroutine: allocate frame, get coroutine handle.
   // ------------------------------------------------------------------------ //
 
   // Constants for initializing coroutine frame.
   auto constZero =
       builder.create<LLVM::ConstantOp>(loc, i32, builder.getI32IntegerAttr(0));
   auto constFalse =
static CoroMachinery setupCoroMachinery(FuncOp func) {
   // async token.
   // ------------------------------------------------------------------------ //
   builder.setInsertionPointToStart(suspendBlock);
 
   // Mark the end of a coroutine: @llvm.coro.end.
   builder.create<LLVM::CallOp>(loc, i1, builder.getSymbolRefAttr(kCoroEnd),
                                ValueRange({coroHdl.getResult(0), constFalse}));
 
-  // Return created `async.token` from the suspend block. This will be the
-  // return value of a coroutine ramp function.
-  builder.create<ReturnOp>(loc, createToken.getResult(0));
+  // 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);
 
   // Branch from the entry block to the cleanup block to create a valid CFG.
   builder.setInsertionPointToEnd(entryBlock);
 
   builder.create<BranchOp>(loc, cleanupBlock);
 
   // `async.await` op lowering will create resume blocks for async
   // continuations, and will conditionally branch to cleanup or suspend blocks.
 
-  return {createToken.getResult(0), coroHdl.getResult(0), cleanupBlock,
-          suspendBlock};
-}
-
-// 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: ...
-//
+  CoroMachinery machinery;
+  machinery.asyncToken = createToken.getResult(0);
+  machinery.returnValues = returnValues;
+  machinery.coroHandle = coroHdl.getResult(0);
+  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 = LLVM::LLVMType::getInt1Ty(ctx);
   auto i8 = LLVM::LLVMType::getInt8Ty(ctx);
 
static void addSuspensionPoint(CoroMachinery coro, Value coroState,
   // ... 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.
+/// 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>();
 
   MLIRContext *ctx = module.getContext();
   Location loc = execute.getLoc();
 
   OpBuilder moduleBuilder(module.getBody()->getTerminator());
 
   // Collect all outlined function inputs.
   llvm::SetVector<mlir::Value> functionInputs(execute.dependencies().begin(),
                                               execute.dependencies().end());
+  assert(execute.operands().empty() && "operands are not supported");
   getUsedValuesDefinedAbove(execute.body(), functionInputs);
 
   // Collect types for the outlined function inputs and outputs.
   auto typesRange = llvm::map_range(
       functionInputs, [](Value value) { return value.getType(); });
   SmallVector<Type, 4> inputTypes(typesRange.begin(), typesRange.end());
   auto outputTypes = execute.getResultTypes();
 
for (size_t i = 0; i < execute.dependencies().size(); ++i)
     builder.create<AwaitOp>(loc, func.getArgument(i));
 
   // Map from function inputs defined above the execute op to the function
   // arguments.
   BlockAndValueMapping valueMapping;
   valueMapping.map(functionInputs, func.getArguments());
 
   // Clone all operations from the execute operation body into the outlined
-  // function body, and replace all `async.yield` operations with a call
-  // to async runtime to emplace the result token.
-  for (Operation &op : execute.body().getOps()) {
-    if (isa<async::YieldOp>(op)) {
-      builder.create<CallOp>(loc, kEmplaceToken, TypeRange(), coro.asyncToken);
-      continue;
-    }
+  // function body.
+  for (Operation &op : execute.body().getOps())
     builder.clone(op, valueMapping);
-  }
 
   // Replace the original `async.execute` with a call to outlined function.
   OpBuilder callBuilder(execute);
   auto callOutlinedFunc =
       callBuilder.create<CallOp>(loc, func.getName(), execute.getResultTypes(),
                                  functionInputs.getArrayRef());
   execute.replaceAllUsesWith(callOutlinedFunc.getResults());
   execute.erase();
 
   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(convertType); }
+  AsyncRuntimeTypeConverter() {
+    addConversion([](Type type) { return type; });
+    addConversion(convertAsyncTypes);
+  }
 
-  static Type convertType(Type type) {
-    MLIRContext *ctx = type.getContext();
-    // Convert async tokens and groups to opaque pointers.
-    if (type.isa<TokenType, GroupType>())
-      return LLVM::LLVMType::getInt8PtrTy(ctx);
-    return type;
+  static Optional<Type> convertAsyncTypes(Type type) {
+    if (type.isa<TokenType, GroupType, ValueType>())
+      return LLVM::LLVMType::getInt8PtrTy(type.getContext());
+    return llvm::None;
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// Convert types for all call operations to lowered async types.
+// Convert return operations that return async values from async regions.
 //===----------------------------------------------------------------------===//
 
 namespace {
-class CallOpOpConversion : public ConversionPattern {
+class ReturnOpOpConversion : public ConversionPattern {
 public:
-  explicit CallOpOpConversion(MLIRContext *ctx)
-      : ConversionPattern(CallOp::getOperationName(), 1, ctx) {}
+  explicit ReturnOpOpConversion(TypeConverter &converter, MLIRContext *ctx)
+      : ConversionPattern(ReturnOp::getOperationName(), 1, converter, ctx) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    AsyncRuntimeTypeConverter converter;
-
-    SmallVector<Type, 5> resultTypes;
-    converter.convertTypes(op->getResultTypes(), resultTypes);
-
-    CallOp call = cast<CallOp>(op);
-    rewriter.replaceOpWithNewOp<CallOp>(op, resultTypes, call.callee(),
-                                        operands);
-
+    rewriter.replaceOpWithNewOp<ReturnOp>(op, operands);
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Async reference counting ops lowering (`async.add_ref` and `async.drop_ref`
 // to the corresponding API calls).
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 template <typename RefCountingOp>
 class RefCountingOpLowering : public ConversionPattern {
 public:
-  explicit RefCountingOpLowering(MLIRContext *ctx, StringRef apiFunctionName)
-      : ConversionPattern(RefCountingOp::getOperationName(), 1, ctx),
+  explicit RefCountingOpLowering(TypeConverter &converter, MLIRContext *ctx,
+                                 StringRef apiFunctionName)
+      : ConversionPattern(RefCountingOp::getOperationName(), 1, converter, ctx),
         apiFunctionName(apiFunctionName) {}
 
   LogicalResult
   matchAndRewrite(Operation *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()));
 
     rewriter.replaceOpWithNewOp<CallOp>(op, TypeRange(), apiFunctionName,
                                         ValueRange({operands[0], count}));
 
     return success();
   }
 
 private:
   StringRef apiFunctionName;
 };
 
-// async.drop_ref op lowering to mlirAsyncRuntimeDropRef function call.
+/// async.drop_ref op lowering to mlirAsyncRuntimeDropRef function call.
 class AddRefOpLowering : public RefCountingOpLowering<AddRefOp> {
 public:
-  explicit AddRefOpLowering(MLIRContext *ctx)
-      : RefCountingOpLowering(ctx, kAddRef) {}
+  explicit AddRefOpLowering(TypeConverter &converter, MLIRContext *ctx)
+      : RefCountingOpLowering(converter, ctx, kAddRef) {}
 };
 
-// async.create_group op lowering to mlirAsyncRuntimeCreateGroup function call.
+/// async.create_group op lowering to mlirAsyncRuntimeCreateGroup function call.
 class DropRefOpLowering : public RefCountingOpLowering<DropRefOp> {
 public:
-  explicit DropRefOpLowering(MLIRContext *ctx)
-      : RefCountingOpLowering(ctx, kDropRef) {}
+  explicit DropRefOpLowering(TypeConverter &converter, MLIRContext *ctx)
+      : RefCountingOpLowering(converter, ctx, kDropRef) {}
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // async.create_group op lowering to mlirAsyncRuntimeCreateGroup function call.
 //===----------------------------------------------------------------------===//
 
 namespace {
 class CreateGroupOpLowering : public ConversionPattern {
 public:
-  explicit CreateGroupOpLowering(MLIRContext *ctx)
-      : ConversionPattern(CreateGroupOp::getOperationName(), 1, ctx) {}
+  explicit CreateGroupOpLowering(TypeConverter &converter, MLIRContext *ctx)
+      : ConversionPattern(CreateGroupOp::getOperationName(), 1, converter,
+                          ctx) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto retTy = GroupType::get(op->getContext());
     rewriter.replaceOpWithNewOp<CallOp>(op, kCreateGroup, retTy);
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // async.add_to_group op lowering to runtime function call.
 //===----------------------------------------------------------------------===//
 
 namespace {
 class AddToGroupOpLowering : public ConversionPattern {
 public:
-  explicit AddToGroupOpLowering(MLIRContext *ctx)
-      : ConversionPattern(AddToGroupOp::getOperationName(), 1, ctx) {}
+  explicit AddToGroupOpLowering(TypeConverter &converter, MLIRContext *ctx)
+      : ConversionPattern(AddToGroupOp::getOperationName(), 1, converter, ctx) {
+  }
 
   LogicalResult
   matchAndRewrite(Operation *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();

 //===----------------------------------------------------------------------===//
 
 namespace {
 
 template <typename AwaitType, typename AwaitableType>
 class AwaitOpLoweringBase : public ConversionPattern {
 protected:
   explicit AwaitOpLoweringBase(
-      MLIRContext *ctx,
+      TypeConverter &converter, MLIRContext *ctx,
       const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions,
       StringRef blockingAwaitFuncName, StringRef coroAwaitFuncName)
-      : ConversionPattern(AwaitType::getOperationName(), 1, ctx),
+      : ConversionPattern(AwaitType::getOperationName(), 1, converter, ctx),
         outlinedFunctions(outlinedFunctions),
         blockingAwaitFuncName(blockingAwaitFuncName),
         coroAwaitFuncName(coroAwaitFuncName) {}
 
 public:
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     // We can only await on one the `AwaitableType` (for `await` it can be
-    // only a `token`, for `await_all` it is a `group`).
+    // 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();
 
     // Check if await operation is inside the outlined coroutine function.
     auto func = await->template getParentOfType<FuncOp>();
     auto outlined = outlinedFunctions.find(func);
     const bool isInCoroutine = outlined != outlinedFunctions.end();
if (isInCoroutine) {
                              awaitAndExecuteArgs);
 
       Block *suspended = op->getBlock();
 
       // 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);
+
+      // Make sure that replacement value will be constructed in resume block.
+      rewriter.setInsertionPointToStart(resume);
     }
 
-    // Original operation was replaced by function call or suspension point.
+    // Replace or erase the await operation with the new value.
+    if (Value replaceWith = getReplacementValue(op, operands[0], rewriter))
+      rewriter.replaceOp(op, replaceWith);
+    else
       rewriter.eraseOp(op);
 
     return success();
   }
 
+  virtual Value getReplacementValue(Operation *op, Value operand,
+                                    ConversionPatternRewriter &rewriter) const {
+    return Value();
+  }
+
 private:
   const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions;
   StringRef blockingAwaitFuncName;
   StringRef coroAwaitFuncName;
 };
 
-// Lowering for `async.await` operation (only token operands are supported).
-class AwaitOpLowering : public AwaitOpLoweringBase<AwaitOp, TokenType> {
+/// Lowering for `async.await` with a token operand.
+class AwaitTokenOpLowering : public AwaitOpLoweringBase<AwaitOp, TokenType> {
   using Base = AwaitOpLoweringBase<AwaitOp, TokenType>;
 
 public:
-  explicit AwaitOpLowering(
-      MLIRContext *ctx,
+  explicit AwaitTokenOpLowering(
+      TypeConverter &converter, MLIRContext *ctx,
       const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions)
-      : Base(ctx, outlinedFunctions, kAwaitToken, kAwaitAndExecute) {}
+      : Base(converter, ctx, outlinedFunctions, kAwaitToken,
+             kAwaitTokenAndExecute) {}
 };
 
-// Lowering for `async.await_all` operation.
+/// 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) {}
+
+  Value
+  getReplacementValue(Operation *op, Value operand,
+                      ConversionPatternRewriter &rewriter) const override {
+    Location loc = op->getLoc();
+    auto i8Ptr = LLVM::LLVMType::getInt8PtrTy(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, llvmValueType.cast<LLVM::LLVMType>().getPointerTo(),
+        storage.getResult(0));
+
+    // Load from the async value storage.
+    auto loaded = rewriter.create<LLVM::LoadOp>(loc, castedStorage.getResult());
+
+    // Cast from LLVM type to the expected value type. This cast will become
+    // no-op after lowering to LLVM.
+    return rewriter.create<LLVM::DialectCastOp>(loc, valueType, loaded);
+  }
+};
+
+/// Lowering for `async.await_all` operation.
 class AwaitAllOpLowering : public AwaitOpLoweringBase<AwaitAllOp, GroupType> {
   using Base = AwaitOpLoweringBase<AwaitAllOp, GroupType>;
 
 public:
   explicit AwaitAllOpLowering(
-      MLIRContext *ctx,
+      TypeConverter &converter, MLIRContext *ctx,
       const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions)
-      : Base(ctx, outlinedFunctions, kAwaitGroup, kAwaitAllAndExecute) {}
+      : Base(converter, ctx, outlinedFunctions, kAwaitGroup,
+             kAwaitAllAndExecute) {}
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
+// async.yield op lowerings to the corresponding async runtime function calls.
+//===----------------------------------------------------------------------===//
+
+class YieldOpLowering : public ConversionPattern {
+public:
+  explicit YieldOpLowering(
+      TypeConverter &converter, MLIRContext *ctx,
+      const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions)
+      : ConversionPattern(async::YieldOp::getOperationName(), 1, converter,
+                          ctx),
+        outlinedFunctions(outlinedFunctions) {}
+
+  LogicalResult
+  matchAndRewrite(Operation *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");
+
+    Location loc = op->getLoc();
+    const CoroMachinery &coro = outlined->getSecond();
+
+    // Store yielded values into the async values storage and emplace them.
+    auto i8Ptr = LLVM::LLVMType::getInt8PtrTy(rewriter.getContext());
+
+    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, yieldValue.getType().cast<LLVM::LLVMType>().getPointerTo(),
+          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);
+    }
+
+    // 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);
+
+    return success();
+  }
+
+private:
+  const llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions;
+};
+
+//===----------------------------------------------------------------------===//
 
 namespace {
 struct ConvertAsyncToLLVMPass
     : public ConvertAsyncToLLVMBase<ConvertAsyncToLLVMPass> {
   void runOnOperation() override;
 };
 
 void ConvertAsyncToLLVMPass::runOnOperation() {
   ModuleOp module = getOperation();
   SymbolTable symbolTable(module);
 
+  MLIRContext *ctx = &getContext();
+
   // Outline all `async.execute` body regions into async functions (coroutines).
   llvm::DenseMap<FuncOp, CoroMachinery> outlinedFunctions;
 
+  // We use conversion to LLVM type to ensure that all `async.value` operands
+  // and results can be lowered to LLVM load and store operations.
+  LLVMTypeConverter llvmConverter(ctx);
+  llvmConverter.addConversion(AsyncRuntimeTypeConverter::convertAsyncTypes);
+
+  // Returns true if the `async.value` payload is convertible to LLVM.
+  auto isConvertibleToLlvm = [&](Type type) -> bool {
+    auto valueType = type.cast<ValueType>().getValueType();
+    return static_cast<bool>(llvmConverter.convertType(valueType));
+  };
+
   WalkResult outlineResult = module.walk([&](ExecuteOp execute) {
+    // All operands and results must be convertible to LLVM.
+    if (!llvm::all_of(execute.operands().getTypes(), isConvertibleToLlvm)) {
+      execute.emitOpError("operands payload must be convertible to LLVM type");
+      return WalkResult::interrupt();
+    }
+    if (!llvm::all_of(execute.results().getTypes(), isConvertibleToLlvm)) {
+      execute.emitOpError("results payload must be convertible to LLVM type");
+      return WalkResult::interrupt();
+    }
+
     // We currently do not support execute operations that have async value
     // operands or produce async results.
-    if (!execute.operands().empty() || !execute.results().empty()) {
-      execute.emitOpError("can't outline async.execute op with async value "
-                          "operands or returned async results");
+    if (!execute.operands().empty()) {
+      execute.emitOpError(
+          "can't outline async.execute op with async value operands");
       return WalkResult::interrupt();
     }
 
     outlinedFunctions.insert(outlineExecuteOp(symbolTable, execute));
 
     return WalkResult::advance();
   });
 
void ConvertAsyncToLLVMPass::runOnOperation() {
   });
 
   // Add declarations for all functions required by the coroutines lowering.
   addResumeFunction(module);
   addAsyncRuntimeApiDeclarations(module);
   addCoroutineIntrinsicsDeclarations(module);
   addCRuntimeDeclarations(module);
 
-  MLIRContext *ctx = &getContext();
-
   // 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);
-  patterns.insert<CallOpOpConversion>(ctx);
-  patterns.insert<AddRefOpLowering, DropRefOpLowering>(ctx);
-  patterns.insert<CreateGroupOpLowering, AddToGroupOpLowering>(ctx);
-  patterns.insert<AwaitOpLowering, AwaitAllOpLowering>(ctx, outlinedFunctions);
+  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);
 
   ConversionTarget target(*ctx);
   target.addLegalOp<ConstantOp>();
   target.addLegalDialect<LLVM::LLVMDialect>();
+
+  // All operations from Async dialect must be lowered to the runtime API 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<CallOp>(
-      [&](CallOp op) { return converter.isLegal(op.getResultTypes()); });
+  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
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertAsyncToLLVMPass() {
   return std::make_unique<ConvertAsyncToLLVMPass>();
 }

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

struct AsyncToken : public RefCounted {
   // Internal state below guarded by a mutex.
   std::mutex mu;
   std::condition_variable cv;
 
   bool ready = false;
   std::vector<std::function<void()>> awaiters;
 };
 
+struct AsyncValue : public RefCounted {
+  // AsyncValue similar to an AsyncToken created with a reference count of 2.
+  AsyncValue(AsyncRuntime *runtime, int32_t size)
+      : RefCounted(runtime, /*count=*/2), storage(size) {}
+
+  // Internal state below guarded by a mutex.
+  std::mutex mu;
+  std::condition_variable cv;
+
+  bool ready = false;
+  std::vector<std::function<void()>> awaiters;
+
+  // Use vector of bytes to store async value payload.
+  std::vector<int8_t> storage;
+};
+
 struct AsyncGroup : public RefCounted {
   AsyncGroup(AsyncRuntime *runtime)
       : RefCounted(runtime), pendingTokens(0), rank(0) {}
 
   std::atomic<int> pendingTokens;
   std::atomic<int> rank;
 
   // Internal state below guarded by a mutex.

 }
 
 // Drops references from reference counted runtime object.
 extern "C" void mlirAsyncRuntimeDropRef(RefCountedObjPtr ptr, int32_t count) {
   RefCounted *refCounted = static_cast<RefCounted *>(ptr);
   refCounted->dropRef(count);
 }
 
-// Create a new `async.token` in not-ready state.
+// Creates a new `async.token` in not-ready state.
 extern "C" AsyncToken *mlirAsyncRuntimeCreateToken() {
   AsyncToken *token = new AsyncToken(getDefaultAsyncRuntimeInstance());
   return token;
 }
 
+// Creates a new `async.value` in not-ready state.
+extern "C" AsyncValue *mlirAsyncRuntimeCreateValue(int32_t size) {
+  AsyncValue *value = new AsyncValue(getDefaultAsyncRuntimeInstance(), size);
+  return value;
+}
+
 // Create a new `async.group` in empty state.
 extern "C" AsyncGroup *mlirAsyncRuntimeCreateGroup() {
   AsyncGroup *group = new AsyncGroup(getDefaultAsyncRuntimeInstance());
   return group;
 }
 
 extern "C" int64_t mlirAsyncRuntimeAddTokenToGroup(AsyncToken *token,
                                                    AsyncGroup *group) {
for (auto &awaiter : token->awaiters)
       awaiter();
   }
 
   // Async tokens created with a ref count `2` to keep token alive until the
   // async task completes. Drop this reference explicitly when token emplaced.
   token->dropRef();
 }
 
+// Switches `async.value` to ready state and runs all awaiters.
+extern "C" void mlirAsyncRuntimeEmplaceValue(AsyncValue *value) {
+  // Make sure that `dropRef` does not destroy the mutex owned by the lock.
+  {
+    std::unique_lock<std::mutex> lock(value->mu);
+    value->ready = true;
+    value->cv.notify_all();
+    for (auto &awaiter : value->awaiters)
+      awaiter();
+  }
+
+  // Async values created with a ref count `2` to keep value alive until the
+  // async task completes. Drop this reference explicitly when value emplaced.
+  value->dropRef();
+}
+
 extern "C" void mlirAsyncRuntimeAwaitToken(AsyncToken *token) {
   std::unique_lock<std::mutex> lock(token->mu);
   if (!token->ready)
     token->cv.wait(lock, [token] { return token->ready; });
 }
 
+extern "C" void mlirAsyncRuntimeAwaitValue(AsyncValue *value) {
+  std::unique_lock<std::mutex> lock(value->mu);
+  if (!value->ready)
+    value->cv.wait(lock, [value] { return value->ready; });
+}
+
 extern "C" void mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *group) {
   std::unique_lock<std::mutex> lock(group->mu);
   if (group->pendingTokens != 0)
     group->cv.wait(lock, [group] { return group->pendingTokens == 0; });
 }
 
+// Returns a pointer to the storage owned by the async value.
+extern "C" ValueStorage mlirAsyncRuntimeGetValueStorage(AsyncValue *value) {
+  return value->storage.data();
+}
+
 extern "C" void mlirAsyncRuntimeExecute(CoroHandle handle, CoroResume resume) {
   (*resume)(handle);
 }
 
 extern "C" void mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *token,
                                                      CoroHandle handle,
                                                      CoroResume resume) {
   std::unique_lock<std::mutex> lock(token->mu);
   auto execute = [handle, resume]() { (*resume)(handle); };
   if (token->ready)
     execute();
   else
     token->awaiters.push_back([execute]() { execute(); });
 }
 
+extern "C" void mlirAsyncRuntimeAwaitValueAndExecute(AsyncValue *value,
+                                                     CoroHandle handle,
+                                                     CoroResume resume) {
+  std::unique_lock<std::mutex> lock(value->mu);
+  auto execute = [handle, resume]() { (*resume)(handle); };
+  if (value->ready)
+    execute();
+  else
+    value->awaiters.push_back([execute]() { execute(); });
+}
+
 extern "C" void mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *group,
                                                           CoroHandle handle,
                                                           CoroResume resume) {
   std::unique_lock<std::mutex> lock(group->mu);
   auto execute = [handle, resume]() { (*resume)(handle); };
   if (group->pendingTokens == 0)
     execute();
   else

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

 
 // Suspend coroutine second time waiting for the group.
 // CHECK: llvm.call @llvm.coro.save
 // CHECK: call @mlirAsyncRuntimeAwaitAllInGroupAndExecute(%arg0, %[[HDL_1]],
 // CHECK: llvm.call @llvm.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
+  %token, %result = async.execute -> !async.value<f32> {
+    %c0 = constant 123.0 : f32
+    async.yield %c0 : f32
+  }
+
+  // CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[RET]]#1)
+  // CHECK: %[[ST_F32:.*]] = llvm.bitcast %[[STORAGE]]
+  // CHECK: %[[LOADED:.*]] = llvm.load %[[ST_F32]] :  !llvm.ptr<float>
+  // CHECK: %[[CASTED:.*]] = llvm.mlir.cast %[[LOADED]] : !llvm.float to f32
+  %0 = async.await %result : !async.value<f32>
+
+  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
+
+// Suspend coroutine in the beginning.
+// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]],
+// CHECK: llvm.call @llvm.coro.suspend
+
+// Emplace result value.
+// CHECK: %[[CST:.*]] = constant 1.230000e+02 : f32
+// CHECK: %[[LLVM_CST:.*]] = llvm.mlir.cast %[[CST]] : f32 to !llvm.float
+// CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[VALUE]])
+// CHECK: %[[ST_F32:.*]] = llvm.bitcast %[[STORAGE]]
+// CHECK: llvm.store %[[LLVM_CST]], %[[ST_F32]] : !llvm.ptr<float>
+// CHECK: call @mlirAsyncRuntimeEmplaceValue(%[[VALUE]])
+
+// Emplace result token.
+// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]])
+

mlir/test/mlir-cpu-runner/async-value.mlir

+// RUN:   mlir-opt %s -async-ref-counting                                      \
+// RUN:               -convert-async-to-llvm                                   \
+// RUN:               -convert-vector-to-llvm                                  \
+// RUN:               -convert-std-to-llvm                                     \
+// RUN: | mlir-cpu-runner                                                      \
+// RUN:     -e main -entry-point-result=void -O0                               \
+// RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_c_runner_utils%shlibext  \
+// RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext    \
+// RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_async_runtime%shlibext   \
+// RUN: | FileCheck %s --dump-input=always
+
+func @main() {
+
+  // ------------------------------------------------------------------------ //
+  // Blocking async.await outside of the async.execute.
+  // ------------------------------------------------------------------------ //
+  %token, %result = async.execute -> !async.value<f32> {
+    %0 = constant 123.456 : f32
+    async.yield %0 : f32
+  }
+  %1 = async.await %result : !async.value<f32>
+
+  // CHECK: 123.456
+  vector.print %1 : f32
+
+  // ------------------------------------------------------------------------ //
+  // Non-blocking async.await inside the async.execute
+  // ------------------------------------------------------------------------ //
+  %token0, %result0 = async.execute -> !async.value<f32> {
+    %token1, %result2 = async.execute -> !async.value<f32> {
+      %2 = constant 456.789 : f32
+      async.yield %2 : f32
+    }
+    %3 = async.await %result2 : !async.value<f32>
+    async.yield %3 : f32
+  }
+  %4 = async.await %result0 : !async.value<f32>
+
+  // CHECK: 456.789
+  vector.print %4 : f32
+
+  // ------------------------------------------------------------------------ //
+  // Memref allocated inside async.execute region.
+  // ------------------------------------------------------------------------ //
+  %token2, %result2 = async.execute[%token0] -> !async.value<memref<f32>> {
+    %5 = alloc() : memref<f32>
+    %c0 = constant 987.654 : f32
+    store %c0, %5[]: memref<f32>
+    async.yield %5 : memref<f32>
+  }
+  %6 = async.await %result2 : !async.value<memref<f32>>
+  %7 = memref_cast %6 :  memref<f32> to memref<*xf32>
+
+  // CHECK: Unranked Memref
+  // CHECK-SAME: rank = 0 offset = 0 sizes = [] strides = []
+  // CHECK-NEXT: [987.654]
+  call @print_memref_f32(%7): (memref<*xf32>) -> ()
+  dealloc %6 : memref<f32>
+
+  return
+}
+
+func private @print_memref_f32(memref<*xf32>)
+  attributes { llvm.emit_c_interface }