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

@@ -59,9 +59,22 @@ bool PyGlobals::loadDialectModule(llvm::StringRef dialectNamespace) {
 
   if (loaded.is_none())
     return false;
+
+  // 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;

mlir/test/python/multithreaded_tests.py

@@ -16,7 +16,7 @@
 2) Tests generation: we use existing tests: test/python/ir/*.py,
 test/python/dialects/*.py, etc to generate multi-threaded tests.
 In details, we perform the following:
-a) we define a list of source tests to be used to generate multi-threaded tests, see `test_modules`.
+a) we define a list of source tests to be used to generate multi-threaded tests, see `TEST_MODULES`.
 b) we define `TestAllMultiThreaded` class and add existing tests to the class. See `add_existing_tests` method.
 c) for each test file, we copy and modify it: test/python/ir/affine_expr.py -> /tmp/ir/affine_expr.py.
 In order to import the test file as python module, we remove all executing functions, like
@@ -376,10 +376,9 @@ def run_construct_and_print_in_module(f):
     return f
 
 
-test_modules = [
+TEST_MODULES = [
     ("execution_engine", run),  # Fail,
     ("pass_manager", run),  # Fail
-
     ("dialects/affine", run_with_insertion_point_v2),  # Pass
     ("dialects/func", run_with_insertion_point_v2),  # Pass
     ("dialects/arith_dialect", run),  # Pass
@@ -410,32 +409,39 @@ def run_construct_and_print_in_module(f):
     ("dialects/transform_bufferization_ext", run_with_insertion_point_v2),  # Pass
     # ("dialects/transform_extras", ),  # Needs a more complicated execution schema
     ("dialects/transform_gpu_ext", run_transform_tensor_ext),  # Pass
-    ("dialects/transform_interpreter", run_with_context_and_location, ["print_", "transform_options", "failed", "include"]),  # Fail
-    ("dialects/transform_loop_ext", run_with_insertion_point_v2, ["loopOutline"]),  # Pass
+    (
+        "dialects/transform_interpreter",
+        run_with_context_and_location,
+        ["print_", "transform_options", "failed", "include"],
+    ),  # Fail
+    (
+        "dialects/transform_loop_ext",
+        run_with_insertion_point_v2,
+        ["loopOutline"],
+    ),  # Pass
     ("dialects/transform_memref_ext", run_with_insertion_point_v2),  # Pass
     ("dialects/transform_nvgpu_ext", run_with_insertion_point_v2),  # Pass
     ("dialects/transform_sparse_tensor_ext", run_transform_tensor_ext),  # Pass
     ("dialects/transform_structured_ext", run_transform_structured_ext),  # Pass
     ("dialects/transform_tensor_ext", run_transform_tensor_ext),  # Pass
-    ("dialects/transform_vector_ext", run_apply_patterns, ["configurable_patterns"]),  # Pass
+    (
+        "dialects/transform_vector_ext",
+        run_apply_patterns,
+        ["configurable_patterns"],
+    ),  # Pass
     ("dialects/transform", run_with_insertion_point_v3),  # Pass
     ("dialects/vector", run_with_context_and_location),  # Pass
-
     ("dialects/gpu/dialect", run_with_context_and_location),  # Pass
     ("dialects/gpu/module-to-binary-nvvm", run_with_context_and_location),  # Pass
     ("dialects/gpu/module-to-binary-rocdl", run_with_context_and_location),  # Fail
-
     ("dialects/linalg/ops", run),  # Pass
     # TO ADD: No proper tests in this dialects/linalg/opsdsl/*
     # ("dialects/linalg/opsdsl/*", ...),  #
-
     ("dialects/sparse_tensor/dialect", run),  # Pass
     ("dialects/sparse_tensor/passes", run),  # Pass
-
     ("integration/dialects/pdl", run_construct_and_print_in_module),  # Pass
     ("integration/dialects/transform", run_construct_and_print_in_module),  # Pass
     ("integration/dialects/linalg/opsrun", run),  # Fail
-
     ("ir/affine_expr", run),  # Pass
     ("ir/affine_map", run),  # Pass
     ("ir/array_attributes", run),  # Pass
@@ -456,16 +462,18 @@ def run_construct_and_print_in_module(f):
     ("ir/value", run),  # Pass
 ]
 
-tests_to_skip = [
+TESTS_TO_SKIP = [
     "test_execution_engine__testNanoTime_multi_threaded",  # testNanoTime can't run in multiple threads, even with GIL
     "test_execution_engine__testSharedLibLoad_multi_threaded",  # testSharedLibLoad can't run in multiple threads, even with GIL
     "test_dialects_arith_dialect__testArithValue_multi_threaded",  # RuntimeError: Value caster is already registered: <class 'dialects/arith_dialect.testArithValue.<locals>.ArithValue'>, even with GIL
     "test_ir_dialects__testAppendPrefixSearchPath_multi_threaded",  # PyGlobals::setDialectSearchPrefixes is not thread-safe, even with GIL. Strange usage of static PyGlobals vs python exposed _cext.globals
     "test_ir_value__testValueCasters_multi_threaded_multi_threaded",  # RuntimeError: Value caster is already registered: <function testValueCasters.<locals>.dont_cast_int, even with GIL
+    "test_ir_operation_testKnownOpView_multi_threaded_multi_threaded",  # uses register_operation method in the test
+    "test_dialects_transform_structured_ext__testMatchInterfaceEnumReplaceAttributeBuilder_multi_threaded",  # uses register_attribute_builder method in the test
 ]
 
 
-tests_to_xfail = [
+TESTS_TO_XFAIL = [
     # execution_engine tests, ctypes related data-races, may be false-positive as libffi was not instrumented with TSAN
     "test_execution_engine__testBF16Memref_multi_threaded",
     "test_execution_engine__testBasicCallback_multi_threaded",
@@ -476,31 +484,28 @@ def run_construct_and_print_in_module(f):
     "test_execution_engine__testF8E5M2Memref_multi_threaded",
     "test_execution_engine__testInvalidModule_multi_threaded",
     "test_execution_engine__testInvokeFloatAdd_multi_threaded",
+    "test_execution_engine__testInvokeVoid_multi_threaded",
     "test_execution_engine__testMemrefAdd_multi_threaded",
     "test_execution_engine__testRankedMemRefCallback_multi_threaded",
     "test_execution_engine__testRankedMemRefWithOffsetCallback_multi_threaded",
     "test_execution_engine__testUnrankedMemRefCallback_multi_threaded",
     "test_execution_engine__testUnrankedMemRefWithOffsetCallback_multi_threaded",
-
     # pass_manager tests
     "test_pass_manager__testPrintIrAfterAll_multi_threaded",  # IRPrinterInstrumentation::runAfterPass is not thread-safe
     "test_pass_manager__testPrintIrBeforeAndAfterAll_multi_threaded",  # IRPrinterInstrumentation::runBeforePass is not thread-safe
     "test_pass_manager__testPrintIrLargeLimitElements_multi_threaded",  # IRPrinterInstrumentation::runAfterPass is not thread-safe
     "test_pass_manager__testPrintIrTree_multi_threaded",  # IRPrinterInstrumentation::runAfterPass is not thread-safe
     "test_pass_manager__testRunPipeline_multi_threaded",  # PrintOpStatsPass::printSummary is not thread-safe
-
     # dialects tests
     "test_dialects_memref__testSubViewOpInferReturnTypeExtensiveSlicing_multi_threaded",  # Related to ctypes data races
     "test_dialects_transform_interpreter__include_multi_threaded",  # mlir::transform::PrintOp::apply(mlir::transform::TransformRewriter...) is not thread-safe
     "test_dialects_transform_interpreter__print_other_multi_threaded",  # Fatal Python error: Aborted or mlir::transform::PrintOp::apply(mlir::transform::TransformRewriter...) is not thread-safe
     "test_dialects_transform_interpreter__print_self_multi_threaded",  # mlir::transform::PrintOp::apply(mlir::transform::TransformRewriter...) is not thread-safe
     "test_dialects_transform_interpreter__transform_options_multi_threaded",  # mlir::transform::PrintOp::apply(mlir::transform::TransformRewriter...) is not thread-safe
     "test_dialects_gpu_module-to-binary-rocdl__testGPUToASMBin_multi_threaded",  # Due to global llvm-project/llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp::GCNTrackers variable mutation
-
     # integration tests
     "test_integration_dialects_linalg_opsrun__test_elemwise_builtin_multi_threaded",  # Related to ctypes data races
     "test_integration_dialects_linalg_opsrun__test_elemwise_generic_multi_threaded",  # Related to ctypes data races
-
     # IR tests
     "test_ir_diagnostic_handler__testDiagnosticCallbackException_multi_threaded",  # mlirEmitError is not thread-safe
     "test_ir_module__testParseSuccess_multi_threaded",  #  mlirOperationDump is not thread-safe
@@ -509,7 +514,7 @@ def run_construct_and_print_in_module(f):
 ]
 
 
-def add_existing_tests(test_prefix: str = "_original_test"):
+def add_existing_tests(test_modules, test_prefix: str = "_original_test"):
     def decorator(test_cls):
         this_folder = Path(__file__).parent.absolute()
         test_cls.output_folder = tempfile.TemporaryDirectory()
@@ -531,16 +536,22 @@ def decorator(test_cls):
             test_mod = import_from_path(test_module_name, dst_filepath)
             for attr_name in dir(test_mod):
                 is_test_fn = test_pattern is None and attr_name.startswith("test")
-                is_test_fn |= test_pattern is not None and any([p in attr_name for p in test_pattern])
+                is_test_fn |= test_pattern is not None and any(
+                    [p in attr_name for p in test_pattern]
+                )
                 if is_test_fn:
                     obj = getattr(test_mod, attr_name)
                     if callable(obj):
                         test_name = f"{test_prefix}_{test_module_name.replace('/', '_')}__{attr_name}"
-                        def wrapped_test_fn(self, *args, __test_fn__=obj, __exec_fn__=exec_fn, **kwargs):
+
+                        def wrapped_test_fn(
+                            self, *args, __test_fn__=obj, __exec_fn__=exec_fn, **kwargs
+                        ):
                             __exec_fn__(__test_fn__)
 
                         setattr(test_cls, test_name, wrapped_test_fn)
         return test_cls
+
     return decorator
 
 
@@ -553,6 +564,7 @@ def multi_threaded(
     multithreaded_test_postfix: str = "_multi_threaded",
 ):
     """Decorator that runs a test in a multi-threaded environment."""
+
     def decorator(test_cls):
         for name, test_fn in test_cls.__dict__.copy().items():
             if not (name.startswith(test_prefix) and callable(test_fn)):
@@ -566,7 +578,9 @@ def decorator(test_cls):
                 ):
                     continue
 
-            def multi_threaded_test_fn(self, capfd, *args, __test_fn__=test_fn, **kwargs):
+            def multi_threaded_test_fn(
+                self, capfd, *args, __test_fn__=test_fn, **kwargs
+            ):
                 barrier = threading.Barrier(num_workers)
 
                 def closure():
@@ -590,7 +604,9 @@ def closure():
                 captured = capfd.readouterr()
                 if len(captured.err) > 0:
                     if "ThreadSanitizer" in captured.err:
-                        raise RuntimeError(f"ThreadSanitizer reported warnings:\n{captured.err}")
+                        raise RuntimeError(
+                            f"ThreadSanitizer reported warnings:\n{captured.err}"
+                        )
                     else:
                         pass
                         # There are tests that write to stderr, we should ignore them
@@ -603,18 +619,19 @@ def closure():
             setattr(test_cls, test_new_name, multi_threaded_test_fn)
 
         return test_cls
+
     return decorator
 
 
 @multi_threaded(
-    num_workers=6,
+    num_workers=8,
     num_runs=20,
-    skip_tests=tests_to_skip,
-    xfail_tests=tests_to_xfail,
+    skip_tests=TESTS_TO_SKIP,
+    xfail_tests=TESTS_TO_XFAIL,
 )
-@add_existing_tests(test_prefix="_original_test")
+@add_existing_tests(test_modules=TEST_MODULES, test_prefix="_original_test")
 class TestAllMultiThreaded:
-    @pytest.fixture(scope='class')
+    @pytest.fixture(scope="class")
     def teardown(self):
         if hasattr(self, "output_folder"):
             self.output_folder.cleanup()