[mlir] Add async.await operation to async dialect

[ezhulenev]

Summary

Add async.await operation to "unwrap" async.values

[mehdi_amini]

mlir/lib/Dialect/Async/IR/Async.cpp

277

Can't this use the declarative assembly?

309

Typo: unwraps

[ezhulenev]

Fix typo

mlir/lib/Dialect/Async/IR/Async.cpp

277

Declarative assembly will require typing unwrapped type, something like %0 = async.await %av :!async.value<type> -> type? Or is it possible to register just the type inference for parser?

[csigg]

mlir/lib/Dialect/Async/IR/Async.cpp

277

custom<> might work.

[ezhulenev]

Use assembly format with custom directives.

[ezhulenev]

mlir/lib/Dialect/Async/IR/Async.cpp

277

Thanks, didn't see it before, it indeed works quite nice in this case.

[herhut]

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

99

nit: becomes

114

Would Optional<AnyType>:$value work?

mlir/lib/Dialect/Async/IR/Async.cpp

263

Would it make sense to have one builder that either extracts the result type if it is an async.value or otherwise builds an await without result?

312

These errors are missing tests.

[ezhulenev]

Change return type to optional + add verification tests.

[ezhulenev]

Update function parameter names to match declaration.

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

let arguments = (ins Variadic<Async_TokenType>:$dependencies,
                        Variadic<Async_AnyValueOrTokenType>:$operands);
 
   let results = (outs Async_TokenType:$token,
                       Variadic<Async_AnyValueType>:$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); }];
 }
 
 def Async_YieldOp :
     Async_Op<"yield", [HasParent<"ExecuteOp">, NoSideEffect, Terminator]> {
   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 = "attr-dict ($operands^ `:` type($operands))?";
 
   let verifier = [{ return ::verify(*this); }];
 }
 
+def Async_AwaitOp : Async_Op<"await", [NoSideEffect]> {
+  let summary = "waits for the argument to become ready";
+  let description = [{
+    The `async.await` operation waits until the argument become 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 Variadic<AnyType>:$value);
+
+  let skipDefaultBuilders = 1;
+
+  let builders = [
+    OpBuilder<"mlir::OpBuilder &builder, OperationState &result,"
+              "Value token, ArrayRef<NamedAttribute> attrs = {}">,
+    OpBuilder<"mlir::OpBuilder &builder, OperationState &result,"
+              "Type type, Value value, ArrayRef<NamedAttribute> attrs = {}">,
+  ];
+
+  let extraClassDeclaration = [{
+    Optional<Type> getResultType() {
+      if (getResultTypes().empty()) return None;
+      return getResultTypes()[0];
+    }
+  }];
+
+  let assemblyFormat = [{
+    attr-dict $operand `:` custom<AwaitValueType>(type($operand), type($value))
+  }];
+
+  let verifier = [{ return ::verify(*this); }];
+}
+
 #endif // ASYNC_OPS

mlir/lib/Dialect/Async/IR/Async.cpp

static LogicalResult verify(ExecuteOp op) {
   if (!std::equal(unwrappedTypes.begin(), unwrappedTypes.end(),
                   op.body().getArgumentTypes().begin()))
     return op.emitOpError("async body region argument types do not match the "
                           "execute operation arguments types");
 
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+/// AwaitOp
+//===----------------------------------------------------------------------===//
+
+void AwaitOp::build(OpBuilder &, OperationState &result, Value token,
+                    ArrayRef<NamedAttribute> attrs) {
+  result.addOperands({token});
+  result.attributes.append(attrs.begin(), attrs.end());
+}
+
+void AwaitOp::build(OpBuilder &, OperationState &result, Type type, Value value,
+                    ArrayRef<NamedAttribute> attrs) {
+  result.addOperands({value});
+  result.addTypes({type});
+  result.attributes.append(attrs.begin(), attrs.end());
+}
+
+static ParseResult parseAwaitValueType(OpAsmParser &parser, Type &operandType,
+                                       SmallVectorImpl<Type> &valueTypes) {
+  if (parser.parseType(operandType))
+    return failure();
+
+  // Add unwrapped async.value type to the returned values types.
+  if (auto valueType = operandType.dyn_cast<ValueType>())
+    valueTypes.push_back(valueType.getValueType());
+
+  return success();
+}
+
+static void printAwaitValueType(OpAsmPrinter &p, Type operandType,
+                                TypeRange valueTypes) {
+  p << operandType;
+}
+
+static LogicalResult verify(AwaitOp op) {
+  Type argType = op.operand().getType();
+
+  // Awaiting on a token does not have any results.
+  if (argType.isa<TokenType>() && !op.getResultTypes().empty())
+    return op.emitOpError("awaiting on a token must have empty result");
+
+  // Awaiting on a value unwraps the async value type.
+  if (auto value = argType.dyn_cast<ValueType>()) {
+    if (op.getResultTypes().size() != 1)
+      return op.emitOpError("awaiting on a value must have one result");
+    if (*op.getResultType() != value.getValueType())
+      return op.emitOpError("result type does not match async value type");
+  }
+
+  return success();
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc"

mlir/test/Dialect/Async/ops.mlir

func @empty_tokens_or_values_operands() {
   // CHECK: async.execute {
   %token2 = async.execute -> () { async.yield }
   // CHECK: async.execute {
   %token3 = async.execute () { async.yield }
   // CHECK: async.execute {
   %token4 = async.execute [] { async.yield }
   return
 }
+
+// CHECK-LABEL: @await_token
+func @await_token(%arg0: !async.token) {
+  // CHECK: async.await %arg0
+  async.await %arg0 : !async.token
+  return
+}
+
+// CHECK-LABEL: @await_value
+func @await_value(%arg0: !async.value<f32>) -> f32 {
+  // CHECK: async.await %arg0
+  %0 = async.await %arg0 : !async.value<f32>
+  return %0 : f32
+}

mlir/test/Dialect/Async/verify.mlir

+// RUN: mlir-opt %s -split-input-file -verify-diagnostics | FileCheck %s
+
+// FileCheck test must have at least one CHECK statement.
+// CHECK-LABEL: @no_op
+func @no_op(%arg0: !async.token) {
+  return
+}
+
+// -----
+
+func @wrong_async_await_arg_type(%arg0: f32) {
+  // expected-error @+1 {{'async.await' op operand #0 must be async value type or token type, but got 'f32'}}
+  async.await %arg0 : f32
+}