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

[llvmlistbot at llvm.org]

Author: Jacques Pienaar
Date: 2024-11-25T15:39:55-08:00
New Revision: 1ea7ced7eebf26aba1938d84ba99e2e6421c503c

URL: https://github.com/llvm/llvm-project/commit/1ea7ced7eebf26aba1938d84ba99e2e6421c503c
DIFF: https://github.com/llvm/llvm-project/commit/1ea7ced7eebf26aba1938d84ba99e2e6421c503c.diff

LOG: [mlir][py] Enable disabling loading all registered (#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.

Added: 
    

Modified: 
    mlir/python/mlir/_mlir_libs/__init__.py

Removed: 
    


################################################################################
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