Add basic Python bindings for the PassManager and bind libTransforms

[mehdi_amini]

Summary

This only exposes the ability to round-trip a textual pipeline at the
moment.
To exercise it, we also bind the libTransforms in a new Python extension. This
does not include any interesting bindings, but it includes all the
mechanism to add separate native extensions and load them dynamically.
As such passes in libTransforms are only registered after `import
mlir.transforms`.
To support this global registration, the TableGen backend is also
extended to bind to the C API the group registration for passes.

[mehdi_amini]

mlir/lib/Bindings/Python/CMakeLists.txt

10

This is a bit unfortunate, but the way the tree is setup right now isn't very convenient for modularity

mlir/test/CMakeLists.txt

103

I've been wondering if we shouldn't:

1. Rename MLIRBindingsPythonExtension to MLIRCoreBindingsPythonExtension
2. Introduce MLIRBindingsPythonExtensions (Plural) to catch all of the Python extensions

Thoughts?

[mehdi_amini]

Minor tweak

[stellaraccident]

mlir/lib/Bindings/Python/Pass.h

2

Stale comment.

mlir/lib/Bindings/Python/Transforms/Transforms.cpp

2

Stale comment.

21

Stale docstring.

mlir/lib/Bindings/Python/mlir/transforms/__init__.py

8

I kind of suspect that we are going to want to just generate the .so file directly into this directory (probably with some kind of naming convention that lets us put other peer transform-extensions in here) instead of maintaining it at the top-level.

Aside from not cluttering up the global namespace when we add more, there are some other build/deploy environments that I have to support where it is extremely obnoxious to have a lot of stuff at the top level.

If it was nested here, you would just import ._builtinTransforms (or whatever you chose to call it).

(I also know you'd probably not like to deal with the rpath issues right now)

mlir/test/Bindings/Python/pass_manager.py

40

Indent.

mlir/test/CMakeLists.txt

103

Makes sense to me and the names sgtm.

[mehdi_amini]

Address the minor comments

[mehdi_amini]

mlir/lib/Bindings/Python/mlir/transforms/__init__.py

8

I don't quite get what is the desired structure here?

For example right now we have:

$ find python/ | grep -v cache
python/
python/_mlir.cpython-38-x86_64-linux-gnu.so
python/_mlirTransforms.cpython-38-x86_64-linux-gnu.so
python/mlir
python/mlir/__init__.py
python/mlir/ir.py
python/mlir/transforms
python/mlir/transforms/__init__.py
python/mlir/dialects
python/mlir/dialects/std.py
python/mlir/dialects/__init__.py
python/mlir/passmanager.py

How do you suggest we structure the tree?

[mehdi_amini]

mlir/lib/Bindings/Python/mlir/transforms/__init__.py

8

Ping here @stellaraccident ?

[stellaraccident]

mlir/lib/Bindings/Python/mlir/transforms/__init__.py

8

The python/mlir/... tree should make sense to the user as a package tree of public APIs.

What I was trying to do (which may be a fool's errand) was to preserve the ability for everything below the top-level python/mlir/__init__.py file to be relocatable: if I had a downstream project that wanted to privately package a variant of the MLIR Python API, they could just copy the python files and replace the top-level "loader" init__.py, which is responsible for the minutae of how to find the backing extensions and dependencies. I was also leaving room here in case if we wanted fully static builds or such to just have the root init__.py be the point of customization vs having absolute imports everywhere. I will mention that aside from future projects, this is not just a hypothetical for you/me: Google's internal deployments of these things require loader customizations in order to associate with the c-extension, and I've worked in enough other weird, static-linking scenarios to know that this flexibility and relocatability pay off eventually (and is a pain to retrofit/sed-your-way-through downstream when needed).

We can keep it the way you have it for now. What we may want to do is expose a top-level _import_extension(name) that imports the _mlir or _mlirTransforms (or whatever other) extensions in an installation specific way. Then instead of an absolute import here, you would:

from . import _import_extension
_cextTransforms = _import_extension("_mlirTransforms")

That is easy enough to do in a followup.

[mehdi_amini]

mlir/lib/Bindings/Python/mlir/transforms/__init__.py

8

Would we avoid the need to set a PYTHONPATH and get a simpler loader with a layout that would look like the following?

python/mlir/_mlir.cpython-38-x86_64-linux-gnu.so
python/mlir/__init__.py
python/mlir/transforms/_mlirTransforms.cpython-38-x86_64-linux-gnu.so
python/mlir/transforms/__init__.py

[mehdi_amini]

mlir/lib/Bindings/Python/mlir/transforms/__init__.py

8

Actually that does not prevent the need for PYTHONPATH to find the mlir folder, but it makes it more "self-contained".
There is still a dependency on libMLIRPublicAPI.so though.

[stellaraccident]

mlir/lib/Bindings/Python/mlir/transforms/__init__.py

8

We went back and forth on this a long time ago. There are some nice characteristics that come from *not* making it self contained. Namely, once you "open" a package with an __init__.py, you lose the ability to treat any part of it as a "Namespace Package", which further restricts your ability to have a separate wheel for it. This stems from the fact that the Python import machinery stops at the first directory it finds that is a package, regardless of whether what you are looking for is located there (i.e. it does not "overlay" PYTHONPATH directories in the same way that C++ include paths are).

Again, this can be done either way and the preference of one vs the other stems mostly from experience of having been burned by packaging concerns. It remains unclear whether we will use the flexibility provided by not deeply co-mingling the extensions with the Python code, but I would advocate for keeping it as separable as possible when the choice has no clear benefits either way.

mlir/lib/Bindings/Python/CMakeLists.txt

 include(AddMLIRPythonExtension)
 add_custom_target(MLIRBindingsPythonExtension)
 ################################################################################
 # Copy python source tree.
 ################################################################################
 
 set(PY_SRC_FILES
   mlir/__init__.py
-  mlir/ir.py
   mlir/dialects/__init__.py
+  mlir/ir.py
+  mlir/passmanager.py
+  mlir/transforms/__init__.py
 )
 
 add_custom_target(MLIRBindingsPythonSources ALL
   DEPENDS ${PY_SRC_FILES}
 )
 add_dependencies(MLIRBindingsPythonExtension MLIRBindingsPythonSources)
 
 add_dependencies(MLIRBindingsPythonExtension MLIRBindingsPythonIncGen)

 ################################################################################
 add_mlir_python_extension(MLIRCoreBindingsPythonExtension _mlir
   INSTALL_DIR
     python
   SOURCES
     MainModule.cpp
     IRModules.cpp
     PybindUtils.cpp
+    Pass.cpp
 )
 add_dependencies(MLIRBindingsPythonExtension MLIRCoreBindingsPythonExtension)
 
 # Note that we copy from the source tree just like for headers because
 # it will not be polluted with py_cache runtime artifacts (from testing and
 # such).
 install(
   DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/mlir
   DESTINATION python
   COMPONENT MLIRBindingsPythonSources
   FILES_MATCHING PATTERN "*.py"
 )
 
 if (NOT LLVM_ENABLE_IDE)
   add_llvm_install_targets(
     install-MLIRBindingsPythonSources
     DEPENDS MLIRBindingsPythonSources
     COMPONENT MLIRBindingsPythonSources)
 endif()
+
+add_subdirectory(Transforms)

mlir/lib/Bindings/Python/MainModule.cpp

 //===- MainModule.cpp - Main pybind module --------------------------------===//
 //
 // 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 <tuple>
 
 #include "PybindUtils.h"
 
 #include "Globals.h"
 #include "IRModules.h"
+#include "Pass.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 // -----------------------------------------------------------------------------
 // PyGlobals
 // -----------------------------------------------------------------------------
m.def(
               return opClass;
             });
       },
       "Class decorator for registering a custom Operation wrapper");
 
   // Define and populate IR submodule.
   auto irModule = m.def_submodule("ir", "MLIR IR Bindings");
   populateIRSubmodule(irModule);
+
+  // Define and populate PassManager submodule.
+  auto passModule =
+      m.def_submodule("passmanager", "MLIR Pass Management Bindings");
+  populatePassManagerSubmodule(passModule);
 }

mlir/lib/Bindings/Python/Pass.h

+//===- Pass.h - PassManager Submodules of pybind module -------------------===//
+//
+// 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_BINDINGS_PYTHON_PASS_H
+#define MLIR_BINDINGS_PYTHON_PASS_H
+
+#include "PybindUtils.h"
+
+namespace mlir {
+namespace python {
+
+void populatePassManagerSubmodule(pybind11::module &m);
+
+} // namespace python
+} // namespace mlir
+
+#endif // MLIR_BINDINGS_PYTHON_PASS_H
 No newline at end of file

mlir/lib/Bindings/Python/Pass.cpp

+//===- Pass.cpp - Pass Management -----------------------------------------===//
+//
+// 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 "Pass.h"
+
+#include "IRModules.h"
+#include "mlir-c/Pass.h"
+
+namespace py = pybind11;
+using namespace mlir;
+using namespace mlir::python;
+
+namespace {
+
+/// Owning Wrapper around a PassManager.
+class PyPassManager {
+public:
+  PyPassManager(MlirPassManager passManager) : passManager(passManager) {}
+  ~PyPassManager() { mlirPassManagerDestroy(passManager); }
+  MlirPassManager get() { return passManager; }
+
+private:
+  MlirPassManager passManager;
+};
+
+} // anonymous namespace
+
+/// Create the `mlir.passmanager` here.
+void mlir::python::populatePassManagerSubmodule(py::module &m) {
+  //----------------------------------------------------------------------------
+  // Mapping of the top-level PassManager
+  //----------------------------------------------------------------------------
+  py::class_<PyPassManager>(m, "PassManager")
+      .def(py::init<>([](DefaultingPyMlirContext context) {
+             MlirPassManager passManager =
+                 mlirPassManagerCreate(context->get());
+             return new PyPassManager(passManager);
+           }),
+           py::arg("context") = py::none(),
+           "Create a new PassManager for the current (or provided) Context.")
+      .def_static(
+          "parse",
+          [](const std::string pipeline, DefaultingPyMlirContext context) {
+            MlirPassManager passManager = mlirPassManagerCreate(context->get());
+            MlirLogicalResult status = mlirParsePassPipeline(
+                mlirPassManagerGetAsOpPassManager(passManager),
+                mlirStringRefCreate(pipeline.data(), pipeline.size()));
+            if (mlirLogicalResultIsFailure(status))
+              throw SetPyError(PyExc_ValueError,
+                               llvm::Twine("invalid pass pipeline '") +
+                                   pipeline + "'.");
+            return new PyPassManager(passManager);
+          },
+          py::arg("pipeline"), py::arg("context") = py::none(),
+          "Parse a textual pass-pipeline and return a top-level PassManager "
+          "that can be applied on a Module. Throw a ValueError if the pipeline "
+          "can't be parsed")
+      .def(
+          "__str__",
+          [](PyPassManager &self) {
+            MlirPassManager passManager = self.get();
+            PyPrintAccumulator printAccum;
+            mlirPrintPassPipeline(
+                mlirPassManagerGetAsOpPassManager(passManager),
+                printAccum.getCallback(), printAccum.getUserData());
+            return printAccum.join();
+          },
+          "Print the textual representation for this PassManager, suitable to "
+          "be passed to `parse` for round-tripping.");
+}

mlir/lib/Bindings/Python/Transforms/CMakeLists.txt

+################################################################################
+# Build python extension
+################################################################################
+
+add_mlir_python_extension(MLIRTransformsBindingsPythonExtension _mlirTransforms
+  INSTALL_DIR
+    python
+  SOURCES
+  Transforms.cpp
+)
 No newline at end of file

mlir/lib/Bindings/Python/Transforms/Transforms.cpp

+//===- Transforms.cpp - Pybind module for the Transforms library ----------===//
+//
+// 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/Transforms.h"
+
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+// -----------------------------------------------------------------------------
+// Module initialization.
+// -----------------------------------------------------------------------------
+
+PYBIND11_MODULE(_mlirTransforms, m) {
+  m.doc() = "MLIR Transforms library";
+
+  // Register all the passes in the Transforms library on load.
+  mlirRegisterTransformsPasses();
+}

mlir/lib/Bindings/Python/mlir/__init__.py

 #  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
 
 # Note that the only function of this module is currently to load the
 # native module and re-export its symbols. In the future, this file is
 # reserved as a trampoline to handle environment specific loading needs
 # and arbitrate any one-time initialization needed in various shared-library
 # scenarios.
 
 __all__ = [
   "ir",
+  "passmanager",
 ]
 
 # Expose the corresponding C-Extension module with a well-known name at this
 # top-level module. This allows relative imports like the following to
 # function:
 #   from .. import _cext
 # This reduces coupling, allowing embedding of the python sources into another
 # project that can just vary based on this top-level loader module.
def _reexport_cext(cext_module_name, target_module_name):
   source_module = getattr(_cext, cext_module_name)
   for attr_name in dir(source_module):
     if not attr_name.startswith("__"):
       setattr(target_module, attr_name, getattr(source_module, attr_name))
 
 
 # Import sub-modules. Since these may import from here, this must come after
 # any exported definitions.
-from . import ir
+from . import ir, passmanager
 
 # Add our 'dialects' parent module to the search path for implementations.
 _cext.globals.append_dialect_search_prefix("mlir.dialects")

mlir/lib/Bindings/Python/mlir/passmanager.py

+#  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
+
+# Simply a wrapper around the extension module of the same name.
+from . import _reexport_cext
+_reexport_cext("passmanager", __name__)
+del _reexport_cext

mlir/lib/Bindings/Python/mlir/transforms/__init__.py

+#  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
+
+# Expose the corresponding C-Extension module with a well-known name at this
+# level.
+import _mlirTransforms as _cextTransforms
+

mlir/test/Bindings/Python/pass_manager.py

+# RUN: %PYTHON %s 2>&1 | FileCheck %s
+
+import gc, sys
+from mlir.ir import *
+from mlir.passmanager import *
+
+# Log everything to stderr and flush so that we have a unified stream to match
+# errors emitted by MLIR to stderr. TODO: this shouldn't be needed when
+# everything is plumbed.
+def log(*args):
+  print(*args, file=sys.stderr)
+  sys.stderr.flush()
+
+def run(f):
+  log("\nTEST:", f.__name__)
+  f()
+  gc.collect()
+  assert Context._get_live_count() == 0
+
+# Verify successful round-trip.
+# CHECK-LABEL: TEST: testParseSuccess
+def testParseSuccess():
+  with Context():
+    # A first import is expected to fail because the pass isn't registered
+    # until we import mlir.transforms
+    try:
+      pm = PassManager.parse("module(func(print-op-stats))")
+      # TODO: this error should be propagate to Python but the C API does not help right now.
+      # CHECK: error: 'print-op-stats' does not refer to a registered pass or pass pipeline
+    except ValueError as e:
+      # CHECK: ValueError exception: invalid pass pipeline 'module(func(print-op-stats))'.
+      log("ValueError exception:", e)
+    else:
+      log("Exception not produced")
+
+    # This will register the pass and round-trip should be possible now.
+    import mlir.transforms
+    pm = PassManager.parse("module(func(print-op-stats))")
+    # CHECK: Roundtrip: module(func(print-op-stats))
+    log("Roundtrip: ", pm)
+run(testParseSuccess)
+
+# Verify failure on unregistered pass.
+# CHECK-LABEL: TEST: testParseFail
+def testParseFail():
+  with Context():
+    try:
+      pm = PassManager.parse("unknown-pass")
+    except ValueError as e:
+      # CHECK: ValueError exception: invalid pass pipeline 'unknown-pass'.
+      log("ValueError exception:", e)
+    else:
+      log("Exception not produced")
+run(testParseFail)

mlir/test/CMakeLists.txt

if(MLIR_VULKAN_RUNNER_ENABLED)
   list(APPEND MLIR_TEST_DEPENDS
     mlir-vulkan-runner
   )
 endif()
 
 if(MLIR_BINDINGS_PYTHON_ENABLED)
   list(APPEND MLIR_TEST_DEPENDS
     MLIRBindingsPythonExtension
+    MLIRTransformsBindingsPythonExtension
   )
 endif()
 
 add_lit_testsuite(check-mlir "Running the MLIR regression tests"
   ${CMAKE_CURRENT_BINARY_DIR}
   DEPENDS ${MLIR_TEST_DEPENDS}
   )
 set_target_properties(check-mlir PROPERTIES FOLDER "Tests")
 
 add_lit_testsuites(MLIR ${CMAKE_CURRENT_SOURCE_DIR}
   DEPENDS ${MLIR_TEST_DEPENDS}
 )

mlir/tools/mlir-tblgen/PassCAPIGen.cpp

 #ifdef __cplusplus
 }
 #endif
 )";
 
 /// Emit TODO
 static bool emitCAPIHeader(const llvm::RecordKeeper &records, raw_ostream &os) {
   os << fileHeader;
+  os << "// Registration for the entire group\n";
+  os << "MLIR_CAPI_EXPORTED void mlirRegister" << groupName << "Passes();\n\n";
   for (const auto *def : records.getAllDerivedDefinitions("PassBase")) {
     Pass pass(def);
     StringRef defName = pass.getDef()->getName();
     os << llvm::formatv(passDecl, groupName, defName);
   }
   os << fileFooter;
   return false;
 }
 
 const char *const passCreateDef = R"(
 MlirPass mlirCreate{0}{1}() {
   return wrap({2}.release());
 }
 void mlirRegister{0}{1}() {
   register{1}Pass();
 }
 
 )";
 
+/// {0}: The name of the pass group.
+const char *const passGroupRegistrationCode = R"(
+//===----------------------------------------------------------------------===//
+// {0} Group Registration
+//===----------------------------------------------------------------------===//
+
+void mlirRegister{0}Passes() {{
+  register{0}Passes();
+}
+)";
+
 static bool emitCAPIImpl(const llvm::RecordKeeper &records, raw_ostream &os) {
   os << "/* Autogenerated by mlir-tblgen; don't manually edit. */";
+  os << llvm::formatv(passGroupRegistrationCode, groupName);
+
   for (const auto *def : records.getAllDerivedDefinitions("PassBase")) {
     Pass pass(def);
     StringRef defName = pass.getDef()->getName();
     os << llvm::formatv(passCreateDef, groupName, defName,
                         pass.getConstructor());
   }
   return false;
 }
 
 static mlir::GenRegistration genCAPIHeader("gen-pass-capi-header",
                                            "Generate pass C API header",
                                            &emitCAPIHeader);
 
 static mlir::GenRegistration genCAPIImpl("gen-pass-capi-impl",
                                          "Generate pass C API implementation",
                                          &emitCAPIImpl);