[mlir] Overhaul C/Python registration APIs to properly scope registration/loading activities.

Since the very first commits, the Python and C MLIR APIs have had mis-placed registration/load functionality for dialects, extensions, etc. This was done pragmatically in order to get bootstrapped and then just grew in. Downstreams largely bypass and do their own thing by providing various APIs to register things they need. Meanwhile, the C++ APIs have stabilized around this and it would make sense to follow suit.

The thing we have observed in canonical usage by downstreams is that each downstream tends to have native entry points that configure its installation to its preferences with one-stop APIs. This patch leans in to this approach with `RegisterEverything.h` and `mlir._mlir_libs._mlirRegisterEverything` being the one-stop entry points for the "upstream packages". The `_mlir_libs.__init__.py` now allows customization of the environment and Context by adding "initialization modules" to the `_mlir_libs` package. If present, `_mlirRegisterEverything` is treated as such a module. Others can be added by downstreams by adding a `_site_initialize_{i}.py` module, where '{i}' is a number starting with zero. The number will be incremented and corresponding module loaded until one is not found. Initialization modules can:

* Perform load time customization to the global environment (i.e. registering passes, hooks, etc).
* Define a `register_dialects(registry: DialectRegistry)` function that can extend the `DialectRegistry` that will be used to bootstrap the `Context`.
* Define a `context_init_hook(context: Context)` function that will be added to a list of callbacks which will be invoked after dialect registration during `Context` initialization.

Note that the `MLIRPythonExtension.RegisterEverything` is not included by default when building a downstream (its corresponding behavior was prior). For downstreams which need the default MLIR initialization to take place, they must add this back in to their Python CMake build just like they add their own components (i.e. to `add_mlir_python_common_capi_library` and `add_mlir_python_modules`). It is perfectly valid to not do this, in which case, only the things explicitly depended on and initialized by downstreams will be built/packaged. If the downstream has not been set up for this, it is recommended to simply add this back for the time being and pay the build time/package size cost.

CMake changes:
* `MLIRCAPIRegistration` -> `MLIRCAPIRegisterEverything` (renamed to signify what it does and force an evaluation: a number of places were incidentally linking this very expensive target)
* `MLIRPythonSoure.Passes` removed (without replacement: just drop)
* `MLIRPythonExtension.AllPassesRegistration` removed (without replacement: just drop)
* `MLIRPythonExtension.Conversions` removed (without replacement: just drop)
* `MLIRPythonExtension.Transforms` removed (without replacement: just drop)

Header changes:
* `mlir-c/Registration.h` is deleted. Dialect registration functionality is now in `IR.h`. Registration of upstream features are in `mlir-c/RegisterEverything.h`. When updating MLIR and a couple of downstreams, I found that proper usage was commingled so required making a choice vs just blind S&R.

Python APIs removed:
  * mlir.transforms and mlir.conversions (previously only had an __init__.py which indirectly triggered `mlirRegisterTransformsPasses()` and `mlirRegisterConversionPasses()` respectively). Downstream impact: Remove these imports if present (they now happen as part of default initialization).
  * mlir._mlir_libs._all_passes_registration, mlir._mlir_libs._mlirTransforms, mlir._mlir_libs._mlirConversions. Downstream impact: None expected (these were internally used).

C-APIs changed:
  * mlirRegisterAllDialects(MlirContext) now takes an MlirDialectRegistry instead. It also used to trigger loading of all dialects, which was already marked with a TODO to remove -- it no longer does, and for direct use, dialects must be explicitly loaded. Downstream impact: Direct C-API users must ensure that needed dialects are loaded or call `mlirContextLoadAllAvailableDialects(MlirContext)` to emulate the prior behavior. Also see the `ir.c` test case (e.g. `  mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("func"));`).
  * mlirDialectHandle* APIs were moved from Registration.h (which now is restricted to just global/upstream registration) to IR.h, arguably where it should have been. Downstream impact: include correct header (likely already doing so).

C-APIs added:
  * mlirContextLoadAllAvailableDialects(MlirContext): Corresponds to C++ API with the same purpose.

Python APIs added:
  * mlir.ir.DialectRegistry: Mapping for an MlirDialectRegistry.
  * mlir.ir.Context.append_dialect_registry(MlirDialectRegistry)
  * mlir.ir.Context.load_all_available_dialects()
  * mlir._mlir_libs._mlirAllRegistration: New native extension that exposes a `register_dialects(MlirDialectRegistry)` entry point and performs all upstream pass/conversion/transforms registration on init. In this first step, we eagerly load this as part of the __init__.py and use it to monkey patch the Context to emulate prior behavior.
  * Type caster and capsule support for MlirDialectRegistry

This should make it possible to build downstream Python dialects that only depend on a subset of MLIR. See: https://github.com/llvm/llvm-project/issues/56037

Here is an example PR, minimally adapting IREE to these changes: https://github.com/iree-org/iree/pull/9638/files In this situation, IREE is opting to not link everything, since it is already configuring the Context to its liking. For projects that would just like to not think about it and pull in everything, add `MLIRPythonExtension.RegisterEverything` to the list of Python sources getting built, and the old behavior will continue.

Reviewed By: mehdi_amini, ftynse

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

6 files changed, 102 insertions, 74 deletions

diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index dc831c9..1fbbadf 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -37,17 +37,8 @@ declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
     runtime/*.py
 )
 
-declare_mlir_python_sources(MLIRPythonSources.Passes
-  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
-  ADD_TO_PARENT MLIRPythonSources
-  SOURCES_GLOB
-    all_passes_registration/*.py
-    conversions/*.py
-    transforms/*.py
-)
-
 declare_mlir_python_sources(MLIRPythonCAPI.HeaderSources
-  ROOT_DIR "${MLIR_MAIN_INCLUDE_DIR}"
+  ROOT_DIR "${MLIR_SOURCE_DIR}/include"
   SOURCES_GLOB "mlir-c/*.h"
 )
 
@@ -283,12 +274,31 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
     MLIRCAPIDebug
     MLIRCAPIIR
     MLIRCAPIInterfaces
-    MLIRCAPIRegistration  # TODO: See about dis-aggregating
 
     # Dialects
     MLIRCAPIFunc
 )
 
+# This extension exposes an API to register all dialects, extensions, and passes
+# packaged in upstream MLIR and it is used for the upstream "mlir" Python
+# package. Downstreams will likely want to provide their own and not depend
+# on this one, since it links in the world.
+# Note that this is not added to any top-level source target for transitive
+# inclusion: It must be included explicitly by downstreams if desired. Note that
+# this has a very large impact on what gets built/packaged.
+declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
+  MODULE_NAME _mlirRegisterEverything
+  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  SOURCES
+    RegisterEverything.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIConversion
+    MLIRCAPITransforms
+    MLIRCAPIRegisterEverything
+)
+
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
   MODULE_NAME _mlirDialectsLinalg
   ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
@@ -341,18 +351,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
     MLIRCAPISparseTensor
 )
 
-declare_mlir_python_extension(MLIRPythonExtension.AllPassesRegistration
-  MODULE_NAME _mlirAllPassesRegistration
-  ROOT_DIR "${PYTHON_SOURCE_DIR}"
-  SOURCES
-    AllPassesRegistration.cpp
-  PRIVATE_LINK_LIBS
-    LLVMSupport
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPIConversion
-    MLIRCAPITransforms
-)
-
 declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
   MODULE_NAME _mlirAsyncPasses
   ADD_TO_PARENT MLIRPythonSources.Dialects.async_dialect
@@ -365,18 +363,6 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
     MLIRCAPIAsync
 )
 
-declare_mlir_python_extension(MLIRPythonExtension.Conversions
-  MODULE_NAME _mlirConversions
-  ADD_TO_PARENT MLIRPythonSources.Passes
-  ROOT_DIR "${PYTHON_SOURCE_DIR}"
-  SOURCES
-    Conversions/Conversions.cpp
-  PRIVATE_LINK_LIBS
-    LLVMSupport
-  EMBED_CAPI_LINK_LIBS
-  MLIRCAPIConversion
-)
-
 # Only enable the ExecutionEngine if the native target is configured in.
 if(TARGET ${LLVM_NATIVE_ARCH})
   declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
@@ -428,18 +414,6 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
     MLIRCAPISparseTensor
 )
 
-declare_mlir_python_extension(MLIRPythonExtension.Transforms
-  MODULE_NAME _mlirTransforms
-  ADD_TO_PARENT MLIRPythonSources.Passes
-  ROOT_DIR "${PYTHON_SOURCE_DIR}"
-  SOURCES
-    Transforms/Transforms.cpp
-  PRIVATE_LINK_LIBS
-    LLVMSupport
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPITransforms
-)
-
 # TODO: Figure out how to put this in the test tree.
 # This should not be included in the main Python extension. However,
 # putting it into MLIRPythonTestSources along with the dialect declaration
@@ -506,7 +480,7 @@ add_mlir_python_common_capi_library(MLIRPythonCAPI
     MLIRPythonCAPI.HeaderSources
   DECLARED_SOURCES
     MLIRPythonSources
-    MLIRPythonExtension.AllPassesRegistration
+    MLIRPythonExtension.RegisterEverything
     ${_ADDL_TEST_SOURCES}
 )
 
@@ -520,7 +494,7 @@ add_mlir_python_modules(MLIRPythonModules
   INSTALL_PREFIX "python_packages/mlir_core/mlir"
   DECLARED_SOURCES
     MLIRPythonSources
-    MLIRPythonExtension.AllPassesRegistration
+    MLIRPythonExtension.RegisterEverything
     ${_ADDL_TEST_SOURCES}
   COMMON_CAPI_LINK_LIBS
     MLIRPythonCAPI
diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py
index 23bc502..add8d92 100644
--- a/mlir/python/mlir/_mlir_libs/__init__.py
+++ b/mlir/python/mlir/_mlir_libs/__init__.py
@@ -9,12 +9,6 @@ import os
 _this_dir = os.path.dirname(__file__)
 
 
-# These submodules have no type stubs and are thus opaque to the type checker.
-_mlirConversions: Any
-_mlirTransforms: Any
-_mlirAllPassesRegistration: Any
-
-
 def get_lib_dirs() -> Sequence[str]:
   """Gets the lib directory for linking to shared libraries.
 
@@ -31,3 +25,77 @@ def get_include_dirs() -> Sequence[str]:
   not be present.
   """
   return [os.path.join(_this_dir, "include")]
+
+
+# Perform Python level site initialization. This involves:
+#   1. Attempting to load initializer modules, specific to the distribution.
+#   2. Defining the concrete mlir.ir.Context that does site specific 
+#      initialization.
+#
+# Aside from just being far more convenient to do this at the Python level,
+# it is actually quite hard/impossible to have such __init__ hooks, given
+# the pybind memory model (i.e. there is not a Python reference to the object
+# in the scope of the base class __init__).
+#
+# For #1, we:
+#   a. Probe for modules named '_mlirRegisterEverything' and 
+#     '_site_initialize_{i}', where 'i' is a number starting at zero and 
+#     proceeding so long as a module with the name is found.
+#   b. If the module has a 'register_dialects' attribute, it will be called
+#     immediately with a DialectRegistry to populate.
+#   c. If the module has a 'context_init_hook', it will be added to a list
+#     of callbacks that are invoked as the last step of Context 
+#     initialization (and passed the Context under construction).
+#
+# This facility allows downstreams to customize Context creation to their
+# needs.
+def _site_initialize():
+  import importlib
+  import itertools
+  import logging
+  from ._mlir import ir
+  registry = ir.DialectRegistry()
+  post_init_hooks = []
+
+  def process_initializer_module(module_name):
+    try:
+      m = importlib.import_module(f".{module_name}", __name__)
+    except ModuleNotFoundError:
+      return False
+
+    logging.debug("Initializing MLIR with module: %s", module_name)
+    if hasattr(m, "register_dialects"):
+      logging.debug("Registering dialects from initializer %r", m)
+      m.register_dialects(registry)
+    if hasattr(m, "context_init_hook"):
+      logging.debug("Adding context init hook from %r", m)
+      post_init_hooks.append(m.context_init_hook)
+    return True
+
+
+  # If _mlirRegisterEverything is built, then include it as an initializer
+  # module.
+  process_initializer_module("_mlirRegisterEverything")
+
+  # Load all _site_initialize_{i} modules, where 'i' is a number starting
+  # at 0.
+  for i in itertools.count():
+    module_name = f"_site_initialize_{i}"
+    if not process_initializer_module(module_name):
+      break
+
+  class Context(ir._BaseContext):
+    def __init__(self, *args, **kwargs):
+      super().__init__(*args, **kwargs)
+      self.append_dialect_registry(registry)
+      for hook in post_init_hooks:
+        hook(self)
+      # TODO: There is some debate about whether we should eagerly load
+      # all dialects. It is being done here in order to preserve existing
+      # behavior. See: https://github.com/llvm/llvm-project/issues/56037
+      self.load_all_available_dialects()
+
+  ir.Context = Context
+
+
+_site_initialize()
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
index cd7eb0a..60bc367 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
@@ -479,6 +479,11 @@ class Context:
     def d(self) -> Dialects: ...
     @property
     def dialects(self) -> Dialects: ...
+    def append_dialect_registry(self, registry: "DialectRegistry") -> None: ...
+    def load_all_available_dialects(self) -> None: ...
+
+class DialectRegistry:
+    def __init__(self) -> None: ...
 
 # TODO: Auto-generated. Audit and fix.
 class DenseElementsAttr(Attribute):
diff --git a/mlir/python/mlir/all_passes_registration/__init__.py b/mlir/python/mlir/all_passes_registration/__init__.py
deleted file mode 100644
index aca557a..0000000
--- a/mlir/python/mlir/all_passes_registration/__init__.py
+++ /dev/null
@@ -1,5 +0,0 @@
-#  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 .._mlir_libs import _mlirAllPassesRegistration as _cextAllPasses
diff --git a/mlir/python/mlir/conversions/__init__.py b/mlir/python/mlir/conversions/__init__.py
deleted file mode 100644
index a6a9eb8..0000000
--- a/mlir/python/mlir/conversions/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-#  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.
-from .._mlir_libs import _mlirConversions as _cextConversions
diff --git a/mlir/python/mlir/transforms/__init__.py b/mlir/python/mlir/transforms/__init__.py
deleted file mode 100644
index 71ea17d..0000000
--- a/mlir/python/mlir/transforms/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-#  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.
-from .._mlir_libs import _mlirTransforms as _cextTransforms