[Mlir-commits] [mlir] [mlir][py] Enable disabling loading all registered (PR #117643)

[Jacques Pienaar]

https://github.com/jpienaar created https://github.com/llvm/llvm-project/pull/117643

There is a pending todo about always eagerly loading or not. Make this behavior optional and give the control to the user in a backwards compatible manner. This is made optional as there were arguments for both forms, kept it in form that is backwards compatible.

>From 2acbdc7127aba497754b3db43b1f9700efb780e5 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Mon, 25 Nov 2024 23:20:37 +0000
Subject: [PATCH] [mlir][py] Enable disabling loading all registered

There is a pending todo about always eagerly loading or not. Make this behavior optional and give the control to the user in a backwards compatible manner. This is made optional as there were arguments for both forms.
---
 mlir/python/mlir/_mlir_libs/__init__.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py
index 98dbbc6adf9ce5..c5cb22c6dccb8f 100644
--- a/mlir/python/mlir/_mlir_libs/__init__.py
+++ b/mlir/python/mlir/_mlir_libs/__init__.py
@@ -80,9 +80,16 @@ def _site_initialize():
     logger = logging.getLogger(__name__)
     post_init_hooks = []
     disable_multithreading = False
+    # This flag disables eagerly loading all dialects. Eagerly loading is often
+    # not the desired behavior (see
+    # https://github.com/llvm/llvm-project/issues/56037), and the logic is that
+    # if any module has this attribute set, then we don't load all (e.g., it's
+    # being used in a solution where the loading is controlled).
+    disable_load_all_available_dialects = False
 
     def process_initializer_module(module_name):
         nonlocal disable_multithreading
+        nonlocal disable_load_all_available_dialects
         try:
             m = importlib.import_module(f".{module_name}", __name__)
         except ModuleNotFoundError:
@@ -107,6 +114,8 @@ def process_initializer_module(module_name):
             if bool(m.disable_multithreading):
                 logger.debug("Disabling multi-threading for context")
                 disable_multithreading = True
+        if hasattr(m, "disable_load_all_available_dialects"):
+            disable_load_all_available_dialects = True
         return True
 
     # If _mlirRegisterEverything is built, then include it as an initializer
@@ -130,10 +139,8 @@ def __init__(self, *args, **kwargs):
                 hook(self)
             if not disable_multithreading:
                 self.enable_multithreading(True)
-            # 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()
+            if not disable_load_all_available_dialects:
+                self.load_all_available_dialects()
             if init_module:
                 logger.debug(
                     "Registering translations from initializer %r", init_module