diff --git a/mlir/include/mlir/Dialect/Async/IR/AsyncDialect.td b/mlir/include/mlir/Dialect/Async/IR/AsyncDialect.td --- a/mlir/include/mlir/Dialect/Async/IR/AsyncDialect.td +++ b/mlir/include/mlir/Dialect/Async/IR/AsyncDialect.td @@ -32,9 +32,8 @@ let extraClassDeclaration = [{ /// The name of a unit attribute on funcs that are allowed to have a - /// blocking async.runtime.await ops. Only useful in combination with - /// 'eliminate-blocking-await-ops' option, which in absence of this - /// attribute might convert a func to a coroutine. + /// blocking async.runtime.await ops. In absence of this attribute the + /// asyncification pass might convert a func to a coroutine. static constexpr StringRef kAllowedToBlockAttrName = "async.allowed_to_block"; }]; diff --git a/mlir/include/mlir/Dialect/Async/Passes.td b/mlir/include/mlir/Dialect/Async/Passes.td --- a/mlir/include/mlir/Dialect/Async/Passes.td +++ b/mlir/include/mlir/Dialect/Async/Passes.td @@ -44,13 +44,6 @@ 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 = [ - // Temporary for bringup, should become the default. - 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", "func::FuncDialect"]; } diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp --- a/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp +++ b/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp @@ -610,137 +610,6 @@ llvm::DenseMap &outlinedFunctions; }; -//===----------------------------------------------------------------------===// - -/// Rewrite a func as a coroutine by: -/// 1) Wrapping the results into `async.value`. -/// 2) Prepending the results with `async.token`. -/// 3) Setting up coroutine blocks. -/// 4) Rewriting return ops as yield op and branch op into the suspend block. -static CoroMachinery rewriteFuncAsCoroutine(func::FuncOp func) { - auto *ctx = func->getContext(); - auto loc = func.getLoc(); - SmallVector resultTypes; - resultTypes.reserve(func.getCallableResults().size()); - llvm::transform(func.getCallableResults(), std::back_inserter(resultTypes), - [](Type type) { return ValueType::get(type); }); - func.setType( - FunctionType::get(ctx, func.getFunctionType().getInputs(), resultTypes)); - func.insertResult(0, TokenType::get(ctx), {}); - for (Block &block : func.getBlocks()) { - Operation *terminator = block.getTerminator(); - if (auto returnOp = dyn_cast(*terminator)) { - ImplicitLocOpBuilder builder(loc, returnOp); - builder.create(returnOp.getOperands()); - returnOp.erase(); - } - } - return setupCoroMachinery(func); -} - -/// Rewrites a call into a function that has been rewritten 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`. -static void rewriteCallsiteForCoroutine(func::CallOp oldCall, - func::FuncOp func) { - auto loc = func.getLoc(); - ImplicitLocOpBuilder callBuilder(loc, oldCall); - auto newCall = callBuilder.create( - func.getName(), func.getCallableResults(), oldCall.getArgOperands()); - - // Await on the async token and all the value results and unwrap the latter. - callBuilder.create(loc, newCall.getResults().front()); - SmallVector unwrappedResults; - unwrappedResults.reserve(newCall->getResults().size() - 1); - for (Value result : newCall.getResults().drop_front()) - unwrappedResults.push_back( - callBuilder.create(loc, result).getResult()); - // Careful, when result of a call is piped into another call this could lead - // to a dangling pointer. - oldCall.replaceAllUsesWith(unwrappedResults); - oldCall.erase(); -} - -static bool isAllowedToBlock(func::FuncOp func) { - return !!func->getAttrOfType(AsyncDialect::kAllowedToBlockAttrName); -} - -static LogicalResult funcsToCoroutines( - ModuleOp module, - llvm::DenseMap &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. - - SmallVector 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 = [&](func::FuncOp func) { - if (isAllowedToBlock(func)) - return; - // 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 (isAllowedToBlock(func) || - outlinedFunctions.find(func) == outlinedFunctions.end()) { - for (Operation &op : func.getBody().getOps()) { - if (isa(op)) { - funcWorklist.push_back(func); - break; - } - } - } - }; - - // Traverse in post-order collecting for each func op the await ops it has. - for (func::FuncOp func : module.getOps()) - addToWorklist(func); - - SymbolTableCollection symbolTable; - SymbolUserMap symbolUserMap(symbolTable, module); - - // Rewrite funcs, while updating call sites and adding them to the worklist. - while (!funcWorklist.empty()) { - auto func = funcWorklist.pop_back_val(); - auto insertion = 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); - SmallVector 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, [](Operation *a, Operation *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 (func::CallOp call = dyn_cast(*op)) { - func::FuncOp caller = call->getParentOfType(); - rewriteCallsiteForCoroutine(call, func); // Careful, erases the call op. - addToWorklist(caller); - } else { - op->emitError("Unexpected reference to func referenced by symbol"); - return failure(); - } - } - } - return success(); -} - //===----------------------------------------------------------------------===// void AsyncToAsyncRuntimePass::runOnOperation() { ModuleOp module = getOperation(); @@ -764,12 +633,6 @@ return outlinedFunctions.find(parentFunc) != outlinedFunctions.end(); }; - if (eliminateBlockingAwaitOps && - failed(funcsToCoroutines(module, outlinedFunctions))) { - signalPassFailure(); - return; - } - // Lower async operations to async.runtime operations. MLIRContext *ctx = module->getContext(); RewritePatternSet asyncPatterns(ctx); @@ -815,12 +678,6 @@ return outlinedFunctions.find(func) == outlinedFunctions.end(); }); - if (eliminateBlockingAwaitOps) - runtimeTarget.addDynamicallyLegalOp( - [&](RuntimeAwaitOp op) -> bool { - return isAllowedToBlock(op->getParentOfType()); - }); - if (failed(applyPartialConversion(module, runtimeTarget, std::move(asyncPatterns)))) { signalPassFailure(); diff --git a/mlir/test/Dialect/Async/async-to-async-runtime-eliminate-blocking.mlir b/mlir/test/Dialect/Async/async-to-async-runtime-eliminate-blocking.mlir deleted file mode 100644 --- a/mlir/test/Dialect/Async/async-to-async-runtime-eliminate-blocking.mlir +++ /dev/null @@ -1,324 +0,0 @@ -// 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 {builtin.foo = "bar"}) -func.func @simple_callee(%arg0: f32) -> (f32 {builtin.foo = "bar"}) { -// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token -// CHECK: %[[RETURNED_STORAGE:.*]] = async.runtime.create : !async.value -// CHECK: %[[ID:.*]] = async.coro.id -// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]] -// CHECK: cf.br ^[[ORIGINAL_ENTRY:.*]] -// CHECK ^[[ORIGINAL_ENTRY]]: -// CHECK: %[[VAL:.*]] = arith.addf %[[ARG]], %[[ARG]] : f32 - %0 = arith.addf %arg0, %arg0 : f32 -// CHECK: %[[VAL_STORAGE:.*]] = async.runtime.create : !async.value - %1 = async.runtime.create: !async.value -// CHECK: async.runtime.store %[[VAL]], %[[VAL_STORAGE]] : - async.runtime.store %0, %1: !async.value -// CHECK: async.runtime.set_available %[[VAL_STORAGE]] : !async.value - async.runtime.set_available %1: !async.value - -// 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 - -// CHECK: ^[[RESUME]]: -// CHECK: %[[IS_ERROR:.*]] = async.runtime.is_error %[[VAL_STORAGE]] : !async.value -// CHECK: cf.cond_br %[[IS_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_OK:.*]] - -// CHECK: ^[[BRANCH_OK]]: -// CHECK: %[[LOADED:.*]] = async.runtime.load %[[VAL_STORAGE]] : -// CHECK: %[[RETURNED:.*]] = arith.mulf %[[ARG]], %[[LOADED]] : f32 -// CHECK: async.runtime.store %[[RETURNED]], %[[RETURNED_STORAGE]] : -// CHECK: async.runtime.set_available %[[RETURNED_STORAGE]] -// CHECK: async.runtime.set_available %[[TOKEN]] -// CHECK: cf.br ^[[CLEANUP]] - %3 = arith.mulf %arg0, %2 : f32 - return %3: f32 - -// CHECK: ^[[BRANCH_ERROR]]: -// CHECK: async.runtime.set_error %[[TOKEN]] -// CHECK: async.runtime.set_error %[[RETURNED_STORAGE]] -// CHECK: cf.br ^[[CLEANUP]] - - -// CHECK: ^[[CLEANUP]]: -// CHECK: async.coro.free %[[ID]], %[[HDL]] -// CHECK: cf.br ^[[SUSPEND]] - -// CHECK: ^[[SUSPEND]]: -// CHECK: async.coro.end %[[HDL]] -// CHECK: return %[[TOKEN]], %[[RETURNED_STORAGE]] : !async.token, !async.value -} - -// CHECK-LABEL: func @simple_caller() -// CHECK-SAME: -> (!async.token, !async.value) -func.func @simple_caller() -> f32 { -// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token -// CHECK: %[[RETURNED_STORAGE:.*]] = async.runtime.create : !async.value -// CHECK: %[[ID:.*]] = async.coro.id -// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]] -// CHECK: cf.br ^[[ORIGINAL_ENTRY:.*]] -// CHECK ^[[ORIGINAL_ENTRY]]: - -// CHECK: %[[CONSTANT:.*]] = arith.constant - %c = arith.constant 1.0 : f32 -// CHECK: %[[RETURNED_TO_CALLER:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value) -// 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: cf.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 -// CHECK: cf.cond_br %[[IS_VALUE_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_VALUE_OK:.*]] - -// CHECK: ^[[BRANCH_VALUE_OK]]: -// CHECK: %[[LOADED:.*]] = async.runtime.load %[[RETURNED_TO_CALLER]]#1 : -// CHECK: async.runtime.store %[[LOADED]], %[[RETURNED_STORAGE]] : -// CHECK: async.runtime.set_available %[[RETURNED_STORAGE]] -// CHECK: async.runtime.set_available %[[TOKEN]] -// CHECK: cf.br ^[[CLEANUP]] - return %r: f32 -// CHECK: ^[[BRANCH_ERROR]]: -// CHECK: async.runtime.set_error %[[TOKEN]] -// CHECK: async.runtime.set_error %[[RETURNED_STORAGE]] -// CHECK: cf.br ^[[CLEANUP]] - - -// CHECK: ^[[CLEANUP]]: -// CHECK: async.coro.free %[[ID]], %[[HDL]] -// CHECK: cf.br ^[[SUSPEND]] - -// CHECK: ^[[SUSPEND]]: -// CHECK: async.coro.end %[[HDL]] -// CHECK: return %[[TOKEN]], %[[RETURNED_STORAGE]] : !async.token, !async.value -} - -// CHECK-LABEL: func @double_caller() -// CHECK-SAME: -> (!async.token, !async.value) -func.func @double_caller() -> f32 { -// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token -// CHECK: %[[RETURNED_STORAGE:.*]] = async.runtime.create : !async.value -// CHECK: %[[ID:.*]] = async.coro.id -// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]] -// CHECK: cf.br ^[[ORIGINAL_ENTRY:.*]] -// CHECK ^[[ORIGINAL_ENTRY]]: - -// CHECK: %[[CONSTANT:.*]] = arith.constant - %c = arith.constant 1.0 : f32 -// CHECK: %[[RETURNED_TO_CALLER_1:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value) -// 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: cf.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 -// CHECK: cf.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 : -// CHECK: %[[RETURNED_TO_CALLER_2:.*]]:2 = call @simple_callee(%[[LOADED_1]]) : (f32) -> (!async.token, !async.value) -// 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: cf.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 -// CHECK: cf.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 : -// CHECK: async.runtime.store %[[LOADED_2]], %[[RETURNED_STORAGE]] : -// CHECK: async.runtime.set_available %[[RETURNED_STORAGE]] -// CHECK: async.runtime.set_available %[[TOKEN]] -// CHECK: cf.br ^[[CLEANUP]] - return %s: f32 -// CHECK: ^[[BRANCH_ERROR]]: -// CHECK: async.runtime.set_error %[[TOKEN]] -// CHECK: async.runtime.set_error %[[RETURNED_STORAGE]] -// CHECK: cf.br ^[[CLEANUP]] - -// CHECK: ^[[CLEANUP]]: -// CHECK: async.coro.free %[[ID]], %[[HDL]] -// CHECK: cf.br ^[[SUSPEND]] - -// CHECK: ^[[SUSPEND]]: -// CHECK: async.coro.end %[[HDL]] -// CHECK: return %[[TOKEN]], %[[RETURNED_STORAGE]] : !async.token, !async.value -} - -// CHECK-LABEL: func @recursive -// CHECK-SAME: (%[[ARG:.*]]: !async.token) -> !async.token -func.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: cf.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: cf.br ^[[CLEANUP]] - -// CHECK: ^[[BRANCH_ERROR]]: -// CHECK: async.runtime.set_error %[[TOKEN]] -// CHECK: cf.br ^[[CLEANUP]] -return - -// CHECK: ^[[CLEANUP]]: -// CHECK: async.coro.free %[[ID]], %[[HDL]] -// CHECK: cf.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.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: cf.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: cf.br ^[[CLEANUP]] - -// CHECK: ^[[BRANCH_ERROR]]: -// CHECK: async.runtime.set_error %[[TOKEN]] -// CHECK: cf.br ^[[CLEANUP]] -return - -// CHECK: ^[[CLEANUP]]: -// CHECK: async.coro.free %[[ID]], %[[HDL]] -// CHECK: cf.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.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: cf.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: cf.br ^[[CLEANUP]] - -// CHECK: ^[[BRANCH_ERROR]]: -// CHECK: async.runtime.set_error %[[TOKEN]] -// CHECK: cf.br ^[[CLEANUP]] -return - -// CHECK: ^[[CLEANUP]]: -// CHECK: async.coro.free %[[ID]], %[[HDL]] -// CHECK: cf.br ^[[SUSPEND]] - -// CHECK: ^[[SUSPEND]]: -// CHECK: async.coro.end %[[HDL]] -// CHECK: return %[[TOKEN]] : !async.token -} - -// CHECK-LABEL: func @caller_allowed_to_block -// CHECK-SAME: () -> f32 -func.func @caller_allowed_to_block() -> f32 attributes { async.allowed_to_block } { -// CHECK: %[[CONSTANT:.*]] = arith.constant - %c = arith.constant 1.0 : f32 -// CHECK: %[[RETURNED_TO_CALLER:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value) -// CHECK: async.runtime.await %[[RETURNED_TO_CALLER]]#0 -// CHECK: async.runtime.await %[[RETURNED_TO_CALLER]]#1 -// CHECK: %[[RETURNED:.*]] = async.runtime.load %[[RETURNED_TO_CALLER]]#1 - %r = call @simple_callee(%c): (f32) -> f32 - -// CHECK: return %[[RETURNED]] : f32 - return %r: f32 -}