[mlir] Overhaul C/Python registration APIs to properly scope registration/loading...
authorStella Laurenzo <stellaraccident@gmail.com>
Sat, 16 Jul 2022 23:09:03 +0000 (16:09 -0700)
committerStella Laurenzo <stellaraccident@gmail.com>
Sun, 17 Jul 2022 00:27:50 +0000 (17:27 -0700)
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

52 files changed:
mlir/examples/standalone/include/Standalone-c/Dialects.h
mlir/examples/standalone/python/CMakeLists.txt
mlir/examples/standalone/test/CAPI/CMakeLists.txt
mlir/examples/standalone/test/CAPI/standalone-capi-test.c
mlir/include/mlir-c/Bindings/Python/Interop.h
mlir/include/mlir-c/Dialect/Async.h
mlir/include/mlir-c/Dialect/ControlFlow.h
mlir/include/mlir-c/Dialect/Func.h
mlir/include/mlir-c/Dialect/GPU.h
mlir/include/mlir-c/Dialect/LLVM.h
mlir/include/mlir-c/Dialect/Linalg.h
mlir/include/mlir-c/Dialect/PDL.h
mlir/include/mlir-c/Dialect/Quant.h
mlir/include/mlir-c/Dialect/SCF.h
mlir/include/mlir-c/Dialect/Shape.h
mlir/include/mlir-c/Dialect/SparseTensor.h
mlir/include/mlir-c/Dialect/Tensor.h
mlir/include/mlir-c/IR.h
mlir/include/mlir-c/Pass.h
mlir/include/mlir-c/RegisterEverything.h [new file with mode: 0644]
mlir/include/mlir-c/Registration.h [deleted file]
mlir/include/mlir/Bindings/Python/PybindAdaptors.h
mlir/include/mlir/CAPI/Registration.h
mlir/lib/Bindings/Python/AllPassesRegistration.cpp [deleted file]
mlir/lib/Bindings/Python/Conversions/Conversions.cpp [deleted file]
mlir/lib/Bindings/Python/IRCore.cpp
mlir/lib/Bindings/Python/IRModule.h
mlir/lib/Bindings/Python/RegisterEverything.cpp [new file with mode: 0644]
mlir/lib/Bindings/Python/Transforms/Transforms.cpp [deleted file]
mlir/lib/CAPI/CMakeLists.txt
mlir/lib/CAPI/Conversion/CMakeLists.txt
mlir/lib/CAPI/IR/IR.cpp
mlir/lib/CAPI/RegisterEverything/CMakeLists.txt [moved from mlir/lib/CAPI/Registration/CMakeLists.txt with 76% similarity]
mlir/lib/CAPI/RegisterEverything/RegisterEverything.cpp [moved from mlir/lib/CAPI/Registration/Registration.cpp with 69% similarity]
mlir/python/CMakeLists.txt
mlir/python/mlir/_mlir_libs/__init__.py
mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
mlir/python/mlir/all_passes_registration/__init__.py [deleted file]
mlir/python/mlir/conversions/__init__.py [deleted file]
mlir/python/mlir/transforms/__init__.py [deleted file]
mlir/test/CAPI/CMakeLists.txt
mlir/test/CAPI/execution_engine.c
mlir/test/CAPI/ir.c
mlir/test/CAPI/pass.c
mlir/test/CAPI/sparse_tensor.c
mlir/test/python/execution_engine.py
mlir/test/python/integration/dialects/linalg/opsrun.py
mlir/test/python/lib/CMakeLists.txt
mlir/test/python/lib/PythonTestCAPI.h
mlir/test/python/pass_manager.py
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel

index 6b98611..b3e4775 100644 (file)
@@ -9,7 +9,7 @@
 #ifndef STANDALONE_C_DIALECTS_H
 #define STANDALONE_C_DIALECTS_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
index 014d606..a8c4382 100644 (file)
@@ -40,6 +40,9 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_SOURCES
     StandalonePythonSources
+    # TODO: Remove this in favor of showing fine grained registration once
+    # available.
+    MLIRPythonExtension.RegisterEverything
     MLIRPythonSources.Core
 )
 
@@ -52,6 +55,9 @@ add_mlir_python_modules(StandalonePythonModules
   INSTALL_PREFIX "python_packages/standalone/mlir_standalone"
   DECLARED_SOURCES
     StandalonePythonSources
+    # TODO: Remove this in favor of showing fine grained registration once
+    # available.
+    MLIRPythonExtension.RegisterEverything
     MLIRPythonSources
   COMMON_CAPI_LINK_LIBS
     StandalonePythonCAPI
index 8958431..eaa6cfc 100644 (file)
@@ -6,7 +6,9 @@ add_mlir_aggregate(StandaloneCAPITestLib
   SHARED
   EMBED_LIBS
   MLIRCAPIIR
-  MLIRCAPIRegistration
+  # TODO: Remove this in favor of showing fine grained dialect registration
+  # (once available).
+  MLIRCAPIRegisterEverything
   StandaloneCAPI
 )
 
index 6478864..54f3ca7 100644 (file)
 
 #include <stdio.h>
 
-#include "mlir-c/IR.h"
 #include "Standalone-c/Dialects.h"
+#include "mlir-c/IR.h"
+#include "mlir-c/RegisterEverything.h"
+
+static void registerAllUpstreamDialects(MlirContext ctx) {
+  MlirDialectRegistry registry = mlirDialectRegistryCreate();
+  mlirRegisterAllDialects(registry);
+  mlirContextAppendDialectRegistry(ctx, registry);
+  mlirDialectRegistryDestroy(registry);
+}
 
 int main(int argc, char **argv) {
   MlirContext ctx = mlirContextCreate();
   // TODO: Create the dialect handles for the builtin dialects and avoid this.
   // This adds dozens of MB of binary size over just the standalone dialect.
-  mlirRegisterAllDialects(ctx);
+  registerAllUpstreamDialects(ctx);
   mlirDialectHandleRegisterDialect(mlirGetDialectHandle__standalone__(), ctx);
 
   MlirModule module = mlirModuleCreateParse(
index 86b57cd..b877f94 100644 (file)
@@ -64,6 +64,8 @@
   MAKE_MLIR_PYTHON_QUALNAME("ir.Attribute._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_CONTEXT                                            \
   MAKE_MLIR_PYTHON_QUALNAME("ir.Context._CAPIPtr")
+#define MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY                                   \
+  MAKE_MLIR_PYTHON_QUALNAME("ir.DialectRegistry._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_EXECUTION_ENGINE                                   \
   MAKE_MLIR_PYTHON_QUALNAME("execution_engine.ExecutionEngine._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_INTEGER_SET                                        \
@@ -172,6 +174,28 @@ static inline MlirContext mlirPythonCapsuleToContext(PyObject *capsule) {
   return context;
 }
 
+/** Creates a capsule object encapsulating the raw C-API MlirDialectRegistry.
+ * The returned capsule does not extend or affect ownership of any Python
+ * objects that reference the context in any way.
+ */
+static inline PyObject *
+mlirPythonDialectRegistryToCapsule(MlirDialectRegistry registry) {
+  return PyCapsule_New(registry.ptr, MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY,
+                       NULL);
+}
+
+/** Extracts an MlirDialectRegistry from a capsule as produced from
+ * mlirPythonDialectRegistryToCapsule. If the capsule is not of the right type,
+ * then a null context is returned (as checked via mlirContextIsNull). In such a
+ * case, the Python APIs will have already set an error. */
+static inline MlirDialectRegistry
+mlirPythonCapsuleToDialectRegistry(PyObject *capsule) {
+  void *ptr =
+      PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY);
+  MlirDialectRegistry registry = {ptr};
+  return registry;
+}
+
 /** Creates a capsule object encapsulating the raw C-API MlirLocation.
  * The returned capsule does not extend or affect ownership of any Python
  * objects that reference the location in any way. */
index 50b6413..e4e32f8 100644 (file)
@@ -10,7 +10,7 @@
 #ifndef MLIR_C_DIALECT_ASYNC_H
 #define MLIR_C_DIALECT_ASYNC_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 #include "mlir-c/Support.h"
 
 #ifdef __cplusplus
index 1ca7054..6d5ff8c 100644 (file)
@@ -10,7 +10,7 @@
 #ifndef MLIR_C_DIALECT_CONTROLFLOW_H
 #define MLIR_C_DIALECT_CONTROLFLOW_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
index 4bdac42..eeb6dfe 100644 (file)
@@ -18,7 +18,7 @@
 #ifndef MLIR_C_DIALECT_FUNC_H
 #define MLIR_C_DIALECT_FUNC_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
index cf4e899..1a18d82 100644 (file)
@@ -10,7 +10,7 @@
 #ifndef MLIR_C_DIALECT_GPU_H
 #define MLIR_C_DIALECT_GPU_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 #include "mlir-c/Support.h"
 
 #ifdef __cplusplus
index 2cf73a3..ba98c33 100644 (file)
@@ -11,7 +11,6 @@
 #define MLIR_C_DIALECT_LLVM_H
 
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
 
 #ifdef __cplusplus
 extern "C" {
index 2fe1872..0ab201e 100644 (file)
@@ -10,7 +10,7 @@
 #ifndef MLIR_C_DIALECT_LINALG_H
 #define MLIR_C_DIALECT_LINALG_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 #include "mlir-c/Support.h"
 
 #ifdef __cplusplus
index 8bd7976..6ad2e2d 100644 (file)
@@ -11,7 +11,6 @@
 #define MLIR_C_DIALECT_PDL_H
 
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
 
 #ifdef __cplusplus
 extern "C" {
index eb529b8..39a1731 100644 (file)
@@ -11,7 +11,6 @@
 #define MLIR_C_DIALECT_QUANT_H
 
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
 
 #ifdef __cplusplus
 extern "C" {
index c1b2567..75f1b68 100644 (file)
@@ -10,7 +10,7 @@
 #ifndef MLIR_C_DIALECT_SCF_H
 #define MLIR_C_DIALECT_SCF_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
index f64da80..3fe3ddf 100644 (file)
@@ -10,7 +10,7 @@
 #ifndef MLIR_C_DIALECT_SHAPE_H
 #define MLIR_C_DIALECT_SHAPE_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
index 68d72b9..252ec68 100644 (file)
@@ -11,7 +11,7 @@
 #define MLIR_C_DIALECT_SPARSETENSOR_H
 
 #include "mlir-c/AffineMap.h"
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
index f749782..74cbc5a 100644 (file)
@@ -10,7 +10,7 @@
 #ifndef MLIR_C_DIALECT_TENSOR_H
 #define MLIR_C_DIALECT_TENSOR_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
index edd0e44..2d38700 100644 (file)
@@ -130,6 +130,11 @@ MLIR_CAPI_EXPORTED MlirDialect mlirContextGetOrLoadDialect(MlirContext context,
 MLIR_CAPI_EXPORTED void mlirContextEnableMultithreading(MlirContext context,
                                                         bool enable);
 
+/// Eagerly loads all available dialects registered with a context, making
+/// them available for use for IR construction.
+MLIR_CAPI_EXPORTED void
+mlirContextLoadAllAvailableDialects(MlirContext context);
+
 /// Returns whether the given fully-qualified operation (i.e.
 /// 'dialect.operation') is registered with the context. This will return true
 /// if the dialect is loaded and the operation is registered within the
@@ -158,6 +163,47 @@ MLIR_CAPI_EXPORTED bool mlirDialectEqual(MlirDialect dialect1,
 MLIR_CAPI_EXPORTED MlirStringRef mlirDialectGetNamespace(MlirDialect dialect);
 
 //===----------------------------------------------------------------------===//
+// DialectHandle API.
+// Registration entry-points for each dialect are declared using the common
+// MLIR_DECLARE_DIALECT_REGISTRATION_CAPI macro, which takes the dialect
+// API name (i.e. "Func", "Tensor", "Linalg") and namespace (i.e. "func",
+// "tensor", "linalg"). The following declarations are produced:
+//
+//   /// Gets the above hook methods in struct form for a dialect by namespace.
+//   /// This is intended to facilitate dynamic lookup and registration of
+//   /// dialects via a plugin facility based on shared library symbol lookup.
+//   const MlirDialectHandle *mlirGetDialectHandle__{NAMESPACE}__();
+//
+// This is done via a common macro to facilitate future expansion to
+// registration schemes.
+//===----------------------------------------------------------------------===//
+
+struct MlirDialectHandle {
+  const void *ptr;
+};
+typedef struct MlirDialectHandle MlirDialectHandle;
+
+#define MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Name, Namespace)                \
+  MLIR_CAPI_EXPORTED MlirDialectHandle mlirGetDialectHandle__##Namespace##__()
+
+/// Returns the namespace associated with the provided dialect handle.
+MLIR_CAPI_EXPORTED
+MlirStringRef mlirDialectHandleGetNamespace(MlirDialectHandle);
+
+/// Inserts the dialect associated with the provided dialect handle into the
+/// provided dialect registry
+MLIR_CAPI_EXPORTED void mlirDialectHandleInsertDialect(MlirDialectHandle,
+                                                       MlirDialectRegistry);
+
+/// Registers the dialect associated with the provided dialect handle.
+MLIR_CAPI_EXPORTED void mlirDialectHandleRegisterDialect(MlirDialectHandle,
+                                                         MlirContext);
+
+/// Loads the dialect associated with the provided dialect handle.
+MLIR_CAPI_EXPORTED MlirDialect mlirDialectHandleLoadDialect(MlirDialectHandle,
+                                                            MlirContext);
+
+//===----------------------------------------------------------------------===//
 // DialectRegistry API.
 //===----------------------------------------------------------------------===//
 
index cb61129..b66bdfe 100644 (file)
@@ -15,7 +15,6 @@
 #define MLIR_C_PASS_H
 
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
 #include "mlir-c/Support.h"
 
 #ifdef __cplusplus
diff --git a/mlir/include/mlir-c/RegisterEverything.h b/mlir/include/mlir-c/RegisterEverything.h
new file mode 100644 (file)
index 0000000..b98ce15
--- /dev/null
@@ -0,0 +1,38 @@
+//===-- mlir-c/RegisterEverything.h - Register all MLIR entities --*- 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
+//
+//===----------------------------------------------------------------------===//
+// This header contains registration entry points for MLIR upstream dialects
+// and passes. Downstream projects typically will not want to use this unless
+// if they don't care about binary size or build bloat and just wish access
+// to the entire set of upstream facilities. For those that do care, they
+// should use registration functions specific to their project.
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_REGISTER_EVERYTHING_H
+#define MLIR_C_REGISTER_EVERYTHING_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// Appends all upstream dialects and extensions to the dialect registry.
+MLIR_CAPI_EXPORTED void mlirRegisterAllDialects(MlirDialectRegistry registry);
+
+/// 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_REGISTER_EVERYTHING_H
diff --git a/mlir/include/mlir-c/Registration.h b/mlir/include/mlir-c/Registration.h
deleted file mode 100644 (file)
index ab37866..0000000
+++ /dev/null
@@ -1,75 +0,0 @@
-//===-- mlir-c/Registration.h - Registration functions for MLIR ---*- 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_REGISTRATION_H
-#define MLIR_C_REGISTRATION_H
-
-#include "mlir-c/IR.h"
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-//===----------------------------------------------------------------------===//
-// Dialect registration declarations.
-// Registration entry-points for each dialect are declared using the common
-// MLIR_DECLARE_DIALECT_REGISTRATION_CAPI macro, which takes the dialect
-// API name (i.e. "Func", "Tensor", "Linalg") and namespace (i.e. "func",
-// "tensor", "linalg"). The following declarations are produced:
-//
-//   /// Gets the above hook methods in struct form for a dialect by namespace.
-//   /// This is intended to facilitate dynamic lookup and registration of
-//   /// dialects via a plugin facility based on shared library symbol lookup.
-//   const MlirDialectHandle *mlirGetDialectHandle__{NAMESPACE}__();
-//
-// This is done via a common macro to facilitate future expansion to
-// registration schemes.
-//===----------------------------------------------------------------------===//
-
-struct MlirDialectHandle {
-  const void *ptr;
-};
-typedef struct MlirDialectHandle MlirDialectHandle;
-
-#define MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Name, Namespace)                \
-  MLIR_CAPI_EXPORTED MlirDialectHandle mlirGetDialectHandle__##Namespace##__()
-
-/// Returns the namespace associated with the provided dialect handle.
-MLIR_CAPI_EXPORTED
-MlirStringRef mlirDialectHandleGetNamespace(MlirDialectHandle);
-
-/// Inserts the dialect associated with the provided dialect handle into the
-/// provided dialect registry
-MLIR_CAPI_EXPORTED void mlirDialectHandleInsertDialect(MlirDialectHandle,
-                                                       MlirDialectRegistry);
-
-/// Registers the dialect associated with the provided dialect handle.
-MLIR_CAPI_EXPORTED void mlirDialectHandleRegisterDialect(MlirDialectHandle,
-                                                         MlirContext);
-
-/// Loads the dialect associated with the provided dialect handle.
-MLIR_CAPI_EXPORTED MlirDialect mlirDialectHandleLoadDialect(MlirDialectHandle,
-                                                            MlirContext);
-
-/// 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
index 588721c..351fb96 100644 (file)
@@ -124,6 +124,25 @@ struct type_caster<MlirContext> {
   }
 };
 
+/// Casts object <-> MlirDialectRegistry.
+template <>
+struct type_caster<MlirDialectRegistry> {
+  PYBIND11_TYPE_CASTER(MlirDialectRegistry, _("MlirDialectRegistry"));
+  bool load(handle src, bool) {
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToDialectRegistry(capsule.ptr());
+    return !mlirDialectRegistryIsNull(value);
+  }
+  static handle cast(MlirDialectRegistry v, return_value_policy, handle) {
+    py::object capsule = py::reinterpret_steal<py::object>(
+        mlirPythonDialectRegistryToCapsule(v));
+    return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+        .attr("DialectRegistry")
+        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+        .release();
+  }
+};
+
 /// Casts object <-> MlirLocation.
 template <>
 struct type_caster<MlirLocation> {
index e57023d..355c4bc 100644 (file)
@@ -10,7 +10,6 @@
 #define MLIR_CAPI_REGISTRATION_H
 
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
 #include "mlir/CAPI/IR.h"
 #include "mlir/CAPI/Support.h"
 
diff --git a/mlir/lib/Bindings/Python/AllPassesRegistration.cpp b/mlir/lib/Bindings/Python/AllPassesRegistration.cpp
deleted file mode 100644 (file)
index f595b20..0000000
+++ /dev/null
@@ -1,22 +0,0 @@
-//===- 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/Conversions/Conversions.cpp b/mlir/lib/Bindings/Python/Conversions/Conversions.cpp
deleted file mode 100644 (file)
index c9d3801..0000000
+++ /dev/null
@@ -1,22 +0,0 @@
-//===- Conversions.cpp - Pybind module for the Conversionss 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/Conversion.h"
-
-#include <pybind11/pybind11.h>
-
-// -----------------------------------------------------------------------------
-// Module initialization.
-// -----------------------------------------------------------------------------
-
-PYBIND11_MODULE(_mlirConversions, m) {
-  m.doc() = "MLIR Conversions library";
-
-  // Register all the passes in the Conversions library on load.
-  mlirRegisterConversionPasses();
-}
index 98566dd..9738351 100644 (file)
@@ -16,7 +16,7 @@
 #include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/Debug.h"
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
+//#include "mlir-c/Registration.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
@@ -474,7 +474,6 @@ py::object PyMlirContext::createFromCapsule(py::object capsule) {
 
 PyMlirContext *PyMlirContext::createNewContextForInit() {
   MlirContext context = mlirContextCreate();
-  mlirRegisterAllDialects(context);
   return new PyMlirContext(context);
 }
 
@@ -793,7 +792,7 @@ py::tuple PyDiagnostic::getNotes() {
 }
 
 //------------------------------------------------------------------------------
-// PyDialect, PyDialectDescriptor, PyDialects
+// PyDialect, PyDialectDescriptor, PyDialects, PyDialectRegistry
 //------------------------------------------------------------------------------
 
 MlirDialect PyDialects::getDialectForKey(const std::string &key,
@@ -807,6 +806,19 @@ MlirDialect PyDialects::getDialectForKey(const std::string &key,
   return dialect;
 }
 
+py::object PyDialectRegistry::getCapsule() {
+  return py::reinterpret_steal<py::object>(
+      mlirPythonDialectRegistryToCapsule(*this));
+}
+
+PyDialectRegistry PyDialectRegistry::createFromCapsule(py::object capsule) {
+  MlirDialectRegistry rawRegistry =
+      mlirPythonCapsuleToDialectRegistry(capsule.ptr());
+  if (mlirDialectRegistryIsNull(rawRegistry))
+    throw py::error_already_set();
+  return PyDialectRegistry(rawRegistry);
+}
+
 //------------------------------------------------------------------------------
 // PyLocation
 //------------------------------------------------------------------------------
@@ -2207,8 +2219,11 @@ void mlir::python::populateIRCore(py::module &m) {
 
   //----------------------------------------------------------------------------
   // Mapping of MlirContext.
+  // Note that this is exported as _BaseContext. The containing, Python level
+  // __init__.py will subclass it with site-specific functionality and set a
+  // "Context" attribute on this module.
   //----------------------------------------------------------------------------
-  py::class_<PyMlirContext>(m, "Context", py::module_local())
+  py::class_<PyMlirContext>(m, "_BaseContext", py::module_local())
       .def(py::init<>(&PyMlirContext::createNewContextForInit))
       .def_static("_get_live_count", &PyMlirContext::getLiveCount)
       .def("_get_context_again",
@@ -2276,7 +2291,16 @@ void mlir::python::populateIRCore(py::module &m) {
             return mlirContextIsRegisteredOperation(
                 self.get(), MlirStringRef{name.data(), name.size()});
           },
-          py::arg("operation_name"));
+          py::arg("operation_name"))
+      .def(
+          "append_dialect_registry",
+          [](PyMlirContext &self, PyDialectRegistry &registry) {
+            mlirContextAppendDialectRegistry(self.get(), registry);
+          },
+          py::arg("registry"))
+      .def("load_all_available_dialects", [](PyMlirContext &self) {
+        mlirContextLoadAllAvailableDialects(self.get());
+      });
 
   //----------------------------------------------------------------------------
   // Mapping of PyDialectDescriptor
@@ -2332,6 +2356,15 @@ void mlir::python::populateIRCore(py::module &m) {
       });
 
   //----------------------------------------------------------------------------
+  // Mapping of PyDialectRegistry
+  //----------------------------------------------------------------------------
+  py::class_<PyDialectRegistry>(m, "DialectRegistry", py::module_local())
+      .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
+                             &PyDialectRegistry::getCapsule)
+      .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyDialectRegistry::createFromCapsule)
+      .def(py::init<>());
+
+  //----------------------------------------------------------------------------
   // Mapping of Location
   //----------------------------------------------------------------------------
   py::class_<PyLocation>(m, "Location", py::module_local())
index 371157a..2e2ebaa 100644 (file)
@@ -390,6 +390,32 @@ private:
   pybind11::object descriptor;
 };
 
+/// Wrapper around an MlirDialectRegistry.
+/// Upon construction, the Python wrapper takes ownership of the
+/// underlying MlirDialectRegistry.
+class PyDialectRegistry {
+public:
+  PyDialectRegistry() : registry(mlirDialectRegistryCreate()) {}
+  PyDialectRegistry(MlirDialectRegistry registry) : registry(registry) {}
+  ~PyDialectRegistry() {
+    if (!mlirDialectRegistryIsNull(registry))
+      mlirDialectRegistryDestroy(registry);
+  }
+  PyDialectRegistry(PyDialectRegistry &) = delete;
+  PyDialectRegistry(PyDialectRegistry &&other) : registry(other.registry) {
+    other.registry = {nullptr};
+  }
+
+  operator MlirDialectRegistry() const { return registry; }
+  MlirDialectRegistry get() const { return registry; }
+
+  pybind11::object getCapsule();
+  static PyDialectRegistry createFromCapsule(pybind11::object capsule);
+
+private:
+  MlirDialectRegistry registry;
+};
+
 /// Wrapper around an MlirLocation.
 class PyLocation : public BaseContextObject {
 public:
diff --git a/mlir/lib/Bindings/Python/RegisterEverything.cpp b/mlir/lib/Bindings/Python/RegisterEverything.cpp
new file mode 100644 (file)
index 0000000..fed5c36
--- /dev/null
@@ -0,0 +1,26 @@
+//===- RegisterEverything.cpp - API to register all dialects/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/RegisterEverything.h"
+#include "mlir-c/Conversion.h"
+#include "mlir-c/Transforms.h"
+
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+
+PYBIND11_MODULE(_mlirRegisterEverything, m) {
+  m.doc() = "MLIR All Upstream Dialects and Passes Registration";
+
+  m.def("register_dialects", [](MlirDialectRegistry registry) {
+    mlirRegisterAllDialects(registry);
+  });
+
+  // Register all passes on load.
+  mlirRegisterAllPasses();
+  mlirRegisterConversionPasses();
+  mlirRegisterTransformsPasses();
+}
diff --git a/mlir/lib/Bindings/Python/Transforms/Transforms.cpp b/mlir/lib/Bindings/Python/Transforms/Transforms.cpp
deleted file mode 100644 (file)
index 944b191..0000000
+++ /dev/null
@@ -1,22 +0,0 @@
-//===- 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>
-
-// -----------------------------------------------------------------------------
-// Module initialization.
-// -----------------------------------------------------------------------------
-
-PYBIND11_MODULE(_mlirTransforms, m) {
-  m.doc() = "MLIR Transforms library";
-
-  // Register all the passes in the Transforms library on load.
-  mlirRegisterTransformsPasses();
-}
index 393b49e..ffb04c2 100644 (file)
@@ -12,7 +12,7 @@ add_subdirectory(Dialect)
 add_subdirectory(Conversion)
 add_subdirectory(Interfaces)
 add_subdirectory(IR)
-add_subdirectory(Registration)
+add_subdirectory(RegisterEverything)
 add_subdirectory(Transforms)
 
 # Only enable the ExecutionEngine if the native target is configured in.
index 166e799..8cafc09 100644 (file)
@@ -2,6 +2,9 @@ get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
 add_mlir_upstream_c_api_library(MLIRCAPIConversion
   Passes.cpp
 
+  DEPENDS
+  MLIRConversionPassIncGen
+
   LINK_LIBS PUBLIC
   ${conversion_libs}
 )
index 6959ad3..c931ea7 100644 (file)
@@ -77,6 +77,10 @@ void mlirContextEnableMultithreading(MlirContext context, bool enable) {
   return unwrap(context)->enableMultithreading(enable);
 }
 
+void mlirContextLoadAllAvailableDialects(MlirContext context) {
+  unwrap(context)->loadAllAvailableDialects();
+}
+
 //===----------------------------------------------------------------------===//
 // Dialect API.
 //===----------------------------------------------------------------------===//
similarity index 76%
rename from mlir/lib/CAPI/Registration/CMakeLists.txt
rename to mlir/lib/CAPI/RegisterEverything/CMakeLists.txt
index 67e26a5..942bba8 100644 (file)
@@ -2,13 +2,15 @@
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 get_property(translation_libs GLOBAL PROPERTY MLIR_TRANSLATION_LIBS)
 get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
-add_mlir_upstream_c_api_library(MLIRCAPIRegistration
-  Registration.cpp
+add_mlir_upstream_c_api_library(MLIRCAPIRegisterEverything
+  RegisterEverything.cpp
 
   LINK_LIBS PUBLIC
-  MLIRCAPIIR
-  MLIRLLVMToLLVMIRTranslation
   ${dialect_libs}
   ${translation_libs}
   ${conversion_libs}
+
+  MLIRCAPIIR
+  MLIRLLVMToLLVMIRTranslation
+  MLIRCAPITransforms
 )
@@ -1,4 +1,4 @@
-//===- Registration.cpp - C Interface for MLIR Registration ---------------===//
+//===- RegisterEverything.cpp - Register all MLIR entities ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -6,17 +6,15 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/RegisterEverything.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 mlirRegisterAllDialects(MlirDialectRegistry registry) {
+  mlir::registerAllDialects(*unwrap(registry));
 }
 
 void mlirRegisterAllLLVMTranslations(MlirContext context) {
index dc831c9..1fbbadf 100644 (file)
@@ -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
index 23bc502..add8d92 100644 (file)
@@ -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()
index cd7eb0a..60bc367 100644 (file)
@@ -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 (file)
index aca557a..0000000
+++ /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 (file)
index a6a9eb8..0000000
+++ /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 (file)
index 71ea17d..0000000
+++ /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
index a4e2f02..b089a37 100644 (file)
@@ -26,8 +26,8 @@ if(TARGET ${LLVM_NATIVE_ARCH})
   LINK_LIBS PRIVATE
     MLIRCAPIConversion
     MLIRCAPIExecutionEngine
-    MLIRCAPIRegistration
-  )
+    MLIRCAPIRegisterEverything
+)
 endif()
 
 _add_capi_test_executable(mlir-capi-ir-test
@@ -35,7 +35,7 @@ _add_capi_test_executable(mlir-capi-ir-test
   LINK_LIBS PRIVATE
     MLIRCAPIIR
     MLIRCAPIFunc
-    MLIRCAPIRegistration
+    MLIRCAPIRegisterEverything
 )
 
 _add_capi_test_executable(mlir-capi-llvm-test
@@ -43,7 +43,7 @@ _add_capi_test_executable(mlir-capi-llvm-test
   LINK_LIBS PRIVATE
     MLIRCAPIIR
     MLIRCAPILLVM
-    MLIRCAPIRegistration
+    MLIRCAPIRegisterEverything
 )
 
 _add_capi_test_executable(mlir-capi-pass-test
@@ -51,7 +51,7 @@ _add_capi_test_executable(mlir-capi-pass-test
   LINK_LIBS PRIVATE
     MLIRCAPIFunc
     MLIRCAPIIR
-    MLIRCAPIRegistration
+    MLIRCAPIRegisterEverything
     MLIRCAPITransforms
 )
 
@@ -59,7 +59,7 @@ _add_capi_test_executable(mlir-capi-sparse-tensor-test
   sparse_tensor.c
   LINK_LIBS PRIVATE
     MLIRCAPIIR
-    MLIRCAPIRegistration
+    MLIRCAPIRegisterEverything
     MLIRCAPISparseTensor
 )
 
@@ -67,7 +67,7 @@ _add_capi_test_executable(mlir-capi-quant-test
   quant.c
   LINK_LIBS PRIVATE
     MLIRCAPIIR
-    MLIRCAPIRegistration
+    MLIRCAPIRegisterEverything
     MLIRCAPIQuant
 )
 
@@ -75,6 +75,6 @@ _add_capi_test_executable(mlir-capi-pdl-test
   pdl.c
   LINK_LIBS PRIVATE
     MLIRCAPIIR
-    MLIRCAPIRegistration
+    MLIRCAPIRegisterEverything
     MLIRCAPIPDL
 )
index 7420af4..662f9f2 100644 (file)
@@ -15,7 +15,7 @@
 #include "mlir-c/Conversion.h"
 #include "mlir-c/ExecutionEngine.h"
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
+#include "mlir-c/RegisterEverything.h"
 
 #include <assert.h>
 #include <math.h>
 #include <stdlib.h>
 #include <string.h>
 
+static void registerAllUpstreamDialects(MlirContext ctx) {
+  MlirDialectRegistry registry = mlirDialectRegistryCreate();
+  mlirRegisterAllDialects(registry);
+  mlirContextAppendDialectRegistry(ctx, registry);
+  mlirDialectRegistryDestroy(registry);
+}
+
 void lowerModuleToLLVM(MlirContext ctx, MlirModule module) {
   MlirPassManager pm = mlirPassManagerCreate(ctx);
   MlirOpPassManager opm = mlirPassManagerGetNestedUnder(
@@ -41,7 +48,8 @@ void lowerModuleToLLVM(MlirContext ctx, MlirModule module) {
 // CHECK-LABEL: Running test 'testSimpleExecution'
 void testSimpleExecution() {
   MlirContext ctx = mlirContextCreate();
-  mlirRegisterAllDialects(ctx);
+  registerAllUpstreamDialects(ctx);
+
   MlirModule module = mlirModuleCreateParse(
       ctx, mlirStringRefCreateFromCString(
                // clang-format off
index 49e6ff9..9c1f95c 100644 (file)
@@ -18,7 +18,7 @@
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/Dialect/Func.h"
 #include "mlir-c/IntegerSet.h"
-#include "mlir-c/Registration.h"
+#include "mlir-c/RegisterEverything.h"
 #include "mlir-c/Support.h"
 
 #include <assert.h>
 #include <stdlib.h>
 #include <string.h>
 
+static void registerAllUpstreamDialects(MlirContext ctx) {
+  MlirDialectRegistry registry = mlirDialectRegistryCreate();
+  mlirRegisterAllDialects(registry);
+  mlirContextAppendDialectRegistry(ctx, registry);
+  mlirDialectRegistryDestroy(registry);
+}
+
 void populateLoopBody(MlirContext ctx, MlirBlock loopBody,
                       MlirLocation location, MlirBlock funcBody) {
   MlirValue iv = mlirBlockGetArgument(loopBody, 0);
@@ -1646,7 +1653,9 @@ int testOperands() {
   // CHECK-LABEL: @testOperands
 
   MlirContext ctx = mlirContextCreate();
-  mlirRegisterAllDialects(ctx);
+  registerAllUpstreamDialects(ctx);
+
+  mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("arith"));
   mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("test"));
   MlirLocation loc = mlirLocationUnknownGet(ctx);
   MlirType indexType = mlirIndexTypeGet(ctx);
@@ -1714,7 +1723,8 @@ int testClone() {
   // CHECK-LABEL: @testClone
 
   MlirContext ctx = mlirContextCreate();
-  mlirRegisterAllDialects(ctx);
+  registerAllUpstreamDialects(ctx);
+
   mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("func"));
   MlirLocation loc = mlirLocationUnknownGet(ctx);
   MlirType indexType = mlirIndexTypeGet(ctx);
@@ -2036,7 +2046,12 @@ void testDiagnostics() {
 
 int main() {
   MlirContext ctx = mlirContextCreate();
-  mlirRegisterAllDialects(ctx);
+  registerAllUpstreamDialects(ctx);
+  mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("func"));
+  mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("memref"));
+  mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("shape"));
+  mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("scf"));
+
   if (constructAndTraverseIr(ctx))
     return 1;
   buildWithInsertionsAndPrint(ctx);
index f73398e..e48045a 100644 (file)
@@ -13,7 +13,7 @@
 #include "mlir-c/Pass.h"
 #include "mlir-c/Dialect/Func.h"
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
+#include "mlir-c/RegisterEverything.h"
 #include "mlir-c/Transforms.h"
 
 #include <assert.h>
 #include <stdlib.h>
 #include <string.h>
 
+static void registerAllUpstreamDialects(MlirContext ctx) {
+  MlirDialectRegistry registry = mlirDialectRegistryCreate();
+  mlirRegisterAllDialects(registry);
+  mlirContextAppendDialectRegistry(ctx, registry);
+  mlirDialectRegistryDestroy(registry);
+}
+
 void testRunPassOnModule() {
   MlirContext ctx = mlirContextCreate();
-  mlirRegisterAllDialects(ctx);
+  registerAllUpstreamDialects(ctx);
 
   MlirModule module = mlirModuleCreateParse(
       ctx,
@@ -62,7 +69,7 @@ void testRunPassOnModule() {
 
 void testRunPassOnNestedModule() {
   MlirContext ctx = mlirContextCreate();
-  mlirRegisterAllDialects(ctx);
+  registerAllUpstreamDialects(ctx);
 
   MlirModule module = mlirModuleCreateParse(
       ctx,
@@ -264,7 +271,7 @@ MlirExternalPassCallbacks makeTestExternalPassCallbacks(
 
 void testExternalPass() {
   MlirContext ctx = mlirContextCreate();
-  mlirRegisterAllDialects(ctx);
+  registerAllUpstreamDialects(ctx);
 
   MlirModule module = mlirModuleCreateParse(
       ctx,
index 78d2f86..10cd812 100644 (file)
@@ -11,7 +11,7 @@
 
 #include "mlir-c/Dialect/SparseTensor.h"
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
+#include "mlir-c/RegisterEverything.h"
 
 #include <assert.h>
 #include <math.h>
index 6eed53c..c46327a 100644 (file)
@@ -62,7 +62,6 @@ run(testInvalidModule)
 
 
 def lowerToLLVM(module):
-  import mlir.conversions
   pm = PassManager.parse(
       "convert-complex-to-llvm,convert-memref-to-llvm,convert-func-to-llvm,reconcile-unrealized-casts")
   pm.run(module)
index 899c355..e22f5a0 100644 (file)
@@ -186,10 +186,6 @@ func.func @main() -> i32 attributes {llvm.emit_c_interface} {
 
 
 def transform(module, boilerplate):
-  import mlir.conversions
-  import mlir.all_passes_registration
-  import mlir.transforms
-
   # TODO: Allow cloning functions from one module to another.
   # Atm we have to resort to string concatenation.
   ops = module.operation.regions[0].blocks[0].operations
index cd45eec..d7cbbfb 100644 (file)
@@ -27,7 +27,6 @@ add_mlir_public_c_api_library(MLIRCAPIPythonTestDialect
   LINK_LIBS PUBLIC
   MLIRCAPIInterfaces
   MLIRCAPIIR
-  MLIRCAPIRegistration
   MLIRPythonTestDialect
 )
 
index dd49102..f72a3c1 100644 (file)
@@ -9,7 +9,7 @@
 #ifndef MLIR_TEST_PYTHON_LIB_PYTHONTESTCAPI_H
 #define MLIR_TEST_PYTHON_LIB_PYTHONTESTCAPI_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
index 6cc627d..1730426 100644 (file)
@@ -33,20 +33,18 @@ run(testCapsule)
 # 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
+    # An unregistered pass should not parse.
     try:
-      pm = PassManager.parse("builtin.module(func.func(print-op-stats{json=false}))")
+      pm = PassManager.parse("builtin.module(func.func(not-existing-pass{json=false}))")
       # 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
+      # CHECK: error: 'not-existing-pass' does not refer to a registered pass or pass pipeline
     except ValueError as e:
-      # CHECK: ValueError exception: invalid pass pipeline 'builtin.module(func.func(print-op-stats{json=false}))'.
+      # CHECK: ValueError exception: invalid pass pipeline 'builtin.module(func.func(not-existing-pass{json=false}))'.
       log("ValueError exception:", e)
     else:
       log("Exception not produced")
 
-    # This will register the pass and round-trip should be possible now.
-    import mlir.transforms
+    # A registered pass should parse successfully.
     pm = PassManager.parse("builtin.module(func.func(print-op-stats{json=false}))")
     # CHECK: Roundtrip: builtin.module(func.func(print-op-stats{json=false}))
     log("Roundtrip: ", pm)
@@ -71,7 +69,6 @@ run(testParseFail)
 def testInvalidNesting():
   with Context():
     try:
-      import mlir.all_passes_registration
       pm = PassManager.parse("func.func(normalize-memrefs)")
     except ValueError as e:
       # CHECK: Can't add pass 'NormalizeMemRefs' restricted to 'builtin.module' on a PassManager intended to run on 'func.func', did you intend to nest?
index 23b3cd2..5058e45 100644 (file)
@@ -384,7 +384,7 @@ mlir_c_api_cc_library(
         "include/mlir-c/IntegerSet.h",
         "include/mlir-c/Interfaces.h",
         "include/mlir-c/Pass.h",
-        "include/mlir-c/Registration.h",
+        "include/mlir-c/RegisterEverything.h",
         "include/mlir-c/Support.h",
         "include/mlir/CAPI/AffineExpr.h",
         "include/mlir/CAPI/AffineMap.h",
@@ -644,9 +644,9 @@ mlir_c_api_cc_library(
 )
 
 mlir_c_api_cc_library(
-    name = "CAPIRegistration",
-    srcs = ["lib/CAPI/Registration/Registration.cpp"],
-    hdrs = ["include/mlir-c/Registration.h"],
+    name = "CAPIRegisterEverything",
+    srcs = ["lib/CAPI/RegisterEverything/RegisterEverything.cpp"],
+    hdrs = ["include/mlir-c/RegisterEverything.h"],
     capi_deps = [
         ":CAPIIR",
     ],
@@ -744,7 +744,6 @@ cc_library(
         ":CAPIGPU",
         ":CAPIIR",
         ":CAPIInterfaces",
-        ":CAPIRegistration",
         ":MLIRBindingsPythonHeadersAndDeps",
         "//llvm:Support",
         "@pybind11",
index ee58abd..abf1662 100644 (file)
@@ -20,13 +20,6 @@ package(
 ##---------------------------------------------------------------------------##
 
 filegroup(
-    name = "ConversionsPyFiles",
-    srcs = glob([
-        "mlir/conversions/*.py",
-    ]),
-)
-
-filegroup(
     name = "DialectCorePyFiles",
     srcs = [
         "mlir/dialects/_ods_common.py",
@@ -83,20 +76,6 @@ filegroup(
     ]),
 )
 
-filegroup(
-    name = "TransformsPyFiles",
-    srcs = glob([
-        "mlir/transforms/*.py",
-    ]),
-)
-
-filegroup(
-    name = "AllPassesRegistrationPyFiles",
-    srcs = glob([
-        "mlir/all_passes_registration/*.py",
-    ]),
-)
-
 ##---------------------------------------------------------------------------##
 # Builtin dialect.
 ##---------------------------------------------------------------------------##