[MLIR] Add async token/value arguments to async.execute op

[ezhulenev]

Summary

Async execute operation can take async arguments as dependencies.

Change async.execute custom parser/printer format to use %value as %unwrapped: !async.value<!type> sytax.

[rriddle]

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

193–194

Functions should not be placed in namespaces, they should be at the top-level. namespaces in .cpp files should only really be for defining local classes.

[rriddle]

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

193–194

Relevant coding style links:
https://llvm.org/docs/CodingStandards.html#use-namespace-qualifiers-to-implement-previously-declared-functions
https://llvm.org/docs/CodingStandards.html#anonymous-namespaces

[ezhulenev]

Use namespace qualifiers to implement previously declared functions in Async.cpp

[rriddle]

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

108

Is the async:: here necessary?

176

You need to check the result of this as it can fail.

200

nit: Use isa instead of dyn_cast if you don't need to use the result.

mlir/test/Dialect/Async/ops.mlir

88

Can you just infer the argument types of the region during parsing? That would remove the need to define the entry block signature.

[ezhulenev]

Cleanup includes + use isa for type matching

[ezhulenev]

Handle parser.resolveOperands failure in async.execute op parsing

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

108

Cleaned up includes, and it's no longer ambiguous with std dialect.

mlir/test/Dialect/Async/ops.mlir

88

Yes, types could be inferred from the execute op operands, but how would I define SSA name for the unwrapped operands then? In this example %arg1: !async.value<f32> unwrapped to f32 and gets %0 ssa name in the attached region.

[herhut]

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

69

Reading the code below, op.done() does not suggest a value. Maybe completionToken?

70

Maybe results better reflects what this is? Or resultValues?

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

128

nit: op.getName()

130–137

Giving this a function like syntax makes it look like a function, which it is not. It does implicit capture, etc. What is the motivation for this?

134

nit: left-over tab

mlir/test/Dialect/Async/ops.mlir

88

The gpu.launch does this with extra syntax. You could have something like

async.execute(%operand as %unwrapped : async.value<foo>) -> ...

but then you would only have the as for unwrapped arguments and not others. We could also require them to be in some order. Something like

async.execute (%operand as %unwrapped : type, ...) deps (%token1, %token2)

lots of room for syntax bike-shedding. I am not a big fan of this looking too function like, as it does not have function behavior. There are no call sites for example.

mlir/test/Dialect/Async/verify.mlir

7

// CHECK-LABEL: @no_op

14

CHECK-LABEL: @wrong_execute_arg also below.

[ezhulenev]

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

130–137

Motivation is your examples in the discourse RFC topic :) Maybe use function call syntax for it?

%token = async.execute (%arg0, %arg1): (!async.token, !async.token) -> !async.token {
}

It's a bit confusing though to see function call syntax with explicit args + implicit capture of "regular" values.

EDIT: after reading comments below I think it makes sense to separate values from tokens, how about this syntax:

%val = ... : !async.value<!type>
%tok = ... : !async.token

%token, %results  = async.execute [%token, ...](%val as %unwrapped): (!async.value<!type>) -> (!async.token, !async.value<!type>>) {
  async.yield %unwrapped : !type
}

1. token dependencies are explicit and types of tokens are not mentioned in the signature
2. syntax is sufficiently different (is it?) from function call

But I'm not sure how to parse %val as %unwrapped part, do you have any examples?

[mehdi_amini]

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

130–137

This syntax looks much nicer to me!
If we always return a token as first result, we can also leave it out from the functional type

[ezhulenev]

Update async.execute op syntax

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

69

I changed to just token, I think that completion part is clear from the context.

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

130–137

Updated printer/parser to use async.execute [%tokens] (%values as %unwrapped) syntax.

mlir/test/Dialect/Async/verify.mlir

14

Functions that fail verification do not show up in the output.

[ezhulenev]

Cleanup async.execute parsing

[ezhulenev]

Use entry block arg names in async.execute op printer

[herhut]

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

54

Mega-nit: I think this indentation is off

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

130–137

Beautiful!

Sorry for creating confusion with my sloppy RFC syntax.

185

You need to check the result here, as well, in case this fails.

[ezhulenev]

Don't forget to check ParseResult and fix indentation.

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

130–137

I actually quite liked it when I saw it first time, but %value as %unwrapped is definitely a readability improvement.

[herhut]

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

192

nit: left-over tab?

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

let typeDescription = [{
     `async.value` represents a value returned by asynchronous operations,
     which may or may not be available currently, but will be available at some
     point in the future.
   }];
 
   Type valueType = type;
 }
 
-def Async_AnyValueType : Type<CPred<"$_self.isa<::mlir::async::ValueType>()">,
-                                    "async value type">;
+def Async_AnyValueType : DialectType<AsyncDialect,
+                           CPred<"$_self.isa<::mlir::async::ValueType>()">,
 
+def Async_AnyValueOrTokenType : AnyTypeOf<[Async_AnyValueType,
+                                           Async_TokenType]>;
+
 #endif // ASYNC_BASE_TD

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

let description = [{
     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 unwraps asynchronous value arguments, and passes
+    values of the underlying type to the attached body region.
+
+    Example:
+
     ```mlir
-    %done, %values = async.execute {
+    %token = ... : !async.token
+    %value = ... : !async.value<f32>
+
+    %done, %values =
+      async.execute (%token: !async.token, %value: !async.value<f32>)
+                 -> (!async.token, !async.value<!some.type>)
+      {
+        ^bb0(%arg0: f32):
           %0 = "compute0"(...) : !some.type
-      async.yield %1 : f32
-    } : !async.token, !async.value<!some.type>
+          async.yield %0 : !some.type
+      }
 
     %1 = "compute1"(...) : !some.type
     ```
+
+    In the example above asynchronous execution starts only after token and
+    value arguments become ready.
   }];
 
-  // TODO: Take async.tokens/async.values as arguments.
-  let arguments = (ins );
+  let arguments = (ins Variadic<Async_AnyValueOrTokenType>:$arguments);
   let results = (outs Async_TokenType:$done,
                       Variadic<Async_AnyValueType>:$values);
   let regions = (region SizedRegion<1>:$body);
 
-  let printer = [{ return ::mlir::async::print(p, *this); }];
-  let parser = [{ return ::mlir::async::parse$cppClass(parser, result); }];
+  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 ::mlir::async::verify(*this); }];
+  let verifier = [{ return ::verify(*this); }];
 }
 
 #endif // ASYNC_OPS

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

 //===- Async.cpp - MLIR Async Operations ----------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Async/IR/Async.h"
 
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/PatternMatch.h"
-#include "mlir/IR/StandardTypes.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/raw_ostream.h"
 
-namespace mlir {
-namespace async {
+using namespace mlir;
+using namespace mlir::async;
 
 void AsyncDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc"
       >();
   addTypes<TokenType>();
   addTypes<ValueType>();
TypeSwitch<Type>(type)
       })
       .Default([](Type) { llvm_unreachable("unexpected 'async' type kind"); });
 }
 
 //===----------------------------------------------------------------------===//
 /// ValueType
 //===----------------------------------------------------------------------===//
 
+namespace mlir {
+namespace async {
 namespace detail {
 
 // Storage for `async.value<T>` type, the only member is the wrapped type.
 struct ValueTypeStorage : public TypeStorage {
   ValueTypeStorage(Type valueType) : valueType(valueType) {}
 
   /// The hash key used for uniquing.
   using KeyTy = Type;
   bool operator==(const KeyTy &key) const { return key == valueType; }
 
   /// Construction.
   static ValueTypeStorage *construct(TypeStorageAllocator &allocator,
                                      Type valueType) {
     return new (allocator.allocate<ValueTypeStorage>())
         ValueTypeStorage(valueType);
   }
 
   Type valueType;
 };
 
 } // namespace detail
+} // namespace async
+} // namespace mlir
 
 ValueType ValueType::get(Type valueType) {
   return Base::get(valueType.getContext(), valueType);
 }
 
 Type ValueType::getValueType() { return getImpl()->valueType; }
 
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(YieldOp op) {
+static LogicalResult verify(async::YieldOp op) {
   // Get the underlying value types from async values returned from the
   // parent `async.execute` operation.
   auto executeOp = op.getParentOfType<ExecuteOp>();
   auto types = llvm::map_range(executeOp.values(), [](const OpResult &result) {
     return result.getType().cast<ValueType>().getValueType();
   });
 
   if (!std::equal(types.begin(), types.end(), op.getOperandTypes().begin()))
     return op.emitOpError("Operand types do not match the types returned from "
                           "the parent ExecuteOp");
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 /// ExecuteOp
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, ExecuteOp op) {
-  p << "async.execute ";
-  p.printRegion(op.body());
+  p << "async.execute";
   p.printOptionalAttrDict(op.getAttrs());
-  p << " : ";
-  p.printType(op.done().getType());
-  if (!op.values().empty())
-    p << ", ";
-  llvm::interleaveComma(op.values(), p, [&](const OpResult &result) {
-    p.printType(result.getType());
+  p << " (";
+  llvm::interleaveComma(llvm::zip(op.arguments(), op.getOperandTypes()), p,
+                        [&](const auto &tuple) {
+                          p << std::get<0>(tuple) << ": " << std::get<1>(tuple);
                         });
+  p << ")";
+  p.printArrowTypeList(op.getResultTypes());
+  p.printRegion(op.body());
 }
 
 static ParseResult parseExecuteOp(OpAsmParser &parser, OperationState &result) {
   MLIRContext *ctx = result.getContext();
 
-  // Parse asynchronous region.
-  Region *body = result.addRegion();
-  if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{},
-                         /*enableNameShadowing=*/false))
-    return failure();
-
   // Parse operation attributes.
   NamedAttrList attrs;
   if (parser.parseOptionalAttrDict(attrs))
     return failure();
   result.addAttributes(attrs);
 
-  // Parse result types.
+  SmallVector<OpAsmParser::OperandType, 4> entryArgs;
+  SmallVector<NamedAttrList, 4> argAttrs;
+  SmallVector<NamedAttrList, 4> resultAttrs;
+  SmallVector<Type, 4> argTypes;
   SmallVector<Type, 4> resultTypes;
-  if (parser.parseColonTypeList(resultTypes))
+
+  // Parse the async.execute signature which is similar to function signature.
+  auto signatureLocation = parser.getCurrentLocation();
+  bool allowVariadic = false;
+  bool isVariadic = false;
+  if (impl::parseFunctionSignature(parser, allowVariadic, entryArgs, argTypes,
+                                   argAttrs, isVariadic, resultTypes,
+                                   resultAttrs))
     return failure();
 
+  // async.execute does not support arg or result attributes, verify that they
+  // are empty.
+  auto emptyAttrs = [](const NamedAttrList &list) { return list.empty(); };
+
+  if (!llvm::all_of(argAttrs, emptyAttrs))
+    return parser.emitError(signatureLocation)
+           << "argument attributes are not supported";
+
+  if (!llvm::all_of(resultAttrs, emptyAttrs))
+    return parser.emitError(signatureLocation)
+           << "result attributes are not supported";
+
+  parser.resolveOperands(entryArgs, argTypes, signatureLocation,
+                         result.operands);
+
   // First result type must be an async token type.
   if (resultTypes.empty() || resultTypes.front() != TokenType::get(ctx))
-    return failure();
+    return parser.emitError(signatureLocation)
+           << "first return type must be an async.token";
   parser.addTypesToList(resultTypes, result.types);
 
+  // Parse asynchronous region.
+  Region *body = result.addRegion();
+  if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{},
+                         /*enableNameShadowing=*/false))
+    return failure();
+
   return success();
 }
 
-} // namespace async
-} // namespace mlir
+static LogicalResult verify(ExecuteOp op) {
+  // Unwrap async.execute operation arguments:
+  //   - `async.token` arguments are skipped
+  //   - `async.value<T>` arguments unwrapped to `T`
+  SmallVector<Type, 4> unwrappedTypes;
+  for (Type type : op.getOperandTypes()) {
+    if (auto token = type.dyn_cast<TokenType>())
+      continue;
+
+    if (auto value = type.dyn_cast<ValueType>()) {
+      unwrappedTypes.push_back(value.getValueType());
+      continue;
+    }
+
+    llvm_unreachable("unexpected 'async' type kind");
+  }
+
+  // Verify that unwrapped argument types matches the body region arguments.
+  if (llvm::size(unwrappedTypes) != llvm::size(op.body().getArgumentTypes()))
+    return op.emitOpError("the number of async body region arguments does not "
+                          "match the number of execute operation arguments");
+
+  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();
+}
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc"

mlir/test/Dialect/Async/ops.mlir

-// RUN: mlir-opt  %s | FileCheck %s
+// RUN: mlir-opt %s | FileCheck %s
 
 // CHECK-LABEL: @identity_token
-func @identity_token(%arg0 : !async.token) -> !async.token {
+func @identity_token(%arg0: !async.token) -> !async.token {
   // CHECK: return %arg0 : !async.token
   return %arg0 : !async.token
 }
 
 // CHECK-LABEL: @identity_value
 func @identity_value(%arg0 : !async.value<f32>) -> !async.value<f32> {
   // CHECK: return %arg0 : !async.value<f32>
   return %arg0 : !async.value<f32>
 }
 
 // CHECK-LABEL: @empty_async_execute
 func @empty_async_execute() -> !async.token {
-  %done = async.execute {
+  %done = async.execute () -> !async.token {
     async.yield
-  } : !async.token
+  }
 
   // CHECK: return %done : !async.token
   return %done : !async.token
 }
 
 // CHECK-LABEL: @return_async_value
 func @return_async_value() -> !async.value<f32> {
-  %done, %values = async.execute {
+  %done, %values = async.execute () -> (!async.token, !async.value<f32>) {
     %cst = constant 1.000000e+00 : f32
     async.yield %cst : f32
-  } : !async.token, !async.value<f32>
+  }
 
   // CHECK: return %values : !async.value<f32>
   return %values : !async.value<f32>
 }
 
+// CHECK-LABEL: @return_captured_value
+func @return_captured_value() -> !async.token {
+  %cst = constant 1.000000e+00 : f32
+
+  %done, %value = async.execute () -> (!async.token, !async.value<f32>) {
+    async.yield %cst : f32
+  }
+
+  // CHECK: return %done : !async.token
+  return %done : !async.token
+}
+
 // CHECK-LABEL: @return_async_values
 func @return_async_values() -> (!async.value<f32>, !async.value<f32>) {
-  %done, %values:2 = async.execute {
+  %done, %values:2 = async.execute () -> (!async.token, !async.value<f32>, !async.value<f32>) {
     %cst1 = constant 1.000000e+00 : f32
     %cst2 = constant 2.000000e+00 : f32
     async.yield %cst1, %cst2 : f32, f32
-  } : !async.token, !async.value<f32>, !async.value<f32>
+  }
 
   // CHECK: return %values#0, %values#1 : !async.value<f32>, !async.value<f32>
   return %values#0, %values#1 : !async.value<f32>, !async.value<f32>
 }
+
+// CHECK-LABEL: @async_token_operands
+func @async_token_operands(%arg0: !async.token) -> !async.token {
+  // CHECK: async.execute (%arg0: !async.token) -> !async.token {
+  %done = async.execute (%arg0 : !async.token) -> !async.token {
+    async.yield
+  }
+
+  // CHECK: return %done : !async.token
+  return %done : !async.token
+}
+
+// CHECK-LABEL: @async_value_operands
+func @async_value_operands(%arg0: !async.value<f32>) -> !async.token {
+  // CHECK: async.execute (%arg0: !async.value<f32>) -> !async.token
+  %done = async.execute (%arg0: !async.value<f32>) -> !async.token {
+    ^bb0(%0: f32):
+      async.yield
+  }
+
+  // CHECK: return %done : !async.token
+  return %done : !async.token
+}
+
+// CHECK-LABEL: @async_token_and_value_operands
+func @async_token_and_value_operands(%arg0: !async.token, %arg1: !async.value<f32>) -> !async.token {
+  // CHECK: async.execute (%arg0: !async.token, %arg1: !async.value<f32>)
+  // CHECK-SAME: -> (!async.token, !async.value<f32>)
+  %done, %value = async.execute (%arg0: !async.token, %arg1: !async.value<f32>) -> (!async.token, !async.value<f32>) {
+    ^bb0(%0: f32):
+      async.yield %0 : f32
+  }
+
+  // CHECK: return %done : !async.token
+  return %done : !async.token
+}
+

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.
+func @no_op(%arg0: !async.token) {
+  // CHECK: @no_op
+  return
+}
+
+// -----
+
+func @wrong_execute_arg(%arg0: f32) {
+  // expected-error @+1 {{'async.execute' op operand #0 must be async value type or token type, but got 'f32'}}
+  %done = async.execute (%arg0: f32) -> !async.token {
+    ^bb0(%0: f32):
+      async.yield
+  }
+}
+
+// -----
+
+func @wrong_async_body_num_args(%arg0: !async.value<f32>) {
+  // expected-error @+1 {{'async.execute' op the number of async body region arguments does not match the number of execute operation arguments}}
+  %done = async.execute (%arg0: !async.value<f32>) -> !async.token {
+    ^bb0(%0: f32, %1: f32):
+      async.yield
+  }
+}
+
+// -----
+
+func @wrong_async_body_arg_type(%arg0: !async.value<f32>) {
+  // expected-error @+1 {{'async.execute' op async body region argument types do not match the execute operation arguments types}}
+  %done = async.execute (%arg0: !async.value<f32>) -> !async.token {
+    ^bb0(%0 : f64):
+      async.yield
+  }
+}