[mlir][python] enable registering dialects with the default Context (#72488)

Author: makslevental

mlir/python/mlir/_mlir_libs/__init__.py

@@ -56,14 +56,28 @@ def get_include_dirs() -> Sequence[str]:
 #
 # This facility allows downstreams to customize Context creation to their
 # needs.
+
+_dialect_registry = None
+
+
+def get_dialect_registry():
+    global _dialect_registry
+
+    if _dialect_registry is None:
+        from ._mlir import ir
+
+        _dialect_registry = ir.DialectRegistry()
+
+    return _dialect_registry
+
+
 def _site_initialize():
     import importlib
     import itertools
     import logging
     from ._mlir import ir
 
     logger = logging.getLogger(__name__)
-    registry = ir.DialectRegistry()
     post_init_hooks = []
     disable_multithreading = False
 
@@ -84,7 +98,7 @@ def process_initializer_module(module_name):
         logger.debug("Initializing MLIR with module: %s", module_name)
         if hasattr(m, "register_dialects"):
             logger.debug("Registering dialects from initializer %r", m)
-            m.register_dialects(registry)
+            m.register_dialects(get_dialect_registry())
         if hasattr(m, "context_init_hook"):
             logger.debug("Adding context init hook from %r", m)
             post_init_hooks.append(m.context_init_hook)
@@ -110,7 +124,7 @@ def process_initializer_module(module_name):
     class Context(ir._BaseContext):
         def __init__(self, *args, **kwargs):
             super().__init__(*args, **kwargs)
-            self.append_dialect_registry(registry)
+            self.append_dialect_registry(get_dialect_registry())
             for hook in post_init_hooks:
                 hook(self)
             if not disable_multithreading:

mlir/python/mlir/dialects/python_test.py

@@ -11,7 +11,7 @@
 )
 
 
-def register_python_test_dialect(context, load=True):
+def register_python_test_dialect(registry):
     from .._mlir_libs import _mlirPythonTest
 
-    _mlirPythonTest.register_python_test_dialect(context, load)
+    _mlirPythonTest.register_dialect(registry)

mlir/python/mlir/ir.py

@@ -5,6 +5,7 @@
 from ._mlir_libs._mlir.ir import *
 from ._mlir_libs._mlir.ir import _GlobalDebug
 from ._mlir_libs._mlir import register_type_caster, register_value_caster
+from ._mlir_libs import get_dialect_registry
 
 
 # Convenience decorator for registering user-friendly Attribute builders.