| //===- AsyncOps.td - Async operations definition -----------*- tablegen -*-===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // This is the operation definition file for Async dialect operations. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #ifndef ASYNC_OPS |
| #define ASYNC_OPS |
| |
| include "mlir/Dialect/Async/IR/AsyncDialect.td" |
| include "mlir/Dialect/Async/IR/AsyncTypes.td" |
| include "mlir/Interfaces/ControlFlowInterfaces.td" |
| include "mlir/Interfaces/SideEffectInterfaces.td" |
| |
| //===----------------------------------------------------------------------===// |
| // Async op definitions |
| //===----------------------------------------------------------------------===// |
| |
| // Base class for the operation in this dialect |
| class Async_Op<string mnemonic, list<OpTrait> traits = []> : |
| Op<AsyncDialect, mnemonic, traits>; |
| |
| def Async_ExecuteOp : |
| Async_Op<"execute", [SingleBlockImplicitTerminator<"YieldOp">, |
| DeclareOpInterfaceMethods<RegionBranchOpInterface, |
| ["getSuccessorEntryOperands", |
| "getNumRegionInvocations"]>, |
| AttrSizedOperandSegments]> { |
| let summary = "Asynchronous execute operation"; |
| let description = [{ |
| The `body` region attached to the `async.execute` operation semantically |
| can be executed concurrently with the successor operation. In the followup |
| example "compute0" can be executed concurrently with "compute1". |
| |
| The actual concurrency semantics depends on the dialect lowering to the |
| executable format. Fully sequential execution ("compute0" completes before |
| "compute1" starts) is a completely legal execution. |
| |
| Because concurrent execution is not guaranteed, it is illegal to create an |
| implicit dependency from "compute1" to "compute0" (e.g. via shared global |
| state). All dependencies must be made explicit with async execute arguments |
| (`async.token` or `async.value`). |
| |
| `async.execute` operation takes `async.token` dependencies and `async.value` |
| operands separately, and starts execution of the attached body region only |
| when all tokens and values become ready. |
| |
| Example: |
| |
| ```mlir |
| %dependency = ... : !async.token |
| %value = ... : !async.value<f32> |
| |
| %token, %results = |
| async.execute [%dependency](%value as %unwrapped: !async.value<f32>) |
| -> !async.value<!some.type> |
| { |
| %0 = "compute0"(%unwrapped): (f32) -> !some.type |
| async.yield %0 : !some.type |
| } |
| |
| %1 = "compute1"(...) : !some.type |
| ``` |
| |
| In the example above asynchronous execution starts only after dependency |
| token and value argument become ready. Unwrapped value passed to the |
| attached body region as an %unwrapped value of f32 type. |
| }]; |
| |
| let arguments = (ins Variadic<Async_TokenType>:$dependencies, |
| Variadic<Async_AnyValueOrTokenType>:$operands); |
| |
| let results = (outs Async_TokenType:$token, |
| Variadic<Async_ValueType>:$results); |
| let regions = (region SizedRegion<1>:$body); |
| |
| let printer = [{ return ::print(p, *this); }]; |
| let parser = [{ return ::parse$cppClass(parser, result); }]; |
| let verifier = [{ return ::verify(*this); }]; |
| |
| let skipDefaultBuilders = 1; |
| let builders = [ |
| OpBuilder<(ins "TypeRange":$resultTypes, "ValueRange":$dependencies, |
| "ValueRange":$operands, |
| CArg<"function_ref<void(OpBuilder &, Location, ValueRange)>", |
| "nullptr">:$bodyBuilder)>, |
| ]; |
| |
| let extraClassDeclaration = [{ |
| using BodyBuilderFn = |
| function_ref<void(OpBuilder &, Location, ValueRange)>; |
| |
| }]; |
| } |
| |
| def Async_YieldOp : |
| Async_Op<"yield", [ |
| HasParent<"ExecuteOp">, NoSideEffect, Terminator, |
| DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface>]> { |
| let summary = "terminator for Async execute operation"; |
| let description = [{ |
| The `async.yield` is a special terminator operation for the block inside |
| `async.execute` operation. |
| }]; |
| |
| let arguments = (ins Variadic<AnyType>:$operands); |
| |
| let assemblyFormat = "($operands^ `:` type($operands))? attr-dict"; |
| |
| let verifier = [{ return ::verify(*this); }]; |
| } |
| |
| def Async_AwaitOp : Async_Op<"await"> { |
| let summary = "waits for the argument to become ready"; |
| let description = [{ |
| The `async.await` operation waits until the argument becomes ready, and for |
| the `async.value` arguments it unwraps the underlying value |
| |
| Example: |
| |
| ```mlir |
| %0 = ... : !async.token |
| async.await %0 : !async.token |
| |
| %1 = ... : !async.value<f32> |
| %2 = async.await %1 : !async.value<f32> |
| ``` |
| }]; |
| |
| let arguments = (ins Async_AnyValueOrTokenType:$operand); |
| let results = (outs Optional<AnyType>:$result); |
| |
| let skipDefaultBuilders = 1; |
| |
| let builders = [ |
| OpBuilder<(ins "Value":$operand, |
| CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>, |
| ]; |
| |
| let extraClassDeclaration = [{ |
| Optional<Type> getResultType() { |
| if (getResultTypes().empty()) return None; |
| return getResultTypes()[0]; |
| } |
| }]; |
| |
| let assemblyFormat = [{ |
| $operand `:` custom<AwaitResultType>( |
| type($operand), type($result) |
| ) attr-dict |
| }]; |
| |
| let verifier = [{ return ::verify(*this); }]; |
| } |
| |
| def Async_CreateGroupOp : Async_Op<"create_group", [NoSideEffect]> { |
| let summary = "creates an empty async group"; |
| let description = [{ |
| The `async.create_group` allocates an empty async group. Async tokens or |
| values can be added to this group later. The size of the group must be |
| specified at construction time, and `await_all` operation will first |
| wait until the number of added tokens or values reaches the group size. |
| |
| Example: |
| |
| ```mlir |
| %size = ... : index |
| %group = async.create_group %size : !async.group |
| ... |
| async.await_all %group |
| ``` |
| }]; |
| |
| let arguments = (ins Index:$size); |
| let results = (outs Async_GroupType:$result); |
| |
| let hasCanonicalizeMethod = 1; |
| |
| let assemblyFormat = "$size `:` type($result) attr-dict"; |
| } |
| |
| def Async_AddToGroupOp : Async_Op<"add_to_group", []> { |
| let summary = "adds and async token or value to the group"; |
| let description = [{ |
| The `async.add_to_group` adds an async token or value to the async group. |
| Returns the rank of the added element in the group. This rank is fixed |
| for the group lifetime. |
| |
| Example: |
| |
| ```mlir |
| %0 = async.create_group %size : !async.group |
| %1 = ... : !async.token |
| %2 = async.add_to_group %1, %0 : !async.token |
| ``` |
| }]; |
| |
| let arguments = (ins Async_AnyValueOrTokenType:$operand, |
| Async_GroupType:$group); |
| let results = (outs Index:$rank); |
| |
| let assemblyFormat = "$operand `,` $group `:` type($operand) attr-dict"; |
| } |
| |
| def Async_AwaitAllOp : Async_Op<"await_all", []> { |
| let summary = "waits for the all async tokens or values in the group to " |
| "become ready"; |
| let description = [{ |
| The `async.await_all` operation waits until all the tokens or values in the |
| group become ready. |
| |
| Example: |
| |
| ```mlir |
| %0 = async.create_group %size : !async.group |
| |
| %1 = ... : !async.token |
| %2 = async.add_to_group %1, %0 : !async.token |
| |
| %3 = ... : !async.token |
| %4 = async.add_to_group %2, %0 : !async.token |
| |
| async.await_all %0 |
| ``` |
| }]; |
| |
| let arguments = (ins Async_GroupType:$operand); |
| let results = (outs); |
| |
| let assemblyFormat = "$operand attr-dict"; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // Async Dialect LLVM Coroutines Operations. |
| //===----------------------------------------------------------------------===// |
| |
| // Async to LLVM dialect lowering converts async tasks (regions inside async |
| // execute operations) to LLVM coroutines [1], and relies on switched-resume |
| // lowering [2] to produce an asynchronous executable. |
| // |
| // We define LLVM coro intrinsics in the async dialect to facilitate progressive |
| // lowering with verifiable and type-safe IR during the multi-step lowering |
| // pipeline. First we convert from high level async operations (e.g. execute) to |
| // the explicit calls to coro intrinsics and runtime API, and then finalize |
| // lowering to LLVM with a simple dialect conversion pass. |
| // |
| // [1] https://llvm.org/docs/Coroutines.html |
| // [2] https://llvm.org/docs/Coroutines.html#switched-resume-lowering |
| |
| def Async_CoroIdOp : Async_Op<"coro.id"> { |
| let summary = "returns a switched-resume coroutine identifier"; |
| let description = [{ |
| The `async.coro.id` returns a switched-resume coroutine identifier. |
| }]; |
| |
| let results = (outs Async_CoroIdType:$id); |
| let assemblyFormat = "attr-dict"; |
| } |
| |
| def Async_CoroBeginOp : Async_Op<"coro.begin"> { |
| let summary = "returns a handle to the coroutine"; |
| let description = [{ |
| The `async.coro.begin` allocates a coroutine frame and returns a handle to |
| the coroutine. |
| }]; |
| |
| let arguments = (ins Async_CoroIdType:$id); |
| let results = (outs Async_CoroHandleType:$handle); |
| let assemblyFormat = "$id attr-dict"; |
| } |
| |
| def Async_CoroFreeOp : Async_Op<"coro.free"> { |
| let summary = "deallocates the coroutine frame"; |
| let description = [{ |
| The `async.coro.free` deallocates the coroutine frame created by the |
| async.coro.begin operation. |
| }]; |
| |
| let arguments = (ins Async_CoroIdType:$id, |
| Async_CoroHandleType:$handle); |
| let assemblyFormat = "$id `,` $handle attr-dict"; |
| } |
| |
| def Async_CoroEndOp : Async_Op<"coro.end"> { |
| let summary = "marks the end of the coroutine in the suspend block"; |
| let description = [{ |
| The `async.coro.end` marks the point where a coroutine needs to return |
| control back to the caller if it is not an initial invocation of the |
| coroutine. It the start part of the coroutine is is no-op. |
| }]; |
| |
| let arguments = (ins Async_CoroHandleType:$handle); |
| let assemblyFormat = "$handle attr-dict"; |
| } |
| |
| def Async_CoroSaveOp : Async_Op<"coro.save"> { |
| let summary = "saves the coroutine state"; |
| let description = [{ |
| The `async.coro.saves` saves the coroutine state. |
| }]; |
| |
| let arguments = (ins Async_CoroHandleType:$handle); |
| let results = (outs Async_CoroStateType:$state); |
| let assemblyFormat = "$handle attr-dict"; |
| } |
| |
| def Async_CoroSuspendOp : Async_Op<"coro.suspend", [Terminator]> { |
| let summary = "suspends the coroutine"; |
| let description = [{ |
| The `async.coro.suspend` suspends the coroutine and transfers control to the |
| `suspend` successor. If suspended coroutine later resumed it will transfer |
| control to the `resume` successor. If it is destroyed it will transfer |
| control to the the `cleanup` successor. |
| |
| In switched-resume lowering coroutine can be already in resumed state when |
| suspend operation is called, in this case control will be transferred to the |
| `resume` successor skipping the `suspend` successor. |
| }]; |
| |
| let arguments = (ins Async_CoroStateType:$state); |
| let successors = (successor AnySuccessor:$suspendDest, |
| AnySuccessor:$resumeDest, |
| AnySuccessor:$cleanupDest); |
| let assemblyFormat = |
| "$state `,` $suspendDest `,` $resumeDest `,` $cleanupDest attr-dict"; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // Async Dialect Runtime Operations. |
| //===----------------------------------------------------------------------===// |
| |
| // The following operations are intermediate async dialect operations to help |
| // lowering from high level async operation like `async.execute` to the Async |
| // Runtime API defined in the `ExecutionEngine/AsyncRuntime.h`. |
| |
| def Async_RuntimeCreateOp : Async_Op<"runtime.create"> { |
| let summary = "creates an async runtime token or value"; |
| let description = [{ |
| The `async.runtime.create` operation creates an async dialect token or |
| value. Tokens and values are created in the non-ready state. |
| }]; |
| |
| let results = (outs Async_AnyValueOrTokenType:$result); |
| let assemblyFormat = "attr-dict `:` type($result)"; |
| } |
| |
| def Async_RuntimeCreateGroupOp : Async_Op<"runtime.create_group"> { |
| let summary = "creates an async runtime group"; |
| let description = [{ |
| The `async.runtime.create_group` operation creates an async dialect group |
| of the given size. Group created in the empty state. |
| }]; |
| |
| let arguments = (ins Index:$size); |
| let results = (outs Async_GroupType:$result); |
| let assemblyFormat = "$size `:` type($result) attr-dict "; |
| } |
| |
| def Async_RuntimeSetAvailableOp : Async_Op<"runtime.set_available"> { |
| let summary = "switches token or value to available state"; |
| let description = [{ |
| The `async.runtime.set_available` operation switches async token or value |
| state to available. |
| }]; |
| |
| let arguments = (ins Async_AnyValueOrTokenType:$operand); |
| let assemblyFormat = "$operand attr-dict `:` type($operand)"; |
| } |
| |
| def Async_RuntimeSetErrorOp : Async_Op<"runtime.set_error"> { |
| let summary = "switches token or value to error state"; |
| let description = [{ |
| The `async.runtime.set_error` operation switches async token or value |
| state to error. |
| }]; |
| |
| let arguments = (ins Async_AnyValueOrTokenType:$operand); |
| let assemblyFormat = "$operand attr-dict `:` type($operand)"; |
| } |
| |
| def Async_RuntimeIsErrorOp : Async_Op<"runtime.is_error"> { |
| let summary = "returns true if token, value or group is in error state"; |
| let description = [{ |
| The `async.runtime.is_error` operation returns true if the token, value or |
| group (any of the async runtime values) is in the error state. It is the |
| caller responsibility to check error state after the call to `await` or |
| resuming after `await_and_resume`. |
| }]; |
| |
| let arguments = (ins Async_AnyAsyncType:$operand); |
| let results = (outs I1:$is_error); |
| |
| let assemblyFormat = "$operand attr-dict `:` type($operand)"; |
| } |
| |
| def Async_RuntimeAwaitOp : Async_Op<"runtime.await"> { |
| let summary = "blocks the caller thread until the operand becomes available"; |
| let description = [{ |
| The `async.runtime.await` operation blocks the caller thread until the |
| operand becomes available or error. |
| }]; |
| |
| let arguments = (ins Async_AnyAsyncType:$operand); |
| let assemblyFormat = "$operand attr-dict `:` type($operand)"; |
| } |
| |
| def Async_RuntimeResumeOp : Async_Op<"runtime.resume"> { |
| let summary = "resumes the coroutine on a thread managed by the runtime"; |
| let description = [{ |
| The `async.runtime.resume` operation resumes the coroutine on a thread |
| managed by the runtime. |
| }]; |
| |
| let arguments = (ins Async_CoroHandleType:$handle); |
| let assemblyFormat = "$handle attr-dict"; |
| } |
| |
| def Async_RuntimeAwaitAndResumeOp : Async_Op<"runtime.await_and_resume"> { |
| let summary = "awaits the async operand and resumes the coroutine"; |
| let description = [{ |
| The `async.runtime.await_and_resume` operation awaits for the operand to |
| become available or error and resumes the coroutine on a thread managed by |
| the runtime. |
| }]; |
| |
| let arguments = (ins Async_AnyAsyncType:$operand, |
| Async_CoroHandleType:$handle); |
| let assemblyFormat = "$operand `,` $handle attr-dict `:` type($operand)"; |
| } |
| |
| def Async_RuntimeStoreOp : Async_Op<"runtime.store", |
| [TypesMatchWith<"type of 'value' matches element type of 'storage'", |
| "storage", "value", |
| "$_self.cast<ValueType>().getValueType()">]> { |
| let summary = "stores the value into the runtime async.value"; |
| let description = [{ |
| The `async.runtime.store` operation stores the value into the runtime |
| async.value storage. |
| }]; |
| |
| let arguments = (ins AnyType:$value, |
| Async_ValueType:$storage); |
| let assemblyFormat = "$value `,` $storage attr-dict `:` type($storage)"; |
| } |
| |
| def Async_RuntimeLoadOp : Async_Op<"runtime.load", |
| [TypesMatchWith<"type of 'value' matches element type of 'storage'", |
| "storage", "result", |
| "$_self.cast<ValueType>().getValueType()">]> { |
| let summary = "loads the value from the runtime async.value"; |
| let description = [{ |
| The `async.runtime.load` operation loads the value from the runtime |
| async.value storage. |
| }]; |
| |
| let arguments = (ins Async_ValueType:$storage); |
| let results = (outs AnyType:$result); |
| let assemblyFormat = "$storage attr-dict `:` type($storage)"; |
| } |
| |
| def Async_RuntimeAddToGroupOp : Async_Op<"runtime.add_to_group", []> { |
| let summary = "adds and async token or value to the group"; |
| let description = [{ |
| The `async.runtime.add_to_group` adds an async token or value to the async |
| group. Returns the rank of the added element in the group. |
| }]; |
| |
| let arguments = (ins Async_AnyValueOrTokenType:$operand, |
| Async_GroupType:$group); |
| let results = (outs Index:$rank); |
| |
| let assemblyFormat = "$operand `,` $group attr-dict `:` type($operand)"; |
| } |
| |
| // All async values (values, tokens, groups) are reference counted at runtime |
| // and automatically destructed when reference count drops to 0. |
| // |
| // All values are semantically created with a reference count of +1 and it is |
| // the responsibility of the last async value user to drop reference count. |
| // |
| // Async values created when: |
| // 1. Operation returns async result (e.g. the result of an `async.execute`). |
| // 2. Async value passed in as a block argument. |
| // |
| // It is the responsibility of the async value user to extend the lifetime by |
| // adding a +1 reference, if the reference counted value captured by the |
| // asynchronously executed region (`async.execute` operation), and drop it after |
| // the last nested use. |
| // |
| // Reference counting operations can be added to the IR using automatic |
| // reference count pass, that relies on liveness analysis to find the last uses |
| // of all reference counted values and automatically inserts |
| // `drop_ref` operations. |
| // |
| // See `AsyncRefCountingPass` documentation for the implementation details. |
| |
| def Async_RuntimeAddRefOp : Async_Op<"runtime.add_ref"> { |
| let summary = "adds a reference to async value"; |
| let description = [{ |
| The `async.runtime.add_ref` operation adds a reference(s) to async value |
| (token, value or group). |
| }]; |
| |
| let arguments = (ins Async_AnyAsyncType:$operand, |
| Confined<I64Attr, [IntPositive]>:$count); |
| |
| let assemblyFormat = [{ |
| $operand attr-dict `:` type($operand) |
| }]; |
| } |
| |
| def Async_RuntimeDropRefOp : Async_Op<"runtime.drop_ref"> { |
| let summary = "drops a reference to async value"; |
| let description = [{ |
| The `async.runtime.drop_ref` operation drops a reference(s) to async value |
| (token, value or group). |
| }]; |
| |
| let arguments = (ins Async_AnyAsyncType:$operand, |
| Confined<I64Attr, [IntPositive]>:$count); |
| |
| let assemblyFormat = [{ |
| $operand attr-dict `:` type($operand) |
| }]; |
| } |
| |
| #endif // ASYNC_OPS |