[Mlir-commits] [mlir] 51c1632 - [mlir] Remove eliminateBlockingAwaitOps option in AsyncToAsyncRuntime pass
Eugene Zhulenev
llvmlistbot at llvm.org
Wed Nov 2 13:58:59 PDT 2022
Author: yijiagu
Date: 2022-11-02T13:58:54-07:00
New Revision: 51c1632f8d500a8439707153439650945c446dfa
URL: https://github.com/llvm/llvm-project/commit/51c1632f8d500a8439707153439650945c446dfa
DIFF: https://github.com/llvm/llvm-project/commit/51c1632f8d500a8439707153439650945c446dfa.diff
LOG: [mlir] Remove eliminateBlockingAwaitOps option in AsyncToAsyncRuntime pass
Remove the eliminateBlockingAwaitOps option in AsyncToAsyncRuntime pass
Today the AsyncToAsyncRuntime pass does two things: one is converting normal funcs with async ops to coroutine cfg; the other is lowering high level async operations to async.coro and async.runtime operations. This patch removes the converting step from AsyncToAsyncRuntime pass.
In the next step we will create a new asyncfication pass for converting normal funcs to the newly added async.func operation.
Reviewed By: ezhulenev
Differential Revision: https://reviews.llvm.org/D137282
Added:
Modified:
mlir/include/mlir/Dialect/Async/IR/AsyncDialect.td
mlir/include/mlir/Dialect/Async/Passes.td
mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
Removed:
mlir/test/Dialect/Async/async-to-async-runtime-eliminate-blocking.mlir
################################################################################
diff --git a/mlir/include/mlir/Dialect/Async/IR/AsyncDialect.td b/mlir/include/mlir/Dialect/Async/IR/AsyncDialect.td
index e3b4db42204b8..eb1d76a180fe2 100644
--- a/mlir/include/mlir/Dialect/Async/IR/AsyncDialect.td
+++ b/mlir/include/mlir/Dialect/Async/IR/AsyncDialect.td
@@ -32,9 +32,8 @@ def AsyncDialect : Dialect {
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
index 16fb8626c0c0e..aed5b4ff7865a 100644
--- a/mlir/include/mlir/Dialect/Async/Passes.td
+++ b/mlir/include/mlir/Dialect/Async/Passes.td
@@ -44,13 +44,6 @@ 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 = [
- // 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
index b4880c0e3b3f5..38f3717c70f9b 100644
--- a/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
+++ b/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
@@ -610,137 +610,6 @@ class AssertOpLowering : public OpConversionPattern<cf::AssertOp> {
llvm::DenseMap<func::FuncOp, CoroMachinery> &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<Type> 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<func::ReturnOp>(*terminator)) {
- ImplicitLocOpBuilder builder(loc, returnOp);
- builder.create<YieldOp>(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::CallOp>(
- func.getName(), func.getCallableResults(), oldCall.getArgOperands());
-
- // Await on the async token and all the value results and unwrap the latter.
- callBuilder.create<AwaitOp>(loc, newCall.getResults().front());
- SmallVector<Value> unwrappedResults;
- unwrappedResults.reserve(newCall->getResults().size() - 1);
- for (Value result : newCall.getResults().drop_front())
- unwrappedResults.push_back(
- callBuilder.create<AwaitOp>(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<UnitAttr>(AsyncDialect::kAllowedToBlockAttrName);
-}
-
-static LogicalResult funcsToCoroutines(
- ModuleOp module,
- llvm::DenseMap<func::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.
-
- SmallVector<func::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 = [&](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<AwaitOp, AwaitAllOp>(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<func::FuncOp>())
- 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<Operation *> 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<func::CallOp>(*op)) {
- func::FuncOp caller = call->getParentOfType<func::FuncOp>();
- 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 @@ void AsyncToAsyncRuntimePass::runOnOperation() {
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 @@ void AsyncToAsyncRuntimePass::runOnOperation() {
return outlinedFunctions.find(func) == outlinedFunctions.end();
});
- if (eliminateBlockingAwaitOps)
- runtimeTarget.addDynamicallyLegalOp<RuntimeAwaitOp>(
- [&](RuntimeAwaitOp op) -> bool {
- return isAllowedToBlock(op->getParentOfType<func::FuncOp>());
- });
-
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
index 5a85c3d8974e6..0000000000000
--- 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<f32> {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<f32>
-// 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<f32>
- %1 = async.runtime.create: !async.value<f32>
-// CHECK: async.runtime.store %[[VAL]], %[[VAL_STORAGE]] : <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: cf.cond_br %[[IS_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_OK:.*]]
-
-// CHECK: ^[[BRANCH_OK]]:
-// CHECK: %[[LOADED:.*]] = async.runtime.load %[[VAL_STORAGE]] : <f32>
-// CHECK: %[[RETURNED:.*]] = arith.mulf %[[ARG]], %[[LOADED]] : f32
-// CHECK: async.runtime.store %[[RETURNED]], %[[RETURNED_STORAGE]] : <f32>
-// 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<f32>
-}
-
-// CHECK-LABEL: func @simple_caller()
-// CHECK-SAME: -> (!async.token, !async.value<f32>)
-func.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: 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<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: 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<f32>
-// 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 : <f32>
-// CHECK: async.runtime.store %[[LOADED]], %[[RETURNED_STORAGE]] : <f32>
-// 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<f32>
-}
-
-// CHECK-LABEL: func @double_caller()
-// CHECK-SAME: -> (!async.token, !async.value<f32>)
-func.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: 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<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: 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<f32>
-// 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 : <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: 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<f32>
-// 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 : <f32>
-// CHECK: async.runtime.store %[[LOADED_2]], %[[RETURNED_STORAGE]] : <f32>
-// 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<f32>
-}
-
-// 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<f32>)
-// 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
-}
More information about the Mlir-commits
mailing list