Updated locks and added docs

Author: vfdev-5

mlir/docs/Bindings/Python.md

@@ -1187,3 +1187,43 @@ or nanobind and
 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.
+
+## Free-threading (No-GIL) support
+
+Free-threading or no-GIL support refers to CPython interpreter (>=3.13) with Global Interpreter Lock made optional. For details on the topic, please check [PEP-703](https://peps.python.org/pep-0703/) and the this [link](https://py-free-threading.github.io/).
+
+MLIR Python PyBind11 bindings are made free-threading compatible with exceptions (discussed below) in the following sense: it is safe to work in multiple threads with **independent** contexts/modules. Below we show an example code of safe usage:
+
+```python
+# python3.13t example.py
+import concurrent.futures
+
+import mlir.dialects.arith as arith
+from mlir.ir import Context, Location, Module, IntegerType, InsertionPoint
+
+
+def func(py_value):
+    with Context() as ctx:
+        module = Module.create(loc=Location.file("foo.txt", 0, 0))
+
+        dtype = IntegerType.get_signless(64)
+        with InsertionPoint(module.body), Location.name("a"):
+            arith.constant(dtype, py_value)
+
+    return module
+
+
+num_workers = 8
+with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
+    futures = []
+    for i in range(num_workers):
+        futures.append(executor.submit(func, i))
+    assert len(list(f.result() for f in futures)) == num_workers
+```
+
+The exceptions to the free-threading compatibility:
+- registration methods and decorators, e.g. `register_dialect`, `register_operation`, `register_dialect`, `register_attribute_builder`, ...
+- `ctypes` is unsafe
+- IR printing is unsafe
+
+For details, please see the list of xfailed tests in `mlir/test/python/multithreaded_tests.py`.

mlir/lib/Bindings/Python/Globals.h

@@ -36,20 +36,6 @@ 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() {

mlir/lib/Bindings/Python/IRCore.cpp

@@ -198,9 +198,7 @@ py::object classmethod(Func f, Args... args) {
 static py::object
 createCustomDialectWrapper(const std::string &dialectNamespace,
                            py::object dialectDescriptor) {
-  auto dialectClass = PyGlobals::withInstance([&](PyGlobals& instance) {
-    return instance.lookupDialectClass(dialectNamespace);
-  });
+  auto dialectClass = PyGlobals::get().lookupDialectClass(dialectNamespace);
   if (!dialectClass) {
     // Use the base class.
     return py::cast(PyDialect(std::move(dialectDescriptor)));
@@ -309,20 +307,15 @@ struct PyAttrBuilderMap {
     return PyGlobals::get().lookupAttributeBuilder(attributeKind).has_value();
   }
   static py::function dunderGetItemNamed(const std::string &attributeKind) {
-    auto builder = PyGlobals::withInstance([&](PyGlobals& instance) {
-        return instance.lookupAttributeBuilder(attributeKind);
-    });
+    auto builder = PyGlobals::get().lookupAttributeBuilder(attributeKind);
     if (!builder)
       throw py::key_error(attributeKind);
     return *builder;
   }
   static void dunderSetItemNamed(const std::string &attributeKind,
                                 py::function func, bool replace) {
-    PyGlobals::withInstance([&](PyGlobals& instance) {
-        instance.registerAttributeBuilder(attributeKind, std::move(func),
-                                          replace);
-        return 0;
-    });
+    PyGlobals::get().registerAttributeBuilder(attributeKind, std::move(func),
+                                              replace);
   }
 
   static void bind(py::module &m) {
@@ -1613,10 +1606,8 @@ py::object PyOperation::createOpView() {
   checkValid();
   MlirIdentifier ident = mlirOperationGetName(get());
   MlirStringRef identStr = mlirIdentifierStr(ident);
-  auto operationCls = PyGlobals::withInstance([&](PyGlobals& instance){
-      return instance.lookupOperationClass(
-          StringRef(identStr.data, identStr.length));
-  });
+  auto operationCls = PyGlobals::get().lookupOperationClass(
+      StringRef(identStr.data, identStr.length));
   if (operationCls)
     return PyOpView::constructDerived(*operationCls, *getRef().get());
   return py::cast(PyOpView(getRef().getObject()));
@@ -2067,9 +2058,7 @@ pybind11::object PyValue::maybeDownCast() {
   assert(!mlirTypeIDIsNull(mlirTypeID) &&
          "mlirTypeID was expected to be non-null.");
   std::optional<pybind11::function> valueCaster =
-      PyGlobals::withInstance([&](PyGlobals& instance) {
-        return instance.lookupValueCaster(mlirTypeID, mlirTypeGetDialect(type));
-  });
+      PyGlobals::get().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);
@@ -3548,10 +3537,8 @@ void mlir::python::populateIRCore(py::module &m) {
         assert(!mlirTypeIDIsNull(mlirTypeID) &&
                "mlirTypeID was expected to be non-null.");
         std::optional<pybind11::function> typeCaster =
-            PyGlobals::withInstance([&](PyGlobals& instance){
-              return instance.lookupTypeCaster(mlirTypeID,
-                                               mlirAttributeGetDialect(self));
-            });
+            PyGlobals::get().lookupTypeCaster(mlirTypeID,
+                                              mlirAttributeGetDialect(self));
         if (!typeCaster)
           return py::cast(self);
         return typeCaster.value()(self);
@@ -3649,10 +3636,8 @@ void mlir::python::populateIRCore(py::module &m) {
              assert(!mlirTypeIDIsNull(mlirTypeID) &&
                     "mlirTypeID was expected to be non-null.");
             std::optional<pybind11::function> typeCaster =
-                PyGlobals::withInstance([&](PyGlobals& instance){
-                  return instance.lookupTypeCaster(mlirTypeID,
-                                                   mlirTypeGetDialect(self));
-                });
+                PyGlobals::get().lookupTypeCaster(mlirTypeID,
+                                                  mlirTypeGetDialect(self));
              if (!typeCaster)
                return py::cast(self);
              return typeCaster.value()(self);

mlir/lib/Bindings/Python/IRModule.cpp

@@ -37,7 +37,21 @@ PyGlobals::PyGlobals() {
 PyGlobals::~PyGlobals() { instance = nullptr; }
 
 bool PyGlobals::loadDialectModule(llvm::StringRef dialectNamespace) {
-  if (loadedDialectModules.contains(dialectNamespace))
+  auto already_loaded = false;
+#ifdef Py_GIL_DISABLED
+  {
+    auto &lock = getLock();
+    PyMutex_Lock(&lock);
+#endif
+
+    already_loaded = loadedDialectModules.contains(dialectNamespace);
+
+#ifdef Py_GIL_DISABLED
+    PyMutex_Unlock(&lock);
+  }
+#endif
+
+  if (already_loaded)
     return true;
   // Since re-entrancy is possible, make a copy of the search prefixes.
   std::vector<std::string> localSearchPrefixes = dialectSearchPrefixes;
@@ -59,9 +73,24 @@ bool PyGlobals::loadDialectModule(llvm::StringRef dialectNamespace) {
 
   if (loaded.is_none())
     return false;
-  // Note: Iterator cannot be shared from prior to loading, since re-entrancy
-  // may have occurred, which may do anything.
-  loadedDialectModules.insert(dialectNamespace);
+
+  // We should use a lock in free-threading as loadDialectModule can be implicitly called by
+  // python functions executed by in multiple threads context (e.g. lookupValueCaster).
+#ifdef Py_GIL_DISABLED
+  {
+    auto &lock = getLock();
+    PyMutex_Lock(&lock);
+#endif
+
+    // Note: Iterator cannot be shared from prior to loading, since re-entrancy
+    // may have occurred, which may do anything.
+    loadedDialectModules.insert(dialectNamespace);
+
+#ifdef Py_GIL_DISABLED
+    PyMutex_Unlock(&lock);
+  }
+#endif
+
   return true;
 }
 

mlir/lib/Bindings/Python/MainModule.cpp

@@ -74,13 +74,7 @@ PYBIND11_MODULE(_mlir, m, py::mod_gil_not_used()) {
               std::string operationName =
                   opClass.attr("OPERATION_NAME").cast<std::string>();
 
-              // Use PyGlobals::withInstance instead of PyGlobals::get()
-              // to prevent data race in multi-threaded context
-              // Error raised in ir/opeation.py testKnownOpView test
-              PyGlobals::withInstance([&](PyGlobals& instance) {
-                  instance.registerOperationImpl(operationName, opClass, replace);
-                  return 0;
-              });
+              PyGlobals::get().registerOperationImpl(operationName, opClass, replace);
               // Dict-stuff the new opClass by name onto the dialect class.
               py::object opClassName = opClass.attr("__name__");
               dialectClass.attr(opClassName) = opClass;