[Mlir-commits] [mlir] Added PyThreadPool as wrapper around MlirLlvmThreadPool in MLIR python bindings (PR #130109)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Sat Mar 8 02:48:44 PST 2025


================
@@ -148,13 +148,18 @@ def process_initializer_module(module_name):
             break
 
     class Context(ir._BaseContext):
-        def __init__(self, load_on_create_dialects=None, *args, **kwargs):
+        def __init__(
+            self, load_on_create_dialects=None, thread_pool=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 thread_pool is None:
+                    self.enable_multithreading(True)
+                else:
+                    self.set_thread_pool(thread_pool)
----------------
vfdev-5 wrote:

Sounds good, I'll commit this:
```
    class Context(ir._BaseContext):
        def __init__(
            self, load_on_create_dialects=None, thread_pool=None, *args, **kwargs
        ):
            super().__init__(*args, **kwargs)
            self.append_dialect_registry(get_dialect_registry())
            for hook in post_init_hooks:
                hook(self)
            if disable_multithreading and thread_pool is not None:
                raise ValueError(
                    "Context constructor has given thread_pool argument, "
                    "but disable_multithreading flag is True. "
                    "Please, set thread_pool argument to None or "
                    "set disable_multithreading flag to False."
                )
            if not disable_multithreading:
```
If there are any other suggestions, please let me know

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


More information about the Mlir-commits mailing list