This is an archive of the discontinued LLVM Phabricator instance.

Differential D137282

[mlir] Remove eliminateBlockingAwaitOps option in AsyncToAsyncRuntime pass
ClosedPublic

Authored by yijia1212 on Nov 2 2022, 12:26 PM.

Details

Reviewers
ezhulenev
nicolasvasilache
Commits
rG51c1632f8d50: [mlir] Remove eliminateBlockingAwaitOps option in AsyncToAsyncRuntime pass
Summary

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.

Diff Detail

Event Timeline

yijia1212 created this revision.Nov 2 2022, 12:26 PM
Herald added a project: Restricted Project. · View Herald TranscriptNov 2 2022, 12:26 PM
Herald added subscribers: Moerafaat, zero9178, bzcheeseman and 21 others. · View Herald Transcript
yijia1212 requested review of this revision.Nov 2 2022, 12:26 PM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptNov 2 2022, 12:26 PM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: stephenneuendorffer, nicolasvasilache. · View Herald Transcript
ezhulenev accepted this revision.Nov 2 2022, 12:30 PM
This revision is now accepted and ready to land.Nov 2 2022, 12:30 PM
ezhulenev retitled this revision from Remove eliminateBlockingAwaitOps option in AsyncToAsyncRuntime pass to [mlir] Remove eliminateBlockingAwaitOps option in AsyncToAsyncRuntime pass.Nov 2 2022, 12:58 PM
Harbormaster completed remote builds in B195781: Diff 472723.Nov 2 2022, 1:50 PM
Closed by commit rG51c1632f8d50: [mlir] Remove eliminateBlockingAwaitOps option in AsyncToAsyncRuntime pass (authored by yijia1212, committed by ezhulenev). · Explain WhyNov 2 2022, 1:59 PM
This revision was automatically updated to reflect the committed changes.
ezhulenev added a commit: rG51c1632f8d50: [mlir] Remove eliminateBlockingAwaitOps option in AsyncToAsyncRuntime pass.

Revision Contents

PathSize
mlir/
include/
mlir/
Dialect/
Async/
IR/
5 lines
7 lines
lib/
Dialect/
Async/
Transforms/
143 lines
test/
Dialect/
Async/

Diff 472751

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

Show All 26 Linesdef AsyncDialect : Dialect {
let description = [{ let description = [{
This dialect contains operations for modeling asynchronous execution. This dialect contains operations for modeling asynchronous execution.
}]; }];
let useDefaultTypePrinterParser = 1; let useDefaultTypePrinterParser = 1;
let extraClassDeclaration = [{ let extraClassDeclaration = [{
/// The name of a unit attribute on funcs that are allowed to have a /// The name of a unit attribute on funcs that are allowed to have a
/// blocking async.runtime.await ops. Only useful in combination with /// blocking async.runtime.await ops. In absence of this attribute the
/// 'eliminate-blocking-await-ops' option, which in absence of this /// asyncification pass might convert a func to a coroutine.
/// attribute might convert a func to a coroutine.
static constexpr StringRef kAllowedToBlockAttrName = static constexpr StringRef kAllowedToBlockAttrName =
"async.allowed_to_block"; "async.allowed_to_block";
}]; }];
} }
#endif // ASYNC_DIALECT_TD #endif // ASYNC_DIALECT_TD

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

Show All 38 Lines let dependentDialects = [
"scf::SCFDialect" "scf::SCFDialect"
]; ];
} }
def AsyncToAsyncRuntime : Pass<"async-to-async-runtime", "ModuleOp"> { def AsyncToAsyncRuntime : Pass<"async-to-async-runtime", "ModuleOp"> {
let summary = "Lower high level async operations (e.g. async.execute) to the" let summary = "Lower high level async operations (e.g. async.execute) to the"
"explicit async.runtime and async.coro operations"; "explicit async.runtime and async.coro operations";
let constructor = "mlir::createAsyncToAsyncRuntimePass()"; 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"]; let dependentDialects = ["async::AsyncDialect", "func::FuncDialect"];
} }
def AsyncRuntimeRefCounting : Pass<"async-runtime-ref-counting"> { def AsyncRuntimeRefCounting : Pass<"async-runtime-ref-counting"> {
let summary = "Automatic reference counting for Async runtime operations"; let summary = "Automatic reference counting for Async runtime operations";
let description = [{ let description = [{
This pass works at the async runtime abtraction level, after all This pass works at the async runtime abtraction level, after all
`async.execute` and `async.await` operations are lowered to the async `async.execute` and `async.await` operations are lowered to the async
▲ Show 20 Lines • Show All 52 Lines • Show Last 20 Lines

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

Show First 20 Lines • Show All 605 Lines • ▼ Show 20 Lines matchAndRewrite(cf::AssertOp op, OpAdaptor adaptor,
return success(); return success();
} }
private: private:
llvm::DenseMap<func::FuncOp, CoroMachinery> &outlinedFunctions; 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() { void AsyncToAsyncRuntimePass::runOnOperation() {
ModuleOp module = getOperation(); ModuleOp module = getOperation();
SymbolTable symbolTable(module); SymbolTable symbolTable(module);
// Outline all `async.execute` body regions into async functions (coroutines). // Outline all `async.execute` body regions into async functions (coroutines).
llvm::DenseMap<func::FuncOp, CoroMachinery> outlinedFunctions; llvm::DenseMap<func::FuncOp, CoroMachinery> outlinedFunctions;
module.walk([&](ExecuteOp execute) { module.walk([&](ExecuteOp execute) {
outlinedFunctions.insert(outlineExecuteOp(symbolTable, execute)); outlinedFunctions.insert(outlineExecuteOp(symbolTable, execute));
}); });
LLVM_DEBUG({ LLVM_DEBUG({
llvm::dbgs() << "Outlined " << outlinedFunctions.size() llvm::dbgs() << "Outlined " << outlinedFunctions.size()
<< " functions built from async.execute operations\n"; << " functions built from async.execute operations\n";
}); });
// Returns true if operation is inside the coroutine. // Returns true if operation is inside the coroutine.
auto isInCoroutine = [&](Operation *op) -> bool { auto isInCoroutine = [&](Operation *op) -> bool {
auto parentFunc = op->getParentOfType<func::FuncOp>(); auto parentFunc = op->getParentOfType<func::FuncOp>();
return outlinedFunctions.find(parentFunc) != outlinedFunctions.end(); return outlinedFunctions.find(parentFunc) != outlinedFunctions.end();
}; };
if (eliminateBlockingAwaitOps &&
failed(funcsToCoroutines(module, outlinedFunctions))) {
signalPassFailure();
return;
}
// Lower async operations to async.runtime operations. // Lower async operations to async.runtime operations.
MLIRContext *ctx = module->getContext(); MLIRContext *ctx = module->getContext();
RewritePatternSet asyncPatterns(ctx); RewritePatternSet asyncPatterns(ctx);
// Conversion to async runtime augments original CFG with the coroutine CFG, // Conversion to async runtime augments original CFG with the coroutine CFG,
// and we have to make sure that structured control flow operations with async // and we have to make sure that structured control flow operations with async
// operations in nested regions will be converted to branch-based control flow // operations in nested regions will be converted to branch-based control flow
// before we add the coroutine basic blocks. // before we add the coroutine basic blocks.
Show All 29 Linesvoid AsyncToAsyncRuntimePass::runOnOperation() {
// Assertions must be converted to runtime errors inside async functions. // Assertions must be converted to runtime errors inside async functions.
runtimeTarget.addDynamicallyLegalOp<cf::AssertOp>( runtimeTarget.addDynamicallyLegalOp<cf::AssertOp>(
[&](cf::AssertOp op) -> bool { [&](cf::AssertOp op) -> bool {
auto func = op->getParentOfType<func::FuncOp>(); auto func = op->getParentOfType<func::FuncOp>();
return outlinedFunctions.find(func) == outlinedFunctions.end(); 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, if (failed(applyPartialConversion(module, runtimeTarget,
std::move(asyncPatterns)))) { std::move(asyncPatterns)))) {
signalPassFailure(); signalPassFailure();
return; return;
} }
} }
std::unique_ptr<OperationPass<ModuleOp>> mlir::createAsyncToAsyncRuntimePass() { std::unique_ptr<OperationPass<ModuleOp>> mlir::createAsyncToAsyncRuntimePass() {
return std::make_unique<AsyncToAsyncRuntimePass>(); return std::make_unique<AsyncToAsyncRuntimePass>();
} }

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

  • This file was deleted.
// 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
}
mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp