[mlir][python] Add python support for async dialect and passes.

since the `async` keyword is reserved in python, the dialect is called async_dialect.

Differential Revision: https://reviews.llvm.org/D101447

11 files changed, 165 insertions, 0 deletions

diff --git a/mlir/include/mlir-c/Dialect/Async.h b/mlir/include/mlir-c/Dialect/Async.h
new file mode 100644
index 000000000000..50b6413efa52
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Async.h
@@ -0,0 +1,28 @@
+//===-- mlir-c/Dialect/Async.h - C API for 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
+//
+//===---------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_ASYNC_H
+#define MLIR_C_DIALECT_ASYNC_H
+
+#include "mlir-c/Registration.h"
+#include "mlir-c/Support.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Async, async);
+
+#ifdef __cplusplus
+}
+#endif
+
+#include "mlir/Dialect/Async/Passes.capi.h.inc"
+
+#endif // MLIR_C_DIALECT_ASYNC_H
diff --git a/mlir/include/mlir/Dialect/Async/CMakeLists.txt b/mlir/include/mlir/Dialect/Async/CMakeLists.txt
index 499e442720cc..cabd5d3087cf 100644
--- a/mlir/include/mlir/Dialect/Async/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/Async/CMakeLists.txt
@@ -2,6 +2,8 @@ add_subdirectory(IR)
 
 set(LLVM_TARGET_DEFINITIONS Passes.td)
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name Async)
+mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix Async)
+mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix Async)
 add_public_tablegen_target(MLIRAsyncPassIncGen)
 
 add_mlir_doc(Passes AsyncPasses ./ -gen-pass-doc)
diff --git a/mlir/lib/Bindings/Python/AsyncOps.td b/mlir/lib/Bindings/Python/AsyncOps.td
new file mode 100644
index 000000000000..b65b9bafdd61
--- /dev/null
+++ b/mlir/lib/Bindings/Python/AsyncOps.td
@@ -0,0 +1,15 @@
+//===-- AsyncOps.td - Entry point async_dialect bindings --*- 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
+//
+//===---------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_ASYNC_OPS
+#define PYTHON_BINDINGS_ASYNC_OPS
+
+include "mlir/Bindings/Python/Attributes.td"
+include "mlir/Dialect/Async/IR/AsyncOps.td"
+
+#endif
diff --git a/mlir/lib/Bindings/Python/AsyncPasses.cpp b/mlir/lib/Bindings/Python/AsyncPasses.cpp
new file mode 100644
index 000000000000..2b83ed40d695
--- /dev/null
+++ b/mlir/lib/Bindings/Python/AsyncPasses.cpp
@@ -0,0 +1,22 @@
+//===- AsyncPasses.cpp - Pybind module for the Async passes -------------===//
+//
+// 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-c/Dialect/Async.h"
+
+#include <pybind11/pybind11.h>
+
+// -----------------------------------------------------------------------------
+// Module initialization.
+// -----------------------------------------------------------------------------
+
+PYBIND11_MODULE(_mlirAsyncPasses, m) {
+  m.doc() = "MLIR Async Dialect Passes";
+
+  // Register all Async passes on load.
+  mlirRegisterAsyncPasses();
+}
diff --git a/mlir/lib/Bindings/Python/CMakeLists.txt b/mlir/lib/Bindings/Python/CMakeLists.txt
index 39192cc54d3c..eba4d2886ec2 100644
--- a/mlir/lib/Bindings/Python/CMakeLists.txt
+++ b/mlir/lib/Bindings/Python/CMakeLists.txt
@@ -31,6 +31,11 @@ endforeach()
 # Generate dialect-specific bindings.
 ################################################################################
 
+add_mlir_dialect_python_bindings(MLIRBindingsPythonAsyncOps
+  TD_FILE AsyncOps.td
+  DIALECT_NAME async_dialect)
+add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonAsyncOps)
+
 add_mlir_dialect_python_bindings(MLIRBindingsPythonBuiltinOps
   TD_FILE BuiltinOps.td
   DIALECT_NAME builtin)
@@ -120,6 +125,14 @@ endif()
 add_subdirectory(Transforms)
 add_subdirectory(Conversions)
 
+add_mlir_python_extension(MLIRAsyncPassesBindingsPythonExtension _mlirAsyncPasses
+  INSTALL_DIR
+    python
+  SOURCES
+    AsyncPasses.cpp
+)
+add_dependencies(MLIRBindingsPythonExtension MLIRAsyncPassesBindingsPythonExtension)
+
 add_mlir_python_extension(MLIRLinalgPassesBindingsPythonExtension _mlirLinalgPasses
   INSTALL_DIR
     python
diff --git a/mlir/lib/Bindings/Python/mlir/dialects/async_dialect/__init__.py b/mlir/lib/Bindings/Python/mlir/dialects/async_dialect/__init__.py
new file mode 100644
index 000000000000..dcf9d6cb2638
--- /dev/null
+++ b/mlir/lib/Bindings/Python/mlir/dialects/async_dialect/__init__.py
@@ -0,0 +1,5 @@
+#  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
+
+from .._async_dialect_ops_gen import *
diff --git a/mlir/lib/Bindings/Python/mlir/dialects/async_dialect/passes/__init__.py b/mlir/lib/Bindings/Python/mlir/dialects/async_dialect/passes/__init__.py
new file mode 100644
index 000000000000..88a7b539c9c7
--- /dev/null
+++ b/mlir/lib/Bindings/Python/mlir/dialects/async_dialect/passes/__init__.py
@@ -0,0 +1,6 @@
+#  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
+
+from ...._cext_loader import _load_extension
+_cextAsyncPasses = _load_extension("_mlirAsyncPasses")
diff --git a/mlir/lib/CAPI/Dialect/Async.cpp b/mlir/lib/CAPI/Dialect/Async.cpp
new file mode 100644
index 000000000000..182cbf9dfd71
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Async.cpp
@@ -0,0 +1,13 @@
+//===- Async.cpp - C Interface for Async dialect --------------------------===//
+//
+// 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-c/Dialect/Async.h"
+#include "mlir/CAPI/Registration.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Async, async, mlir::async::AsyncDialect)
diff --git a/mlir/lib/CAPI/Dialect/AsyncPasses.cpp b/mlir/lib/CAPI/Dialect/AsyncPasses.cpp
new file mode 100644
index 000000000000..aa2074dcd173
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/AsyncPasses.cpp
@@ -0,0 +1,26 @@
+//===- AsyncPasses.cpp - C API for Async Dialect Passes -----------------===//
+//
+// 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/CAPI/Pass.h"
+#include "mlir/Dialect/Async/Passes.h"
+#include "mlir/Pass/Pass.h"
+
+// Must include the declarations as they carry important visibility attributes.
+#include "mlir/Dialect/Async/Passes.capi.h.inc"
+
+using namespace mlir;
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include "mlir/Dialect/Async/Passes.capi.cpp.inc"
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/mlir/lib/CAPI/Dialect/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt
index 41c659d6ab75..3f6265e8a2fe 100644
--- a/mlir/lib/CAPI/Dialect/CMakeLists.txt
+++ b/mlir/lib/CAPI/Dialect/CMakeLists.txt
@@ -1,5 +1,7 @@
 # TODO: Make the check source feature optional as an argument on *_add_library.
 set(LLVM_OPTIONAL_SOURCES
+  Async.cpp
+  AsyncPasses.cpp
   Linalg.cpp
   LinalgPasses.cpp
   SCF.cpp
@@ -8,6 +10,20 @@ set(LLVM_OPTIONAL_SOURCES
   Tensor.cpp
 )
 
+add_mlir_public_c_api_library(MLIRCAPIAsync
+  Async.cpp
+  AsyncPasses.cpp
+
+  DEPENDS
+  MLIRAsyncPassIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRAsync
+  MLIRAsyncTransforms
+  MLIRPass
+)
+
 add_mlir_public_c_api_library(MLIRCAPILinalg
   Linalg.cpp
   LinalgPasses.cpp
diff --git a/mlir/test/Bindings/Python/dialects/async_dialect.py b/mlir/test/Bindings/Python/dialects/async_dialect.py
new file mode 100644
index 000000000000..6a33bd6b6d03
--- /dev/null
+++ b/mlir/test/Bindings/Python/dialects/async_dialect.py
@@ -0,0 +1,19 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+import mlir.dialects.async_dialect
+import mlir.dialects.async_dialect.passes
+from mlir.passmanager import *
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+
+def testAsyncPass():
+  with Context() as context:
+    PassManager.parse('async-to-async-runtime')
+  print('SUCCESS')
+
+# CHECK-LABEL: testAsyncPass
+#       CHECK: SUCCESS
+run(testAsyncPass)