[skip-ci] Added lock on PyGlobals::get and PyMlirContext liveContexts

WIP on adding multithreaded_tests

Author: vfdev-5

mlir/docs/Bindings/Python.md

@@ -1035,7 +1035,7 @@ class ConstantOp(_ods_ir.OpView):
     ...
 ```
 
-expects `value` to be a `TypedAttr` (e.g., `IntegerAttr` or `FloatAttr`). 
+expects `value` to be a `TypedAttr` (e.g., `IntegerAttr` or `FloatAttr`).
 Thus, a natural extension is a builder that accepts a MLIR type and a Python value and instantiates the appropriate `TypedAttr`:
 
 ```python
@@ -1181,9 +1181,9 @@ make the passes available along with the dialect.
 Dialect functionality other than IR objects or passes, such as helper functions,
 can be exposed to Python similarly to attributes and types. C API is expected to
 exist for this functionality, which can then be wrapped using pybind11 and
-`[include/mlir/Bindings/Python/PybindAdaptors.h](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/PybindAdaptors.h)`,
+[`include/mlir/Bindings/Python/PybindAdaptors.h`](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/PybindAdaptors.h),
 or nanobind and
-`[include/mlir/Bindings/Python/NanobindAdaptors.h](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h)`
+[`include/mlir/Bindings/Python/NanobindAdaptors.h`](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h)
 utilities to connect to the rest of Python API. The bindings can be located in a
 separate module or in the same module as attributes and types, and
 loaded along with the dialect.

mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp

@@ -12,9 +12,11 @@
 #include "Standalone-c/Dialects.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 
+namespace py = pybind11;
+
 using namespace mlir::python::adaptors;
 
-PYBIND11_MODULE(_standaloneDialectsPybind11, m) {
+PYBIND11_MODULE(_standaloneDialects, m, py::mod_gil_not_used()) {
   //===--------------------------------------------------------------------===//
   // standalone dialect
   //===--------------------------------------------------------------------===//

mlir/lib/Bindings/Python/Globals.h

@@ -36,6 +36,20 @@ class PyGlobals {
     return *instance;
   }
 
+  template<typename F>
+  static inline auto withInstance(const F& cb) -> decltype(cb(get())) {
+    auto &instance = get();
+#ifdef Py_GIL_DISABLED
+    auto &lock = getLock();
+    PyMutex_Lock(&lock);
+#endif
+    auto result = cb(instance);
+#ifdef Py_GIL_DISABLED
+    PyMutex_Unlock(&lock);
+#endif
+    return result;
+  }
+
   /// Get and set the list of parent modules to search for dialect
   /// implementation classes.
   std::vector<std::string> &getDialectSearchPrefixes() {
@@ -125,6 +139,14 @@ class PyGlobals {
   /// Set of dialect namespaces that we have attempted to import implementation
   /// modules for.
   llvm::StringSet<> loadedDialectModules;
+
+#ifdef Py_GIL_DISABLED
+  static PyMutex &getLock() {
+    static PyMutex lock;
+    return lock;
+  }
+#endif
+
 };
 
 } // namespace python

mlir/lib/Bindings/Python/IRCore.cpp

@@ -198,7 +198,9 @@ py::object classmethod(Func f, Args... args) {
 static py::object
 createCustomDialectWrapper(const std::string &dialectNamespace,
                            py::object dialectDescriptor) {
-  auto dialectClass = PyGlobals::get().lookupDialectClass(dialectNamespace);
+  auto dialectClass = PyGlobals::withInstance([&](PyGlobals& instance) {
+    return instance.lookupDialectClass(dialectNamespace);
+  });
   if (!dialectClass) {
     // Use the base class.
     return py::cast(PyDialect(std::move(dialectDescriptor)));
@@ -601,16 +603,23 @@ class PyOpOperandIterator {
 
 PyMlirContext::PyMlirContext(MlirContext context) : context(context) {
   py::gil_scoped_acquire acquire;
-  auto &liveContexts = getLiveContexts();
-  liveContexts[context.ptr] = this;
+  withLiveContexts([&](LiveContextMap& liveContexts) {
+    liveContexts[context.ptr] = this;
+    return this;
+  });
 }
 
 PyMlirContext::~PyMlirContext() {
   // Note that the only public way to construct an instance is via the
   // forContext method, which always puts the associated handle into
   // liveContexts.
   py::gil_scoped_acquire acquire;
-  getLiveContexts().erase(context.ptr);
+
+  withLiveContexts([&](LiveContextMap& liveContexts) {
+    liveContexts.erase(context.ptr);
+    return this;
+  });
+
   mlirContextDestroy(context);
 }
 
@@ -632,27 +641,32 @@ PyMlirContext *PyMlirContext::createNewContextForInit() {
 
 PyMlirContextRef PyMlirContext::forContext(MlirContext context) {
   py::gil_scoped_acquire acquire;
-  auto &liveContexts = getLiveContexts();
-  auto it = liveContexts.find(context.ptr);
-  if (it == liveContexts.end()) {
-    // Create.
-    PyMlirContext *unownedContextWrapper = new PyMlirContext(context);
-    py::object pyRef = py::cast(unownedContextWrapper);
-    assert(pyRef && "cast to py::object failed");
-    liveContexts[context.ptr] = unownedContextWrapper;
-    return PyMlirContextRef(unownedContextWrapper, std::move(pyRef));
-  }
-  // Use existing.
-  py::object pyRef = py::cast(it->second);
-  return PyMlirContextRef(it->second, std::move(pyRef));
+  return withLiveContexts([&](LiveContextMap& liveContexts) {
+    auto it = liveContexts.find(context.ptr);
+    if (it == liveContexts.end()) {
+      // Create.
+      PyMlirContext *unownedContextWrapper = new PyMlirContext(context);
+      py::object pyRef = py::cast(unownedContextWrapper);
+      assert(pyRef && "cast to py::object failed");
+      liveContexts[context.ptr] = unownedContextWrapper;
+      return PyMlirContextRef(unownedContextWrapper, std::move(pyRef));
+    }
+    // Use existing.
+    py::object pyRef = py::cast(it->second);
+    return PyMlirContextRef(it->second, std::move(pyRef));
+  });
 }
 
 PyMlirContext::LiveContextMap &PyMlirContext::getLiveContexts() {
   static LiveContextMap liveContexts;
   return liveContexts;
 }
 
-size_t PyMlirContext::getLiveCount() { return getLiveContexts().size(); }
+size_t PyMlirContext::getLiveCount() {
+  return withLiveContexts([&](LiveContextMap& liveContexts) {
+    return liveContexts.size();
+  });
+}
 
 size_t PyMlirContext::getLiveOperationCount() { return liveOperations.size(); }
 
@@ -1556,8 +1570,10 @@ py::object PyOperation::createOpView() {
   checkValid();
   MlirIdentifier ident = mlirOperationGetName(get());
   MlirStringRef identStr = mlirIdentifierStr(ident);
-  auto operationCls = PyGlobals::get().lookupOperationClass(
-      StringRef(identStr.data, identStr.length));
+  auto operationCls = PyGlobals::withInstance([&](PyGlobals& instance){
+      return instance.lookupOperationClass(
+          StringRef(identStr.data, identStr.length));
+  });
   if (operationCls)
     return PyOpView::constructDerived(*operationCls, *getRef().get());
   return py::cast(PyOpView(getRef().getObject()));
@@ -2008,7 +2024,9 @@ pybind11::object PyValue::maybeDownCast() {
   assert(!mlirTypeIDIsNull(mlirTypeID) &&
          "mlirTypeID was expected to be non-null.");
   std::optional<pybind11::function> valueCaster =
-      PyGlobals::get().lookupValueCaster(mlirTypeID, mlirTypeGetDialect(type));
+      PyGlobals::withInstance([&](PyGlobals& instance) {
+        return instance.lookupValueCaster(mlirTypeID, mlirTypeGetDialect(type));
+  });
   // py::return_value_policy::move means use std::move to move the return value
   // contents into a new instance that will be owned by Python.
   py::object thisObj = py::cast(this, py::return_value_policy::move);
@@ -3487,8 +3505,10 @@ void mlir::python::populateIRCore(py::module &m) {
         assert(!mlirTypeIDIsNull(mlirTypeID) &&
                "mlirTypeID was expected to be non-null.");
         std::optional<pybind11::function> typeCaster =
-            PyGlobals::get().lookupTypeCaster(mlirTypeID,
-                                              mlirAttributeGetDialect(self));
+            PyGlobals::withInstance([&](PyGlobals& instance){
+              return instance.lookupTypeCaster(mlirTypeID,
+                                               mlirAttributeGetDialect(self));
+            });
         if (!typeCaster)
           return py::cast(self);
         return typeCaster.value()(self);
@@ -3585,9 +3605,11 @@ void mlir::python::populateIRCore(py::module &m) {
              MlirTypeID mlirTypeID = mlirTypeGetTypeID(self);
              assert(!mlirTypeIDIsNull(mlirTypeID) &&
                     "mlirTypeID was expected to be non-null.");
-             std::optional<pybind11::function> typeCaster =
-                 PyGlobals::get().lookupTypeCaster(mlirTypeID,
+            std::optional<pybind11::function> typeCaster =
+                PyGlobals::withInstance([&](PyGlobals& instance){
+                  return instance.lookupTypeCaster(mlirTypeID,
                                                    mlirTypeGetDialect(self));
+                });
              if (!typeCaster)
                return py::cast(self);
              return typeCaster.value()(self);

mlir/lib/Bindings/Python/IRModule.h

@@ -263,6 +263,27 @@ class PyMlirContext {
   using LiveContextMap = llvm::DenseMap<void *, PyMlirContext *>;
   static LiveContextMap &getLiveContexts();
 
+#ifdef Py_GIL_DISABLED
+  static PyMutex &getLock() {
+    static PyMutex lock;
+    return lock;
+  }
+#endif
+
+  template<typename F>
+  static inline auto withLiveContexts(const F& cb) -> decltype(cb(getLiveContexts())) {
+    auto &liveContexts = getLiveContexts();
+#ifdef Py_GIL_DISABLED
+    auto &lock = getLock();
+    PyMutex_Lock(&lock);
+#endif
+    auto result = cb(liveContexts);
+#ifdef Py_GIL_DISABLED
+    PyMutex_Unlock(&lock);
+#endif
+    return result;
+  }
+
   // Interns all live modules associated with this context. Modules tracked
   // in this map are valid. When a module is invalidated, it is removed
   // from this map, and while it still exists as an instance, any

mlir/test/python/execution_engine.py

@@ -306,7 +306,7 @@ def callback(a):
         log(arr)
 
     with Context():
-        # The module takes a subview of the argument memref, casts it to an unranked memref and 
+        # The module takes a subview of the argument memref, casts it to an unranked memref and
         # calls the callback with it.
         module = Module.parse(
             r"""

mlir/test/python/lib/PythonTestModulePybind11.cpp

@@ -23,7 +23,7 @@ static bool mlirTypeIsARankedIntegerTensor(MlirType t) {
          mlirTypeIsAInteger(mlirShapedTypeGetElementType(t));
 }
 
-PYBIND11_MODULE(_mlirPythonTestPybind11, m) {
+PYBIND11_MODULE(_mlirPythonTest, m, py::mod_gil_not_used()) {
   m.def(
       "register_python_test_dialect",
       [](MlirContext context, bool load) {