1//===- AsyncOps.td - Async operations definition -----------*- tablegen -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This is the operation definition file for Async dialect operations.
10//
11//===----------------------------------------------------------------------===//
12
13#ifndef ASYNC_OPS
14#define ASYNC_OPS
15
16include "mlir/Dialect/Async/IR/AsyncBase.td"
17include "mlir/Interfaces/ControlFlowInterfaces.td"
18include "mlir/Interfaces/SideEffectInterfaces.td"
19
20//===----------------------------------------------------------------------===//
21// Async op definitions
22//===----------------------------------------------------------------------===//
23
24// Base class for the operation in this dialect
25class Async_Op<string mnemonic, list<OpTrait> traits = []> :
26    Op<AsyncDialect, mnemonic, traits>;
27
28def Async_ExecuteOp :
29  Async_Op<"execute", [SingleBlockImplicitTerminator<"YieldOp">,
30                       DeclareOpInterfaceMethods<RegionBranchOpInterface,
31                                                 ["getNumRegionInvocations"]>,
32                       AttrSizedOperandSegments]> {
33  let summary = "Asynchronous execute operation";
34  let description = [{
35    The `body` region attached to the `async.execute` operation semantically
36    can be executed concurrently with the successor operation. In the followup
37    example "compute0" can be executed concurrently with "compute1".
38
39    The actual concurrency semantics depends on the dialect lowering to the
40    executable format. Fully sequential execution ("compute0" completes before
41    "compute1" starts) is a completely legal execution.
42
43    Because concurrent execution is not guaranteed, it is illegal to create an
44    implicit dependency from "compute1" to "compute0" (e.g. via shared global
45    state). All dependencies must be made explicit with async execute arguments
46    (`async.token` or `async.value`).
47
48   `async.execute` operation takes `async.token` dependencies and `async.value`
49    operands separately, and starts execution of the attached body region only
50    when all tokens and values become ready.
51
52    Example:
53
54    ```mlir
55    %dependency = ... : !async.token
56    %value = ... : !async.value<f32>
57
58    %token, %results =
59      async.execute [%dependency](%value as %unwrapped: !async.value<f32>)
60                 -> !async.value<!some.type>
61      {
62        %0 = "compute0"(%unwrapped): (f32) -> !some.type
63        async.yield %0 : !some.type
64      }
65
66    %1 = "compute1"(...) : !some.type
67    ```
68
69    In the example above asynchronous execution starts only after dependency
70    token and value argument become ready. Unwrapped value passed to the
71    attached body region as an %unwrapped value of f32 type.
72  }];
73
74  let arguments = (ins Variadic<Async_TokenType>:$dependencies,
75                       Variadic<Async_AnyValueOrTokenType>:$operands);
76
77  let results = (outs Async_TokenType:$token,
78                      Variadic<Async_AnyValueType>:$results);
79  let regions = (region SizedRegion<1>:$body);
80
81  let printer = [{ return ::print(p, *this); }];
82  let parser = [{ return ::parse$cppClass(parser, result); }];
83  let verifier = [{ return ::verify(*this); }];
84
85  let skipDefaultBuilders = 1;
86  let builders = [
87    OpBuilderDAG<(ins "TypeRange":$resultTypes, "ValueRange":$dependencies,
88      "ValueRange":$operands,
89      CArg<"function_ref<void(OpBuilder &, Location, ValueRange)>",
90           "nullptr">:$bodyBuilder)>,
91  ];
92
93  let extraClassDeclaration = [{
94    using BodyBuilderFn =
95        function_ref<void(OpBuilder &, Location, ValueRange)>;
96
97  }];
98}
99
100def Async_YieldOp :
101    Async_Op<"yield", [HasParent<"ExecuteOp">, NoSideEffect, Terminator]> {
102  let summary = "terminator for Async execute operation";
103  let description = [{
104    The `async.yield` is a special terminator operation for the block inside
105    `async.execute` operation.
106  }];
107
108  let arguments = (ins Variadic<AnyType>:$operands);
109
110  let assemblyFormat = "($operands^ `:` type($operands))? attr-dict";
111
112  let verifier = [{ return ::verify(*this); }];
113}
114
115def Async_AwaitOp : Async_Op<"await"> {
116  let summary = "waits for the argument to become ready";
117  let description = [{
118    The `async.await` operation waits until the argument becomes ready, and for
119    the `async.value` arguments it unwraps the underlying value
120
121    Example:
122
123    ```mlir
124    %0 = ... : !async.token
125    async.await %0 : !async.token
126
127    %1 = ... : !async.value<f32>
128    %2 = async.await %1 : !async.value<f32>
129    ```
130  }];
131
132  let arguments = (ins Async_AnyValueOrTokenType:$operand);
133  let results = (outs Optional<AnyType>:$result);
134
135  let skipDefaultBuilders = 1;
136
137  let builders = [
138    OpBuilderDAG<(ins "Value":$operand,
139      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
140  ];
141
142  let extraClassDeclaration = [{
143    Optional<Type> getResultType() {
144      if (getResultTypes().empty()) return None;
145      return getResultTypes()[0];
146    }
147  }];
148
149  let assemblyFormat = [{
150    $operand `:` custom<AwaitResultType>(
151      type($operand), type($result)
152    ) attr-dict
153  }];
154
155  let verifier = [{ return ::verify(*this); }];
156}
157
158def Async_CreateGroupOp : Async_Op<"create_group", [NoSideEffect]> {
159  let summary = "creates an empty async group";
160  let description = [{
161    The `async.create_group` allocates an empty async group. Async tokens or
162    values can be added to this group later.
163
164    Example:
165
166    ```mlir
167    %0 = async.create_group
168    ...
169    async.await_all %0
170    ```
171  }];
172
173  let arguments = (ins );
174  let results = (outs Async_GroupType:$result);
175
176  let assemblyFormat = "attr-dict";
177}
178
179def Async_AddToGroupOp : Async_Op<"add_to_group", []> {
180  let summary = "adds and async token or value to the group";
181  let description = [{
182    The `async.add_to_group` adds an async token or value to the async group.
183    Returns the rank of the added element in the group. This rank is fixed
184    for the group lifetime.
185
186    Example:
187
188    ```mlir
189    %0 = async.create_group
190    %1 = ... : !async.token
191    %2 = async.add_to_group %1, %0 : !async.token
192    ```
193  }];
194
195  let arguments = (ins Async_AnyValueOrTokenType:$operand,
196                       Async_GroupType:$group);
197  let results = (outs Index:$rank);
198
199  let assemblyFormat = "$operand `,` $group `:` type($operand) attr-dict";
200}
201
202def Async_AwaitAllOp : Async_Op<"await_all", []> {
203  let summary = "waits for the all async tokens or values in the group to "
204                "become ready";
205  let description = [{
206    The `async.await_all` operation waits until all the tokens or values in the
207    group become ready.
208
209    Example:
210
211    ```mlir
212    %0 = async.create_group
213
214    %1 = ... : !async.token
215    %2 = async.add_to_group %1, %0 : !async.token
216
217    %3 = ... : !async.token
218    %4 = async.add_to_group %2, %0 : !async.token
219
220    async.await_all %0
221    ```
222  }];
223
224  let arguments = (ins Async_GroupType:$operand);
225  let results = (outs);
226
227  let assemblyFormat = "$operand attr-dict";
228}
229
230//===----------------------------------------------------------------------===//
231// Async Dialect Automatic Reference Counting Operations.
232//===----------------------------------------------------------------------===//
233
234// All async values (values, tokens, groups) are reference counted at runtime
235// and automatically destructed when reference count drops to 0.
236//
237// All values are semantically created with a reference count of +1 and it is
238// the responsibility of the last async value user to drop reference count.
239//
240// Async values created when:
241//   1. Operation returns async result (e.g. the result of an `async.execute`).
242//   2. Async value passed in as a block argument.
243//
244// It is the responsiblity of the async value user to extend the lifetime by
245// adding a +1 reference, if the reference counted value captured by the
246// asynchronously executed region (`async.execute` operation), and drop it after
247// the last nested use.
248//
249// Reference counting operations can be added to the IR using automatic
250// reference count pass, that relies on liveness analysis to find the last uses
251// of all reference counted values and automatically inserts
252// `drop_ref` operations.
253//
254// See `AsyncRefCountingPass` documentation for the implementation details.
255
256def Async_AddRefOp : Async_Op<"add_ref"> {
257  let summary = "adds a reference to async value";
258  let description = [{
259    The `async.add_ref` operation adds a reference(s) to async value (token,
260    value or group).
261  }];
262
263  let arguments = (ins Async_AnyAsyncType:$operand,
264                       Confined<I32Attr, [IntPositive]>:$count);
265  let results = (outs );
266
267  let assemblyFormat = [{
268    $operand attr-dict `:` type($operand)
269  }];
270}
271
272def Async_DropRefOp : Async_Op<"drop_ref"> {
273  let summary = "drops a reference to async value";
274  let description = [{
275    The `async.drop_ref` operation drops a reference(s) to async value (token,
276    value or group).
277  }];
278
279  let arguments = (ins Async_AnyAsyncType:$operand,
280                       Confined<I32Attr, [IntPositive]>:$count);
281  let results = (outs );
282
283  let assemblyFormat = [{
284    $operand attr-dict `:` type($operand)
285  }];
286}
287
288#endif // ASYNC_OPS
289