[Mlir-commits] [mlir] 97f15ed - [mlir][python] Provide "all passes" registration module in Python

[Aart Bik]

Author: Aart Bik
Date: 2021-05-26T15:14:57-07:00
New Revision: 97f15eda4f26eb18c914884ce808b4e366e29c34

URL: https://github.com/llvm/llvm-project/commit/97f15eda4f26eb18c914884ce808b4e366e29c34
DIFF: https://github.com/llvm/llvm-project/commit/97f15eda4f26eb18c914884ce808b4e366e29c34.diff

LOG: [mlir][python] Provide "all passes" registration module in Python

Currently, passes are registered on a per-dialect basis, which
provides the smallest footprint obviously. But for prototyping
and experimentation, a convenience "all passes" module is provided,
which registers all known MLIR passes in one run.

Usage in Python:

import mlir.all_passes_registration

Reviewed By: mehdi_amini

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

Added: 
    mlir/lib/Bindings/Python/AllPassesRegistration.cpp
    mlir/python/mlir/all_passes_registration/__init__.py

Modified: 
    mlir/include/mlir-c/Registration.h
    mlir/lib/Bindings/Python/CMakeLists.txt
    mlir/lib/CAPI/Registration/Registration.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir-c/Registration.h b/mlir/include/mlir-c/Registration.h
index 4cfc967192c9f..e8604329fd381 100644
--- a/mlir/include/mlir-c/Registration.h
+++ b/mlir/include/mlir-c/Registration.h
@@ -60,6 +60,9 @@ MLIR_CAPI_EXPORTED void mlirRegisterAllDialects(MlirContext context);
 /// Register all translations to LLVM IR for dialects that can support it.
 MLIR_CAPI_EXPORTED void mlirRegisterAllLLVMTranslations(MlirContext context);
 
+/// Register all compiler passes of MLIR.
+MLIR_CAPI_EXPORTED void mlirRegisterAllPasses();
+
 #ifdef __cplusplus
 }
 #endif

diff  --git a/mlir/lib/Bindings/Python/AllPassesRegistration.cpp b/mlir/lib/Bindings/Python/AllPassesRegistration.cpp
new file mode 100644
index 0000000000000..f595b20ba8d25
--- /dev/null
+++ b/mlir/lib/Bindings/Python/AllPassesRegistration.cpp
@@ -0,0 +1,22 @@
+//===- AllPassesRegistration.cpp - Pybind module to register all 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/Registration.h"
+
+#include <pybind11/pybind11.h>
+
+// -----------------------------------------------------------------------------
+// Module initialization.
+// -----------------------------------------------------------------------------
+
+PYBIND11_MODULE(_mlirAllPassesRegistration, m) {
+  m.doc() = "MLIR All Passes Convenience Module";
+
+  // Register all passes on load.
+  mlirRegisterAllPasses();
+}

diff  --git a/mlir/lib/Bindings/Python/CMakeLists.txt b/mlir/lib/Bindings/Python/CMakeLists.txt
index 575b9dbbd4ed7..173cf48c0d5c9 100644
--- a/mlir/lib/Bindings/Python/CMakeLists.txt
+++ b/mlir/lib/Bindings/Python/CMakeLists.txt
@@ -25,6 +25,14 @@ add_dependencies(MLIRBindingsPythonExtension MLIRCoreBindingsPythonExtension)
 add_subdirectory(Transforms)
 add_subdirectory(Conversions)
 
+add_mlir_python_extension(MLIRAllPassesRegistrationBindingsPythonExtension _mlirAllPassesRegistration
+  INSTALL_DIR
+    python
+  SOURCES
+    AllPassesRegistration.cpp
+)
+add_dependencies(MLIRBindingsPythonExtension MLIRAllPassesRegistrationBindingsPythonExtension)
+
 add_mlir_python_extension(MLIRAsyncPassesBindingsPythonExtension _mlirAsyncPasses
   INSTALL_DIR
     python
@@ -37,7 +45,7 @@ add_mlir_python_extension(MLIRSparseTensorPassesBindingsPythonExtension _mlirSpa
   INSTALL_DIR
     python
   SOURCES
-  SparseTensorPasses.cpp
+    SparseTensorPasses.cpp
 )
 add_dependencies(MLIRBindingsPythonExtension MLIRSparseTensorPassesBindingsPythonExtension)
 

diff  --git a/mlir/lib/CAPI/Registration/Registration.cpp b/mlir/lib/CAPI/Registration/Registration.cpp
index dea7824531522..4ac300d1fac70 100644
--- a/mlir/lib/CAPI/Registration/Registration.cpp
+++ b/mlir/lib/CAPI/Registration/Registration.cpp
@@ -10,6 +10,7 @@
 
 #include "mlir/CAPI/IR.h"
 #include "mlir/InitAllDialects.h"
+#include "mlir/InitAllPasses.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 
 void mlirRegisterAllDialects(MlirContext context) {
@@ -21,3 +22,5 @@ void mlirRegisterAllDialects(MlirContext context) {
 void mlirRegisterAllLLVMTranslations(MlirContext context) {
   mlir::registerLLVMDialectTranslation(*unwrap(context));
 }
+
+void mlirRegisterAllPasses() { mlir::registerAllPasses(); }

diff  --git a/mlir/python/mlir/all_passes_registration/__init__.py b/mlir/python/mlir/all_passes_registration/__init__.py
new file mode 100644
index 0000000000000..cf3367cfe92ff
--- /dev/null
+++ b/mlir/python/mlir/all_passes_registration/__init__.py
@@ -0,0 +1,8 @@
+#  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
+
+_cextAllPasses = _load_extension("_mlirAllPassesRegistration")
+del _load_extension