From 1ea7ced7eebf26aba1938d84ba99e2e6421c503c Mon Sep 17 00:00:00 2001 From: Jacques Pienaar Date: Mon, 25 Nov 2024 15:39:55 -0800 Subject: [PATCH] [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. --- 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 98dbbc6adf9c..c5cb22c6dccb 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 _site_initialize(): 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 _site_initialize(): 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