Add Async Function to the Async Dialect

[yijia1212]

Summary

Add Async Function to the Async Dialect

Today async.execute operation semantics requires attached region to be executed in a thread managed by the runtime, and always returns an !async.token result. We need to model async functions that are not necessarily executed in a runtime-managed threads, but eventually lowered to llvm coroutines.

Example:

async.func @foo(%arg0: !async.value<f32>) -> !async.token {
  %0 = async.await %arg0: !async.value<f32>
  "do_something_with_f32"(%0)
  return
}

If arg0 is available this function will be executed in the caller thread. If it's not available it will be suspended and resumed later later on a thread managed by the async runtime. Currently this is not representable with async.execute operations.

The longer term goal is to make async dialect more like https://github.com/lewissbaker/cppcoro to be able to represent structured host concurrency in MLIR.

(1) Add async.func, async.call, and async.return operations in Async Dialect

[rriddle]

The description of this seems extremely light, and doesn't really touch on any rationale.

[rriddle]

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

145

Symbols aren't really expected to have results. These are SSA results, not function results.

[rriddle]

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

155–161

Are these create methods really necessary? func::FuncOp has them only really for historical reasons, using a builder is generally just cleaner/better.

251

277–297

The formatting here looks off.

[ezhulenev]

In D137189#3900108, @rriddle wrote:

The description of this seems extremely light, and doesn't really touch on any rationale.

I added a bit more context to the description.

[yijia1212]

Address review.

[ezhulenev]

mlir/test/Dialect/Async/verify.mlir

48

missing new line

[yijia1212]

Address comments.

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

 
 #ifndef MLIR_DIALECT_ASYNC_IR_ASYNC_H
 #define MLIR_DIALECT_ASYNC_IR_ASYNC_H
 
 #include "mlir/Dialect/Async/IR/AsyncTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
-#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // Async Dialect
 //===----------------------------------------------------------------------===//
 

 /// Returns true if the type is reference counted at runtime.
 inline bool isRefCounted(Type type) {
   return type.isa<TokenType, ValueType, GroupType>();
 }
 
 } // namespace async
 } // namespace mlir
 
+namespace llvm {
+
+/// Allow stealing the low bits of async::FuncOp.
+template <>
+struct PointerLikeTypeTraits<mlir::async::FuncOp> {
+  static inline void *getAsVoidPointer(mlir::async::FuncOp val) {
+    return const_cast<void *>(val.getAsOpaquePointer());
+  }
+  static inline mlir::async::FuncOp getFromVoidPointer(void *p) {
+    return mlir::async::FuncOp::getFromOpaquePointer(p);
+  }
+  static constexpr int numLowBitsAvailable = 3;
+};
+} // namespace llvm
+
 #endif // MLIR_DIALECT_ASYNC_IR_ASYNC_H

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

 #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/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/CallInterfaces.td"
+include "mlir/IR/SymbolInterfaces.td"
+include "mlir/IR/FunctionInterfaces.td"
+include "mlir/IR/OpAsmInterface.td"
+
 
 //===----------------------------------------------------------------------===//
 // Async op definitions
 //===----------------------------------------------------------------------===//
 
 // Base class for the operation in this dialect
 class Async_Op<string mnemonic, list<Trait> traits = []> :
     Op<AsyncDialect, mnemonic, traits>;
def Async_ExecuteOp :
 
   let extraClassDeclaration = [{
     using BodyBuilderFn =
         function_ref<void(OpBuilder &, Location, ValueRange)>;
 
   }];
 }
 
+def Async_FuncOp : Async_Op<"func",
+    [CallableOpInterface, FunctionOpInterface,
+     IsolatedFromAbove, OpAsmOpInterface, Symbol]> {
+  let summary = "async function operation";
+  let description = [{
+    An async function is like a normal function, but supports non-blocking
+    await. Internally, async function is lowered to the LLVM coroutinue with
+    async runtime intrinsic. It can return an async token and/or async values. 
+    The token represents the execution state of async function and can be used 
+    when users want to express dependencies on some side effects, e.g.,
+    the token becomes available once every thing in the func body is executed.
+
+    Example:
+
+    ```mlir
+    // Async function can't return void, it always must be some async thing.
+    async.func @async.0() -> !async.token {
+      return
+    }
+
+    // Function returns only async value.
+    async.func @async.1() -> !async.value<i32> {
+      %0 = arith.constant 42 : i32
+      return %0 : i32
+    }
+
+    // Implicit token can be added to return types.
+    async.func @async.2() -> !async.token, !async.value<i32> {
+      %0 = arith.constant 42 : i32
+      return %0 : i32
+    }
+    ```
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name,
+                       TypeAttrOf<FunctionType>:$function_type,
+                       OptionalAttr<StrAttr>:$sym_visibility);
+
+  let regions = (region AnyRegion:$body);
+
+  let builders = [
+    OpBuilder<(ins "StringRef":$name, "FunctionType":$type,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs,
+      CArg<"ArrayRef<DictionaryAttr>", "{}">:$argAttrs)>
+  ];
+
+  let extraClassDeclaration = [{
+    //===------------------------------------------------------------------===//
+    // CallableOpInterface
+    //===------------------------------------------------------------------===//
+
+    /// Returns the region on the current operation that is callable. This may
+    /// return null in the case of an external callable object, e.g. an external
+    /// function.
+    ::mlir::Region *getCallableRegion() { return isExternal() ? nullptr
+                                                              : &getBody(); }
+
+    /// Returns the results types that the callable region produces when
+    /// executed.
+    ArrayRef<Type> getCallableResults() { return getFunctionType()
+                                                    .getResults(); }
+
+    //===------------------------------------------------------------------===//
+    // FunctionOpInterface Methods
+    //===------------------------------------------------------------------===//
+
+    /// Returns the argument types of this async function.
+    ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }
+
+    /// Returns the result types of this async function.
+    ArrayRef<Type> getResultTypes() { return getFunctionType().getResults(); }
+
+    /// Returns the number of results of this async function
+    unsigned getNumResults() {return getResultTypes().size();}
+
+    /// Is the async func stateful
+    bool isStateful() { return isa<TokenType>(getFunctionType().getResult(0));}
+
+    //===------------------------------------------------------------------===//
+    // OpAsmOpInterface Methods
+    //===------------------------------------------------------------------===//
+
+    /// Allow the dialect prefix to be omitted.
+    static StringRef getDefaultDialect() { return "async"; }
+
+    //===------------------------------------------------------------------===//
+    // SymbolOpInterface Methods
+    //===------------------------------------------------------------------===//
+
+    bool isDeclaration() { return isExternal(); }
+  }];
+  let hasCustomAssemblyFormat = 1;
+
+  let hasVerifier = 1;
+}
+
+def Async_CallOp : Async_Op<"call",
+    [CallOpInterface, DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
+  let summary = "async call operation";
+  let description = [{
+    The `async.call` operation represents a direct call to an async function
+    that is within the same symbol scope as the call. The operands and result
+    types of the call must match the specified async function type. The callee
+    is encoded as a symbol reference attribute named "callee".
+
+    Example:
+
+    ```mlir
+    %2 = async.call @my_add(%0, %1) : (f32, f32) -> !async.value<f32>
+    ```
+  }];
+
+  let arguments = (ins FlatSymbolRefAttr:$callee, Variadic<AnyType>:$operands);
+  let results = (outs Variadic<Async_AnyValueOrTokenType>);
+
+  let builders = [
+    OpBuilder<(ins "FuncOp":$callee, CArg<"ValueRange", "{}">:$operands), [{
+      $_state.addOperands(operands);
+      $_state.addAttribute("callee", SymbolRefAttr::get(callee));
+      $_state.addTypes(callee.getFunctionType().getResults());
+    }]>,
+    OpBuilder<(ins "SymbolRefAttr":$callee, "TypeRange":$results,
+      CArg<"ValueRange", "{}">:$operands), [{
+      $_state.addOperands(operands);
+      $_state.addAttribute("callee", callee);
+      $_state.addTypes(results);
+    }]>,
+    OpBuilder<(ins "StringAttr":$callee, "TypeRange":$results,
+      CArg<"ValueRange", "{}">:$operands), [{
+      build($_builder, $_state, SymbolRefAttr::get(callee), results, operands);
+    }]>,
+    OpBuilder<(ins "StringRef":$callee, "TypeRange":$results,
+      CArg<"ValueRange", "{}">:$operands), [{
+      build($_builder, $_state, StringAttr::get($_builder.getContext(), callee),
+            results, operands);
+    }]>
+  ];
+
+  let extraClassDeclaration = [{
+    FunctionType getCalleeType();
+
+    /// Get the argument operands to the called function.
+    operand_range getArgOperands() {
+      return {arg_operand_begin(), arg_operand_end()};
+    }
+
+    operand_iterator arg_operand_begin() { return operand_begin(); }
+    operand_iterator arg_operand_end() { return operand_end(); }
+
+    /// Return the callee of this operation.
+    CallInterfaceCallable getCallableForCallee() {
+      return (*this)->getAttrOfType<SymbolRefAttr>("callee");
+    }
+  }];
+
+  let assemblyFormat = [{
+    $callee `(` $operands `)` attr-dict `:` functional-type($operands, results)
+  }];
+}
+
+def Async_ReturnOp : Async_Op<"return",
+    [Pure, HasParent<"FuncOp">, ReturnLike, Terminator]> {
+  let summary = "Async function return operation";
+  let description = [{
+    The `async.return` is a special terminator operation for Async function.
+
+    Example:
+
+    ```mlir
+    async.func @foo() : !async.token {
+      return
+    }
+    ```
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$operands);
+
+  let builders = [OpBuilder<(ins), [{build($_builder, $_state, llvm::None);}]>];
+
+  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
+  let hasVerifier = 1;
+}
+
 def Async_YieldOp :
     Async_Op<"yield", [
       HasParent<"ExecuteOp">, Pure, 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 hasVerifier = 1;
 }
 
 def Async_AwaitOp : Async_Op<"await"> {

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/IR/BlockAndValueMapping.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/FunctionImplementation.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::async;
 
 #include "mlir/Dialect/Async/IR/AsyncOpsDialect.cpp.inc"
 
 constexpr StringRef AsyncDialect::kAllowedToBlockAttrName;
if (*getResultType() != value.getValueType())
                            << " does not match async value type "
                            << value.getValueType();
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// FuncOp
+//===----------------------------------------------------------------------===//
+
+void FuncOp::build(OpBuilder &builder, OperationState &state, StringRef name,
+                   FunctionType type, ArrayRef<NamedAttribute> attrs,
+                   ArrayRef<DictionaryAttr> argAttrs) {
+  state.addAttribute(SymbolTable::getSymbolAttrName(),
+                     builder.getStringAttr(name));
+  state.addAttribute(FunctionOpInterface::getTypeAttrName(),
+                     TypeAttr::get(type));
+
+  state.attributes.append(attrs.begin(), attrs.end());
+  state.addRegion();
+
+  if (argAttrs.empty())
+    return;
+  assert(type.getNumInputs() == argAttrs.size());
+  function_interface_impl::addArgAndResultAttrs(builder, state, argAttrs,
+                                                /*resultAttrs=*/llvm::None);
+}
+
+ParseResult FuncOp::parse(OpAsmParser &parser, OperationState &result) {
+  auto buildFuncType =
+      [](Builder &builder, ArrayRef<Type> argTypes, ArrayRef<Type> results,
+         function_interface_impl::VariadicFlag,
+         std::string &) { return builder.getFunctionType(argTypes, results); };
+
+  return function_interface_impl::parseFunctionOp(
+      parser, result, /*allowVariadic=*/false, buildFuncType);
+}
+
+void FuncOp::print(OpAsmPrinter &p) {
+  function_interface_impl::printFunctionOp(p, *this, /*isVariadic=*/false);
+}
+
+/// Check that the result type of async.func is not void and must be
+/// some async token or async values.
+LogicalResult FuncOp::verify() {
+  auto resultTypes = getResultTypes();
+  if (resultTypes.empty())
+    return emitOpError()
+           << "result is expected to be at least of size 1, but got "
+           << resultTypes.size();
+
+  for (unsigned i = 0, e = resultTypes.size(); i != e; ++i) {
+    auto type = resultTypes[i];
+    if (!type.isa<TokenType>() && !type.isa<ValueType>())
+      return emitOpError() << "result type must be async value type or async "
+                              "token type, but got "
+                           << type;
+    // We only allow AsyncToken appear as the first return value
+    if (type.isa<TokenType>() && i != 0) {
+      return emitOpError()
+             << " results' (optional) async token type is expected "
+                "to appear as the 1st return value, but got "
+             << i + 1;
+    }
+  }
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+/// CallOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult CallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  // Check that the callee attribute was specified.
+  auto fnAttr = (*this)->getAttrOfType<FlatSymbolRefAttr>("callee");
+  if (!fnAttr)
+    return emitOpError("requires a 'callee' symbol reference attribute");
+  FuncOp fn = symbolTable.lookupNearestSymbolFrom<FuncOp>(*this, fnAttr);
+  if (!fn)
+    return emitOpError() << "'" << fnAttr.getValue()
+                         << "' does not reference a valid async function";
+
+  // Verify that the operand and result types match the callee.
+  auto fnType = fn.getFunctionType();
+  if (fnType.getNumInputs() != getNumOperands())
+    return emitOpError("incorrect number of operands for callee");
+
+  for (unsigned i = 0, e = fnType.getNumInputs(); i != e; ++i)
+    if (getOperand(i).getType() != fnType.getInput(i))
+      return emitOpError("operand type mismatch: expected operand type ")
+             << fnType.getInput(i) << ", but provided "
+             << getOperand(i).getType() << " for operand number " << i;
+
+  if (fnType.getNumResults() != getNumResults())
+    return emitOpError("incorrect number of results for callee");
+
+  for (unsigned i = 0, e = fnType.getNumResults(); i != e; ++i)
+    if (getResult(i).getType() != fnType.getResult(i)) {
+      auto diag = emitOpError("result type mismatch at index ") << i;
+      diag.attachNote() << "      op result types: " << getResultTypes();
+      diag.attachNote() << "function result types: " << fnType.getResults();
+      return diag;
+    }
+
+  return success();
+}
+
+FunctionType CallOp::getCalleeType() {
+  return FunctionType::get(getContext(), getOperandTypes(), getResultTypes());
+}
+
+//===----------------------------------------------------------------------===//
+/// ReturnOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult ReturnOp::verify() {
+  auto funcOp = (*this)->getParentOfType<FuncOp>();
+  ArrayRef<Type> resultTypes = funcOp.isStateful()
+                                   ? funcOp.getResultTypes().drop_front()
+                                   : funcOp.getResultTypes();
+  // Get the underlying value types from async types returned from the
+  // parent `async.func` operation.
+  auto types = llvm::map_range(resultTypes, [](const Type &result) {
+    return result.cast<ValueType>().getValueType();
+  });
+
+  if (getOperandTypes() != types)
+    return emitOpError("operand types do not match the types returned from "
+                       "the parent FuncOp");
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // TableGen'd type method definitions

mlir/test/Dialect/Async/ops.mlir

func.func @create_group_and_await_all(%arg0: !async.token,
   // CHECK: async.add_to_group %arg1
   %1 = async.add_to_group %arg0, %0 : !async.token
   %2 = async.add_to_group %arg1, %0 : !async.value<f32>
   async.await_all %0
 
   %3 = arith.addi %1, %2 : index
   return %3 : index
 }
+
+// CHECK-LABEL: @async_func_return_token
+async.func @async_func_return_token() -> !async.token {
+  // CHECK: return
+  return
+}
+
+// CHECK-LABEL: @async_func_return_value
+async.func @async_func_return_value() -> !async.value<i32> {
+  %0 = arith.constant 42 : i32
+  // CHECK: return %[[value:.*]] : i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: @async_func_return_optional_token
+async.func @async_func_return_optional_token() -> (!async.token, !async.value<i32>) {
+  %0 = arith.constant 42 : i32
+  // CHECK: return %[[value:.*]] : i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: @async_call
+func.func @async_call() {
+  // CHECK: async.call @async_func_return_token
+  // CHECK: async.call @async_func_return_value
+  // CHECK: async.call @async_func_return_optional_token
+  %0 = async.call @async_func_return_token() : () -> !async.token
+  %1 = async.call @async_func_return_value() : () -> !async.value<i32>
+  %2, %3 = async.call @async_func_return_optional_token() : () -> (!async.token, !async.value<i32>)
+  return
+}

mlir/test/Dialect/Async/verify.mlir

 }
 
 // -----
 
 func.func @wrong_async_await_result_type(%arg0: !async.value<f32>) {
   // expected-error @+1 {{'async.await' op result type 'f64' does not match async value type 'f32'}}
   %0 = "async.await"(%arg0): (!async.value<f32>) -> f64
 }
+
+
+// -----
+// expected-error @+1 {{'async.func' op result is expected to be at least of size 1, but got 0}}
+async.func @wrong_async_func_void_result_type(%arg0: f32) {
+  return
+}
+
+
+// -----
+// expected-error @+1 {{'async.func' op result type must be async value type or async token type, but got 'f32'}}
+async.func @wrong_async_func_result_type(%arg0: f32) -> f32 {
+  return %arg0 : f32
+}
+
+// -----
+// expected-error @+1 {{'async.func' op  results' (optional) async token type is expected to appear as the 1st return value, but got 2}}
+async.func @wrong_async_func_token_type_placement(%arg0: f32) -> (!async.value<f32>, !async.token) {
+  return %arg0 : f32
+}
+
+// -----
+async.func @wrong_async_func_return_type(%arg0: f32) -> (!async.token, !async.value<i32>) {
+  // expected-error @+1 {{'async.return' op operand types do not match the types returned from the parent FuncOp}}
+  return %arg0 : f32
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

td_library(
     name = "AsyncOpsTdFiles",
     srcs = [
         "include/mlir/Dialect/Async/IR/AsyncDialect.td",
         "include/mlir/Dialect/Async/IR/AsyncOps.td",
         "include/mlir/Dialect/Async/IR/AsyncTypes.td",
     ],
     includes = ["include"],
     deps = [
+        ":CallInterfacesTdFiles",
         ":ControlFlowInterfacesTdFiles",
+        ":FunctionInterfacesTdFiles",
         ":InferTypeOpInterfaceTdFiles",
         ":OpBaseTdFiles",
         ":SideEffectInterfacesTdFiles",
     ],
 )
 
 gentbl_cc_library(
     name = "AsyncOpsIncGen",