[MLIR] Add async.value type to Async dialect

[ezhulenev]

Summary

Return values from async regions as !async.value<...>.

[mehdi_amini]

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

51

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

74

Can you add a class doc (something simple like "Storage for async.value<T> type, the only member is the wrapped type.")

103

Please fix :)

106

This can fail, you need to check in the verifier. Also I rather be more strict here: we really want to map the immediate parent only:

114

Nit: remote trivial braces.

131

interleaveComma should do it here.

This can't be done in the assembly format by the way?

You're also missing attribute printing

[ezhulenev]

Update parse/print + fix linter warning

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

106

Wouldn't this be checked by HasParent trait defined on .td file?

131

Changed to interleaveComma.

Assembly format does not seem to work for variadic results.

Added attrs parsing/printing.

[ezhulenev]

Verify async.yield op parent

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

106

Didn't find anything like that in the generated file, added the check to this function.

[mehdi_amini]

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

106

You're right: the HasParent trait will verify it (there isn't a specific verification in the TableGen generated file, it is dispatched to the trait by the simple fact of inheriting from it)

[ezhulenev]

Do not check async.yield parent op in verifier, it is already checked by HasParent trait.

[csigg]

Looks good!

Just name bike shedding: I assume you considered both async.value and async.future. What's your reason to pick async.value?

[ezhulenev]

In D88510#2304208, @csigg wrote:

What's your reason to pick async.value?

When I was writing this PR I didn't see the async.future proposal ;) Although I'm not a fan of async.future because both async and future have very similar semantics, and having them both in the same type looks weird to me. I'm fine changing it to future, if people have strong opinion about that. I'll prepare the next PR with async arguments, and then we can change the type if needed.

[csigg]

In D88510#2304241, @ezhulenev wrote:

Although I'm not a fan of async.future because both async and future have very similar semantics...

I see where you are coming from. I look at !async.<type> as <type> in the async dialect, and that way future makes more sense than value.

mlir/include/mlir/Dialect/Async/IR/Async.h

 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {
 namespace async {
 
+namespace detail {
+struct ValueTypeStorage;
+} // namespace detail
+
 /// The token type to represent asynchronous operation completion.
 class TokenType : public Type::TypeBase<TokenType, Type, TypeStorage> {
 public:
   using Base::Base;
 };
 
+/// The value type to represent values returned from asynchronous operations.
+class ValueType
+    : public Type::TypeBase<ValueType, Type, detail::ValueTypeStorage> {
+public:
+  using Base::Base;
+
+  /// Get or create an async ValueType with the provided value type.
+  static ValueType get(Type valueType);
+
+  Type getValueType();
+};
+
 } // namespace async
 } // namespace mlir
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Async/IR/AsyncOps.h.inc"
 
 #include "mlir/Dialect/Async/IR/AsyncOpsDialect.h.inc"
 
 #endif // MLIR_DIALECT_ASYNC_IR_ASYNC_H

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

def Async_TokenType : DialectType<AsyncDialect,
     CPred<"$_self.isa<::mlir::async::TokenType>()">, "token type">,
     BuildableType<"$_builder.getType<::mlir::async::TokenType>()"> {
   let typeDescription = [{
     `async.token` is a type returned by asynchronous operations, and it becomes
     `ready` when the asynchronous operations that created it is completed.
   }];
 }
 
+class Async_ValueType<Type type>
+    : DialectType<AsyncDialect,
+        And<[
+          CPred<"$_self.isa<::mlir::async::ValueType>()">,
+          SubstLeaves<"$_self",
+                      "$_self.cast<::mlir::async::ValueType>().getValueType()",
+                      type.predicate>
+       ]>, "async value type with " # type.description # " underlying type"> {
+  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">;
+
 #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`).
 
-    Example:
-
     ```mlir
-    %0 = async.execute {
-      "compute0"(...)
-      async.yield
-    } : !async.token
+    %done, %values = async.execute {
+      %0 = "compute0"(...) : !some.type
+      async.yield %1 : f32
+    } : !async.token, !async.value<!some.type>
 
-    %1 = "compute1"(...)
+    %1 = "compute1"(...) : !some.type
     ```
   }];
 
   // TODO: Take async.tokens/async.values as arguments.
   let arguments = (ins );
-  let results = (outs Async_TokenType:$done);
+  let results = (outs Async_TokenType:$done,
+                      Variadic<Async_AnyValueType>:$values);
   let regions = (region SizedRegion<1>:$body);
 
-  let assemblyFormat = "$body attr-dict `:` type($done)";
+  let printer = [{ return ::mlir::async::print(p, *this); }];
+  let parser = [{ return ::mlir::async::parse$cppClass(parser, result); }];
 }
 
 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); }];
 }
 
 #endif // ASYNC_OPS

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

 #include "mlir/IR/DialectImplementation.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"
 
-using namespace mlir;
-using namespace mlir::async;
+namespace mlir {
+namespace async {
 
 void AsyncDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc"
       >();
   addTypes<TokenType>();
+  addTypes<ValueType>();
 }
 
 /// Parse a type registered to this dialect.
 Type AsyncDialect::parseType(DialectAsmParser &parser) const {
   StringRef keyword;
   if (parser.parseKeyword(&keyword))
     return Type();
 
   if (keyword == "token")
     return TokenType::get(getContext());
 
+  if (keyword == "value") {
+    Type ty;
+    if (parser.parseLess() || parser.parseType(ty) || parser.parseGreater()) {
+      parser.emitError(parser.getNameLoc(), "failed to parse async value type");
+      return Type();
+    }
+    return ValueType::get(ty);
+  }
+
   parser.emitError(parser.getNameLoc(), "unknown async type: ") << keyword;
   return Type();
 }
 
 /// Print a type registered to this dialect.
 void AsyncDialect::printType(Type type, DialectAsmPrinter &os) const {
   TypeSwitch<Type>(type)
-      .Case<TokenType>([&](Type) { os << "token"; })
+      .Case<TokenType>([&](TokenType) { os << "token"; })
+      .Case<ValueType>([&](ValueType valueTy) {
+        os << "value<";
+        os.printType(valueTy.getValueType());
+        os << '>';
+      })
       .Default([](Type) { llvm_unreachable("unexpected 'async' type kind"); });
 }
 
+//===----------------------------------------------------------------------===//
+/// ValueType
+//===----------------------------------------------------------------------===//
+
+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
+
+ValueType ValueType::get(Type valueType) {
+  return Base::get(valueType.getContext(), valueType);
+}
+
+Type ValueType::getValueType() { return getImpl()->valueType; }
+
+//===----------------------------------------------------------------------===//
+// YieldOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(YieldOp op) {
+  // Immediate parent must be an ExecuteOp.
+  auto executeOp = dyn_cast_or_null<ExecuteOp>(op.getParentOp());
+  if (!executeOp)
+    return op.emitOpError() << "Expect parent operation to be `async.execute`";
+
+  // Get the underlying value types from async values returned from the
+  // parent `async.execute` operation.
+  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.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());
+  });
+}
+
+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<Type, 4> resultTypes;
+  if (parser.parseColonTypeList(resultTypes))
+    return failure();
+
+  // First result type must be an async token type.
+  if (resultTypes.empty() || resultTypes.front() != TokenType::get(ctx))
+    return failure();
+  parser.addTypesToList(resultTypes, result.types);
+
+  return success();
+}
+
+} // namespace async
+} // namespace mlir
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc"

mlir/test/Dialect/Async/ops.mlir

 // RUN: mlir-opt  %s | FileCheck %s
 
-// CHECK-LABEL: @identity
-func @identity(%arg0 : !async.token) -> !async.token {
+// CHECK-LABEL: @identity_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 {
-  %0 = async.execute {
+  %done = async.execute {
     async.yield
   } : !async.token
 
-  return %0 : !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 {
+    %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_async_values
+func @return_async_values() -> (!async.value<f32>, !async.value<f32>) {
+  %done, %values:2 = async.execute {
+    %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>
 }