[Mlir-commits] [mlir] Added free-threading CPython mode support in MLIR Python bindings (PR #107103)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Dec 13 10:07:26 PST 2024
https://github.com/vfdev-5 updated https://github.com/llvm/llvm-project/pull/107103
>From 8ff7314c6a72624b3e7460760e336d17e57efe31 Mon Sep 17 00:00:00 2001
From: vfdev-5 <vfdev.5 at gmail.com>
Date: Tue, 3 Sep 2024 15:02:12 +0200
Subject: [PATCH 1/5] Added free-threading CPython mode support in Python
bindings - temporarily updated requirements
---
mlir/lib/Bindings/Python/AsyncPasses.cpp | 4 +++-
mlir/lib/Bindings/Python/DialectGPU.cpp | 2 +-
mlir/lib/Bindings/Python/DialectLLVM.cpp | 2 +-
mlir/lib/Bindings/Python/DialectLinalg.cpp | 2 +-
mlir/lib/Bindings/Python/DialectNVGPU.cpp | 2 +-
mlir/lib/Bindings/Python/DialectPDL.cpp | 2 +-
mlir/lib/Bindings/Python/DialectQuant.cpp | 2 +-
mlir/lib/Bindings/Python/DialectSparseTensor.cpp | 2 +-
mlir/lib/Bindings/Python/DialectTransform.cpp | 2 +-
mlir/lib/Bindings/Python/ExecutionEngineModule.cpp | 2 +-
mlir/lib/Bindings/Python/GPUPasses.cpp | 4 +++-
mlir/lib/Bindings/Python/LinalgPasses.cpp | 4 +++-
mlir/lib/Bindings/Python/MainModule.cpp | 2 +-
mlir/lib/Bindings/Python/RegisterEverything.cpp | 2 +-
mlir/lib/Bindings/Python/SparseTensorPasses.cpp | 4 +++-
mlir/lib/Bindings/Python/TransformInterpreter.cpp | 2 +-
mlir/python/requirements.txt | 6 ++++--
17 files changed, 28 insertions(+), 18 deletions(-)
diff --git a/mlir/lib/Bindings/Python/AsyncPasses.cpp b/mlir/lib/Bindings/Python/AsyncPasses.cpp
index b611a758dbbb37..d34a164b6e30c2 100644
--- a/mlir/lib/Bindings/Python/AsyncPasses.cpp
+++ b/mlir/lib/Bindings/Python/AsyncPasses.cpp
@@ -11,11 +11,13 @@
#include <pybind11/detail/common.h>
#include <pybind11/pybind11.h>
+namespace py = pybind11;
+
// -----------------------------------------------------------------------------
// Module initialization.
// -----------------------------------------------------------------------------
-PYBIND11_MODULE(_mlirAsyncPasses, m) {
+PYBIND11_MODULE(_mlirAsyncPasses, m, py::mod_gil_not_used()) {
m.doc() = "MLIR Async Dialect Passes";
// Register all Async passes on load.
diff --git a/mlir/lib/Bindings/Python/DialectGPU.cpp b/mlir/lib/Bindings/Python/DialectGPU.cpp
index 560a54bcd15919..5acfad007c3055 100644
--- a/mlir/lib/Bindings/Python/DialectGPU.cpp
+++ b/mlir/lib/Bindings/Python/DialectGPU.cpp
@@ -23,7 +23,7 @@ using namespace mlir::python::adaptors;
// Module initialization.
// -----------------------------------------------------------------------------
-PYBIND11_MODULE(_mlirDialectsGPU, m) {
+PYBIND11_MODULE(_mlirDialectsGPU, m, py::mod_gil_not_used()) {
m.doc() = "MLIR GPU Dialect";
//===-------------------------------------------------------------------===//
// AsyncTokenType
diff --git a/mlir/lib/Bindings/Python/DialectLLVM.cpp b/mlir/lib/Bindings/Python/DialectLLVM.cpp
index cccf1370b8cc87..6981e5ed6b8427 100644
--- a/mlir/lib/Bindings/Python/DialectLLVM.cpp
+++ b/mlir/lib/Bindings/Python/DialectLLVM.cpp
@@ -136,7 +136,7 @@ void populateDialectLLVMSubmodule(const pybind11::module &m) {
});
}
-PYBIND11_MODULE(_mlirDialectsLLVM, m) {
+PYBIND11_MODULE(_mlirDialectsLLVM, m, py::mod_gil_not_used()) {
m.doc() = "MLIR LLVM Dialect";
populateDialectLLVMSubmodule(m);
diff --git a/mlir/lib/Bindings/Python/DialectLinalg.cpp b/mlir/lib/Bindings/Python/DialectLinalg.cpp
index 2e54ebeb61fb10..118c4ab3f2f573 100644
--- a/mlir/lib/Bindings/Python/DialectLinalg.cpp
+++ b/mlir/lib/Bindings/Python/DialectLinalg.cpp
@@ -21,7 +21,7 @@ static void populateDialectLinalgSubmodule(py::module m) {
"op.");
}
-PYBIND11_MODULE(_mlirDialectsLinalg, m) {
+PYBIND11_MODULE(_mlirDialectsLinalg, m, py::mod_gil_not_used()) {
m.doc() = "MLIR Linalg dialect.";
populateDialectLinalgSubmodule(m);
diff --git a/mlir/lib/Bindings/Python/DialectNVGPU.cpp b/mlir/lib/Bindings/Python/DialectNVGPU.cpp
index 754e0a75b0abc7..4c962c403082cb 100644
--- a/mlir/lib/Bindings/Python/DialectNVGPU.cpp
+++ b/mlir/lib/Bindings/Python/DialectNVGPU.cpp
@@ -34,7 +34,7 @@ static void populateDialectNVGPUSubmodule(const pybind11::module &m) {
py::arg("ctx") = py::none());
}
-PYBIND11_MODULE(_mlirDialectsNVGPU, m) {
+PYBIND11_MODULE(_mlirDialectsNVGPU, m, py::mod_gil_not_used()) {
m.doc() = "MLIR NVGPU dialect.";
populateDialectNVGPUSubmodule(m);
diff --git a/mlir/lib/Bindings/Python/DialectPDL.cpp b/mlir/lib/Bindings/Python/DialectPDL.cpp
index 8d3f9a7ab1d6ac..e8542d5e777a65 100644
--- a/mlir/lib/Bindings/Python/DialectPDL.cpp
+++ b/mlir/lib/Bindings/Python/DialectPDL.cpp
@@ -100,7 +100,7 @@ void populateDialectPDLSubmodule(const pybind11::module &m) {
py::arg("context") = py::none());
}
-PYBIND11_MODULE(_mlirDialectsPDL, m) {
+PYBIND11_MODULE(_mlirDialectsPDL, m, py::mod_gil_not_used()) {
m.doc() = "MLIR PDL dialect.";
populateDialectPDLSubmodule(m);
}
diff --git a/mlir/lib/Bindings/Python/DialectQuant.cpp b/mlir/lib/Bindings/Python/DialectQuant.cpp
index 9a871f2c122d12..e974ecb66effb8 100644
--- a/mlir/lib/Bindings/Python/DialectQuant.cpp
+++ b/mlir/lib/Bindings/Python/DialectQuant.cpp
@@ -309,7 +309,7 @@ static void populateDialectQuantSubmodule(const py::module &m) {
});
}
-PYBIND11_MODULE(_mlirDialectsQuant, m) {
+PYBIND11_MODULE(_mlirDialectsQuant, m, py::mod_gil_not_used()) {
m.doc() = "MLIR Quantization dialect";
populateDialectQuantSubmodule(m);
diff --git a/mlir/lib/Bindings/Python/DialectSparseTensor.cpp b/mlir/lib/Bindings/Python/DialectSparseTensor.cpp
index a730bf500be98c..00d8482a91df7a 100644
--- a/mlir/lib/Bindings/Python/DialectSparseTensor.cpp
+++ b/mlir/lib/Bindings/Python/DialectSparseTensor.cpp
@@ -143,7 +143,7 @@ static void populateDialectSparseTensorSubmodule(const py::module &m) {
});
}
-PYBIND11_MODULE(_mlirDialectsSparseTensor, m) {
+PYBIND11_MODULE(_mlirDialectsSparseTensor, m, py::mod_gil_not_used()) {
m.doc() = "MLIR SparseTensor dialect.";
populateDialectSparseTensorSubmodule(m);
}
diff --git a/mlir/lib/Bindings/Python/DialectTransform.cpp b/mlir/lib/Bindings/Python/DialectTransform.cpp
index 6b57e652aa9d8b..df665dd66bdc28 100644
--- a/mlir/lib/Bindings/Python/DialectTransform.cpp
+++ b/mlir/lib/Bindings/Python/DialectTransform.cpp
@@ -117,7 +117,7 @@ void populateDialectTransformSubmodule(const pybind11::module &m) {
"Get the type this ParamType is associated with.");
}
-PYBIND11_MODULE(_mlirDialectsTransform, m) {
+PYBIND11_MODULE(_mlirDialectsTransform, m, py::mod_gil_not_used()) {
m.doc() = "MLIR Transform dialect.";
populateDialectTransformSubmodule(m);
}
diff --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
index b3df30583fc963..ddd81d1e7d592e 100644
--- a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
+++ b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
@@ -64,7 +64,7 @@ class PyExecutionEngine {
} // namespace
/// Create the `mlir.execution_engine` module here.
-PYBIND11_MODULE(_mlirExecutionEngine, m) {
+PYBIND11_MODULE(_mlirExecutionEngine, m, py::mod_gil_not_used()) {
m.doc() = "MLIR Execution Engine";
//----------------------------------------------------------------------------
diff --git a/mlir/lib/Bindings/Python/GPUPasses.cpp b/mlir/lib/Bindings/Python/GPUPasses.cpp
index e276a3ce3a56a0..bfc43e99946bb9 100644
--- a/mlir/lib/Bindings/Python/GPUPasses.cpp
+++ b/mlir/lib/Bindings/Python/GPUPasses.cpp
@@ -11,11 +11,13 @@
#include <pybind11/detail/common.h>
#include <pybind11/pybind11.h>
+namespace py = pybind11;
+
// -----------------------------------------------------------------------------
// Module initialization.
// -----------------------------------------------------------------------------
-PYBIND11_MODULE(_mlirGPUPasses, m) {
+PYBIND11_MODULE(_mlirGPUPasses, m, py::mod_gil_not_used()) {
m.doc() = "MLIR GPU Dialect Passes";
// Register all GPU passes on load.
diff --git a/mlir/lib/Bindings/Python/LinalgPasses.cpp b/mlir/lib/Bindings/Python/LinalgPasses.cpp
index 3f230207a42114..e3d8f237e2bab3 100644
--- a/mlir/lib/Bindings/Python/LinalgPasses.cpp
+++ b/mlir/lib/Bindings/Python/LinalgPasses.cpp
@@ -10,11 +10,13 @@
#include <pybind11/pybind11.h>
+namespace py = pybind11;
+
// -----------------------------------------------------------------------------
// Module initialization.
// -----------------------------------------------------------------------------
-PYBIND11_MODULE(_mlirLinalgPasses, m) {
+PYBIND11_MODULE(_mlirLinalgPasses, m, py::mod_gil_not_used()) {
m.doc() = "MLIR Linalg Dialect Passes";
// Register all Linalg passes on load.
diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp
index 7c27021902de31..168a939ff3cd8a 100644
--- a/mlir/lib/Bindings/Python/MainModule.cpp
+++ b/mlir/lib/Bindings/Python/MainModule.cpp
@@ -22,7 +22,7 @@ using namespace mlir::python;
// Module initialization.
// -----------------------------------------------------------------------------
-PYBIND11_MODULE(_mlir, m) {
+PYBIND11_MODULE(_mlir, m, py::mod_gil_not_used()) {
m.doc() = "MLIR Python Native Extension";
py::class_<PyGlobals>(m, "_Globals", py::module_local())
diff --git a/mlir/lib/Bindings/Python/RegisterEverything.cpp b/mlir/lib/Bindings/Python/RegisterEverything.cpp
index 6b2f6b0a6a3b86..5c5c6e32102712 100644
--- a/mlir/lib/Bindings/Python/RegisterEverything.cpp
+++ b/mlir/lib/Bindings/Python/RegisterEverything.cpp
@@ -9,7 +9,7 @@
#include "mlir-c/RegisterEverything.h"
#include "mlir/Bindings/Python/PybindAdaptors.h"
-PYBIND11_MODULE(_mlirRegisterEverything, m) {
+PYBIND11_MODULE(_mlirRegisterEverything, m, py::mod_gil_not_used()) {
m.doc() = "MLIR All Upstream Dialects, Translations and Passes Registration";
m.def("register_dialects", [](MlirDialectRegistry registry) {
diff --git a/mlir/lib/Bindings/Python/SparseTensorPasses.cpp b/mlir/lib/Bindings/Python/SparseTensorPasses.cpp
index 2a8e2b802df9c4..1bbdf2f3ccf4e5 100644
--- a/mlir/lib/Bindings/Python/SparseTensorPasses.cpp
+++ b/mlir/lib/Bindings/Python/SparseTensorPasses.cpp
@@ -10,11 +10,13 @@
#include <pybind11/pybind11.h>
+namespace py = pybind11;
+
// -----------------------------------------------------------------------------
// Module initialization.
// -----------------------------------------------------------------------------
-PYBIND11_MODULE(_mlirSparseTensorPasses, m) {
+PYBIND11_MODULE(_mlirSparseTensorPasses, m, py::mod_gil_not_used()) {
m.doc() = "MLIR SparseTensor Dialect Passes";
// Register all SparseTensor passes on load.
diff --git a/mlir/lib/Bindings/Python/TransformInterpreter.cpp b/mlir/lib/Bindings/Python/TransformInterpreter.cpp
index 0c8c0e0a965aa7..d55efe13319b79 100644
--- a/mlir/lib/Bindings/Python/TransformInterpreter.cpp
+++ b/mlir/lib/Bindings/Python/TransformInterpreter.cpp
@@ -100,7 +100,7 @@ static void populateTransformInterpreterSubmodule(py::module &m) {
py::arg("target"), py::arg("other"));
}
-PYBIND11_MODULE(_mlirTransformInterpreter, m) {
+PYBIND11_MODULE(_mlirTransformInterpreter, m, py::mod_gil_not_used()) {
m.doc() = "MLIR Transform dialect interpreter functionality.";
populateTransformInterpreterSubmodule(m);
}
diff --git a/mlir/python/requirements.txt b/mlir/python/requirements.txt
index ab8a9122919e19..449748bb02cc22 100644
--- a/mlir/python/requirements.txt
+++ b/mlir/python/requirements.txt
@@ -1,5 +1,7 @@
nanobind>=2.0, <3.0
numpy>=1.19.5, <=2.1.2
-pybind11>=2.10.0, <=2.13.6
+# pybind11>=2.14.0, <2.15.0
+# Temporarily set pybind11 version to master waiting the next release to 2.13.6
+pybind11 @ git+https://github.com/pybind/pybind11@master
PyYAML>=5.4.0, <=6.0.1
-ml_dtypes>=0.1.0, <=0.5.0 # provides several NumPy dtype extensions, including the bf16
+ml_dtypes>=0.5.0, <=0.6.0 # provides several NumPy dtype extensions, including the bf16
>From 6303c5b4dea86d5c0df63711899a1f4a7cc679a9 Mon Sep 17 00:00:00 2001
From: vfdev-5 <vfdev.5 at gmail.com>
Date: Fri, 20 Sep 2024 22:30:49 +0200
Subject: [PATCH 2/5] [skip-ci] Added lock on PyGlobals::get and PyMlirContext
liveContexts WIP on adding multithreaded_tests
---
mlir/docs/Bindings/Python.md | 6 +-
.../python/StandaloneExtensionPybind11.cpp | 4 +-
mlir/lib/Bindings/Python/Globals.h | 22 +++
mlir/lib/Bindings/Python/IRCore.cpp | 72 +++++---
mlir/lib/Bindings/Python/IRModule.h | 21 +++
mlir/test/python/execution_engine.py | 2 +-
.../python/lib/PythonTestModulePybind11.cpp | 2 +-
mlir/test/python/multithreaded_tests.py | 154 ++++++++++++++++++
mlir/test/python/test_to_remove.py | 46 ++++++
9 files changed, 298 insertions(+), 31 deletions(-)
create mode 100644 mlir/test/python/multithreaded_tests.py
create mode 100644 mlir/test/python/test_to_remove.py
diff --git a/mlir/docs/Bindings/Python.md b/mlir/docs/Bindings/Python.md
index a0bd1cac118bad..32df3310d811d7 100644
--- a/mlir/docs/Bindings/Python.md
+++ b/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.
diff --git a/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
index 397db4c20e7432..0bf8e150ee354c 100644
--- a/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
+++ b/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
//===--------------------------------------------------------------------===//
diff --git a/mlir/lib/Bindings/Python/Globals.h b/mlir/lib/Bindings/Python/Globals.h
index a022067f5c7e57..05400608ba9ffa 100644
--- a/mlir/lib/Bindings/Python/Globals.h
+++ b/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
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 3e96f8c60ba7cd..727cbc2e106d5b 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/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,8 +603,10 @@ 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() {
@@ -610,7 +614,12 @@ PyMlirContext::~PyMlirContext() {
// 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,19 +641,20 @@ 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() {
@@ -652,7 +662,11 @@ PyMlirContext::LiveContextMap &PyMlirContext::getLiveContexts() {
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);
diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h
index 172898cfda0c52..cbafed86c3ea85 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/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
diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py
index 7c375ce81de0eb..1435ae0f7569ed 100644
--- a/mlir/test/python/execution_engine.py
+++ b/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"""
diff --git a/mlir/test/python/lib/PythonTestModulePybind11.cpp b/mlir/test/python/lib/PythonTestModulePybind11.cpp
index 94a5f5178d16e8..9b63ea07a0eaba 100644
--- a/mlir/test/python/lib/PythonTestModulePybind11.cpp
+++ b/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) {
diff --git a/mlir/test/python/multithreaded_tests.py b/mlir/test/python/multithreaded_tests.py
new file mode 100644
index 00000000000000..fa545861cb0d7d
--- /dev/null
+++ b/mlir/test/python/multithreaded_tests.py
@@ -0,0 +1,154 @@
+import concurrent.futures
+import functools
+import importlib.util
+import sys
+import threading
+import tempfile
+
+from collections import defaultdict
+from contextlib import redirect_stderr, redirect_stdout
+from pathlib import Path
+from typing import Optional
+
+import pytest
+
+import mlir.dialects.arith as arith
+from mlir.ir import Context, Location, Module, IntegerType, F64Type, InsertionPoint
+
+
+
+def import_from_path(module_name: str, file_path: Path):
+ spec = importlib.util.spec_from_file_location(module_name, file_path)
+ module = importlib.util.module_from_spec(spec)
+ sys.modules[module_name] = module
+ spec.loader.exec_module(module)
+ return module
+
+
+def copy_and_update(src_filepath: Path, dst_filepath: Path):
+ # We should remove all calls like `run(testMethod)`
+ with open(src_filepath, "r") as reader, open(dst_filepath, "w") as writer:
+ while True:
+ src_line = reader.readline()
+ if len(src_line) == 0:
+ break
+ skip_lines = [
+ "run(",
+ "@run",
+ "@constructAndPrintInModule",
+ ]
+ if any(src_line.startswith(line) for line in skip_lines):
+ continue
+ writer.write(src_line)
+
+
+test_modules = [
+ "execution_engine",
+ # "pass_manager",
+]
+
+
+def add_existing_tests(test_prefix: str = "_original_test"):
+ def decorator(test_cls):
+ this_folder = Path(__file__).parent.absolute()
+ test_cls.output_folder = tempfile.TemporaryDirectory()
+ output_folder = Path(test_cls.output_folder.name)
+
+ for test_module_name in test_modules:
+ src_filepath = this_folder / f"{test_module_name}.py"
+ dst_filepath = (output_folder / f"{test_module_name}.py").absolute()
+ if not dst_filepath.parent.exists():
+ dst_filepath.parent.mkdir(parents=True)
+ copy_and_update(src_filepath, dst_filepath)
+ test_mod = import_from_path(test_module_name, dst_filepath)
+ for attr_name in dir(test_mod):
+ if attr_name.startswith("test"):
+ obj = getattr(test_mod, attr_name)
+ if callable(obj):
+ test_name = f"{test_prefix}_{test_module_name.replace('/', '_')}__{attr_name}"
+ def wrapped_test_fn(*args, __test_fn__=obj, **kwargs):
+ __test_fn__()
+
+ setattr(test_cls, test_name, wrapped_test_fn)
+ return test_cls
+ return decorator
+
+
+def multi_threaded(
+ num_workers: int,
+ num_runs: int = 5,
+ skip_tests: Optional[list[str]] = None,
+ test_prefix: str = "_original_test",
+):
+ """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)):
+ continue
+
+ name = f"test{name[len(test_prefix):]}"
+ if skip_tests is not None:
+ if any(test_name in name for test_name in skip_tests):
+ continue
+
+ def multi_threaded_test_fn(self, capfd, *args, __test_fn__=test_fn, **kwargs):
+ barrier = threading.Barrier(num_workers)
+
+ def closure():
+ barrier.wait()
+ for _ in range(num_runs):
+ __test_fn__(self, *args, **kwargs)
+
+ with concurrent.futures.ThreadPoolExecutor(
+ max_workers=num_workers
+ ) as executor:
+ futures = []
+ for _ in range(num_workers):
+ futures.append(executor.submit(closure))
+ # We should call future.result() to re-raise an exception if test has
+ # failed
+ list(f.result() for f in futures)
+
+ captured = capfd.readouterr()
+ if len(captured.err) > 0:
+ if "ThreadSanitizer" in captured.err:
+ raise RuntimeError(f"ThreadSanitizer reported warnings:\n{captured.err}")
+ else:
+ raise RuntimeError(f"Other error:\n{captured.err}")
+
+ setattr(test_cls, f"{name}_multi_threaded", multi_threaded_test_fn)
+
+ return test_cls
+ return decorator
+
+
+ at multi_threaded(num_workers=4, num_runs=10)
+ at add_existing_tests(test_prefix="_original_test")
+class TestAllMultiThreaded:
+ @pytest.fixture(scope='class')
+ def teardown(self):
+ self.output_folder.cleanup()
+
+ def _original_test_create_context(self):
+ with Context() as ctx:
+ print(ctx._get_live_count())
+ print(ctx._get_live_module_count())
+ print(ctx._get_live_operation_count())
+ print(ctx._get_live_operation_objects())
+ print(ctx._get_context_again() is ctx)
+ print(ctx._clear_live_operations())
+
+ def _original_test_create_module_with_consts(self):
+ py_values = [123, 234, 345]
+ 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_values[0])
+
+ with InsertionPoint(module.body), Location.name("b"):
+ arith.constant(dtype, py_values[1])
+
+ with InsertionPoint(module.body), Location.name("c"):
+ arith.constant(dtype, py_values[2])
diff --git a/mlir/test/python/test_to_remove.py b/mlir/test/python/test_to_remove.py
new file mode 100644
index 00000000000000..b2b2cc4b552fe9
--- /dev/null
+++ b/mlir/test/python/test_to_remove.py
@@ -0,0 +1,46 @@
+import concurrent.futures
+import threading
+import inspect
+
+
+def decorator(f):
+ # Introspect the callable for optional features.
+ sig = inspect.signature(f)
+ for param in sig.parameters.values():
+ pass
+
+ def emit_call_op(*call_args):
+ pass
+
+ wrapped = emit_call_op
+ return wrapped
+
+
+def test_dialects_vector_repro_3():
+ num_workers = 6
+ num_runs = 10
+ barrier = threading.Barrier(num_workers)
+
+ def closure():
+ barrier.wait()
+ for _ in range(num_runs):
+
+ @decorator
+ def print_vector(arg):
+ return 0
+
+ barrier.wait()
+
+ with concurrent.futures.ThreadPoolExecutor(
+ max_workers=num_workers
+ ) as executor:
+ futures = []
+ for _ in range(num_workers):
+ futures.append(executor.submit(closure))
+ # We should call future.result() to re-raise an exception if test has
+ # failed
+ assert len(list(f.result() for f in futures)) == num_workers
+
+
+if __name__ == "__main__":
+ test_dialects_vector_repro_3()
\ No newline at end of file
>From f1c1a9164b36f65740f4628c9326f89d60c3916b Mon Sep 17 00:00:00 2001
From: vfdev-5 <vfdev.5 at gmail.com>
Date: Thu, 14 Nov 2024 13:57:04 +0100
Subject: [PATCH 3/5] [skip-ci] More tests and added a lock to
_cext.register_operation
---
mlir/lib/Bindings/Python/MainModule.cpp | 9 +++-
mlir/test/python/multithreaded_tests.py | 65 ++++++++++++++++++++++---
2 files changed, 66 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp
index 168a939ff3cd8a..590d0590f51e4a 100644
--- a/mlir/lib/Bindings/Python/MainModule.cpp
+++ b/mlir/lib/Bindings/Python/MainModule.cpp
@@ -73,9 +73,14 @@ PYBIND11_MODULE(_mlir, m, py::mod_gil_not_used()) {
[dialectClass, replace](py::type opClass) -> py::type {
std::string operationName =
opClass.attr("OPERATION_NAME").cast<std::string>();
- PyGlobals::get().registerOperationImpl(operationName, opClass,
- replace);
+ // 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;
+ });
// Dict-stuff the new opClass by name onto the dialect class.
py::object opClassName = opClass.attr("__name__");
dialectClass.attr(opClassName) = opClass;
diff --git a/mlir/test/python/multithreaded_tests.py b/mlir/test/python/multithreaded_tests.py
index fa545861cb0d7d..87d779ee303454 100644
--- a/mlir/test/python/multithreaded_tests.py
+++ b/mlir/test/python/multithreaded_tests.py
@@ -1,5 +1,6 @@
import concurrent.futures
import functools
+import gc
import importlib.util
import sys
import threading
@@ -42,9 +43,55 @@ def copy_and_update(src_filepath: Path, dst_filepath: Path):
writer.write(src_line)
+def run(f):
+ f()
+
+
+def constructAndPrintInModule(f):
+ print("\nTEST:", f.__name__)
+ with Context(), Location.unknown():
+ module = Module.create()
+ with InsertionPoint(module.body):
+ f()
+ print(module)
+
+
+def run_with_context_and_location(f):
+ print("\nTEST:", f.__name__)
+ with Context(), Location.unknown():
+ f()
+ return f
+
+
test_modules = [
- "execution_engine",
- # "pass_manager",
+ ("execution_engine", run), # Fail
+ ("pass_manager", run), # Fail
+
+ # Dialects tests
+ ("dialects/affine", constructAndPrintInModule), # Fail
+ ("dialects/vector", run_with_context_and_location), # Fail
+
+ # IR tests
+ ("ir/affine_expr", run), # Pass
+ ("ir/affine_map", run), # Pass
+ ("ir/array_attributes", run), # Pass
+ ("ir/attributes", run), # Pass
+ ("ir/blocks", run), # Pass
+ ("ir/builtin_types", run), # Pass
+ ("ir/context_managers", run), # Pass
+ ("ir/debug", run), # Fail
+ ("ir/diagnostic_handler", run), # Fail
+ ("ir/dialects", run), # Fail
+ ("ir/exception", run), # Fail
+ ("ir/insertion_point", run), # Pass
+ ("ir/insertion_point", run), # Pass
+ ("ir/integer_set", run), # Pass
+ ("ir/location", run), # Pass
+ ("ir/module", run), # Pass but may fail randomly on mlirOperationDump in testParseSuccess
+ ("ir/operation", run), # Pass
+ ("ir/symbol_table", run), # Pass
+ ("ir/value", run), # Fail/Crash
+
]
@@ -54,7 +101,7 @@ def decorator(test_cls):
test_cls.output_folder = tempfile.TemporaryDirectory()
output_folder = Path(test_cls.output_folder.name)
- for test_module_name in test_modules:
+ for test_module_name, exec_fn in test_modules:
src_filepath = this_folder / f"{test_module_name}.py"
dst_filepath = (output_folder / f"{test_module_name}.py").absolute()
if not dst_filepath.parent.exists():
@@ -66,8 +113,8 @@ def decorator(test_cls):
obj = getattr(test_mod, attr_name)
if callable(obj):
test_name = f"{test_prefix}_{test_module_name.replace('/', '_')}__{attr_name}"
- def wrapped_test_fn(*args, __test_fn__=obj, **kwargs):
- __test_fn__()
+ 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
@@ -99,6 +146,10 @@ def closure():
for _ in range(num_runs):
__test_fn__(self, *args, **kwargs)
+ barrier.wait()
+ gc.collect()
+ assert Context._get_live_count() == 0
+
with concurrent.futures.ThreadPoolExecutor(
max_workers=num_workers
) as executor:
@@ -114,7 +165,9 @@ def closure():
if "ThreadSanitizer" in captured.err:
raise RuntimeError(f"ThreadSanitizer reported warnings:\n{captured.err}")
else:
- raise RuntimeError(f"Other error:\n{captured.err}")
+ pass
+ # There are tests that write to stderr, we should ignore them
+ # raise RuntimeError(f"Other error:\n{captured.err}")
setattr(test_cls, f"{name}_multi_threaded", multi_threaded_test_fn)
>From 303c87e77e67b46b78f409a33335f8224c839462 Mon Sep 17 00:00:00 2001
From: vfdev-5 <vfdev.5 at gmail.com>
Date: Tue, 10 Dec 2024 19:33:19 +0100
Subject: [PATCH 4/5] [skip ci] Updated tests + labelled passing and failing
tests
---
mlir/lib/Bindings/Python/IRCore.cpp | 65 +++-
mlir/python/requirements.txt | 2 +-
mlir/test/python/multithreaded_tests.py | 485 ++++++++++++++++++++++--
mlir/test/python/test_to_remove.py | 46 ---
4 files changed, 516 insertions(+), 82 deletions(-)
delete mode 100644 mlir/test/python/test_to_remove.py
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 727cbc2e106d5b..ed2d32d6bbe7f5 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -241,9 +241,20 @@ static MlirBlock createBlock(const py::sequence &pyArgTypes,
/// Wrapper for the global LLVM debugging flag.
struct PyGlobalDebugFlag {
- static void set(py::object &o, bool enable) { mlirEnableGlobalDebug(enable); }
+ static void set(py::object &o, bool enable) {
+ withLock([&](){
+ mlirEnableGlobalDebug(enable);
+ return 0;
+ });
+ }
- static bool get(const py::object &) { return mlirIsGlobalDebugEnabled(); }
+ static bool get(const py::object &) {
+ // Use lock in free-threading to avoid data-races,
+ // observed in ir/debug.py, testDebugDlag
+ return withLock([&](){
+ return mlirIsGlobalDebugEnabled();
+ });
+ }
static void bind(py::module &m) {
// Debug flags.
@@ -253,7 +264,10 @@ struct PyGlobalDebugFlag {
.def_static(
"set_types",
[](const std::string &type) {
- mlirSetGlobalDebugType(type.c_str());
+ withLock([&](){
+ mlirSetGlobalDebugType(type.c_str());
+ return 0;
+ });
},
"types"_a, "Sets specific debug types to be produced by LLVM")
.def_static("set_types", [](const std::vector<std::string> &types) {
@@ -261,32 +275,61 @@ struct PyGlobalDebugFlag {
pointers.reserve(types.size());
for (const std::string &str : types)
pointers.push_back(str.c_str());
- mlirSetGlobalDebugTypes(pointers.data(), pointers.size());
+ withLock([&](){
+ mlirSetGlobalDebugTypes(pointers.data(), pointers.size());
+ return 0;
+ });
});
}
+
+private:
+#ifdef Py_GIL_DISABLED
+ static PyMutex &getLock() {
+ static PyMutex lock;
+ return lock;
+ }
+#endif
+
+ template<typename F>
+ static inline auto withLock(const F& cb) -> decltype(cb()) {
+#ifdef Py_GIL_DISABLED
+ auto &lock = getLock();
+ PyMutex_Lock(&lock);
+#endif
+ auto result = cb();
+#ifdef Py_GIL_DISABLED
+ PyMutex_Unlock(&lock);
+#endif
+ return result;
+ }
};
struct PyAttrBuilderMap {
static bool dunderContains(const std::string &attributeKind) {
return PyGlobals::get().lookupAttributeBuilder(attributeKind).has_value();
}
- static py::function dundeGetItemNamed(const std::string &attributeKind) {
- auto builder = PyGlobals::get().lookupAttributeBuilder(attributeKind);
+ static py::function dunderGetItemNamed(const std::string &attributeKind) {
+ auto builder = PyGlobals::withInstance([&](PyGlobals& instance) {
+ return instance.lookupAttributeBuilder(attributeKind);
+ });
if (!builder)
throw py::key_error(attributeKind);
return *builder;
}
- static void dundeSetItemNamed(const std::string &attributeKind,
+ static void dunderSetItemNamed(const std::string &attributeKind,
py::function func, bool replace) {
- PyGlobals::get().registerAttributeBuilder(attributeKind, std::move(func),
- replace);
+ PyGlobals::withInstance([&](PyGlobals& instance) {
+ instance.registerAttributeBuilder(attributeKind, std::move(func),
+ replace);
+ return 0;
+ });
}
static void bind(py::module &m) {
py::class_<PyAttrBuilderMap>(m, "AttrBuilder", py::module_local())
.def_static("contains", &PyAttrBuilderMap::dunderContains)
- .def_static("get", &PyAttrBuilderMap::dundeGetItemNamed)
- .def_static("insert", &PyAttrBuilderMap::dundeSetItemNamed,
+ .def_static("get", &PyAttrBuilderMap::dunderGetItemNamed)
+ .def_static("insert", &PyAttrBuilderMap::dunderSetItemNamed,
"attribute_kind"_a, "attr_builder"_a, "replace"_a = false,
"Register an attribute builder for building MLIR "
"attributes from python values.");
diff --git a/mlir/python/requirements.txt b/mlir/python/requirements.txt
index 449748bb02cc22..66c6168d78540d 100644
--- a/mlir/python/requirements.txt
+++ b/mlir/python/requirements.txt
@@ -3,5 +3,5 @@ numpy>=1.19.5, <=2.1.2
# pybind11>=2.14.0, <2.15.0
# Temporarily set pybind11 version to master waiting the next release to 2.13.6
pybind11 @ git+https://github.com/pybind/pybind11@master
-PyYAML>=5.4.0, <=6.0.1
+PyYAML>=5.4.0, <=6.0.2
ml_dtypes>=0.5.0, <=0.6.0 # provides several NumPy dtype extensions, including the bf16
diff --git a/mlir/test/python/multithreaded_tests.py b/mlir/test/python/multithreaded_tests.py
index 87d779ee303454..ab5b7bf388909c 100644
--- a/mlir/test/python/multithreaded_tests.py
+++ b/mlir/test/python/multithreaded_tests.py
@@ -1,3 +1,223 @@
+"""
+This script generates multi-threaded tests to check free-threading mode using CPython compiled with TSAN.
+Tests can be run using pytest:
+```bash
+python3.13t -mpytest -vvv multithreaded_tests.py
+```
+
+IMPORTANT. Running tests are not checking the correctness, but just the execution of the tests in multi-threaded context
+and passing if no warnings reported by TSAN and failing otherwise.
+
+
+Details on the generated tests and execution:
+1) Multi-threaded execution: all generated tests are executed independently by
+a pool of threads, running each test multiple times, see @multi_threaded for details
+
+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`.
+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
+`@run` or `run(testMethod)`. See `copy_and_update` and `add_existing_tests` methods for details.
+
+
+Observed warnings reported by TSAN.
+- Python 3.13.1 experimental free-threading build (tags/v3.13.1:06714517797, Dec 10 2024, 00:18:06) [Clang 15.0.7 ]
+- numpy 2.3.0.dev0
+- nanobind 2.5.0.dev1
+- pybind11 master
+
+CPython and free-threading known data-races:
+1) ctypes => try to build libffi with TSAN. Reported data races may be false positives. => can't build libffi from source
+```
+WARNING: ThreadSanitizer: data race (pid=99593)
+ Atomic read of size 1 at 0x7f6054c485a8 by thread T3:
+ #0 pthread_mutex_lock <null> (python3.13t+0xe83ca) (BuildId: de51a96f802ffcb0f2dcf5c04836201f1a81133c)
+ #1 ffi_closure_alloc <null> (libffi.so.8+0x5d55) (BuildId: 59c2a6b204f74f358ca7711d2dfd349d88711f6a)
+ #2 PyCFuncPtr_new /tmp/cpython-tsan/./Modules/_ctypes/_ctypes.c:3949:13 (_ctypes.cpython-313t-x86_64-linux-gnu.so+0x19205) (BuildId: 3866c4f0cc959b64602a68236b872ff98967ec7a)
+ #3 type_call /tmp/cpython-tsan/Objects/typeobject.c:1981:11 (python3.13t+0x2d9080) (BuildId: de51a96f802ffcb0f2dcf5c04836201f1a81133c)
+ #4 _PyObject_MakeTpCall /tmp/cpython-tsan/Objects/call.c:242:18 (python3.13t+0x1d6a6c) (BuildId: de51a96f802ffcb0f2dcf5c04836201f1a81133c)
+```
+
+LLVM related data-races
+1) mlir pass manager
+```
+E WARNING: ThreadSanitizer: data race (pid=109173)
+E Write of size 8 at 0x7fef4f12a4a0 by thread T108 (mutexes: write M0):
+E #0 llvm::raw_fd_ostream::write_impl(char const*, unsigned long) /tmp/jax/llvm-project/llvm/lib/Support/raw_ostream.cpp:743:7 (libMLIRPythonCAPI.so.20.0git+0x4897686) (BuildId: 85b3b16da1be79a4)
+E #1 llvm::raw_ostream::write(char const*, unsigned long) /tmp/jax/llvm-project/llvm/lib/Support/raw_ostream.cpp:252:9 (libMLIRPythonCAPI.so.20.0git+0x4895aac) (BuildId: 85b3b16da1be79a4)
+E #2 llvm::raw_ostream::operator<<(llvm::StringRef) /tmp/jax/llvm-project/llvm/include/llvm/Support/raw_ostream.h:230:14 (libMLIRPythonCAPI.so.20.0git+0x478ccbe) (BuildId: 85b3b16da1be79a4)
+E #3 llvm::raw_ostream::operator<<(char const*) /tmp/jax/llvm-project/llvm/include/llvm/Support/raw_ostream.h:257:18 (libMLIRPythonCAPI.so.20.0git+0x478ccbe)
+E #4 (anonymous namespace)::IRPrinterInstrumentation::runAfterPass(mlir::Pass*, mlir::Operation*)::$_1::operator()(llvm::raw_ostream&) const /tmp/jax/llvm-project/mlir/lib/Pass/IRPrinting.cpp:109:9 (libMLIRPythonCAPI.so.20.0git+0x478ccbe)
+E #5 void llvm::function_ref<void (llvm::raw_ostream&)>::callback_fn<(anonymous namespace)::IRPrinterInstrumentation::runAfterPass(mlir::Pass*, mlir::Operation*)::$_1>(long, llvm::raw_ostream&) /tmp/jax/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:46:12 (libMLIRPythonCAPI.so.20.0git+0x478ccbe)
+E #6 llvm::function_ref<void (llvm::raw_ostream&)>::operator()(llvm::raw_ostream&) const /tmp/jax/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69:12 (libMLIRPythonCAPI.so.20.0git+0x478d54e) (BuildId: 85b3b16da1be79a4)
+E #7 (anonymous namespace)::BasicIRPrinterConfig::printAfterIfEnabled(mlir::Pass*, mlir::Operation*, llvm::function_ref<void (llvm::raw_ostream&)>) /tmp/jax/llvm-project/mlir/lib/Pass/IRPrinting.cpp:195:7 (libMLIRPythonCAPI.so.20.0git+0x478d54e)
+E #8 (anonymous namespace)::IRPrinterInstrumentation::runAfterPass(mlir::Pass*, mlir::Operation*) /tmp/jax/llvm-project/mlir/lib/Pass/IRPrinting.cpp:108:11 (libMLIRPythonCAPI.so.20.0git+0x478b7d2) (BuildId: 85b3b16da1be79a4)
+E #9 mlir::PassInstrumentor::runAfterPass(mlir::Pass*, mlir::Operation*) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:1037:12 (libMLIRPythonCAPI.so.20.0git+0x4797282) (BuildId: 85b3b16da1be79a4)
+E #10 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:563:11 (libMLIRPythonCAPI.so.20.0git+0x4797282)
+E #11 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:592:16 (libMLIRPythonCAPI.so.20.0git+0x479b55b) (BuildId: 85b3b16da1be79a4)
+E #12 mlir::PassManager::runPasses(mlir::Operation*, mlir::AnalysisManager) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:905:10 (libMLIRPythonCAPI.so.20.0git+0x479b55b)
+E #13 mlir::PassManager::run(mlir::Operation*) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:885:60 (libMLIRPythonCAPI.so.20.0git+0x479b55b)
+E #14 mlirPassManagerRunOnOp /tmp/jax/llvm-project/mlir/lib/CAPI/IR/Pass.cpp:44:36 (libMLIRPythonCAPI.so.20.0git+0x46c6150) (BuildId: 85b3b16da1be79a4)
+E #15 mlir::python::populatePassManagerSubmodule(pybind11::module_&)::$_5::operator()((anonymous namespace)::PyPassManager&, mlir::python::PyOperationBase&, bool) const /tmp/jax/llvm-project/mlir/lib/Bindings/Python/Pass.cpp:154:40 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x343e3e) (BuildId: 6e516b7f12acec76)
+E #16 void pybind11::detail::argument_loader<(anonymous namespace)::PyPassManager&, mlir::python::PyOperationBase&, bool>::call_impl<void, mlir::python::populatePassManagerSubmodule(pybind11::module_&)::$_5&, 0ul, 1ul, 2ul, pybind11::detail::void_type>(mlir::python::populatePassManagerSubmodule(pybind11::module_&)::$_5&, std::integer_sequence<unsigned long, 0ul, 1ul, 2ul>, pybind11::detail::void_type&&) && /usr/local/include/pybind11/cast.h:1685:16 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x343e3e)
+```
+
+```
+E WARNING: ThreadSanitizer: data race (pid=13402)
+E Write of size 8 at 0x7f5bf342a4a0 by thread T63:
+E #0 llvm::raw_fd_ostream::write_impl(char const*, unsigned long) /tmp/jax/llvm-project/llvm/lib/Support/raw_ostream.cpp:743:7 (libMLIRPythonCAPI.so.20.0git+0x4897686) (BuildId: 85b3b16da1be79a4)
+E #1 llvm::raw_ostream::write(char const*, unsigned long) /tmp/jax/llvm-project/llvm/lib/Support/raw_ostream.cpp:252:9 (libMLIRPythonCAPI.so.20.0git+0x4895aac) (BuildId: 85b3b16da1be79a4)
+E #2 llvm::raw_ostream::operator<<(llvm::StringRef) /tmp/jax/llvm-project/llvm/include/llvm/Support/raw_ostream.h:230:14 (libMLIRPythonCAPI.so.20.0git+0x6996bd2) (BuildId: 85b3b16da1be79a4)
+E #3 llvm::raw_ostream::operator<<(char const*) /tmp/jax/llvm-project/llvm/include/llvm/Support/raw_ostream.h:257:18 (libMLIRPythonCAPI.so.20.0git+0x6996bd2)
+E #4 (anonymous namespace)::PrintOpStatsPass::printSummary() /tmp/jax/llvm-project/mlir/lib/Transforms/OpStats.cpp:62:6 (libMLIRPythonCAPI.so.20.0git+0x6996bd2)
+E #5 (anonymous namespace)::PrintOpStatsPass::runOnOperation() /tmp/jax/llvm-project/mlir/lib/Transforms/OpStats.cpp:57:5 (libMLIRPythonCAPI.so.20.0git+0x6996bd2)
+E #6 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::$_7::operator()() const /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:526:17 (libMLIRPythonCAPI.so.20.0git+0x4796dc4) (BuildId: 85b3b16da1be79a4)
+E #7 void llvm::function_ref<void ()>::callback_fn<mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::$_7>(long) /tmp/jax/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:46:12 (libMLIRPythonCAPI.so.20.0git+0x4796dc4)
+E #8 llvm::function_ref<void ()>::operator()() const /tmp/jax/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69:12 (libMLIRPythonCAPI.so.20.0git+0x4796dc4)
+E #9 void mlir::MLIRContext::executeAction<mlir::PassExecutionAction, mlir::Pass&>(llvm::function_ref<void ()>, llvm::ArrayRef<mlir::IRUnit>, mlir::Pass&) /tmp/jax/llvm-project/mlir/include/mlir/IR/MLIRContext.h:280:7 (libMLIRPythonCAPI.so.20.0git+0x4796dc4)
+E #10 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:520:21 (libMLIRPythonCAPI.so.20.0git+0x4796dc4)
+E #11 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:592:16 (libMLIRPythonCAPI.so.20.0git+0x479b55b) (BuildId: 85b3b16da1be79a4)
+E #12 mlir::PassManager::runPasses(mlir::Operation*, mlir::AnalysisManager) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:905:10 (libMLIRPythonCAPI.so.20.0git+0x479b55b)
+E #13 mlir::PassManager::run(mlir::Operation*) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:885:60 (libMLIRPythonCAPI.so.20.0git+0x479b55b)
+E #14 mlirPassManagerRunOnOp /tmp/jax/llvm-project/mlir/lib/CAPI/IR/Pass.cpp:44:36 (libMLIRPythonCAPI.so.20.0git+0x46c6150) (BuildId: 85b3b16da1be79a4)
+E #15 mlir::python::populatePassManagerSubmodule(pybind11::module_&)::$_5::operator()((anonymous namespace)::PyPassManager&, mlir::python::PyOperationBase&, bool) const /tmp/jax/llvm-project/mlir/lib/Bindings/Python/Pass.cpp:154:40 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x343e3e) (BuildId: 6e516b7f12acec76)
+E #16 void pybind11::detail::argument_loader<(anonymous namespace)::PyPassManager&, mlir::python::PyOperationBase&, bool>::call_impl<void, mlir::python::populatePassManagerSubmodule(pybind11::module_&)::$_5&, 0ul, 1ul, 2ul, pybind11::detail::void_type>(mlir::python::populatePassManagerSubmodule(pybind11::module_&)::$_5&, std::integer_sequence<unsigned long, 0ul, 1ul, 2ul>, pybind11::detail::void_type&&) && /usr/local/include/pybind11/cast.h:1685:16 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x343e3e)
+```
+
+
+```
+E WARNING: ThreadSanitizer: data race (pid=14122)
+E Write of size 8 at 0x7fa88142a4a0 by thread T59 (mutexes: write M0):
+E #0 llvm::raw_fd_ostream::write_impl(char const*, unsigned long) /tmp/jax/llvm-project/llvm/lib/Support/raw_ostream.cpp:743:7 (libMLIRPythonCAPI.so.20.0git+0x4897686) (BuildId: 85b3b16da1be79a4)
+E #1 llvm::raw_ostream::write(char const*, unsigned long) /tmp/jax/llvm-project/llvm/lib/Support/raw_ostream.cpp:252:9 (libMLIRPythonCAPI.so.20.0git+0x4895aac) (BuildId: 85b3b16da1be79a4)
+E #2 llvm::raw_ostream::operator<<(llvm::StringRef) /tmp/jax/llvm-project/llvm/include/llvm/Support/raw_ostream.h:230:14 (libMLIRPythonCAPI.so.20.0git+0x478c5be) (BuildId: 85b3b16da1be79a4)
+E #3 llvm::raw_ostream::operator<<(char const*) /tmp/jax/llvm-project/llvm/include/llvm/Support/raw_ostream.h:257:18 (libMLIRPythonCAPI.so.20.0git+0x478c5be)
+E #4 (anonymous namespace)::IRPrinterInstrumentation::runBeforePass(mlir::Pass*, mlir::Operation*)::$_0::operator()(llvm::raw_ostream&) const /tmp/jax/llvm-project/mlir/lib/Pass/IRPrinting.cpp:78:9 (libMLIRPythonCAPI.so.20.0git+0x478c5be)
+E #5 void llvm::function_ref<void (llvm::raw_ostream&)>::callback_fn<(anonymous namespace)::IRPrinterInstrumentation::runBeforePass(mlir::Pass*, mlir::Operation*)::$_0>(long, llvm::raw_ostream&) /tmp/jax/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:46:12 (libMLIRPythonCAPI.so.20.0git+0x478c5be)
+E #6 llvm::function_ref<void (llvm::raw_ostream&)>::operator()(llvm::raw_ostream&) const /tmp/jax/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69:12 (libMLIRPythonCAPI.so.20.0git+0x478d49e) (BuildId: 85b3b16da1be79a4)
+E #7 (anonymous namespace)::BasicIRPrinterConfig::printBeforeIfEnabled(mlir::Pass*, mlir::Operation*, llvm::function_ref<void (llvm::raw_ostream&)>) /tmp/jax/llvm-project/mlir/lib/Pass/IRPrinting.cpp:189:7 (libMLIRPythonCAPI.so.20.0git+0x478d49e)
+E #8 (anonymous namespace)::IRPrinterInstrumentation::runBeforePass(mlir::Pass*, mlir::Operation*) /tmp/jax/llvm-project/mlir/lib/Pass/IRPrinting.cpp:77:11 (libMLIRPythonCAPI.so.20.0git+0x478b506) (BuildId: 85b3b16da1be79a4)
+E #9 mlir::PassInstrumentor::runBeforePass(mlir::Pass*, mlir::Operation*) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:1030:12 (libMLIRPythonCAPI.so.20.0git+0x4796ca2) (BuildId: 85b3b16da1be79a4)
+E #10 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:517:9 (libMLIRPythonCAPI.so.20.0git+0x4796ca2)
+E #11 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:592:16 (libMLIRPythonCAPI.so.20.0git+0x479b55b) (BuildId: 85b3b16da1be79a4)
+E #12 mlir::PassManager::runPasses(mlir::Operation*, mlir::AnalysisManager) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:905:10 (libMLIRPythonCAPI.so.20.0git+0x479b55b)
+E #13 mlir::PassManager::run(mlir::Operation*) /tmp/jax/llvm-project/mlir/lib/Pass/Pass.cpp:885:60 (libMLIRPythonCAPI.so.20.0git+0x479b55b)
+E #14 mlirPassManagerRunOnOp /tmp/jax/llvm-project/mlir/lib/CAPI/IR/Pass.cpp:44:36 (libMLIRPythonCAPI.so.20.0git+0x46c6150) (BuildId: 85b3b16da1be79a4)
+E #15 mlir::python::populatePassManagerSubmodule(pybind11::module_&)::$_5::operator()((anonymous namespace)::PyPassManager&, mlir::python::PyOperationBase&, bool) const /tmp/jax/llvm-project/mlir/lib/Bindings/Python/Pass.cpp:154:40 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x343e3e) (BuildId: 6e516b7f12acec76)
+E #16 void pybind11::detail::argument_loader<(anonymous namespace)::PyPassManager&, mlir::python::PyOperationBase&, bool>::call_impl<void, mlir::python::populatePassManagerSubmodule(pybind11::module_&)::$_5&, 0ul, 1ul, 2ul, pybind11::detail::void_type>(mlir::python::populatePassManagerSubmodule(pybind11::module_&)::$_5&, std::integer_sequence<unsigned long, 0ul, 1ul, 2ul>, pybind11::detail::void_type&&) && /usr/local/include/pybind11/cast.h:1685:16 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x343e3e)
+```
+
+
+2) dialects/transform_interpreter.py
+
+```
+E WARNING: ThreadSanitizer: data race (pid=15594)
+E Read of size 8 at 0x7fdece62a3f8 by thread T58:
+E #0 llvm::raw_ostream::operator<<(llvm::StringRef) /tmp/jax/llvm-project/llvm/include/llvm/Support/raw_ostream.h:229:25 (libMLIRPythonCAPI.so.20.0git+0x5963297) (BuildId: 85b3b16da1be79a4)
+E #1 llvm::raw_ostream::operator<<(char const*) /tmp/jax/llvm-project/llvm/include/llvm/Support/raw_ostream.h:257:18 (libMLIRPythonCAPI.so.20.0git+0x5963297)
+E #2 mlir::transform::PrintOp::apply(mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/IR/TransformOps.cpp:2787:16 (libMLIRPythonCAPI.so.20.0git+0x5963297)
+E #3 mlir::transform::detail::TransformOpInterfaceInterfaceTraits::Model<mlir::transform::PrintOp>::apply(mlir::transform::detail::TransformOpInterfaceInterfaceTraits::Concept const*, mlir::Operation*, mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/build-tsan/tools/mlir/include/mlir/Dialect/Transform/Interfaces/TransformInterfaces.h.inc:477:56 (libMLIRPythonCAPI.so.20.0git+0x58e1e9a) (BuildId: 85b3b16da1be79a4)
+E #4 mlir::transform::TransformOpInterface::apply(mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/build-tsan/tools/mlir/include/mlir/Dialect/Transform/Interfaces/TransformInterfaces.cpp.inc:61:14 (libMLIRPythonCAPI.so.20.0git+0xb246960) (BuildId: 85b3b16da1be79a4)
+E #5 mlir::transform::TransformState::applyTransform(mlir::transform::TransformOpInterface) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/Interfaces/TransformInterfaces.cpp:951:48 (libMLIRPythonCAPI.so.20.0git+0xb246960)
+E #6 applySequenceBlock(mlir::Block&, mlir::transform::FailurePropagationMode, mlir::transform::TransformState&, mlir::transform::TransformResults&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/IR/TransformOps.cpp:1786:15 (libMLIRPythonCAPI.so.20.0git+0x5958aa7) (BuildId: 85b3b16da1be79a4)
+E #7 mlir::transform::IncludeOp::apply(mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/IR/TransformOps.cpp:1827:40 (libMLIRPythonCAPI.so.20.0git+0x59584b3) (BuildId: 85b3b16da1be79a4)
+E #8 mlir::transform::detail::TransformOpInterfaceInterfaceTraits::Model<mlir::transform::IncludeOp>::apply(mlir::transform::detail::TransformOpInterfaceInterfaceTraits::Concept const*, mlir::Operation*, mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/build-tsan/tools/mlir/include/mlir/Dialect/Transform/Interfaces/TransformInterfaces.h.inc:477:56 (libMLIRPythonCAPI.so.20.0git+0x58d4dfa) (BuildId: 85b3b16da1be79a4)
+E #9 mlir::transform::TransformOpInterface::apply(mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/build-tsan/tools/mlir/include/mlir/Dialect/Transform/Interfaces/TransformInterfaces.cpp.inc:61:14 (libMLIRPythonCAPI.so.20.0git+0xb246960) (BuildId: 85b3b16da1be79a4)
+E #10 mlir::transform::TransformState::applyTransform(mlir::transform::TransformOpInterface) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/Interfaces/TransformInterfaces.cpp:951:48 (libMLIRPythonCAPI.so.20.0git+0xb246960)
+E #11 applySequenceBlock(mlir::Block&, mlir::transform::FailurePropagationMode, mlir::transform::TransformState&, mlir::transform::TransformResults&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/IR/TransformOps.cpp:1786:15 (libMLIRPythonCAPI.so.20.0git+0x5958aa7) (BuildId: 85b3b16da1be79a4)
+E #12 mlir::transform::IncludeOp::apply(mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/IR/TransformOps.cpp:1827:40 (libMLIRPythonCAPI.so.20.0git+0x59584b3) (BuildId: 85b3b16da1be79a4)
+E #13 mlir::transform::detail::TransformOpInterfaceInterfaceTraits::Model<mlir::transform::IncludeOp>::apply(mlir::transform::detail::TransformOpInterfaceInterfaceTraits::Concept const*, mlir::Operation*, mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/build-tsan/tools/mlir/include/mlir/Dialect/Transform/Interfaces/TransformInterfaces.h.inc:477:56 (libMLIRPythonCAPI.so.20.0git+0x58d4dfa) (BuildId: 85b3b16da1be79a4)
+E #14 mlir::transform::TransformOpInterface::apply(mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/build-tsan/tools/mlir/include/mlir/Dialect/Transform/Interfaces/TransformInterfaces.cpp.inc:61:14 (libMLIRPythonCAPI.so.20.0git+0xb246960) (BuildId: 85b3b16da1be79a4)
+E #15 mlir::transform::TransformState::applyTransform(mlir::transform::TransformOpInterface) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/Interfaces/TransformInterfaces.cpp:951:48 (libMLIRPythonCAPI.so.20.0git+0xb246960)
+E #16 applySequenceBlock(mlir::Block&, mlir::transform::FailurePropagationMode, mlir::transform::TransformState&, mlir::transform::TransformResults&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/IR/TransformOps.cpp:1786:15 (libMLIRPythonCAPI.so.20.0git+0x5958aa7) (BuildId: 85b3b16da1be79a4)
+E #17 mlir::transform::NamedSequenceOp::apply(mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/IR/TransformOps.cpp:2153:10 (libMLIRPythonCAPI.so.20.0git+0x595ea87) (BuildId: 85b3b16da1be79a4)
+E #18 mlir::transform::detail::TransformOpInterfaceInterfaceTraits::Model<mlir::transform::NamedSequenceOp>::apply(mlir::transform::detail::TransformOpInterfaceInterfaceTraits::Concept const*, mlir::Operation*, mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/build-tsan/tools/mlir/include/mlir/Dialect/Transform/Interfaces/TransformInterfaces.h.inc:477:56 (libMLIRPythonCAPI.so.20.0git+0x58dda7a) (BuildId: 85b3b16da1be79a4)
+E #19 mlir::transform::TransformOpInterface::apply(mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/build-tsan/tools/mlir/include/mlir/Dialect/Transform/Interfaces/TransformInterfaces.cpp.inc:61:14 (libMLIRPythonCAPI.so.20.0git+0xb246960) (BuildId: 85b3b16da1be79a4)
+E #20 mlir::transform::TransformState::applyTransform(mlir::transform::TransformOpInterface) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/Interfaces/TransformInterfaces.cpp:951:48 (libMLIRPythonCAPI.so.20.0git+0xb246960)
+E #21 mlir::transform::applyTransforms(mlir::Operation*, mlir::transform::TransformOpInterface, mlir::RaggedArray<llvm::PointerUnion<mlir::Operation*, mlir::Attribute, mlir::Value>> const&, mlir::transform::TransformOptions const&, bool, llvm::function_ref<void (mlir::transform::TransformState&)>, llvm::function_ref<llvm::LogicalResult (mlir::transform::TransformState&)>) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/Interfaces/TransformInterfaces.cpp:2018:13 (libMLIRPythonCAPI.so.20.0git+0xb256a9c) (BuildId: 85b3b16da1be79a4)
+E #22 mlir::transform::applyTransformNamedSequence(mlir::RaggedArray<llvm::PointerUnion<mlir::Operation*, mlir::Attribute, mlir::Value>>, mlir::transform::TransformOpInterface, mlir::ModuleOp, mlir::transform::TransformOptions const&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp:234:10 (libMLIRPythonCAPI.so.20.0git+0x59991ab) (BuildId: 85b3b16da1be79a4)
+E #23 mlir::transform::applyTransformNamedSequence(mlir::Operation*, mlir::Operation*, mlir::ModuleOp, mlir::transform::TransformOptions const&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp:196:10 (libMLIRPythonCAPI.so.20.0git+0x59986c4) (BuildId: 85b3b16da1be79a4)
+E #24 mlirTransformApplyNamedSequence /tmp/jax/llvm-project/mlir/lib/CAPI/Dialect/TransformInterpreter.cpp:71:15 (libMLIRPythonCAPI.so.20.0git+0x46d16b4) (BuildId: 85b3b16da1be79a4)
+E #25 populateTransformInterpreterSubmodule(pybind11::module_&)::$_4::operator()(MlirOperation, MlirOperation, MlirOperation, (anonymous namespace)::PyMlirTransformOptions const&) const /tmp/jax/llvm-project/mlir/lib/Bindings/Python/TransformInterpreter.cpp:74:36 (_mlirTransformInterpreter.cpython-313t-x86_64-linux-gnu.so+0x3fc0d) (BuildId: 7d3d2fbf5b54f7b5)
+E #26 void pybind11::detail::argument_loader<MlirOperation, MlirOperation, MlirOperation, (anonymous namespace)::PyMlirTransformOptions const&>::call_impl<void, populateTransformInterpreterSubmodule(pybind11::module_&)::$_4&, 0ul, 1ul, 2ul, 3ul, pybind11::detail::void_type>(populateTransformInterpreterSubmodule(pybind11::module_&)::$_4&, std::integer_sequence<unsigned long, 0ul, 1ul, 2ul, 3ul>, pybind11::detail::void_type&&) && /usr/local/include/pybind11/cast.h:1685:16 (_mlirTransformInterpreter.cpython-313t-x86_64-linux-gnu.so+0x3fc0d)
+```
+
+```
+E WARNING: ThreadSanitizer: data race (pid=16690)
+E Read of size 8 at 0x7fcda1b2a3f8 by thread T62:
+E #0 llvm::raw_ostream::operator<<(llvm::StringRef) /tmp/jax/llvm-project/llvm/include/llvm/Support/raw_ostream.h:229:25 (libMLIRPythonCAPI.so.20.0git+0x5963297) (BuildId: 85b3b16da1be79a4)
+E #1 llvm::raw_ostream::operator<<(char const*) /tmp/jax/llvm-project/llvm/include/llvm/Support/raw_ostream.h:257:18 (libMLIRPythonCAPI.so.20.0git+0x5963297)
+E #2 mlir::transform::PrintOp::apply(mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/IR/TransformOps.cpp:2787:16 (libMLIRPythonCAPI.so.20.0git+0x5963297)
+E #3 mlir::transform::detail::TransformOpInterfaceInterfaceTraits::Model<mlir::transform::PrintOp>::apply(mlir::transform::detail::TransformOpInterfaceInterfaceTraits::Concept const*, mlir::Operation*, mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/build-tsan/tools/mlir/include/mlir/Dialect/Transform/Interfaces/TransformInterfaces.h.inc:477:56 (libMLIRPythonCAPI.so.20.0git+0x58e1e9a) (BuildId: 85b3b16da1be79a4)
+E #4 mlir::transform::TransformOpInterface::apply(mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/build-tsan/tools/mlir/include/mlir/Dialect/Transform/Interfaces/TransformInterfaces.cpp.inc:61:14 (libMLIRPythonCAPI.so.20.0git+0xb246960) (BuildId: 85b3b16da1be79a4)
+E #5 mlir::transform::TransformState::applyTransform(mlir::transform::TransformOpInterface) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/Interfaces/TransformInterfaces.cpp:951:48 (libMLIRPythonCAPI.so.20.0git+0xb246960)
+E #6 applySequenceBlock(mlir::Block&, mlir::transform::FailurePropagationMode, mlir::transform::TransformState&, mlir::transform::TransformResults&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/IR/TransformOps.cpp:1786:15 (libMLIRPythonCAPI.so.20.0git+0x5958aa7) (BuildId: 85b3b16da1be79a4)
+E #7 mlir::transform::NamedSequenceOp::apply(mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/IR/TransformOps.cpp:2153:10 (libMLIRPythonCAPI.so.20.0git+0x595ea87) (BuildId: 85b3b16da1be79a4)
+E #8 mlir::transform::detail::TransformOpInterfaceInterfaceTraits::Model<mlir::transform::NamedSequenceOp>::apply(mlir::transform::detail::TransformOpInterfaceInterfaceTraits::Concept const*, mlir::Operation*, mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/build-tsan/tools/mlir/include/mlir/Dialect/Transform/Interfaces/TransformInterfaces.h.inc:477:56 (libMLIRPythonCAPI.so.20.0git+0x58dda7a) (BuildId: 85b3b16da1be79a4)
+E #9 mlir::transform::TransformOpInterface::apply(mlir::transform::TransformRewriter&, mlir::transform::TransformResults&, mlir::transform::TransformState&) /tmp/jax/llvm-project/build-tsan/tools/mlir/include/mlir/Dialect/Transform/Interfaces/TransformInterfaces.cpp.inc:61:14 (libMLIRPythonCAPI.so.20.0git+0xb246960) (BuildId: 85b3b16da1be79a4)
+E #10 mlir::transform::TransformState::applyTransform(mlir::transform::TransformOpInterface) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/Interfaces/TransformInterfaces.cpp:951:48 (libMLIRPythonCAPI.so.20.0git+0xb246960)
+E #11 mlir::transform::applyTransforms(mlir::Operation*, mlir::transform::TransformOpInterface, mlir::RaggedArray<llvm::PointerUnion<mlir::Operation*, mlir::Attribute, mlir::Value>> const&, mlir::transform::TransformOptions const&, bool, llvm::function_ref<void (mlir::transform::TransformState&)>, llvm::function_ref<llvm::LogicalResult (mlir::transform::TransformState&)>) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/Interfaces/TransformInterfaces.cpp:2018:13 (libMLIRPythonCAPI.so.20.0git+0xb256a9c) (BuildId: 85b3b16da1be79a4)
+E #12 mlir::transform::applyTransformNamedSequence(mlir::RaggedArray<llvm::PointerUnion<mlir::Operation*, mlir::Attribute, mlir::Value>>, mlir::transform::TransformOpInterface, mlir::ModuleOp, mlir::transform::TransformOptions const&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp:234:10 (libMLIRPythonCAPI.so.20.0git+0x59991ab) (BuildId: 85b3b16da1be79a4)
+E #13 mlir::transform::applyTransformNamedSequence(mlir::Operation*, mlir::Operation*, mlir::ModuleOp, mlir::transform::TransformOptions const&) /tmp/jax/llvm-project/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp:196:10 (libMLIRPythonCAPI.so.20.0git+0x59986c4) (BuildId: 85b3b16da1be79a4)
+E #14 mlirTransformApplyNamedSequence /tmp/jax/llvm-project/mlir/lib/CAPI/Dialect/TransformInterpreter.cpp:71:15 (libMLIRPythonCAPI.so.20.0git+0x46d16b4) (BuildId: 85b3b16da1be79a4)
+E #15 populateTransformInterpreterSubmodule(pybind11::module_&)::$_4::operator()(MlirOperation, MlirOperation, MlirOperation, (anonymous namespace)::PyMlirTransformOptions const&) const /tmp/jax/llvm-project/mlir/lib/Bindings/Python/TransformInterpreter.cpp:74:36 (_mlirTransformInterpreter.cpython-313t-x86_64-linux-gnu.so+0x3fc0d) (BuildId: 7d3d2fbf5b54f7b5)
+E #16 void pybind11::detail::argument_loader<MlirOperation, MlirOperation, MlirOperation, (anonymous namespace)::PyMlirTransformOptions const&>::call_impl<void, populateTransformInterpreterSubmodule(pybind11::module_&)::$_4&, 0ul, 1ul, 2ul, 3ul, pybind11::detail::void_type>(populateTransformInterpreterSubmodule(pybind11::module_&)::$_4&, std::integer_sequence<unsigned long, 0ul, 1ul, 2ul, 3ul>, pybind11::detail::void_type&&) && /usr/local/include/pybind11/cast.h:1685:16 (_mlirTransformInterpreter.cpython-313t-x86_64-linux-gnu.so+0x3fc0d)
+```
+
+
+3) ir/diagnostic_handler.py
+```
+E WARNING: ThreadSanitizer: data race (pid=19144)
+E Write of size 8 at 0x7ff1be8294a0 by thread T60 (mutexes: write M0):
+E #0 llvm::raw_fd_ostream::write_impl(char const*, unsigned long) /tmp/jax/llvm-project/llvm/lib/Support/raw_ostream.cpp:743:7 (libMLIRPythonCAPI.so.20.0git+0x4897686) (BuildId: 85b3b16da1be79a4)
+E #1 llvm::raw_ostream::write(char const*, unsigned long) /tmp/jax/llvm-project/llvm/lib/Support/raw_ostream.cpp:252:9 (libMLIRPythonCAPI.so.20.0git+0x4895aac) (BuildId: 85b3b16da1be79a4)
+E #2 llvm::raw_ostream::operator<<(llvm::StringRef) /tmp/jax/llvm-project/llvm/include/llvm/Support/raw_ostream.h:230:14 (libMLIRPythonCAPI.so.20.0git+0x49d130f) (BuildId: 85b3b16da1be79a4)
+E #3 llvm::raw_ostream::operator<<(char const*) /tmp/jax/llvm-project/llvm/include/llvm/Support/raw_ostream.h:257:18 (libMLIRPythonCAPI.so.20.0git+0x49d130f)
+E #4 mlir::detail::DiagnosticEngineImpl::emit(mlir::Diagnostic&&) /tmp/jax/llvm-project/mlir/lib/IR/Diagnostics.cpp:264:6 (libMLIRPythonCAPI.so.20.0git+0x49d130f)
+E #5 mlir::DiagnosticEngine::emit(mlir::Diagnostic&&) /tmp/jax/llvm-project/mlir/lib/IR/Diagnostics.cpp:299:9 (libMLIRPythonCAPI.so.20.0git+0x49d0ff4) (BuildId: 85b3b16da1be79a4)
+E #6 mlir::InFlightDiagnostic::report() /tmp/jax/llvm-project/mlir/lib/IR/Diagnostics.cpp:210:12 (libMLIRPythonCAPI.so.20.0git+0x49d0ff4)
+E #7 mlir::InFlightDiagnostic::~InFlightDiagnostic() /tmp/jax/llvm-project/mlir/include/mlir/IR/Diagnostics.h:325:7 (libMLIRPythonCAPI.so.20.0git+0x46bae43) (BuildId: 85b3b16da1be79a4)
+E #8 mlirEmitError /tmp/jax/llvm-project/mlir/lib/CAPI/IR/Diagnostics.cpp:79:3 (libMLIRPythonCAPI.so.20.0git+0x46bae43)
+E #9 mlir::python::populateIRCore(pybind11::module_&)::$_43::operator()(mlir::python::PyLocation&, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char>>) const /tmp/jax/llvm-project/mlir/lib/Bindings/Python/IRCore.cpp:2939:13 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x26cce9) (BuildId: d60a319200e3423d)
+E #10 void pybind11::detail::argument_loader<mlir::python::PyLocation&, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char>>>::call_impl<void, mlir::python::populateIRCore(pybind11::module_&)::$_43&, 0ul, 1ul, pybind11::detail::void_type>(mlir::python::populateIRCore(pybind11::module_&)::$_43&, std::integer_sequence<unsigned long, 0ul, 1ul>, pybind11::detail::void_type&&) && /usr/local/include/pybind11/cast.h:1685:16 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x26cce9)
+```
+
+5) ir/module.py
+```
+E WARNING: ThreadSanitizer: data race (pid=19407)
+E Write of size 8 at 0x7f1858429468 by thread T62:
+E #0 llvm::raw_ostream::copy_to_buffer(char const*, unsigned long) /tmp/jax/llvm-project/llvm/lib/Support/raw_ostream.cpp:306:13 (libMLIRPythonCAPI.so.20.0git+0x48958bb) (BuildId: 85b3b16da1be79a4)
+E #1 llvm::raw_ostream::write(char const*, unsigned long) /tmp/jax/llvm-project/llvm/lib/Support/raw_ostream.cpp:285:3 (libMLIRPythonCAPI.so.20.0git+0x48958bb)
+E #2 llvm::raw_ostream& write_padding<(char)32>(llvm::raw_ostream&, unsigned int) /tmp/jax/llvm-project/llvm/lib/Support/raw_ostream.cpp:486:15 (libMLIRPythonCAPI.so.20.0git+0x489610e) (BuildId: 85b3b16da1be79a4)
+E #3 llvm::raw_ostream::indent(unsigned int) /tmp/jax/llvm-project/llvm/lib/Support/raw_ostream.cpp:498:10 (libMLIRPythonCAPI.so.20.0git+0x489610e)
+E #4 (anonymous namespace)::OperationPrinter::printFullOpWithIndentAndLoc(mlir::Operation*) /tmp/jax/llvm-project/mlir/lib/IR/AsmPrinter.cpp:3426:6 (libMLIRPythonCAPI.so.20.0git+0x4912af0) (BuildId: 85b3b16da1be79a4)
+E #5 mlir::Operation::print(llvm::raw_ostream&, mlir::AsmState&) /tmp/jax/llvm-project/mlir/lib/IR/AsmPrinter.cpp:3979:13 (libMLIRPythonCAPI.so.20.0git+0x4910238) (BuildId: 85b3b16da1be79a4)
+E #6 mlir::Operation::print(llvm::raw_ostream&, mlir::OpPrintingFlags const&) /tmp/jax/llvm-project/mlir/lib/IR/AsmPrinter.cpp:3971:3 (libMLIRPythonCAPI.so.20.0git+0x4913a05) (BuildId: 85b3b16da1be79a4)
+E #7 mlir::Operation::dump() /tmp/jax/llvm-project/mlir/lib/IR/AsmPrinter.cpp:3984:3 (libMLIRPythonCAPI.so.20.0git+0x4913a05)
+E #8 mlirOperationDump /tmp/jax/llvm-project/mlir/lib/CAPI/IR/IR.cpp:715:63 (libMLIRPythonCAPI.so.20.0git+0x46c05d9) (BuildId: 85b3b16da1be79a4)
+E #9 mlir::python::populateIRCore(pybind11::module_&)::$_50::operator()(mlir::python::PyModule&) const /tmp/jax/llvm-project/mlir/lib/Bindings/Python/IRCore.cpp:3000:13 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x26fc90) (BuildId: d60a319200e3423d)
+E #10 void pybind11::detail::argument_loader<mlir::python::PyModule&>::call_impl<void, mlir::python::populateIRCore(pybind11::module_&)::$_50&, 0ul, pybind11::detail::void_type>(mlir::python::populateIRCore(pybind11::module_&)::$_50&, std::integer_sequence<unsigned long, 0ul>, pybind11::detail::void_type&&) && /usr/local/include/pybind11/cast.h:1685:16 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x26fc90)
+```
+
+```
+E WARNING: ThreadSanitizer: data race (pid=19669)
+E Write of size 8 at 0x7f3882591840 by thread T62:
+E #0 unicode_fill_utf8 /tmp/cpython-tsan/Objects/unicodeobject.c:5444:30 (python3.13t+0x3142a1) (BuildId: 31ded8fac9cea6b44dd4f3246be3853edec8d0dc)
+E #1 PyUnicode_AsUTF8AndSize /tmp/cpython-tsan/Objects/unicodeobject.c:4066:13 (python3.13t+0x3142a1)
+E #2 pybind11::detail::string_caster<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char>>, false>::load(pybind11::handle, bool) /usr/local/include/pybind11/cast.h:444:51 (_mlir.cpython-313t-x86_64-linux-gnu.so+0xd6a25) (BuildId: d60a319200e3423d)
+E #3 bool pybind11::detail::argument_loader<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char>> const&, mlir::python::DefaultingPyMlirContext>::load_impl_sequence<0ul, 1ul>(pybind11::detail::function_call&, std::integer_sequence<unsigned long, 0ul, 1ul>) /usr/local/include/pybind11/cast.h:1670:47 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x26e4a5) (BuildId: d60a319200e3423d)
+E #4 pybind11::detail::argument_loader<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char>> const&, mlir::python::DefaultingPyMlirContext>::load_args(pybind11::detail::function_call&) /usr/local/include/pybind11/cast.h:1648:50 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x26e036) (BuildId: d60a319200e3423d)
+E #5 void pybind11::cpp_function::initialize<mlir::python::populateIRCore(pybind11::module_&)::$_45, pybind11::object, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char>> const&, mlir::python::DefaultingPyMlirContext, pybind11::name, pybind11::scope, pybind11::sibling, pybind11::arg, pybind11::arg_v, char [168]>(mlir::python::populateIRCore(pybind11::module_&)::$_45&&, pybind11::object (*)(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char>> const&, mlir::python::DefaultingPyMlirContext), pybind11::name const&, pybind11::scope const&, pybind11::sibling const&, pybind11::arg const&, pybind11::arg_v const&, char const (&) [168])::'lambda'(pybind11::detail::function_call&)::operator()(pybind11::detail::function_call&) const /usr/local/include/pybind11/pybind11.h:259:33 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x26e036)
+E #6 void pybind11::cpp_function::initialize<mlir::python::populateIRCore(pybind11::module_&)::$_45, pybind11::object, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char>> const&, mlir::python::DefaultingPyMlirContext, pybind11::name, pybind11::scope, pybind11::sibling, pybind11::arg, pybind11::arg_v, char [168]>(mlir::python::populateIRCore(pybind11::module_&)::$_45&&, pybind11::object (*)(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char>> const&, mlir::python::DefaultingPyMlirContext), pybind11::name const&, pybind11::scope const&, pybind11::sibling const&, pybind11::arg const&, pybind11::arg_v const&, char const (&) [168])::'lambda'(pybind11::detail::function_call&)::__invoke(pybind11::detail::function_call&) /usr/local/include/pybind11/pybind11.h:255:21 (_mlir.cpython-313t-x86_64-linux-gnu.so+0x26e036)
+E #7 pybind11::cpp_function::dispatcher(_object*, _object*, _object*) /usr/local/include/pybind11/pybind11.h:986:30 (_mlir.cpython-313t-x86_64-linux-gnu.so+0xe6227) (BuildId: d60a319200e3423d)
+E #8 cfunction_call /tmp/cpython-tsan/Objects/methodobject.c:540:18 (python3.13t+0x2795e7) (BuildId: 31ded8fac9cea6b44dd4f3246be3853edec8d0dc)
+
+```
+
+"""
import concurrent.futures
import functools
import gc
@@ -14,10 +234,10 @@
import pytest
import mlir.dialects.arith as arith
+from mlir.dialects import transform
from mlir.ir import Context, Location, Module, IntegerType, F64Type, InsertionPoint
-
def import_from_path(module_name: str, file_path: Path):
spec = importlib.util.spec_from_file_location(module_name, file_path)
module = importlib.util.module_from_spec(spec)
@@ -37,41 +257,185 @@ def copy_and_update(src_filepath: Path, dst_filepath: Path):
"run(",
"@run",
"@constructAndPrintInModule",
+ "run_apply_patterns(",
+ "@run_apply_patterns",
+ "@test_in_context",
+ "@construct_and_print_in_module",
]
if any(src_line.startswith(line) for line in skip_lines):
continue
writer.write(src_line)
+# Helper run functions
def run(f):
f()
-def constructAndPrintInModule(f):
+def run_with_context_and_location(f):
+ print("\nTEST:", f.__name__)
+ with Context(), Location.unknown():
+ f()
+ return f
+
+
+def run_with_insertion_point(f):
+ print("\nTEST:", f.__name__)
+ with Context() as ctx, Location.unknown():
+ module = Module.create()
+ with InsertionPoint(module.body):
+ f(ctx)
+ print(module)
+
+
+def run_with_insertion_point_v2(f):
print("\nTEST:", f.__name__)
with Context(), Location.unknown():
module = Module.create()
with InsertionPoint(module.body):
f()
print(module)
+ return f
-def run_with_context_and_location(f):
+def run_with_insertion_point_v3(f):
+ with Context(), Location.unknown():
+ module = Module.create()
+ with InsertionPoint(module.body):
+ print("\nTEST:", f.__name__)
+ f(module)
+ print(module)
+ return f
+
+
+def run_with_insertion_point_v4(f):
print("\nTEST:", f.__name__)
+ with Context() as ctx, Location.unknown():
+ ctx.allow_unregistered_dialects = True
+ module = Module.create()
+ with InsertionPoint(module.body):
+ f()
+ return f
+
+
+def run_apply_patterns(f):
with Context(), Location.unknown():
- f()
+ module = Module.create()
+ with InsertionPoint(module.body):
+ sequence = transform.SequenceOp(
+ transform.FailurePropagationMode.Propagate,
+ [],
+ transform.AnyOpType.get(),
+ )
+ with InsertionPoint(sequence.body):
+ apply = transform.ApplyPatternsOp(sequence.bodyTarget)
+ with InsertionPoint(apply.patterns):
+ f()
+ transform.YieldOp()
+ print("\nTEST:", f.__name__)
+ print(module)
+ return f
+
+
+def run_transform_tensor_ext(f):
+ print("\nTEST:", f.__name__)
+ with Context(), Location.unknown():
+ module = Module.create()
+ with InsertionPoint(module.body):
+ sequence = transform.SequenceOp(
+ transform.FailurePropagationMode.Propagate,
+ [],
+ transform.AnyOpType.get(),
+ )
+ with InsertionPoint(sequence.body):
+ f(sequence.bodyTarget)
+ transform.YieldOp()
+ print(module)
+ return f
+
+
+def run_transform_structured_ext(f):
+ with Context(), Location.unknown():
+ module = Module.create()
+ with InsertionPoint(module.body):
+ print("\nTEST:", f.__name__)
+ f()
+ module.operation.verify()
+ print(module)
+ return f
+
+
+def run_construct_and_print_in_module(f):
+ print("\nTEST:", f.__name__)
+ with Context(), Location.unknown():
+ module = Module.create()
+ with InsertionPoint(module.body):
+ module = f(module)
+ if module is not None:
+ print(module)
return f
test_modules = [
- ("execution_engine", run), # Fail
+ ("execution_engine", run), # Fail,
("pass_manager", run), # Fail
- # Dialects tests
- ("dialects/affine", constructAndPrintInModule), # Fail
- ("dialects/vector", run_with_context_and_location), # Fail
+ ("dialects/affine", run_with_insertion_point_v2), # Pass
+ ("dialects/func", run_with_insertion_point_v2), # Pass
+ ("dialects/arith_dialect", run), # Pass
+ ("dialects/arith_llvm", run), # Pass
+ ("dialects/async_dialect", run), # Pass
+ ("dialects/builtin", run), # Pass
+ ("dialects/cf", run_with_insertion_point_v4), # Pass
+ ("dialects/complex_dialect", run), # Pass
+ ("dialects/func", run_with_insertion_point_v2), # Pass
+ ("dialects/index_dialect", run_with_insertion_point), # Pass
+ ("dialects/llvm", run_with_insertion_point_v2), # Pass
+ ("dialects/math_dialect", run), # Pass
+ ("dialects/memref", run), # Fail
+ ("dialects/ml_program", run_with_insertion_point_v2), # Pass
+ ("dialects/nvgpu", run_with_insertion_point_v2), # Pass
+ ("dialects/nvvm", run_with_insertion_point_v2), # Pass
+ ("dialects/ods_helpers", run), # Pass
+ ("dialects/openmp_ops", run_with_insertion_point_v2), # Pass
+ ("dialects/pdl_ops", run_with_insertion_point_v2), # Pass
+ # ("dialects/python_test", run), # TODO: Need to pass pybind11 or nanobind argv
+ ("dialects/quant", run), # Pass
+ ("dialects/rocdl", run_with_insertion_point_v2), # Pass
+ ("dialects/scf", run_with_insertion_point_v2), # Pass
+ ("dialects/shape", run), # Pass
+ ("dialects/spirv_dialect", run), # Pass
+ ("dialects/tensor", run), # Pass
+ # ("dialects/tosa", ), # Nothing to test
+ ("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_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", 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 tests
("ir/affine_expr", run), # Pass
("ir/affine_map", run), # Pass
("ir/array_attributes", run), # Pass
@@ -82,16 +446,66 @@ def run_with_context_and_location(f):
("ir/debug", run), # Fail
("ir/diagnostic_handler", run), # Fail
("ir/dialects", run), # Fail
- ("ir/exception", run), # Fail
- ("ir/insertion_point", run), # Pass
+ ("ir/exception", run), # Pass
("ir/insertion_point", run), # Pass
("ir/integer_set", run), # Pass
("ir/location", run), # Pass
- ("ir/module", run), # Pass but may fail randomly on mlirOperationDump in testParseSuccess
+ ("ir/module", run), # Fail
("ir/operation", run), # Pass
("ir/symbol_table", run), # Pass
- ("ir/value", run), # Fail/Crash
+ ("ir/value", run), # Pass
+]
+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
+]
+
+
+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",
+ "test_execution_engine__testComplexMemrefAdd_multi_threaded",
+ "test_execution_engine__testComplexUnrankedMemrefAdd_multi_threaded",
+ "test_execution_engine__testDynamicMemrefAdd2D_multi_threaded",
+ "test_execution_engine__testF16MemrefAdd_multi_threaded",
+ "test_execution_engine__testF8E5M2Memref_multi_threaded",
+ "test_execution_engine__testInvalidModule_multi_threaded",
+ "test_execution_engine__testInvokeFloatAdd_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
+ "test_ir_module__testRoundtripUnicode_multi_threaded", # randomly fails due to data race in PyUnicode_AsUTF8AndSize
+ "test_ir_module__testRoundtripBinary_multi_threaded", # randomly fails due to data race in PyUnicode_AsUTF8AndSize
]
@@ -101,7 +515,14 @@ def decorator(test_cls):
test_cls.output_folder = tempfile.TemporaryDirectory()
output_folder = Path(test_cls.output_folder.name)
- for test_module_name, exec_fn in test_modules:
+ for test_mod_info in test_modules:
+ assert isinstance(test_mod_info, tuple) and len(test_mod_info) in (2, 3)
+ if len(test_mod_info) == 2:
+ test_module_name, exec_fn = test_mod_info
+ test_pattern = None
+ else:
+ test_module_name, exec_fn, test_pattern = test_mod_info
+
src_filepath = this_folder / f"{test_module_name}.py"
dst_filepath = (output_folder / f"{test_module_name}.py").absolute()
if not dst_filepath.parent.exists():
@@ -109,7 +530,9 @@ def decorator(test_cls):
copy_and_update(src_filepath, dst_filepath)
test_mod = import_from_path(test_module_name, dst_filepath)
for attr_name in dir(test_mod):
- if attr_name.startswith("test"):
+ 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])
+ if is_test_fn:
obj = getattr(test_mod, attr_name)
if callable(obj):
test_name = f"{test_prefix}_{test_module_name.replace('/', '_')}__{attr_name}"
@@ -125,7 +548,9 @@ def multi_threaded(
num_workers: int,
num_runs: int = 5,
skip_tests: Optional[list[str]] = None,
+ xfail_tests: Optional[list[str]] = None,
test_prefix: str = "_original_test",
+ multithreaded_test_postfix: str = "_multi_threaded",
):
"""Decorator that runs a test in a multi-threaded environment."""
def decorator(test_cls):
@@ -135,7 +560,10 @@ def decorator(test_cls):
name = f"test{name[len(test_prefix):]}"
if skip_tests is not None:
- if any(test_name in name for test_name in skip_tests):
+ if any(
+ test_name.replace(multithreaded_test_postfix, "") in name
+ for test_name in skip_tests
+ ):
continue
def multi_threaded_test_fn(self, capfd, *args, __test_fn__=test_fn, **kwargs):
@@ -146,10 +574,6 @@ def closure():
for _ in range(num_runs):
__test_fn__(self, *args, **kwargs)
- barrier.wait()
- gc.collect()
- assert Context._get_live_count() == 0
-
with concurrent.futures.ThreadPoolExecutor(
max_workers=num_workers
) as executor:
@@ -158,7 +582,10 @@ def closure():
futures.append(executor.submit(closure))
# We should call future.result() to re-raise an exception if test has
# failed
- list(f.result() for f in futures)
+ assert len(list(f.result() for f in futures)) == num_workers
+
+ gc.collect()
+ assert Context._get_live_count() == 0
captured = capfd.readouterr()
if len(captured.err) > 0:
@@ -169,18 +596,28 @@ def closure():
# There are tests that write to stderr, we should ignore them
# raise RuntimeError(f"Other error:\n{captured.err}")
- setattr(test_cls, f"{name}_multi_threaded", multi_threaded_test_fn)
+ test_new_name = f"{name}{multithreaded_test_postfix}"
+ if xfail_tests is not None and test_new_name in xfail_tests:
+ multi_threaded_test_fn = pytest.mark.xfail(multi_threaded_test_fn)
+
+ setattr(test_cls, test_new_name, multi_threaded_test_fn)
return test_cls
return decorator
- at multi_threaded(num_workers=4, num_runs=10)
+ at multi_threaded(
+ num_workers=6,
+ num_runs=20,
+ skip_tests=tests_to_skip,
+ xfail_tests=tests_to_xfail,
+)
@add_existing_tests(test_prefix="_original_test")
class TestAllMultiThreaded:
@pytest.fixture(scope='class')
def teardown(self):
- self.output_folder.cleanup()
+ if hasattr(self, "output_folder"):
+ self.output_folder.cleanup()
def _original_test_create_context(self):
with Context() as ctx:
diff --git a/mlir/test/python/test_to_remove.py b/mlir/test/python/test_to_remove.py
deleted file mode 100644
index b2b2cc4b552fe9..00000000000000
--- a/mlir/test/python/test_to_remove.py
+++ /dev/null
@@ -1,46 +0,0 @@
-import concurrent.futures
-import threading
-import inspect
-
-
-def decorator(f):
- # Introspect the callable for optional features.
- sig = inspect.signature(f)
- for param in sig.parameters.values():
- pass
-
- def emit_call_op(*call_args):
- pass
-
- wrapped = emit_call_op
- return wrapped
-
-
-def test_dialects_vector_repro_3():
- num_workers = 6
- num_runs = 10
- barrier = threading.Barrier(num_workers)
-
- def closure():
- barrier.wait()
- for _ in range(num_runs):
-
- @decorator
- def print_vector(arg):
- return 0
-
- barrier.wait()
-
- with concurrent.futures.ThreadPoolExecutor(
- max_workers=num_workers
- ) as executor:
- futures = []
- for _ in range(num_workers):
- futures.append(executor.submit(closure))
- # We should call future.result() to re-raise an exception if test has
- # failed
- assert len(list(f.result() for f in futures)) == num_workers
-
-
-if __name__ == "__main__":
- test_dialects_vector_repro_3()
\ No newline at end of file
>From 0683720e18f164812700e57381b90399b53214d0 Mon Sep 17 00:00:00 2001
From: vfdev-5 <vfdev.5 at gmail.com>
Date: Fri, 13 Dec 2024 18:32:42 +0100
Subject: [PATCH 5/5] Updated locks and added docs
---
mlir/docs/Bindings/Python.md | 40 ++++++++++++++
mlir/lib/Bindings/Python/Globals.h | 14 -----
mlir/lib/Bindings/Python/IRCore.cpp | 37 ++++---------
mlir/lib/Bindings/Python/IRModule.cpp | 13 +++++
mlir/lib/Bindings/Python/MainModule.cpp | 8 +--
mlir/test/python/multithreaded_tests.py | 71 +++++++++++++++----------
6 files changed, 109 insertions(+), 74 deletions(-)
diff --git a/mlir/docs/Bindings/Python.md b/mlir/docs/Bindings/Python.md
index 32df3310d811d7..bf181964b72dc0 100644
--- a/mlir/docs/Bindings/Python.md
+++ b/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`.
diff --git a/mlir/lib/Bindings/Python/Globals.h b/mlir/lib/Bindings/Python/Globals.h
index 05400608ba9ffa..facc1428a1f9d5 100644
--- a/mlir/lib/Bindings/Python/Globals.h
+++ b/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() {
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index ed2d32d6bbe7f5..f3633c0a0c9a3a 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/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);
diff --git a/mlir/lib/Bindings/Python/IRModule.cpp b/mlir/lib/Bindings/Python/IRModule.cpp
index 6727860c094a2a..7c1deb13a0a405 100644
--- a/mlir/lib/Bindings/Python/IRModule.cpp
+++ b/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;
}
diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp
index 590d0590f51e4a..397455d61561b6 100644
--- a/mlir/lib/Bindings/Python/MainModule.cpp
+++ b/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;
diff --git a/mlir/test/python/multithreaded_tests.py b/mlir/test/python/multithreaded_tests.py
index ab5b7bf388909c..6a5307667b87a1 100644
--- a/mlir/test/python/multithreaded_tests.py
+++ b/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,19 +484,18 @@ 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
@@ -496,11 +503,9 @@ def run_construct_and_print_in_module(f):
"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,
)
- at add_existing_tests(test_prefix="_original_test")
+ at 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()
More information about the Mlir-commits
mailing list