Added get_num_threads and _mlir_thread_pool_ptr methods to _BaseContext

Added thread_pool arg to the constructor: `mlir.ir.Context(thread_pool=tp)`

Author: vfdev-5

mlir/include/mlir-c/IR.h

@@ -162,6 +162,15 @@ MLIR_CAPI_EXPORTED bool mlirContextIsRegisteredOperation(MlirContext context,
 MLIR_CAPI_EXPORTED void mlirContextSetThreadPool(MlirContext context,
                                                  MlirLlvmThreadPool threadPool);
 
+/// Gets the number of threads of the thread pool of the context when
+/// multithreading is enabled. Returns 1 if no multithreading.
+MLIR_CAPI_EXPORTED unsigned mlirContextGetNumThreads(MlirContext context);
+
+/// Gets the thread pool of the context when enabled multithreading, otherwise
+/// an assertion is raised.
+MLIR_CAPI_EXPORTED MlirLlvmThreadPool
+mlirContextGetThreadPool(MlirContext context);
+
 //===----------------------------------------------------------------------===//
 // Dialect API.
 //===----------------------------------------------------------------------===//

mlir/lib/Bindings/Python/IRCore.cpp

@@ -2747,7 +2747,8 @@ void mlir::python::populateIRCore(nb::module_ &m) {
   // Expose DefaultThreadPool to python
   nb::class_<PyThreadPool>(m, "ThreadPool")
       .def("__init__", [](PyThreadPool &self) { new (&self) PyThreadPool(); })
-      .def("get_max_concurrency", &PyThreadPool::getMaxConcurrency);
+      .def("get_max_concurrency", &PyThreadPool::getMaxConcurrency)
+      .def("_mlir_thread_pool_ptr", &PyThreadPool::_mlir_thread_pool_ptr);
 
   nb::class_<PyMlirContext>(m, "_BaseContext")
       .def("__init__",
@@ -2822,8 +2823,23 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           nb::arg("enable"))
       .def("set_thread_pool",
            [](PyMlirContext &self, PyThreadPool &pool) {
+             // we should disable multi-threading first before setting
+             // new thread pool otherwise the assert in
+             // MLIRContext::setThreadPool will be raised.
+             mlirContextEnableMultithreading(self.get(), false);
              mlirContextSetThreadPool(self.get(), pool.get());
            })
+      .def("get_num_threads",
+           [](PyMlirContext &self) {
+             return mlirContextGetNumThreads(self.get());
+           })
+      .def("_mlir_thread_pool_ptr",
+           [](PyMlirContext &self) {
+             MlirLlvmThreadPool pool = mlirContextGetThreadPool(self.get());
+             std::stringstream ss;
+             ss << pool.ptr;
+             return ss.str();
+           })
       .def(
           "is_registered_operation",
           [](PyMlirContext &self, std::string &name) {

mlir/lib/Bindings/Python/IRModule.h

@@ -11,6 +11,7 @@
 #define MLIR_BINDINGS_PYTHON_IRMODULES_H
 
 #include <optional>
+#include <sstream>
 #include <utility>
 #include <vector>
 
@@ -172,6 +173,12 @@ class PyThreadPool {
   int getMaxConcurrency() const { return ownedThreadPool->getMaxConcurrency(); }
   MlirLlvmThreadPool get() { return wrap(ownedThreadPool.get()); }
 
+  std::string _mlir_thread_pool_ptr() const {
+    std::stringstream ss;
+    ss << ownedThreadPool.get();
+    return ss.str();
+  }
+
 private:
   std::unique_ptr<llvm::ThreadPoolInterface> ownedThreadPool;
 };

mlir/lib/CAPI/IR/IR.cpp

@@ -114,6 +114,14 @@ void mlirContextSetThreadPool(MlirContext context,
   unwrap(context)->setThreadPool(*unwrap(threadPool));
 }
 
+unsigned mlirContextGetNumThreads(MlirContext context) {
+  return unwrap(context)->getNumThreads();
+}
+
+MlirLlvmThreadPool mlirContextGetThreadPool(MlirContext context) {
+  return wrap(&unwrap(context)->getThreadPool());
+}
+
 //===----------------------------------------------------------------------===//
 // Dialect API.
 //===----------------------------------------------------------------------===//

mlir/python/mlir/_mlir_libs/__init__.py

@@ -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)
             if load_on_create_dialects is not None:
                 logger.debug(
                     "Loading all dialects from load_on_create_dialects arg %r",

mlir/test/python/ir/context_lifecycle.py

@@ -55,9 +55,18 @@
 tp = mlir.ir.ThreadPool()
 assert tp.get_max_concurrency() > 0
 c5 = mlir.ir.Context()
-c5.enable_multithreading(False)
 c5.set_thread_pool(tp)
+assert c5.get_num_threads() == tp.get_max_concurrency()
+assert c5._mlir_thread_pool_ptr() == tp._mlir_thread_pool_ptr()
 c6 = mlir.ir.Context()
-c6.enable_multithreading(False)
 c6.set_thread_pool(tp)
-assert mlir.ir.Context._get_live_count() == 2
+assert c6.get_num_threads() == tp.get_max_concurrency()
+assert c6._mlir_thread_pool_ptr() == tp._mlir_thread_pool_ptr()
+c7 = mlir.ir.Context(thread_pool=tp)
+assert c7.get_num_threads() == tp.get_max_concurrency()
+assert c7._mlir_thread_pool_ptr() == tp._mlir_thread_pool_ptr()
+assert mlir.ir.Context._get_live_count() == 3
+c5 = None
+c6 = None
+c7 = None
+gc.collect()