[Mlir-commits] [mlir] [mlir][py] Enable loading only specified dialects during creation. (PR #121421)

[Maksim Levental]

================
@@ -132,15 +148,28 @@ def process_initializer_module(module_name):
             break
 
     class Context(ir._BaseContext):
-        def __init__(self, *args, **kwargs):
+        def __init__(self, load_on_create_dialects=None, *args, **kwargs):
             super().__init__(*args, **kwargs)
             self.append_dialect_registry(get_dialect_registry())
             for hook in post_init_hooks:
                 hook(self)
             if not disable_multithreading:
                 self.enable_multithreading(True)
-            if not disable_load_all_available_dialects:
-                self.load_all_available_dialects()
+            if load_on_create_dialects is not None:
+                logger.debug("Loading all dialects from load_on_create_dialects arg %r", _load_on_create_dialects)
+                for dialect in load_on_create_dialects:
+                    # Load dialect.
+                    _ = self.dialects[dialect]
+            else:
+                if disable_load_all_available_dialects:
+                    if _load_on_create_dialects:
----------------
makslevental wrote:

another nit: `_load_on_create_dialects` should gotten through `get_load_on_create_dialects`

https://github.com/llvm/llvm-project/pull/121421