[MLIR] Add Async dialect with trivial async.region operation

[ezhulenev]

Summary

Start Async dialect for modeling asynchronous execution.

[rriddle]

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

50

nit: Please do not put usernames in TODOs.

[ezhulenev]

Update TODO

[mehdi_amini]

LG but please wait for @herhut to have a look

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

14

Can you fix the clang-tidy warning?

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

38

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

32

I think it is important to clarify whether it is guaranteed that they will execute concurrently, i.e. is it valid to have in compute0 a spin-lock on a flag that has to be set by compute1.
Or only explicit dependencies are making side-effects visible?

55

Could skip the type entirely here?

[ezhulenev]

Update documentation + fix header guard.

[ezhulenev]

Fix typos

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

32

I think it is important to have sequential execution legal, clarified it in description.

55

I was thinking that with an addition of return values from async regions, it will be clearer to have an explicit token return type.

Example:

%0#2 = async.region {
  %1 = alloc : memref<10xf32>
  async.yield %0 : memref<10xf32>
} : () -> (!async.token, !async.value<memref<10xf32>>)
             ^                       ^
             |                       |  
              region completed    return value is ready

[ezhulenev]

Rename async.region to async.execute

[herhut]

Thanks for starting this!

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

55

Do we need both types here? Does it ever happen that the async.value is ready before or after the async.token becomes ready? I can see that we want a async.value<void> like thing in case the async.region does not produce a value (and we can name it async.token). I am just not sure about requiring it to always be present.

[ezhulenev]

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

55

Not sure about that. Having a token is convenient when you have two async regions that produce values, however there is no "true data dependency" between them.

Example:

%0:3 = async.execute (%arg : memref) {
} : !async.token, !async.value<foo>, !async.value<bar>

%1 = async.execute [%0#1](%arg : memref) {
   ... do not need the value foo here
} : ...

both regions read-write into memref (e.g. scratch buffer) but second op does not need any of the produced values. Token can be constructed from any async value, but then it becomes a question, which value to choose. I feel that explicit token makes semantics clearer.

[mehdi_amini]

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

55

It seems easier to me to handle the region with having always this token there: it makes it more regular.
It is also less ambiguous if a dependency is on the side-effects of the region or on the availability of a particular result. Otherwise a result may have users that aren't interested in the value itself but are just synchronizing on the async task.

mlir/include/mlir/Dialect/Async/CMakeLists.txt

+add_subdirectory(IR)

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

+//===- Async.h - MLIR Async dialect -----------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the async dialect that is used for modeling asynchronous
+// execution.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ASYNC_IR_ASYNC_H
+#define MLIR_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 {
+
+/// The token type to represent asynchronous operation completion.
+class TokenType : public Type::TypeBase<TokenType, Type, TypeStorage> {
+public:
+  using Base::Base;
+};
+
+} // 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_ASYNC_IR_ASYNC_H

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

+//===- AsyncBase.td ----------------------------------------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Base definitions for the `async` dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ASYNC_BASE_TD
+#define ASYNC_BASE_TD
+
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// Async dialect definitions
+//===----------------------------------------------------------------------===//
+
+def AsyncDialect : Dialect {
+  let name = "async";
+
+  let summary = "Types and operations for async dialect";
+  let description = [{
+    This dialect contains operations for modeling asynchronous execution.
+  }];
+
+  let cppNamespace = "::mlir::async";
+}
+
+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 asyncrhonous oeprations that created it is
+    completed.
+  }];
+}
+
+#endif // ASYNC_BASE_TD

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

+//===- AsyncOps.td - Async operations definition -----------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the operation definition file for Async dialect operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ASYNC_OPS
+#define ASYNC_OPS
+
+include "mlir/Dialect/Async/IR/AsyncBase.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+//===----------------------------------------------------------------------===//
+// Async op definitions
+//===----------------------------------------------------------------------===//
+
+// Base class for the operation in this dialect
+class Async_Op<string mnemonic, list<OpTrait> traits = []> :
+    Op<AsyncDialect, mnemonic, traits>;
+
+def Async_RegionOp : Async_Op<"region"> {
+  let summary = "Asynchronous region operation";
+  let description = [{
+    The `body` region attached to the `async.region` operation semantically
+    can be executed concurrently with the successor operation. In the followup
+    example "compute0" can be executed concurrently with "compute1".
+
+    The actual concurrency semantics depends on the dialect lowering to the
+    executable format. Fully sequential execution ("compute0" completes before
+    "compute1" starts) is a completely legal execution.
+
+    Example:
+
+    ```mlir
+    %0 = async.region {
+      "compute0"(...)
+      async.yield
+    } : () -> !async.token
+
+    %1 = "compute1"(...)
+    ```
+  }];
+
+  // TODO(ezhulenev): Take async.tokens/async.values as arguments.
+  let arguments = (ins );
+  let results = (outs Async_TokenType:$done);
+  let regions = (region SizedRegion<1>:$body);
+
+  let assemblyFormat = "$body attr-dict `:` `(` `)` `->` type($done)";
+}
+
+def Async_YieldOp :
+    Async_Op<"yield", [HasParent<"RegionOp">, NoSideEffect, Terminator]> {
+  let summary = "terminator for Async region operation";
+  let description = [{
+    The `async.yield` is a special terminator operation for the block inside
+    `async.region` operation.
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$operands);
+
+  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
+}
+
+#endif // ASYNC_OPS

mlir/include/mlir/Dialect/Async/IR/CMakeLists.txt

+add_mlir_dialect(AsyncOps async)
+add_mlir_doc(AsyncOps -gen-dialect-doc AsyncDialect Dialects/)

mlir/include/mlir/Dialect/CMakeLists.txt

 add_subdirectory(Affine)
+add_subdirectory(Async)
 add_subdirectory(AVX512)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(OpenACC)
 add_subdirectory(OpenMP)
 add_subdirectory(PDL)
 add_subdirectory(PDLInterp)
 add_subdirectory(Quant)
 add_subdirectory(SCF)
 add_subdirectory(Shape)
 add_subdirectory(SPIRV)
 add_subdirectory(StandardOps)
 add_subdirectory(Vector)

mlir/include/mlir/InitAllDialects.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INITALLDIALECTS_H_
 #define MLIR_INITALLDIALECTS_H_
 
 #include "mlir/Dialect/AVX512/AVX512Dialect.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMAVX512Dialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"

 
 namespace mlir {
 
 // Add all the MLIR dialects to the provided registry.
 inline void registerAllDialects(DialectRegistry &registry) {
   // clang-format off
   registry.insert<acc::OpenACCDialect,
                   AffineDialect,
+                  async::AsyncDialect,
                   avx512::AVX512Dialect,
                   gpu::GPUDialect,
                   LLVM::LLVMAVX512Dialect,
                   LLVM::LLVMDialect,
                   linalg::LinalgDialect,
                   scf::SCFDialect,
                   omp::OpenMPDialect,
                   pdl::PDLDialect,

mlir/lib/Dialect/Async/CMakeLists.txt

+add_subdirectory(IR)

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/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;
+
+void AsyncDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc"
+      >();
+  addTypes<TokenType>();
+}
+
+/// 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());
+
+  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"; })
+      .Default([](Type) { llvm_unreachable("unexpected 'async' type kind"); });
+}
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc"

mlir/lib/Dialect/Async/IR/CMakeLists.txt

+add_mlir_dialect_library(MLIRAsync
+  Async.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Async
+
+  DEPENDS
+  MLIRAsyncOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRDialect
+  MLIRIR
+  )

mlir/lib/Dialect/CMakeLists.txt

 add_subdirectory(Affine)
+add_subdirectory(Async)
 add_subdirectory(AVX512)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(OpenACC)
 add_subdirectory(OpenMP)
 add_subdirectory(PDL)
 add_subdirectory(PDLInterp)

mlir/test/Dialect/Async/ops.mlir

+// RUN: mlir-opt  %s | FileCheck %s
+
+// CHECK-LABEL: @identity
+func @identity(%arg0 : !async.token) -> !async.token {
+  // CHECK: return %arg0 : !async.token
+  return %arg0 : !async.token
+}
+
+// CHECK-LABEL: @empty_async_region
+func @empty_async_region() -> !async.token {
+  %0 = async.region {
+    async.yield
+  } : () -> !async.token
+
+  return %0 : !async.token
+}