diff --git a/mlir/include/mlir/Dialect/Async/CMakeLists.txt b/mlir/include/mlir/Dialect/Async/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Async/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/mlir/include/mlir/Dialect/Async/IR/Async.h b/mlir/include/mlir/Dialect/Async/IR/Async.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Async/IR/Async.h @@ -0,0 +1,39 @@ +//===- 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 { +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 diff --git a/mlir/include/mlir/Dialect/Async/IR/AsyncBase.td b/mlir/include/mlir/Dialect/Async/IR/AsyncBase.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Async/IR/AsyncBase.td @@ -0,0 +1,43 @@ +//===- 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()">, "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 diff --git a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td @@ -0,0 +1,71 @@ +//===- 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 traits = []> : + Op; + +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:$operands); + + let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?"; +} + +#endif // ASYNC_OPS diff --git a/mlir/include/mlir/Dialect/Async/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/Async/IR/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Async/IR/CMakeLists.txt @@ -0,0 +1,2 @@ +add_mlir_dialect(AsyncOps async) +add_mlir_doc(AsyncOps -gen-dialect-doc AsyncDialect Dialects/) diff --git a/mlir/include/mlir/Dialect/CMakeLists.txt b/mlir/include/mlir/Dialect/CMakeLists.txt --- a/mlir/include/mlir/Dialect/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/CMakeLists.txt @@ -1,4 +1,5 @@ add_subdirectory(Affine) +add_subdirectory(Async) add_subdirectory(AVX512) add_subdirectory(GPU) add_subdirectory(Linalg) diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h --- a/mlir/include/mlir/InitAllDialects.h +++ b/mlir/include/mlir/InitAllDialects.h @@ -16,6 +16,7 @@ #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" @@ -42,6 +43,7 @@ // clang-format off registry.insert(); + addTypes(); +} + +/// 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) + .Case([&](Type) { os << "token"; }) + .Default([](Type) { llvm_unreachable("unexpected 'async' type kind"); }); +} + +#define GET_OP_CLASSES +#include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc" diff --git a/mlir/lib/Dialect/Async/IR/CMakeLists.txt b/mlir/lib/Dialect/Async/IR/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/Async/IR/CMakeLists.txt @@ -0,0 +1,13 @@ +add_mlir_dialect_library(MLIRAsync + Async.cpp + + ADDITIONAL_HEADER_DIRS + ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Async + + DEPENDS + MLIRAsyncOpsIncGen + + LINK_LIBS PUBLIC + MLIRDialect + MLIRIR + ) diff --git a/mlir/lib/Dialect/CMakeLists.txt b/mlir/lib/Dialect/CMakeLists.txt --- a/mlir/lib/Dialect/CMakeLists.txt +++ b/mlir/lib/Dialect/CMakeLists.txt @@ -1,4 +1,5 @@ add_subdirectory(Affine) +add_subdirectory(Async) add_subdirectory(AVX512) add_subdirectory(GPU) add_subdirectory(Linalg) diff --git a/mlir/test/Dialect/Async/ops.mlir b/mlir/test/Dialect/Async/ops.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Async/ops.mlir @@ -0,0 +1,16 @@ +// 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 +}