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

[stellaraccident]

Summary

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.

[tschuett]

Do you believe that there is a place for a HowTo.md that makes downstream users really happy?

[stellaraccident]

In D128593#3611008, @tschuett wrote:

Do you believe that there is a place for a HowTo.md that makes downstream users really happy?

As in a "How to Integrate into a Downstream Project?" guide. Yes, we should create that.

[tschuett]

In D128593#3611022, @stellaraccident wrote:

In D128593#3611008, @tschuett wrote:

Do you believe that there is a place for a HowTo.md that makes downstream users really happy?

As in a "How to Integrate into a Downstream Project?" guide. Yes, we should create that.

Exactly.

[stellaraccident]

In D128593#3611022, @stellaraccident wrote:

In D128593#3611008, @tschuett wrote:

Do you believe that there is a place for a HowTo.md that makes downstream users really happy?

As in a "How to Integrate into a Downstream Project?" guide. Yes, we should create that.

I will note that the Standalone example does wire this together, but it may be too simple. Since we never got the C/Python registration stuff right (from the get-go), my goal with this patch was to be as NFC as possible to set the stage. Then I was going to make changes to how it is all set up, reflected in the Standalone example. That would probably be the right time to add a guide as well. We have not previously specified how to "do registration right", and my view is that all of the downstreams made pragmatic (and different) choices there. We should establish a norm and use it upstream/Standalone-sample and advocate for downstreams to follow suit, imo.

[ftynse]

Would you mind posting a PSA on the forum copying the change description and the tentative evolution plan?

mlir/include/mlir-c/Registration.h

1–6

Should we rename the file to RegisterEverything.h or something to make it clear?

mlir/lib/Bindings/Python/IRCore.cpp

19

Nit: drop this line

mlir/python/mlir/_mlir_libs/__init__.py

42

This looks unnecessary.

mlir/tools/mlir-pdll/mlir-pdll.cpp

122–126 ↗

(On Diff #440014)

Looks unrelated.

[stellaraccident]

In D128593#3611605, @ftynse wrote:

Would you mind posting a PSA on the forum copying the change description and the tentative evolution plan?

Thanks, Alex. I'm going to take another pass through with an eye to getting things lined up to the final state. Will ping for another look then.

[stellaraccident]

Move closer to end state.

[stellaraccident]

mlir/python/mlir/_mlir_libs/__init__.py

42

Reworked all of this.

[stellaraccident]

Comments.

[ftynse]

I am not a huge fan of magic method / file names, but this is a common thing to do in Python. My only usability concern is about debugging:

with Context():
  ir.parse("custom.func @foo()",)

will work or break depending on where the Context is imported from: from mlir.ir import Context vs from custom.ir import Context. I suppose we can get used to it.

mlir/lib/Bindings/Python/IRCore.cpp

19

Still here...

mlir/lib/Bindings/Python/RegisterEverything.cpp

19–26

Not super-relevant for this patch, but the inconsistency between having to call a function to register dialects and having the passes registered on load feels confusing.

[mehdi_amini]

I have some concerns with:

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.

Because contrary to pass registration for example, there isn't a global context or global dialect registry in MLIR: that means that I should be able to create a context inside my python library and not have it "polluted" by something that was registered by another library for its own purpose.

That said this is still quite an big step forward from where we're at right now, even though I fear we may dig ourselves in a place where it'll be hard to get out of.

mlir/lib/Bindings/Python/RegisterEverything.cpp

19–26

This matches somehow MLIR though in that passes are in a global registries while Dialects aren't.

mlir/lib/CAPI/RegisterEverything/RegisterEverything.cpp

19

Great that you fixed this TODO! :)

[stellaraccident]

In D128593#3614871, @ftynse wrote:

I am not a huge fan of magic method / file names, but this is a common thing to do in Python. My only usability concern is about debugging:

with Context():
  ir.parse("custom.func @foo()",)

will work or break depending on where the Context is imported from: from mlir.ir import Context vs from custom.ir import Context. I suppose we can get used to it.

It is probably best to think of these as completely different libraries: iree.ir.Context, circt.ir.Context are IREE/CIRCT APIs respectively. It is kind of weird if their contexts don't come instantiated to do the job. If they were just open coding a Context, they would just declare their own, but in actuality, this class is very wired in to the core code and not easily modifiable by downstreams. Therefore we provide some hooks for customizing it, so they can make it their own... That part is more mechanics for working around practicalities of interop between the native code and extenders.

I have less of an opinion (after we land all of this) about whether the stock mlir.ir.Context comes pre registered with everything. Aside from matching current behavior, which is an important intermediate state during this transition, we could choose to make the mlirRegisterEverything part of the automatic setup be explicit for upstream in some way. So I see this as digging us halfway out of the hole we're in on that and giving downstreams the ability to climb out completely. Let's discuss the upstream register everything default in a follow-up. Anything we would do there would be an API break and we should discuss that explicitly. I think this patch leaves us the options we need there.

[stellaraccident]

In D128593#3615157, @mehdi_amini wrote:

I have some concerns with:

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.

Because contrary to pass registration for example, there isn't a global context or global dialect registry in MLIR: that means that I should be able to create a context inside my python library and not have it "polluted" by something that was registered by another library for its own purpose.

That said this is still quite an big step forward from where we're at right now, even though I fear we may dig ourselves in a place where it'll be hard to get out of.

Thanks for your analysis. To avoid double typing, I think I addressed some of them in my comment above to Alex. Lmk and happy to keep resolving the default registration case post this patch.

[mikeurbach]

Sorry for the delay, but I finally got around to taking this for a spin.

In terms of implementation details, I think this looks good in general. The _site_initialize_i.py mechanism would be useful for us downstream in CIRCT, and I agree with previous comments about using it for the short term to continue registering everything upstream, until we decide to stop doing that by default.

In CIRCT, I measured a reduction from 379 MB to 84 MB in our add_mlir_python_common_capi_library library (we depend only on MLIRPythonSources.Core, no other dialects or transforms).

So in general, +1 from me. Thanks for cleaning this up @stellaraccident, and for providing a hook for downstreams to do site-specific initialization (and upstream to retain previous behavior).

mlir/lib/Bindings/Python/IRCore.cpp

19

Just bumping this since I'm reviewing right now.

mlir/lib/Bindings/Python/RegisterEverything.cpp

19–26

Agreed that this sort of matches MLIR. In the case of CIRCT, we've done something similar in our own bespoke version of this, where import circt registers passes, but for dialects, you have to call circt.register_dialects.

I think that will change to all happen automatically using the new _site_initialization hook after this patch, but just weighing in that this is how we split it downstream as well.

[stellaraccident]

In D128593#3631463, @mikeurbach wrote:

Sorry for the delay, but I finally got around to taking this for a spin.

In terms of implementation details, I think this looks good in general. The _site_initialize_i.py mechanism would be useful for us downstream in CIRCT, and I agree with previous comments about using it for the short term to continue registering everything upstream, until we decide to stop doing that by default.

In CIRCT, I measured a reduction from 379 MB to 84 MB in our add_mlir_python_common_capi_library library (we depend only on MLIRPythonSources.Core, no other dialects or transforms).

So in general, +1 from me. Thanks for cleaning this up @stellaraccident, and for providing a hook for downstreams to do site-specific initialization (and upstream to retain previous behavior).

Thanke, Mike. Not that it is super relevant to this patch but I'm curious: is that some kind of debug build? 379MiB seems quite hefty otherwise! I expect that maybe you are <10MiB in a stripped/release build?

[mikeurbach]

Not that it is super relevant to this patch but I'm curious: is that some kind of debug build? 379MiB seems quite hefty otherwise! I expect that maybe you are <10MiB in a stripped/release build?

Yes, that was a debug build.

[stellaraccident]

Rebase, fix standalone example, attempt to patch bazel as best I could.

[ftynse]

This broke SparseTensor integration tests. I will fix forward.

[ftynse]

Seems there was a follow-up in beebffa9ab81bba3de91b2cee7a62578ebe1ae00

[stellaraccident]

Sorry, yes - I submitted a follow-up and neglected to note on this thread: https://reviews.llvm.org/rGbeebffa9ab81bba3de91b2cee7a62578ebe1ae00

mlir/examples/standalone/include/Standalone-c/Dialects.h

 //===- Dialects.h - CAPI for dialects -----------------------------*- C -*-===//
 //
 // This file is licensed 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 STANDALONE_C_DIALECTS_H
 #define STANDALONE_C_DIALECTS_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Standalone, standalone);
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // STANDALONE_C_DIALECTS_H

mlir/examples/standalone/python/CMakeLists.txt

 
 add_mlir_python_common_capi_library(StandalonePythonCAPI
   INSTALL_COMPONENT StandalonePythonModules
   INSTALL_DESTINATION python_packages/standalone/mlir_standalone/_mlir_libs
   OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/standalone/mlir_standalone/_mlir_libs"
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_SOURCES
     StandalonePythonSources
+    # TODO: Remove this in favor of showing fine grained registration once
+    # available.
+    MLIRPythonExtension.RegisterEverything
     MLIRPythonSources.Core
 )
 
 ################################################################################
 # Instantiation of all Python modules
 ################################################################################
 
 add_mlir_python_modules(StandalonePythonModules
   ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/standalone/mlir_standalone"
   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
   )

mlir/examples/standalone/test/CAPI/CMakeLists.txt

 # Here we create a single aggregate shared library with the parts of the CAPI
 # that we want to bundle together. Then we link a simple C executable
 # against it to demonstrate that it does have the fully self contained
 # core MLIR library and our own standalone dialect.
 add_mlir_aggregate(StandaloneCAPITestLib
   SHARED
   EMBED_LIBS
   MLIRCAPIIR
-  MLIRCAPIRegistration
+  # TODO: Remove this in favor of showing fine grained dialect registration
+  # (once available).
+  MLIRCAPIRegisterEverything
   StandaloneCAPI
 )
 
 add_llvm_executable(standalone-capi-test
   standalone-capi-test.c
 )
 llvm_update_compile_flags(standalone-capi-test)
 target_link_libraries(standalone-capi-test
   PRIVATE StandaloneCAPITestLib)

mlir/examples/standalone/test/CAPI/standalone-capi-test.c

 //===- standalone-cap-demo.c - Simple demo of C-API -----------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 // RUN: standalone-capi-test 2>&1 | FileCheck %s
 
 #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(
       ctx, mlirStringRefCreateFromCString("%0 = arith.constant 2 : i32\n"
                                           "%1 = standalone.foo %0 : i32\n"));
   if (mlirModuleIsNull(module)) {
     printf("ERROR: Could not parse.\n");
     mlirContextDestroy(ctx);

mlir/include/mlir-c/Bindings/Python/Interop.h

 #define MLIR_PYTHON_CAPSULE_AFFINE_EXPR                                        \
   MAKE_MLIR_PYTHON_QUALNAME("ir.AffineExpr._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_AFFINE_MAP                                         \
   MAKE_MLIR_PYTHON_QUALNAME("ir.AffineMap._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_ATTRIBUTE                                          \
   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                                        \
   MAKE_MLIR_PYTHON_QUALNAME("ir.IntegerSet._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_LOCATION                                           \
   MAKE_MLIR_PYTHON_QUALNAME("ir.Location._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_MODULE                                             \
   MAKE_MLIR_PYTHON_QUALNAME("ir.Module._CAPIPtr")

  * a null context is returned (as checked via mlirContextIsNull). In such a
  * case, the Python APIs will have already set an error. */
 static inline MlirContext mlirPythonCapsuleToContext(PyObject *capsule) {
   void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_CONTEXT);
   MlirContext context = {ptr};
   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. */
 static inline PyObject *mlirPythonLocationToCapsule(MlirLocation loc) {
   return PyCapsule_New(MLIR_PYTHON_GET_WRAPPED_POINTER(loc),
                        MLIR_PYTHON_CAPSULE_LOCATION, NULL);
 }
 

mlir/include/mlir-c/Dialect/Async.h

 //===-- mlir-c/Dialect/Async.h - C API for Async dialect ---------*- 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_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
 extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Async, async);
 
 #ifdef __cplusplus
 }
 #endif
 
 #include "mlir/Dialect/Async/Passes.capi.h.inc"
 
 #endif // MLIR_C_DIALECT_ASYNC_H

mlir/include/mlir-c/Dialect/ControlFlow.h

 //===-- mlir-c/Dialect/ControlFlow.h - C API for ControlFlow ------*- 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_DIALECT_CONTROLFLOW_H
 #define MLIR_C_DIALECT_CONTROLFLOW_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(ControlFlow, cf);
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_DIALECT_CONTROLFLOW_H

mlir/include/mlir-c/Dialect/Func.h

 // before using any of its attributes, operations or types. Parser and pass
 // manager can load registered dialects automatically.
 //
 //===----------------------------------------------------------------------===//
 
 #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" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Func, func);
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_DIALECT_FUNC_H

mlir/include/mlir-c/Dialect/GPU.h

 //===-- mlir-c/Dialect/GPU.h - C API for GPU dialect -------------*- 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_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
 extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(GPU, gpu);
 
 #ifdef __cplusplus
 }
 #endif
 
 #include "mlir/Dialect/GPU/Transforms/Passes.capi.h.inc"
 
 #endif // MLIR_C_DIALECT_GPU_H

mlir/include/mlir-c/Dialect/LLVM.h

 //===-- mlir-c/Dialect/LLVM.h - C API for LLVM --------------------*- 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_DIALECT_LLVM_H
 #define MLIR_C_DIALECT_LLVM_H
 
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(LLVM, llvm);
 
 /// Creates an llvm.ptr type.

mlir/include/mlir-c/Dialect/Linalg.h

 //===-- mlir-c/Dialect/Linalg.h - C API for Linalg dialect -------*- 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_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
 extern "C" {
 #endif
 
 /// Apply the special region builder for the builtin named Linalg op.
 /// Assert that `mlirOp` is a builtin named Linalg op.

mlir/include/mlir-c/Dialect/PDL.h

 //===-- mlir-c/Dialect/PDL.h - C API for PDL Dialect --------------*- 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_DIALECT_PDL_H
 #define MLIR_C_DIALECT_PDL_H
 
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(PDL, pdl);
 
 //===---------------------------------------------------------------------===//

mlir/include/mlir-c/Dialect/Quant.h

 //===-- mlir-c/Dialect/Quant.h - C API for LLVM -------------------*- 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_DIALECT_QUANT_H
 #define MLIR_C_DIALECT_QUANT_H
 
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(quant, quant);
 
 //===---------------------------------------------------------------------===//

mlir/include/mlir-c/Dialect/SCF.h

 //===-- mlir-c/Dialect/SCF.h - C API for SCF dialect --------------*- 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_DIALECT_SCF_H
 #define MLIR_C_DIALECT_SCF_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(SCF, scf);
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_DIALECT_SCF_H

mlir/include/mlir-c/Dialect/Shape.h

 //===-- mlir-c/Dialect/Shape.h - C API for Shape dialect ----------*- 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_DIALECT_SHAPE_H
 #define MLIR_C_DIALECT_SHAPE_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Shape, shape);
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_DIALECT_SHAPE_H

mlir/include/mlir-c/Dialect/SparseTensor.h

 //===-- mlir-c/Dialect/SparseTensor.h - C API for SparseTensor ----*- 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_DIALECT_SPARSETENSOR_H
 #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" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor);
 
 /// Dimension level types that define sparse tensors:

mlir/include/mlir-c/Dialect/Tensor.h

 //===-- mlir-c/Dialect/Tensor.h - C API for Tensor dialect --------*- 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_DIALECT_TENSOR_H
 #define MLIR_C_DIALECT_TENSOR_H
 
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Tensor, tensor);
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_DIALECT_TENSOR_H

mlir/include/mlir-c/IR.h

 /// dialect.
 MLIR_CAPI_EXPORTED MlirDialect mlirContextGetOrLoadDialect(MlirContext context,
                                                            MlirStringRef name);
 
 /// Set threading mode (must be set to false to mlir-print-ir-after-all).
 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
 /// dialect.
 MLIR_CAPI_EXPORTED bool mlirContextIsRegisteredOperation(MlirContext context,
                                                          MlirStringRef name);
 
 //===----------------------------------------------------------------------===//

 /// from different contexts will not compare equal.
 MLIR_CAPI_EXPORTED bool mlirDialectEqual(MlirDialect dialect1,
                                          MlirDialect dialect2);
 
 /// Returns the namespace of the given dialect.
 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.
 //===----------------------------------------------------------------------===//
 
 /// Creates a dialect registry and transfers its ownership to the caller.
 MLIR_CAPI_EXPORTED MlirDialectRegistry mlirDialectRegistryCreate();
 
 /// Checks if the dialect registry is null.
 static inline bool mlirDialectRegistryIsNull(MlirDialectRegistry registry) {

mlir/include/mlir-c/Pass.h

 // This header declares the C interface to MLIR pass manager.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_C_PASS_H
 #define MLIR_C_PASS_H
 
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
 #include "mlir-c/Support.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 //===----------------------------------------------------------------------===//
 // Opaque type declarations.

mlir/include/mlir-c/RegisterEverything.h

+//===-- 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

mlir/include/mlir-c/Registration.h

-//===-- 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

mlir/include/mlir/Bindings/Python/PybindAdaptors.h

if (src.is_none()) {
                 .attr("current");
     }
     py::object capsule = mlirApiObjectToCapsule(src);
     value = mlirPythonCapsuleToContext(capsule.ptr());
     return !mlirContextIsNull(value);
   }
 };
 
+/// 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> {
   PYBIND11_TYPE_CASTER(MlirLocation, _("MlirLocation"));
   bool load(handle src, bool) {
     if (src.is_none()) {
       // Gets the current thread-bound context.
       src = py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))

mlir/include/mlir/CAPI/Registration.h

 //===- Registration.h - C API Registration implementation  ------*- 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_CAPI_REGISTRATION_H
 #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"
 
 //===----------------------------------------------------------------------===//
 // Corrolary to MLIR_DECLARE_CAPI_DIALECT_REGISTRATION that defines an impl.
 // Takes the same name passed to the above and the fully qualified class name
 // of the dialect class.
 //===----------------------------------------------------------------------===//

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/Conversions/Conversions.cpp

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

mlir/lib/Bindings/Python/IRCore.cpp

 #include "Globals.h"
 #include "PybindUtils.h"
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/BuiltinAttributes.h"
 #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"
 
 #include <utility>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
py::object PyMlirContext::createFromCapsule(py::object capsule) {
   MlirContext rawContext = mlirPythonCapsuleToContext(capsule.ptr());
   if (mlirContextIsNull(rawContext))
     throw py::error_already_set();
   return forContext(rawContext).releaseObject();
 }
 
 PyMlirContext *PyMlirContext::createNewContextForInit() {
   MlirContext context = mlirContextCreate();
-  mlirRegisterAllDialects(context);
   return new PyMlirContext(context);
 }
 
 PyMlirContextRef PyMlirContext::forContext(MlirContext context) {
   py::gil_scoped_acquire acquire;
   auto &liveContexts = getLiveContexts();
   auto it = liveContexts.find(context.ptr);
   if (it == liveContexts.end()) {
for (intptr_t i = 0; i < numNotes; ++i) {
     MlirDiagnostic noteDiag = mlirDiagnosticGetNote(diagnostic, i);
     py::object pyNoteDiag = py::cast(PyDiagnostic(noteDiag));
     PyTuple_SET_ITEM(materializedNotes->ptr(), i, pyNoteDiag.ptr());
   }
   return *materializedNotes;
 }
 
 //------------------------------------------------------------------------------
-// PyDialect, PyDialectDescriptor, PyDialects
+// PyDialect, PyDialectDescriptor, PyDialects, PyDialectRegistry
 //------------------------------------------------------------------------------
 
 MlirDialect PyDialects::getDialectForKey(const std::string &key,
                                          bool attrError) {
   MlirDialect dialect = mlirContextGetOrLoadDialect(getContext()->get(),
                                                     {key.data(), key.size()});
   if (mlirDialectIsNull(dialect)) {
     throw SetPyError(attrError ? PyExc_AttributeError : PyExc_IndexError,
                      Twine("Dialect '") + key + "' not found");
   }
   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
 //------------------------------------------------------------------------------
 
 py::object PyLocation::getCapsule() {
   return py::reinterpret_steal<py::object>(mlirPythonLocationToCapsule(*this));
 }
 
py::class_<PyDiagnosticHandler>(m, "DiagnosticHandler", py::module_local())
       .def("detach", &PyDiagnosticHandler::detach)
       .def_property_readonly("attached", &PyDiagnosticHandler::isAttached)
       .def_property_readonly("had_error", &PyDiagnosticHandler::getHadError)
       .def("__enter__", &PyDiagnosticHandler::contextEnter)
       .def("__exit__", &PyDiagnosticHandler::contextExit);
 
   //----------------------------------------------------------------------------
   // 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",
            [](PyMlirContext &self) {
              PyMlirContextRef ref = PyMlirContext::forContext(self.get());
              return ref.releaseObject();
            })
       .def("_get_live_operation_count", &PyMlirContext::getLiveOperationCount)
py::class_<PyMlirContext>(m, "_BaseContext", py::module_local())
           },
           py::arg("enable"))
       .def(
           "is_registered_operation",
           [](PyMlirContext &self, std::string &name) {
             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
   //----------------------------------------------------------------------------
   py::class_<PyDialectDescriptor>(m, "DialectDescriptor", py::module_local())
       .def_property_readonly("namespace",
                              [](PyDialectDescriptor &self) {
                                MlirStringRef ns =
py::class_<PyDialect>(m, "Dialect", py::module_local())
         auto clazz = self.attr("__class__");
         return py::str("<Dialect ") +
                self.attr("descriptor").attr("namespace") + py::str(" (class ") +
                clazz.attr("__module__") + py::str(".") +
                clazz.attr("__name__") + py::str(")>");
       });
 
   //----------------------------------------------------------------------------
+  // 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())
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyLocation::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyLocation::createFromCapsule)
       .def("__enter__", &PyLocation::contextEnter)
       .def("__exit__", &PyLocation::contextExit)
       .def("__eq__",

mlir/lib/Bindings/Python/IRModule.h

public:
   PyDialect(pybind11::object descriptor) : descriptor(std::move(descriptor)) {}
 
   pybind11::object getDescriptor() { return descriptor; }
 
 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:
   PyLocation(PyMlirContextRef contextRef, MlirLocation loc)
       : BaseContextObject(std::move(contextRef)), loc(loc) {}
 
   operator MlirLocation() const { return loc; }
   MlirLocation get() const { return loc; }

mlir/lib/Bindings/Python/RegisterEverything.cpp

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

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>
-
-// -----------------------------------------------------------------------------
-// Module initialization.
-// -----------------------------------------------------------------------------
-
-PYBIND11_MODULE(_mlirTransforms, m) {
-  m.doc() = "MLIR Transforms library";
-
-  // Register all the passes in the Transforms library on load.
-  mlirRegisterTransformsPasses();
-}

mlir/lib/CAPI/CMakeLists.txt

 # For upstream, we accumulate all libraries into the MLIR_CAPI_LIBRARIES
 # property via a custom wrapper function. This is then used to create an
 # aggregate below.
 set_property(GLOBAL APPEND PROPERTY MLIR_CAPI_LIBRARIES)
 function(add_mlir_upstream_c_api_library name)
   add_mlir_public_c_api_library(${name} ${ARGN})
   set_property(GLOBAL APPEND PROPERTY MLIR_CAPI_LIBRARIES ${name})
 endfunction()
 
 add_subdirectory(Debug)
 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.
 if(TARGET ${LLVM_NATIVE_ARCH})
   add_subdirectory(ExecutionEngine)
 endif()
 
 # Build the optional CAPI dylib.

mlir/lib/CAPI/Conversion/CMakeLists.txt

 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}
 )

mlir/lib/CAPI/IR/IR.cpp

 bool mlirContextIsRegisteredOperation(MlirContext context, MlirStringRef name) {
   return unwrap(context)->isOperationRegistered(unwrap(name));
 }
 
 void mlirContextEnableMultithreading(MlirContext context, bool enable) {
   return unwrap(context)->enableMultithreading(enable);
 }
 
+void mlirContextLoadAllAvailableDialects(MlirContext context) {
+  unwrap(context)->loadAllAvailableDialects();
+}
+
 //===----------------------------------------------------------------------===//
 // Dialect API.
 //===----------------------------------------------------------------------===//
 
 MlirContext mlirDialectGetContext(MlirDialect dialect) {
   return wrap(unwrap(dialect)->getContext());
 }
 

mlir/lib/CAPI/RegisterEverything/CMakeLists.txt

 # Dialect registration.
 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
 )

mlir/lib/CAPI/RegisterEverything/RegisterEverything.cpp

-//===- 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.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#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) {
   mlir::registerLLVMDialectTranslation(*unwrap(context));
 }
 
 void mlirRegisterAllPasses() { mlir::registerAllPasses(); }

mlir/python/CMakeLists.txt

declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
   ADD_TO_PARENT MLIRPythonSources
   SOURCES
     execution_engine.py
     _mlir_libs/_mlirExecutionEngine.pyi
   SOURCES_GLOB
     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"
 )
 
 ################################################################################
 # Dialect bindings
 ################################################################################
 
 declare_mlir_dialect_python_bindings(
SOURCES
     Pass.h
     PybindUtils.h
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
     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
   ROOT_DIR "${PYTHON_SOURCE_DIR}"
   SOURCES
     DialectLinalg.cpp
   PRIVATE_LINK_LIBS
     LLVMSupport
SOURCES
     DialectSparseTensor.cpp
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     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
   ROOT_DIR "${PYTHON_SOURCE_DIR}"
   SOURCES
     AsyncPasses.cpp
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
     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
     MODULE_NAME _mlirExecutionEngine
     ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
     ROOT_DIR "${PYTHON_SOURCE_DIR}"
     SOURCES
       ExecutionEngineModule.cpp
declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
   SOURCES
     SparseTensorPasses.cpp
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
     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
 # above confuses Python module loader when running under lit.
 set(_ADDL_TEST_SOURCES)
 if(MLIR_INCLUDE_TESTS)
   set(_ADDL_TEST_SOURCES MLIRPythonTestSources)
   declare_mlir_python_sources(MLIRPythonTestSources)
add_mlir_python_common_capi_library(MLIRPythonCAPI
   INSTALL_COMPONENT MLIRPythonModules
   INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
   OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_HEADERS
     MLIRPythonCAPI.HeaderSources
   DECLARED_SOURCES
     MLIRPythonSources
-    MLIRPythonExtension.AllPassesRegistration
+    MLIRPythonExtension.RegisterEverything
     ${_ADDL_TEST_SOURCES}
 )
 
 ################################################################################
 # The fully assembled package of modules.
 # This must come last.
 ################################################################################
 
 add_mlir_python_modules(MLIRPythonModules
   ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
   INSTALL_PREFIX "python_packages/mlir_core/mlir"
   DECLARED_SOURCES
     MLIRPythonSources
-    MLIRPythonExtension.AllPassesRegistration
+    MLIRPythonExtension.RegisterEverything
     ${_ADDL_TEST_SOURCES}
   COMMON_CAPI_LINK_LIBS
     MLIRPythonCAPI
 )

mlir/python/mlir/_mlir_libs/__init__.py

 # Licensed 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 typing import Any, Sequence
 
 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.
 
   On some platforms, the package may need to be built specially to export
   development libraries.
   """
   return [_this_dir]
 
 
 def get_include_dirs() -> Sequence[str]:
   """Gets the include directory for compiling against exported C libraries.
 
   Depending on how the package was build, development C libraries may or may
   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()

mlir/python/mlir/_mlir_libs/_mlir/ir.pyi

class Context:
     def __enter__(self) -> Context: ...
     def __exit__(self, arg0: object, arg1: object, arg2: object) -> None: ...
     @property
     def _CAPIPtr(self) -> object: ...
     @property
     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):
     def __init__(self, cast_from_attr: Attribute) -> None: ...
     @staticmethod
     def get(*args, **kwargs) -> DenseElementsAttr: ...
     @staticmethod
     def get_splat(*args, **kwargs) -> Any: ...

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 .._mlir_libs import _mlirAllPassesRegistration as _cextAllPasses

mlir/python/mlir/conversions/__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.
-from .._mlir_libs import _mlirConversions as _cextConversions

mlir/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.
-from .._mlir_libs import _mlirTransforms as _cextTransforms

mlir/test/CAPI/CMakeLists.txt

 
 # Only enable the ExecutionEngine if the native target is configured in.
 if(TARGET ${LLVM_NATIVE_ARCH})
   _add_capi_test_executable(mlir-capi-execution-engine-test
     execution_engine.c
   LINK_LIBS PRIVATE
     MLIRCAPIConversion
     MLIRCAPIExecutionEngine
-    MLIRCAPIRegistration
-  )
+    MLIRCAPIRegisterEverything
 endif()
 
 _add_capi_test_executable(mlir-capi-ir-test
   ir.c
   LINK_LIBS PRIVATE
     MLIRCAPIIR
     MLIRCAPIFunc
-    MLIRCAPIRegistration
+    MLIRCAPIRegisterEverything
 )
 
 _add_capi_test_executable(mlir-capi-llvm-test
   llvm.c
   LINK_LIBS PRIVATE
     MLIRCAPIIR
     MLIRCAPILLVM
-    MLIRCAPIRegistration
+    MLIRCAPIRegisterEverything
 )
 
 _add_capi_test_executable(mlir-capi-pass-test
   pass.c
   LINK_LIBS PRIVATE
     MLIRCAPIFunc
     MLIRCAPIIR
-    MLIRCAPIRegistration
+    MLIRCAPIRegisterEverything
     MLIRCAPITransforms
 )
 
 _add_capi_test_executable(mlir-capi-sparse-tensor-test
   sparse_tensor.c
   LINK_LIBS PRIVATE
     MLIRCAPIIR
-    MLIRCAPIRegistration
+    MLIRCAPIRegisterEverything
     MLIRCAPISparseTensor
 )
 
 _add_capi_test_executable(mlir-capi-quant-test
   quant.c
   LINK_LIBS PRIVATE
     MLIRCAPIIR
-    MLIRCAPIRegistration
+    MLIRCAPIRegisterEverything
     MLIRCAPIQuant
 )
 
 _add_capi_test_executable(mlir-capi-pdl-test
   pdl.c
   LINK_LIBS PRIVATE
     MLIRCAPIIR
-    MLIRCAPIRegistration
+    MLIRCAPIRegisterEverything
     MLIRCAPIPDL
 )

mlir/test/CAPI/execution_engine.c

 /* RUN: mlir-capi-execution-engine-test 2>&1 | FileCheck %s
  */
 /* REQUIRES: native
  */
 
 #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 <stdio.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(
       pm, mlirStringRefCreateFromCString("func.func"));
   mlirPassManagerAddOwnedPass(pm, mlirCreateConversionConvertFuncToLLVM());
   mlirOpPassManagerAddOwnedPass(opm,
                                 mlirCreateConversionConvertArithmeticToLLVM());
   MlirLogicalResult status = mlirPassManagerRun(pm, module);
   if (mlirLogicalResultIsFailure(status)) {
     fprintf(stderr, "Unexpected failure running pass pipeline\n");
     exit(2);
   }
   mlirPassManagerDestroy(pm);
 }
 
 // CHECK-LABEL: Running test 'testSimpleExecution'
 void testSimpleExecution() {
   MlirContext ctx = mlirContextCreate();
-  mlirRegisterAllDialects(ctx);
+  registerAllUpstreamDialects(ctx);
+
   MlirModule module = mlirModuleCreateParse(
       ctx, mlirStringRefCreateFromCString(
                // clang-format off
 "module {                                                                    \n"
 "  func.func @add(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {     \n"
 "    %res = arith.addi %arg0, %arg0 : i32                                        \n"
 "    return %res : i32                                                           \n"
 "  }                                                                             \n"

mlir/test/CAPI/ir.c

 #include "mlir-c/IR.h"
 #include "mlir-c/AffineExpr.h"
 #include "mlir-c/AffineMap.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 #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 <inttypes.h>
 #include <math.h>
 #include <stdio.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);
   MlirValue funcArg0 = mlirBlockGetArgument(funcBody, 0);
   MlirValue funcArg1 = mlirBlockGetArgument(funcBody, 1);
   MlirType f32Type =
       mlirTypeParseGet(ctx, mlirStringRefCreateFromCString("f32"));
 

 }
 
 /// Tests operand APIs.
 int testOperands() {
   fprintf(stderr, "@testOperands\n");
   // 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);
 
   // Create some constants to use as operands.
   MlirAttribute indexZeroLiteral =
       mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString("0 : index"));
   MlirNamedAttribute indexZeroValueAttr = mlirNamedAttributeGet(

 }
 
 /// Tests clone APIs.
 int testClone() {
   fprintf(stderr, "@testClone\n");
   // CHECK-LABEL: @testClone
 
   MlirContext ctx = mlirContextCreate();
-  mlirRegisterAllDialects(ctx);
+  registerAllUpstreamDialects(ctx);
+
   mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("func"));
   MlirLocation loc = mlirLocationUnknownGet(ctx);
   MlirType indexType = mlirIndexTypeGet(ctx);
   MlirStringRef valueStringRef = mlirStringRefCreateFromCString("value");
 
   MlirAttribute indexZeroLiteral =
       mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString("0 : index"));
   MlirNamedAttribute indexZeroValueAttr = mlirNamedAttributeGet(
void testDiagnostics() {
   // CHECK: deleting user data (userData: 42)
   // CHECK-NOT: processing diagnostic
   // CHECK:     more test diagnostics
   mlirContextDestroy(ctx);
 }
 
 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);
   if (createOperationWithTypeInference(ctx))
     return 2;
 
   if (printBuiltinTypes(ctx))
     return 3;

mlir/test/CAPI/pass.c

 //===- pass.c - Simple test of C APIs -------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 /* RUN: mlir-capi-pass-test 2>&1 | FileCheck %s
  */
 
 #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 <math.h>
 #include <stdio.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,
       // clang-format off
                             mlirStringRefCreateFromCString(
 "func.func @foo(%arg0 : i32) -> i32 {                                   \n"
 "  %res = arith.addi %arg0, %arg0 : i32                                     \n"
 "  return %res : i32                                                        \n"
// CHECK: func.return        , 1
     mlirPassManagerDestroy(pm);
   }
   mlirModuleDestroy(module);
   mlirContextDestroy(ctx);
 }
 
 void testRunPassOnNestedModule() {
   MlirContext ctx = mlirContextCreate();
-  mlirRegisterAllDialects(ctx);
+  registerAllUpstreamDialects(ctx);
 
   MlirModule module = mlirModuleCreateParse(
       ctx,
       // clang-format off
                             mlirStringRefCreateFromCString(
 "func.func @foo(%arg0 : i32) -> i32 {                                   \n"
 "  %res = arith.addi %arg0, %arg0 : i32                                     \n"
 "  return %res : i32                                                        \n"
MlirExternalPassCallbacks makeTestExternalPassCallbacks(
     void (*runPass)(MlirOperation op, MlirExternalPass, void *userData)) {
   return (MlirExternalPassCallbacks){testConstructExternalPass,
                                      testDestructExternalPass, initializePass,
                                      testCloneExternalPass, runPass};
 }
 
 void testExternalPass() {
   MlirContext ctx = mlirContextCreate();
-  mlirRegisterAllDialects(ctx);
+  registerAllUpstreamDialects(ctx);
 
   MlirModule module = mlirModuleCreateParse(
       ctx,
       // clang-format off
       mlirStringRefCreateFromCString(
 "func.func @foo(%arg0 : i32) -> i32 {                                   \n"
 "  %res = arith.addi %arg0, %arg0 : i32                                     \n"
 "  return %res : i32                                                        \n"

mlir/test/CAPI/sparse_tensor.c

 //===- sparse_tensor.c - Test of sparse_tensor APIs -----------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 // RUN: mlir-capi-sparse-tensor-test 2>&1 | FileCheck %s
 
 #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>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
 
 // CHECK-LABEL: testRoundtripEncoding()

mlir/test/python/execution_engine.py

with Context():
     except RuntimeError as e:
       log("Got RuntimeError: ", e)
 
 
 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)
   return module
 
 
 # Test simple ExecutionEngine execution
 # CHECK-LABEL: TEST: testInvokeVoid

mlir/test/python/integration/dialects/linalg/opsrun.py

func.func @main() -> i32 attributes {llvm.emit_c_interface} {
 
   // TODO: FFI-based solution to allow testing and printing with python code.
   return %0 : i32
 }
 """
 
 
 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
   mod = Module.parse("\n".join([str(op) for op in ops]) + boilerplate)
 
   pm = PassManager.parse(
       "func.func(convert-linalg-to-loops, lower-affine, " +
       "convert-math-to-llvm, convert-scf-to-cf, arith-expand, memref-expand), "

mlir/test/python/lib/CMakeLists.txt

add_mlir_public_c_api_library(MLIRCAPIPythonTestDialect
   PythonTestCAPI.cpp
 
   DEPENDS
   MLIRPythonTestIncGen
 
   LINK_LIBS PUBLIC
   MLIRCAPIInterfaces
   MLIRCAPIIR
-  MLIRCAPIRegistration
   MLIRPythonTestDialect
 )
 

mlir/test/python/lib/PythonTestCAPI.h

 //===- PythonTestCAPI.h - C API for the PythonTest dialect ------*- 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_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" {
 #endif
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(PythonTest, python_test);
 
 MLIR_CAPI_EXPORTED bool

mlir/test/python/pass_manager.py

with Context():
     assert pm1 is not None  # And does not crash.
 run(testCapsule)
 
 
 # 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
+    # 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)
 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)
 
 
 # Verify failure on incorrect level of nesting.
 # CHECK-LABEL: TEST: testInvalidNesting
 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?
       # CHECK: ValueError exception: invalid pass pipeline 'func.func(normalize-memrefs)'.
       log("ValueError exception:", e)
     else:
       log("Exception not produced")
 run(testInvalidNesting)

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

hdrs = [
         "include/mlir-c/BuiltinTypes.h",
         "include/mlir-c/Diagnostics.h",
         "include/mlir-c/Dialect/Func.h",
         "include/mlir-c/ExecutionEngine.h",
         "include/mlir-c/IR.h",
         "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",
         "include/mlir/CAPI/Diagnostics.h",
         "include/mlir/CAPI/IR.h",
         "include/mlir/CAPI/IntegerSet.h",
         "include/mlir/CAPI/Pass.h",
         "include/mlir/CAPI/Registration.h",
mlir_c_api_cc_library(
     includes = ["include"],
     deps = [
         ":Pass",
         ":Transforms",
     ],
 )
 
 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",
     ],
     includes = ["include"],
     deps = [
         ":AllPassesAndDialects",
         ":LLVMToLLVMIRTranslation",
     ],
tags = [
         "nobuildkite",  # TODO(gcmn): Add support for this target
     ],
     deps = [
         ":CAPIAsync",
         ":CAPIDebug",
         ":CAPIGPU",
         ":CAPIIR",
         ":CAPIInterfaces",
-        ":CAPIRegistration",
         ":MLIRBindingsPythonHeadersAndDeps",
         "//llvm:Support",
         "@pybind11",
         "@python_runtime//:headers",
     ],
 )
 
 cc_library(

utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel

package(
     licenses = ["notice"],
 )
 
 ##---------------------------------------------------------------------------##
 # Core IR modules.
 ##---------------------------------------------------------------------------##
 
 filegroup(
-    name = "ConversionsPyFiles",
-    srcs = glob([
-        "mlir/conversions/*.py",
-    ]),
-)
-
-filegroup(
     name = "DialectCorePyFiles",
     srcs = [
         "mlir/dialects/_ods_common.py",
     ],
 )
 
 filegroup(
     name = "ExecutionEnginePyFiles",

 
 filegroup(
     name = "RuntimePyFiles",
     srcs = glob([
         "mlir/runtime/*.py",
     ]),
 )
 
-filegroup(
-    name = "TransformsPyFiles",
-    srcs = glob([
-        "mlir/transforms/*.py",
-    ]),
-)
-
-filegroup(
-    name = "AllPassesRegistrationPyFiles",
-    srcs = glob([
-        "mlir/all_passes_registration/*.py",
-    ]),
-)
-
 ##---------------------------------------------------------------------------##
 # Builtin dialect.
 ##---------------------------------------------------------------------------##
 
 td_library(
     name = "BuiltinOpsPyTdFiles",
     srcs = [
         "mlir/dialects/BuiltinOps.td",