Optionally eliminate blocking runtime.await calls by converting functions to coroutines.

[bakhtiyarneyman]

Summary

Interop parallelism requires needs awaiting on results. Blocking awaits are bad for performance. TFRT supports lightweight resumption on threads, and coroutines are an abstraction than can be used to lower the kernels onto TFRT threads.

[bakhtiyarneyman]

Address clang warnings.

[mehdi_amini]

mlir/include/mlir/Dialect/Async/Passes.td

47

Why is this an option and not just the default?

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

562

This is repeating the function name, can you expand on how it is achieving this instead?

566

Nit: you don't need to specify the size in SmallVector anymore.
(same below in other places)

572

setResultAttrs and getResultAttrs aren't really efficient API. setAllResultAttrs and getAllResultAttrs should be able to do it better.

But also, it seems that what you really want to do here is to say "add a token as the first return for the function", and there is an API for this on FuncOp I believe:

/// Insert a single result of type `resultType` at `resultIndex`.
  void insertResult(unsigned resultIndex, Type resultType,
                    DictionaryAttr resultAttrs) {

576

Nit early-form reduces indentation:

for (Block &block : func.getBlocks()) {
    if (&block != coro.suspend)
       continue;
    Operation *terminator = block.getTerminator();
    ...

584

Where is this invariant verified/enforced before getting into this point?

591

595

Can you make this function (and all the others) static please

608

Please favor for-range loop and avoid indexing when possible:

SmallVector<Value> unwrappedResults;
unwrappedResults.reserve(newCall->getResults().size() - 1);
for (Value result : newCall.getResults().drop_front())
   unwrappedResults.push_back(callBuilder.create<AwaitOp>(loc, result));

610

I don't understand this comment, it may miss some context: dangling pointer where? Is this referring to the comment in front of this function The invocation of this function is safe only when call ops are traversed in reverse order of how they appear in a single block. See funcsToCoroutines.? If so, it seems a bit out-of-place here (we're not careful in these two lines, only in the contract with this function and its caller which is documented on the API itself).

644

walk() is recursive and would traverse all the ops inside the functions.

651–652

(you'll see that LLVM/MLIR uses SmallVector for worklists in general, and because we don't support exceptions we have pop_back_val() :) )

658

You can save double-queries to maps when you know you want the entry inserted regardless:

auto newMapEntry = outlinedFunctions.insert({func, CoroMachinery{});
if (!insertion.second) {
    // This function has already been processed because this is either
    // the corecursive case, or a caller with multiple calls to a newly
    // created corouting. Either way, skip updating the call sites.
    continue;
}
insertion.first->second = rewriteFuncAsCoroutine(func);

664

677

I suspect this can fire on valid IR: can you make this a proper error handling (or if you want to keep an assert here, document the invariant and make it a runtime check and signalPassFailure() at the beginning of runOnOperation.

[ezhulenev]

mlir/include/mlir/Dialect/Async/Passes.td

47

I should be the default once it will be tested, I'm +1 to have it as an option for now, so it can be quickly disabled without a need to push complex fixes.

[mehdi_amini]

mlir/include/mlir/Dialect/Async/Passes.td

47

OK, then please document it if the option is temporary for bringup :)

[bakhtiyarneyman]

Addressing comments.

[mehdi_amini]

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

574

Is this all still useful here? I'd expect the call below to take care of all this.

677

How is this handled now?

[ezhulenev]

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

570

SmallVector<Type> resultTypes = {TokenType::get(ctx)} to skip func.insertResult below?

[mehdi_amini]

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

570

insertResult actually handles also argument attributes remapping.

574

Oh I missed that this is wrapping the results into a "ValueType", forget this comment.

[bakhtiyarneyman]

Addressing comments.

[bakhtiyarneyman]

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

570

Yes, the reason we are using insertResult is to avoid attributes being shifted.

677

I think it was not pushed properly, PTAL.

[mehdi_amini]

Thanks, LGTM

mlir/include/mlir/Dialect/Async/Passes.td

def AsyncParallelFor : Pass<"async-parallel-for", "ModuleOp"> {
 
   let dependentDialects = ["async::AsyncDialect", "scf::SCFDialect"];
 }
 
 def AsyncToAsyncRuntime : Pass<"async-to-async-runtime", "ModuleOp"> {
   let summary = "Lower high level async operations (e.g. async.execute) to the"
                 "explicit async.runtime and async.coro operations";
   let constructor = "mlir::createAsyncToAsyncRuntimePass()";
+  let options = [
+    Option<"eliminateBlockingAwaitOps", "eliminate-blocking-await-ops", "bool",
+            /*default=*/"false",
+           "Rewrite functions with blocking async.runtime.await as coroutines "
+           "with async.runtime.await_and_resume.">
+  ];
   let dependentDialects = ["async::AsyncDialect"];
 }
 
 def AsyncRuntimeRefCounting : Pass<"async-runtime-ref-counting"> {
   let summary = "Automatic reference counting for Async runtime operations";
   let description = [{
     This pass works at the async runtime abtraction level, after all
     `async.execute` and `async.await` operations are lowered to the async

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

struct CoroMachinery {
 
   Value coroHandle; // coroutine handle (!async.coro.handle value)
   Block *setError;  // switch completion token and all values to error state
   Block *cleanup;   // coroutine cleanup block
   Block *suspend;   // coroutine suspension block
 };
 } // namespace
 
-/// Builds an coroutine template compatible with LLVM coroutines switched-resume
-/// lowering using `async.runtime.*` and `async.coro.*` operations.
+/// Utility to partially update the regular function CFG to the coroutine CFG
+/// compatible with LLVM coroutines switched-resume lowering using
+/// `async.runtime.*` and `async.coro.*` operations. Modifies the entry block
+/// by prepending its ops with coroutine setup. Also inserts trailing blocks.
+///
+/// The result types of the passed `func` must start with an `async.token`
+/// and be continued with some number of `async.value`s.
+///
+/// It's up to the caller of this function to fix up the terminators of the
+/// preexisting blocks of the passed func op. If the passed `func` is legal,
+/// this typically means rewriting every return op as a yield op and a branch op
+/// to the suspend block.
 ///
 /// See LLVM coroutines documentation: https://llvm.org/docs/Coroutines.html
 ///
 ///  - `entry` block sets up the coroutine.
 ///  - `set_error` block sets completion token and async values state to error.
 ///  - `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>)
+///   func @some_fn(<function-arguments>) -> (!async.token, !async.value<T>)
 ///   {
 ///     ^entry(<function-arguments>):
 ///       %token = <async token> : !async.token    // create async runtime token
 ///       %value = <async value> : !async.value<T> // create async value
 ///       %id = async.coro.id                      // create a coroutine id
 ///       %hdl = async.coro.begin %id              // create a coroutine handle
-///       br ^cleanup
+///       /* other ops of the preexisting entry block */
+///
+///     /* other preexisting blocks */
 ///
 ///     ^set_error: // this block created lazily only if needed (see code below)
 ///       async.runtime.set_error %token : !async.token
 ///       async.runtime.set_error %value : !async.value<T>
 ///       br ^cleanup
 ///
 ///     ^cleanup:
 ///       async.coro.free %hdl // delete the coroutine state
 ///       br ^suspend
 ///
 ///     ^suspend:
 ///       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");
+  assert(!func.getBlocks().empty() && "Function must have an entry block");
 
   MLIRContext *ctx = func.getContext();
-  Block *entryBlock = func.addEntryBlock();
-
+  Block *entryBlock = &func.getBlocks().front();
   auto builder = ImplicitLocOpBuilder::atBlockBegin(func->getLoc(), entryBlock);
 
   // ------------------------------------------------------------------------ //
   // Allocate async token/values that we will return from a ramp function.
   // ------------------------------------------------------------------------ //
   auto retToken = builder.create<RuntimeCreateOp>(TokenType::get(ctx)).result();
 
   llvm::SmallVector<Value, 4> retValues;
static CoroMachinery setupCoroMachinery(FuncOp func) {
   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{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.func = func;
   machinery.asyncToken = retToken;
   machinery.returnValues = retValues;
   machinery.coroHandle = coroHdlOp.handle();
outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
   // multiple nested async.execute operations).
   FuncOp func = FuncOp::create(loc, kAsyncFnPrefix, funcType, funcAttrs);
   symbolTable.insert(func);
 
   SymbolTable::setSymbolVisibility(func, SymbolTable::Visibility::Private);
 
   // Prepare a function for coroutine lowering by adding entry/cleanup/suspend
   // blocks, adding async.coro operations and setting up control flow.
+  func.addEntryBlock();
   CoroMachinery coro = setupCoroMachinery(func);
 
   // Suspend async function at the end of an entry block, and resume it using
   // 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);
+  auto builder = ImplicitLocOpBuilder::atBlockEnd(loc, entryBlock);
 
   // Save the coroutine state: async.coro.save
   auto coroSaveOp =
       builder.create<CoroSaveOp>(CoroStateType::get(ctx), coro.coroHandle);
 
   // Pass coroutine to the runtime to be resumed on a runtime managed thread.
   builder.create<RuntimeResumeOp>(coro.coroHandle);
+  builder.create<BranchOp>(coro.cleanup);
 
   // Split the entry block before the terminator (branch to suspend block).
   auto *terminatorOp = entryBlock->getTerminator();
   Block *suspended = terminatorOp->getBlock();
   Block *resume = suspended->splitBlock(terminatorOp);
 
   // Add async.coro.suspend as a suspended block terminator.
   builder.setInsertionPointToEnd(suspended);
public:
   }
 
 private:
   llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions;
 };
 
 //===----------------------------------------------------------------------===//
 
+/// Rewrite a func as a coroutine.
+CoroMachinery rewriteFuncAsCoroutine(FuncOp func) {
+  auto *ctx = func->getContext();
+  auto loc = func.getLoc();
+  SmallVector<Type, 4> resultTypes{TokenType::get(ctx)};
+  llvm::transform(func.getCallableResults(), std::back_inserter(resultTypes),
+                  [](Type type) { return ValueType::get(type); });
+  func.setType(FunctionType::get(ctx, func.getType().getInputs(), resultTypes));
+  for (int i = resultTypes.size() - 1; i > 0; --i) {
+    func.setResultAttrs(i, func.getResultAttrs(i - 1));
+  }
+  func.setResultAttrs(0, ArrayRef<NamedAttribute>());
+  CoroMachinery coro = setupCoroMachinery(func);
+  for (Block &block : func.getBlocks()) {
+    if (&block != coro.suspend) {
+      auto *terminator = block.getTerminator();
+      if (auto returnOp = dyn_cast<ReturnOp>(*terminator)) {
+        ImplicitLocOpBuilder builder(loc, returnOp);
+        builder.create<YieldOp>(returnOp.getOperands());
+        builder.create<BranchOp>(coro.cleanup);
+        returnOp.erase();
+      } else {
+        assert(dyn_cast<BranchOp>(*terminator) && "Unexpected terminator.");
+      }
+    }
+  }
+  return coro;
+};
+
+/// Rewrites a call into a function that has been rewiten as a coroutine.
+///
+/// The invocation of this function is safe only when call ops are traversed in
+/// reverse order of how they appear in a single block. See `funcsToCoroutines`.
+void rewriteCallsiteForCoroutine(CallOp oldCall, FuncOp func) {
+  auto loc = func.getLoc();
+  ImplicitLocOpBuilder callBuilder(loc, oldCall);
+  auto newCall = callBuilder.create<CallOp>(
+      func.getName(), func.getCallableResults(), oldCall.getArgOperands());
+
+  // Await on the async token and all the value results and unwrap the latter.
+  int size = newCall->getResults().size() - 1;
+  SmallVector<Value, 4> unwrappedResults(size);
+  callBuilder.create<AwaitOp>(loc, newCall.getResult(0));
+  for (size_t i = 0; i < size; ++i) {
+    Value result = newCall.getResult(i + 1);
+    unwrappedResults[i] = callBuilder.create<AwaitOp>(loc, result).result();
+  }
+  // Careful, when result of a call is piped into another call this could lead
+  // to a dangling pointer.
+  oldCall.replaceAllUsesWith(unwrappedResults);
+  oldCall.erase();
+};
+
+void funcsToCoroutines(
+    ModuleOp module, llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions) {
+  // The following code supports the general case when 2 functions mutually
+  // recurse into each other. Because of this and that we are relying on
+  // SymbolUserMap to find pointers to calling FuncOps, we cannot simply erase
+  // a FuncOp while inserting an equivalent coroutine, because that could lead
+  // to dangling pointers.
+
+  std::vector<FuncOp> funcWorklist;
+
+  // Careful, it's okay to add a func to the worklist multiple times if and only
+  // if the loop processing the worklist will skip the functions that have
+  // already been converted to coroutines.
+  auto addToWorklist = [&outlinedFunctions, &funcWorklist](FuncOp func) {
+    // N.B. To refactor this code into a separate pass the lookup in
+    // outlinedFunctions is the most obvious obstacle. Looking at an arbitrary
+    // func and recognizing if it has a coroutine structure is messy. Passing
+    // this dict between the passes is ugly.
+    if (outlinedFunctions.find(func) == outlinedFunctions.end()) {
+      for (Operation &op : func.body().getOps()) {
+        if (dyn_cast<AwaitOp>(op) || dyn_cast<AwaitAllOp>(op)) {
+          funcWorklist.push_back(func);
+          break;
+        }
+      }
+    }
+  };
+
+  // Traverse in post-order collecting for each func op the await ops it has.
+  module.walk(addToWorklist);
+
+  SymbolTableCollection symbolTable;
+  SymbolUserMap symbolUserMap(symbolTable, module);
+
+  // Rewrite funcs, while updating call sites and adding them to the worklist.
+  while (!funcWorklist.empty()) {
+    auto func = funcWorklist.back();
+    funcWorklist.pop_back();
+    if (outlinedFunctions.find(func) != outlinedFunctions.end())
+      // This function has already been processed because this is either
+      // the corecursive case, or a caller with multiple calls to a newly
+      // created corouting. Either way, skip updating the call sites.
+      continue;
+    outlinedFunctions.insert({func, rewriteFuncAsCoroutine(func)});
+    SmallVector<Operation *, 4> users(symbolUserMap.getUsers(func).begin(),
+                                      symbolUserMap.getUsers(func).end());
+    // If there are multiple calls from the same block they need to be traversed
+    // in reverse order so that symbolUserMap references are not invalidated
+    // when updating the users of the call op which is earlier in the block.
+    llvm::sort(users, [](auto *a, auto *b) {
+      Block *blockA = a->getBlock();
+      Block *blockB = b->getBlock();
+      // Impose arbitrary order on blocks so that there is a well-defined order.
+      return blockA > blockB || (blockA == blockB && !a->isBeforeInBlock(b));
+    });
+    // Rewrite the callsites to await on results of the newly created coroutine.
+    for (Operation *op : users) {
+      if (CallOp call = dyn_cast<mlir::CallOp>(*op)) {
+        FuncOp caller = call->getParentOfType<FuncOp>();
+        rewriteCallsiteForCoroutine(call, func); // Careful, erases the call op.
+        addToWorklist(caller);
+      } else {
+        assert(false && "FuncOp used by an op other than CallOp");
+      }
+    }
+  }
+}
+
+//===----------------------------------------------------------------------===//
 void AsyncToAsyncRuntimePass::runOnOperation() {
   ModuleOp module = getOperation();
   SymbolTable symbolTable(module);
 
   // Outline all `async.execute` body regions into async functions (coroutines).
   llvm::DenseMap<FuncOp, CoroMachinery> outlinedFunctions;
 
   module.walk([&](ExecuteOp execute) {
     outlinedFunctions.insert(outlineExecuteOp(symbolTable, execute));
   });
 
   LLVM_DEBUG({
     llvm::dbgs() << "Outlined " << outlinedFunctions.size()
                  << " functions built from async.execute operations\n";
   });
 
+  if (eliminateBlockingAwaitOps)
+    funcsToCoroutines(module, outlinedFunctions);
+
   // Lower async operations to async.runtime operations.
   MLIRContext *ctx = module->getContext();
   RewritePatternSet asyncPatterns(ctx);
 
   // Async lowering does not use type converter because it must preserve all
   // types for async.runtime operations.
   asyncPatterns.add<CreateGroupOpLowering, AddToGroupOpLowering>(ctx);
   asyncPatterns.add<AwaitTokenOpLowering, AwaitValueOpLowering,
void AsyncToAsyncRuntimePass::runOnOperation() {
 
   // Assertions must be converted to runtime errors inside async functions.
   runtimeTarget.addDynamicallyLegalOp<AssertOp>([&](AssertOp op) -> bool {
     auto func = op->getParentOfType<FuncOp>();
     return outlinedFunctions.find(func) == outlinedFunctions.end();
   });
   runtimeTarget.addLegalOp<CondBranchOp>();
 
+  if (eliminateBlockingAwaitOps)
+    runtimeTarget.addIllegalOp<RuntimeAwaitOp>();
+
   if (failed(applyPartialConversion(module, runtimeTarget,
                                     std::move(asyncPatterns)))) {
     signalPassFailure();
     return;
   }
 }
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createAsyncToAsyncRuntimePass() {
   return std::make_unique<AsyncToAsyncRuntimePass>();
 }

mlir/test/Dialect/Async/async-to-async-runtime-eliminate-blocking.mlir

+// RUN: mlir-opt %s -split-input-file \
+// RUN:    -async-to-async-runtime="eliminate-blocking-await-ops=true" \
+// RUN:   | FileCheck %s --dump-input=always
+
+// CHECK-LABEL: func @simple_callee
+// CHECK-SAME: (%[[ARG:.*]]: f32)
+// CHECK-SAME: -> (!async.token, !async.value<f32> {builtin.foo = "bar"})
+func @simple_callee(%arg0: f32) -> (f32 {builtin.foo = "bar"}) {
+// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
+// CHECK: %[[RETURNED_STORAGE:.*]] = async.runtime.create : !async.value<f32>
+// CHECK: %[[ID:.*]] = async.coro.id
+// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
+
+// CHECK: %[[VAL:.*]] = addf %[[ARG]], %[[ARG]] : f32
+  %0 = addf %arg0, %arg0 : f32
+// CHECK: %[[VAL_STORAGE:.*]] = async.runtime.create : !async.value<f32>
+  %1 = async.runtime.create: !async.value<f32>
+// CHECK: async.runtime.store %[[VAL]], %[[VAL_STORAGE]] : !async.value<f32>
+  async.runtime.store %0, %1: !async.value<f32>
+// CHECK: async.runtime.set_available %[[VAL_STORAGE]] : !async.value<f32>
+  async.runtime.set_available %1: !async.value<f32>
+
+// CHECK: %[[SAVED:.*]] = async.coro.save %[[HDL]]
+// CHECK: async.runtime.await_and_resume %[[VAL_STORAGE]], %[[HDL]]
+// CHECK: async.coro.suspend %[[SAVED]]
+// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]
+  %2 = async.await %1 : !async.value<f32>
+
+// CHECK: ^[[RESUME]]:
+// CHECK:   %[[IS_ERROR:.*]] = async.runtime.is_error %[[VAL_STORAGE]] : !async.value<f32>
+// CHECK:   cond_br %[[IS_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_OK:.*]]
+
+// CHECK: ^[[BRANCH_OK]]:
+// CHECK:   %[[LOADED:.*]] = async.runtime.load %[[VAL_STORAGE]] : !async.value<f32>
+// CHECK:   %[[RETURNED:.*]] = mulf %[[ARG]], %[[LOADED]] : f32
+// CHECK:   async.runtime.store %[[RETURNED]], %[[RETURNED_STORAGE]] : !async.value<f32>
+// CHECK:   async.runtime.set_available %[[RETURNED_STORAGE]]
+// CHECK:   async.runtime.set_available %[[TOKEN]]
+// CHECK:   br ^[[CLEANUP]]
+  %3 = mulf %arg0, %2 : f32
+  return %3: f32
+
+// CHECK: ^[[BRANCH_ERROR]]:
+// CHECK:   async.runtime.set_error %[[TOKEN]]
+// CHECK:   async.runtime.set_error %[[RETURNED_STORAGE]]
+// CHECK:   br ^[[CLEANUP]]
+
+
+// CHECK: ^[[CLEANUP]]:
+// CHECK:   async.coro.free %[[ID]], %[[HDL]]
+// CHECK:   br ^[[SUSPEND]]
+
+// CHECK: ^[[SUSPEND]]:
+// CHECK:   async.coro.end %[[HDL]]
+// CHECK:   return %[[TOKEN]], %[[RETURNED_STORAGE]] : !async.token, !async.value<f32>
+}
+
+// CHECK-LABEL: func @simple_caller()
+// CHECK-SAME: -> (!async.token, !async.value<f32>)
+func @simple_caller() -> f32 {
+// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
+// CHECK: %[[RETURNED_STORAGE:.*]] = async.runtime.create : !async.value<f32>
+// CHECK: %[[ID:.*]] = async.coro.id
+// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
+
+// CHECK: %[[CONSTANT:.*]] = constant
+  %c = constant 1.0 : f32
+// CHECK: %[[RETURNED_TO_CALLER:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value<f32>)
+// CHECK: %[[SAVED:.*]] = async.coro.save %[[HDL]]
+// CHECK: async.runtime.await_and_resume %[[RETURNED_TO_CALLER]]#0, %[[HDL]]
+// CHECK: async.coro.suspend %[[SAVED]]
+// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]
+  %r = call @simple_callee(%c): (f32) -> f32
+
+// CHECK: ^[[RESUME]]:
+// CHECK:   %[[IS_TOKEN_ERROR:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER]]#0 : !async.token
+// CHECK:   cond_br %[[IS_TOKEN_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_TOKEN_OK:.*]]
+
+// CHECK: ^[[BRANCH_TOKEN_OK]]:
+// CHECK:   %[[IS_VALUE_ERROR:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER]]#1 : !async.value<f32>
+// CHECK:   cond_br %[[IS_VALUE_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_VALUE_OK:.*]]
+
+// CHECK: ^[[BRANCH_VALUE_OK]]:
+// CHECK:   %[[LOADED:.*]] = async.runtime.load %[[RETURNED_TO_CALLER]]#1 : !async.value<f32>
+// CHECK:   async.runtime.store %[[LOADED]], %[[RETURNED_STORAGE]] : !async.value<f32>
+// CHECK:   async.runtime.set_available %[[RETURNED_STORAGE]]
+// CHECK:   async.runtime.set_available %[[TOKEN]]
+// CHECK:   br ^[[CLEANUP]]
+  return %r: f32
+// CHECK: ^[[BRANCH_ERROR]]:
+// CHECK:   async.runtime.set_error %[[TOKEN]]
+// CHECK:   async.runtime.set_error %[[RETURNED_STORAGE]]
+// CHECK:   br ^[[CLEANUP]]
+
+
+// CHECK: ^[[CLEANUP]]:
+// CHECK:   async.coro.free %[[ID]], %[[HDL]]
+// CHECK:   br ^[[SUSPEND]]
+
+// CHECK: ^[[SUSPEND]]:
+// CHECK:   async.coro.end %[[HDL]]
+// CHECK:   return %[[TOKEN]], %[[RETURNED_STORAGE]] : !async.token, !async.value<f32>
+}
+
+// CHECK-LABEL: func @double_caller()
+// CHECK-SAME: -> (!async.token, !async.value<f32>)
+func @double_caller() -> f32 {
+// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
+// CHECK: %[[RETURNED_STORAGE:.*]] = async.runtime.create : !async.value<f32>
+// CHECK: %[[ID:.*]] = async.coro.id
+// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
+
+// CHECK: %[[CONSTANT:.*]] = constant
+  %c = constant 1.0 : f32
+// CHECK: %[[RETURNED_TO_CALLER_1:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value<f32>)
+// CHECK: %[[SAVED_1:.*]] = async.coro.save %[[HDL]]
+// CHECK: async.runtime.await_and_resume %[[RETURNED_TO_CALLER_1]]#0, %[[HDL]]
+// CHECK: async.coro.suspend %[[SAVED_1]]
+// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_1:.*]], ^[[CLEANUP:.*]]
+  %r = call @simple_callee(%c): (f32) -> f32
+
+// CHECK: ^[[RESUME_1]]:
+// CHECK:   %[[IS_TOKEN_ERROR_1:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER_1]]#0 : !async.token
+// CHECK:   cond_br %[[IS_TOKEN_ERROR_1]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_TOKEN_OK_1:.*]]
+
+// CHECK: ^[[BRANCH_TOKEN_OK_1]]:
+// CHECK:   %[[IS_VALUE_ERROR_1:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER_1]]#1 : !async.value<f32>
+// CHECK:   cond_br %[[IS_VALUE_ERROR_1]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_VALUE_OK_1:.*]]
+
+// CHECK: ^[[BRANCH_VALUE_OK_1]]:
+// CHECK:   %[[LOADED_1:.*]] = async.runtime.load %[[RETURNED_TO_CALLER_1]]#1 : !async.value<f32>
+// CHECK:   %[[RETURNED_TO_CALLER_2:.*]]:2 = call @simple_callee(%[[LOADED_1]]) : (f32) -> (!async.token, !async.value<f32>)
+// CHECK:   %[[SAVED_2:.*]] = async.coro.save %[[HDL]]
+// CHECK:   async.runtime.await_and_resume %[[RETURNED_TO_CALLER_2]]#0, %[[HDL]]
+// CHECK:   async.coro.suspend %[[SAVED_2]]
+// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_2:.*]], ^[[CLEANUP:.*]]
+  %s = call @simple_callee(%r): (f32) -> f32
+
+// CHECK: ^[[RESUME_2]]:
+// CHECK:   %[[IS_TOKEN_ERROR_2:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER_2]]#0 : !async.token
+// CHECK:   cond_br %[[IS_TOKEN_ERROR_2]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_TOKEN_OK_2:.*]]
+
+// CHECK: ^[[BRANCH_TOKEN_OK_2]]:
+// CHECK:   %[[IS_VALUE_ERROR_2:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER_2]]#1 : !async.value<f32>
+// CHECK:   cond_br %[[IS_VALUE_ERROR_2]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_VALUE_OK_2:.*]]
+
+// CHECK: ^[[BRANCH_VALUE_OK_2]]:
+// CHECK:   %[[LOADED_2:.*]] = async.runtime.load %[[RETURNED_TO_CALLER_2]]#1 : !async.value<f32>
+// CHECK:   async.runtime.store %[[LOADED_2]], %[[RETURNED_STORAGE]] : !async.value<f32>
+// CHECK:   async.runtime.set_available %[[RETURNED_STORAGE]]
+// CHECK:   async.runtime.set_available %[[TOKEN]]
+// CHECK:   br ^[[CLEANUP]]
+  return %s: f32
+// CHECK: ^[[BRANCH_ERROR]]:
+// CHECK:   async.runtime.set_error %[[TOKEN]]
+// CHECK:   async.runtime.set_error %[[RETURNED_STORAGE]]
+// CHECK:   br ^[[CLEANUP]]
+
+// CHECK: ^[[CLEANUP]]:
+// CHECK:   async.coro.free %[[ID]], %[[HDL]]
+// CHECK:   br ^[[SUSPEND]]
+
+// CHECK: ^[[SUSPEND]]:
+// CHECK:   async.coro.end %[[HDL]]
+// CHECK:   return %[[TOKEN]], %[[RETURNED_STORAGE]] : !async.token, !async.value<f32>
+}
+
+// CHECK-LABEL: func @recursive
+// CHECK-SAME: (%[[ARG:.*]]: !async.token) -> !async.token
+func @recursive(%arg: !async.token) {
+// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
+// CHECK: %[[ID:.*]] = async.coro.id
+// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
+// CHECK: %[[SAVED_1:.*]] = async.coro.save %[[HDL]]
+// CHECK: async.runtime.await_and_resume %[[ARG]], %[[HDL]]
+// CHECK: async.coro.suspend %[[SAVED_1]]
+// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_1:.*]], ^[[CLEANUP:.*]]
+
+  async.await %arg : !async.token
+// CHECK: ^[[RESUME_1]]:
+// CHECK:   %[[IS_ERROR:.*]] = async.runtime.is_error %[[ARG]] : !async.token
+// CHECK:   cond_br %[[IS_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_OK:.*]]
+
+// CHECK: ^[[BRANCH_OK]]:
+// CHECK:   %[[GIVEN:.*]] = async.runtime.create : !async.token
+%r = async.runtime.create : !async.token
+// CHECK:   async.runtime.set_available %[[GIVEN]]
+async.runtime.set_available %r: !async.token
+// CHECK:   %[[RETURNED_TO_CALLER:.*]] = call @recursive(%[[GIVEN]]) : (!async.token) -> !async.token
+call @recursive(%r): (!async.token) -> ()
+// CHECK: %[[SAVED_2:.*]] = async.coro.save %[[HDL]]
+// CHECK: async.runtime.await_and_resume %[[RETURNED_TO_CALLER]], %[[HDL]]
+// CHECK: async.coro.suspend %[[SAVED_2]]
+// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_2:.*]], ^[[CLEANUP:.*]]
+
+// CHECK: ^[[RESUME_2]]:
+// CHECK:   async.runtime.set_available %[[TOKEN]]
+// CHECK:   br ^[[CLEANUP]]
+
+// CHECK: ^[[BRANCH_ERROR]]:
+// CHECK:   async.runtime.set_error %[[TOKEN]]
+// CHECK:   br ^[[CLEANUP]]
+return
+
+// CHECK: ^[[CLEANUP]]:
+// CHECK:   async.coro.free %[[ID]], %[[HDL]]
+// CHECK:   br ^[[SUSPEND]]
+
+// CHECK: ^[[SUSPEND]]:
+// CHECK:   async.coro.end %[[HDL]]
+// CHECK:   return %[[TOKEN]] : !async.token
+}
+
+// CHECK-LABEL: func @corecursive1
+// CHECK-SAME: (%[[ARG:.*]]: !async.token) -> !async.token
+func @corecursive1(%arg: !async.token) {
+// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
+// CHECK: %[[ID:.*]] = async.coro.id
+// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
+// CHECK: %[[SAVED_1:.*]] = async.coro.save %[[HDL]]
+// CHECK: async.runtime.await_and_resume %[[ARG]], %[[HDL]]
+// CHECK: async.coro.suspend %[[SAVED_1]]
+// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_1:.*]], ^[[CLEANUP:.*]]
+
+  async.await %arg : !async.token
+// CHECK: ^[[RESUME_1]]:
+// CHECK:   %[[IS_ERROR:.*]] = async.runtime.is_error %[[ARG]] : !async.token
+// CHECK:   cond_br %[[IS_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_OK:.*]]
+
+// CHECK: ^[[BRANCH_OK]]:
+// CHECK:   %[[GIVEN:.*]] = async.runtime.create : !async.token
+%r = async.runtime.create : !async.token
+// CHECK:   async.runtime.set_available %[[GIVEN]]
+async.runtime.set_available %r: !async.token
+// CHECK:   %[[RETURNED_TO_CALLER:.*]] = call @corecursive2(%[[GIVEN]]) : (!async.token) -> !async.token
+call @corecursive2(%r): (!async.token) -> ()
+// CHECK: %[[SAVED_2:.*]] = async.coro.save %[[HDL]]
+// CHECK: async.runtime.await_and_resume %[[RETURNED_TO_CALLER]], %[[HDL]]
+// CHECK: async.coro.suspend %[[SAVED_2]]
+// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_2:.*]], ^[[CLEANUP:.*]]
+
+// CHECK: ^[[RESUME_2]]:
+// CHECK:   async.runtime.set_available %[[TOKEN]]
+// CHECK:   br ^[[CLEANUP]]
+
+// CHECK: ^[[BRANCH_ERROR]]:
+// CHECK:   async.runtime.set_error %[[TOKEN]]
+// CHECK:   br ^[[CLEANUP]]
+return
+
+// CHECK: ^[[CLEANUP]]:
+// CHECK:   async.coro.free %[[ID]], %[[HDL]]
+// CHECK:   br ^[[SUSPEND]]
+
+// CHECK: ^[[SUSPEND]]:
+// CHECK:   async.coro.end %[[HDL]]
+// CHECK:   return %[[TOKEN]] : !async.token
+}
+
+// CHECK-LABEL: func @corecursive2
+// CHECK-SAME: (%[[ARG:.*]]: !async.token) -> !async.token
+func @corecursive2(%arg: !async.token) {
+// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
+// CHECK: %[[ID:.*]] = async.coro.id
+// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
+// CHECK: %[[SAVED_1:.*]] = async.coro.save %[[HDL]]
+// CHECK: async.runtime.await_and_resume %[[ARG]], %[[HDL]]
+// CHECK: async.coro.suspend %[[SAVED_1]]
+// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_1:.*]], ^[[CLEANUP:.*]]
+
+  async.await %arg : !async.token
+// CHECK: ^[[RESUME_1]]:
+// CHECK:   %[[IS_ERROR:.*]] = async.runtime.is_error %[[ARG]] : !async.token
+// CHECK:   cond_br %[[IS_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_OK:.*]]
+
+// CHECK: ^[[BRANCH_OK]]:
+// CHECK:   %[[GIVEN:.*]] = async.runtime.create : !async.token
+%r = async.runtime.create : !async.token
+// CHECK:   async.runtime.set_available %[[GIVEN]]
+async.runtime.set_available %r: !async.token
+// CHECK:   %[[RETURNED_TO_CALLER:.*]] = call @corecursive1(%[[GIVEN]]) : (!async.token) -> !async.token
+call @corecursive1(%r): (!async.token) -> ()
+// CHECK: %[[SAVED_2:.*]] = async.coro.save %[[HDL]]
+// CHECK: async.runtime.await_and_resume %[[RETURNED_TO_CALLER]], %[[HDL]]
+// CHECK: async.coro.suspend %[[SAVED_2]]
+// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_2:.*]], ^[[CLEANUP:.*]]
+
+// CHECK: ^[[RESUME_2]]:
+// CHECK:   async.runtime.set_available %[[TOKEN]]
+// CHECK:   br ^[[CLEANUP]]
+
+// CHECK: ^[[BRANCH_ERROR]]:
+// CHECK:   async.runtime.set_error %[[TOKEN]]
+// CHECK:   br ^[[CLEANUP]]
+return
+
+// CHECK: ^[[CLEANUP]]:
+// CHECK:   async.coro.free %[[ID]], %[[HDL]]
+// CHECK:   br ^[[SUSPEND]]
+
+// CHECK: ^[[SUSPEND]]:
+// CHECK:   async.coro.end %[[HDL]]
+// CHECK:   return %[[TOKEN]] : !async.token
+}