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

[aartbik]

Summary

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

[mehdi_amini]

Can we have "register" somewhere in the name? This submodule won't do anything else right?
I would even consider making it a method instead:

import mlir
mlir.registerAllPasses()

[aartbik]

In D103130#2780990, @mehdi_amini wrote:

Can we have "register" somewhere in the name?

I am following the other registration mechanism to the letter here, not one of them uses that technique or have "register" in the name.
I don't strongly feel either way, but what I did seemed much more consistent with the style used throughout.

PYBIND11_MODULE(_mlirLinalgPasses, m) {
  m.doc() = "MLIR Linalg Dialect Passes";
  // Register all Linalg passes on load.
  mlirRegisterLinalgPasses();
}
....
PYBIND11_MODULE(_mlirAsyncPasses, m) {
  m.doc() = "MLIR Async Dialect Passes";
  // Register all Async passes on load.
  mlirRegisterAsyncPasses();
}
....
etc.

[mehdi_amini]

The others may expose other content than just registration, in particular if/when we develop programmatic API for individual passes. This import is a) top-level which makes it more prominent and b) does not seem like it'll have any other purpose than registration.

[aartbik]

In D103130#2781071, @mehdi_amini wrote:

The others may expose other content than just registration.

The other pass oriented imports don't expose anything but passes either.
For example we have

    import mlir.dialects.linalg.passes -> _load_extension("_mlirLinalgPasses") -> mlirRegisterLinalgPasses(); 
    import mlir.dialects.async         -> _load_extension("_mlirAsyncPasses")  -> mlirRegisterAsyncPasses();
    import mlir.dialects.gpu           -> _load_extension("_mlirGPUPasses")    -> mlirRegisterGPUPasses();
    import mlir.transforms             -> _load_extension("_mlirTransforms")   -> mlirRegisterTransformsPasses();
    import mlir.conversion             -> _load_extension("_mlirConversions")  -> mlirRegisterConversionPasses()

and I simply added

  import mlir.all_passes     ->  _load_extension("_mlirAllPasses") -> mlirRegisterAllPasses();

to this repertoire.This way a developer can simply start with "import mlir.all_passes", but then replace that line with a more restricted import when it turns out not everything is needed.

[mehdi_amini]

In D103130#2782654, @aartbik wrote:

In D103130#2781071, @mehdi_amini wrote:

The others may expose other content than just registration.

The other pass oriented imports don't expose anything but passes either.

"don't expose anything" yet. That's the main difference: the others are design so that they can expose pass classes for example. This one won't.

[aartbik]

Would:

import mlir.all_passes_registration

be acceptable to you (I believe using a noun is preferred over a verb here)?

[mehdi_amini]

I'd prefer a method call instead of a package, but I can live with that.

[aartbik]

add "registration" to the import name

[aartbik]

In D103130#2782734, @mehdi_amini wrote:

I'd prefer a method call instead of a package, but I can live with that.

Thanks Mehdi. PTAL

[aartbik]

I hope you simply clicked the wrong button, but otherwise what changes would you like to see, my favorite sparring partner? ;-)

[mehdi_amini]

Indeed :)

[jpienaar]

Is https://lab.llvm.org/buildbot/#/builders/61/builds/10308/steps/5/logs/stdio related?

[jpienaar]

(build was broken for a few days)

[aartbik]

In D103130#2785440, @jpienaar wrote:

(build was broken for a few days)

Oops. Just saw this. I did not get any notification on that (but after chat I realize failure was masked by other failures)

mlir/include/mlir-c/Registration.h

 /// Registers all dialects known to core MLIR with the provided Context.
 /// This is needed before creating IR for these Dialects.
 /// TODO: Remove this function once the real registration API is finished.
 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
 
 #endif // MLIR_C_REGISTRATION_H

mlir/lib/Bindings/Python/AllPassesRegistration.cpp

+//===- 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();
+}

mlir/lib/Bindings/Python/CMakeLists.txt

SOURCES
     Pass.cpp
     ExecutionEngine.cpp
 )
 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
   SOURCES
     AsyncPasses.cpp
 )
 add_dependencies(MLIRBindingsPythonExtension MLIRAsyncPassesBindingsPythonExtension)
 
 add_mlir_python_extension(MLIRSparseTensorPassesBindingsPythonExtension _mlirSparseTensorPasses
   INSTALL_DIR
     python
   SOURCES
     SparseTensorPasses.cpp
 )
 add_dependencies(MLIRBindingsPythonExtension MLIRSparseTensorPassesBindingsPythonExtension)
 
 add_mlir_python_extension(MLIRGPUPassesBindingsPythonExtension _mlirGPUPasses
   INSTALL_DIR
     python
   SOURCES
     GPUPasses.cpp

mlir/lib/CAPI/Registration/Registration.cpp

 //===- Registration.cpp - C Interface for MLIR Registration ---------------===//
 //
 // 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 "mlir/CAPI/IR.h"
 #include "mlir/InitAllDialects.h"
+#include "mlir/InitAllPasses.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 
 void mlirRegisterAllDialects(MlirContext context) {
   mlir::registerAllDialects(*unwrap(context));
   // TODO: we may not want to eagerly load here.
   unwrap(context)->loadAllAvailableDialects();
 }
 
 void mlirRegisterAllLLVMTranslations(MlirContext context) {
   mlir::registerLLVMDialectTranslation(*unwrap(context));
 }
+
+void mlirRegisterAllPasses() { mlir::registerAllPasses(); }

mlir/python/mlir/all_passes_registration/__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
+
+from .._cext_loader import _load_extension
+
+_cextAllPasses = _load_extension("_mlirAllPassesRegistration")
+del _load_extension