[Mlir-commits] [mlir] [MLIR][Python] Fix stale multithreaded_tests.py (PR #237018)
Maksim Levental
llvmlistbot at llvm.org
Thu Oct 29 16:46:40 PDT 2026
https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/237018
## What was broken
`mlir/test/python/multithreaded_tests.py` is broken on `main` when run on a free-threaded CPython (checked with 3.14t). Every generated test failed, and then the process aborted (SIGABRT) at `test_dialects_builtin__testBuildFuncOp_multi_threaded`. There were several separate problems:
1. **Removed APIs.** `_original_test_create_context` still called `Context._get_live_operation_count`, `_get_live_operation_objects` and `_clear_live_operations`. These were removed in #155114. I dropped the calls.
2. **Live-context check.** Each generated test asserted `Context._get_live_count() == 0`. However, `ir/array_attributes.py` deliberately keeps a module-level `Context` alive (`testDanglingResource`, #150561), so this assertion failed for every test. The check now compares against the number of live contexts taken before the test.
3. **Caster hijack causing the abort.** Importing `_mlirPythonTestNanobind` (which `ir/diagnostic_handler.py` does) runs `PyTestIntegerRankedTensorType::bind`. That registers `TestIntegerRankedTensorType` as the caster for the builtin `RankedTensorType` TypeID, but its isa check only accepts integer element types. After that, every non-integer ranked tensor type fails to downcast with `ValueError`. Inside the `noexcept` `type_caster<MlirType>::from_cpp` this calls `std::terminate`; #237010 fixes that part separately. Here, the test module restores the builtin caster right after binding the test type.
- I kept `getTypeIdFunction`, because `python_test.py` uses `TestIntegerRankedTensorType.static_typeid` and expects "Type caster is already registered".
- `testCustomTypeTypeCaster` registers the custom caster explicitly, so nothing depends on the automatic registration.
4. **Collection bug.** `add_existing_tests` collected every callable attribute whose name starts with `test`, including imported ones such as the C++ helper `test_diagnostics_with_errors_and_notes`. That helper then ran without a Context. Only functions whose `__module__` is the source test module are collected now.
5. **Test-order interference.** `testMatchInterfaceEnumReplaceAttributeBuilder` globally replaces the `MatchInterfaceEnum` attribute builder with one that only handles strings. That broke `testMatchInterfaceEnum`, which sorts after it. Only the test that does the replacing is skipped, so `testMatchInterfaceEnum` keeps running.
6. **Nondeterministic xfails.**
- `TESTS_TO_XFAIL` entries are only reliably expected failures under TSAN, where the harness raises on captured "ThreadSanitizer" output. The original free-threading PRs (#107103 / #122684) assumed "CPython compiled with TSAN" but never detected it.
- Without TSAN, those tests randomly pass (unexpected successes) or fail. They are now skipped unless CPython is built with TSAN. The check mirrors CPython's `test.support.check_sanitizer(thread=True)`: `-fsanitize=thread` in `CFLAGS` or `--with-thread-sanitizer` in `CONFIG_ARGS`.
- Tests that use the `mlir.runtime.np_to_memref` descriptor helpers (`get_ranked_memref_descriptor`, `get_unranked_memref_descriptor`, `make_nd_memref_descriptor`) get the same treatment: xfail under TSAN, skip otherwise. They are found by scanning each test's code object, not by listing names.
- These helpers race on concurrent ctypes type creation (python/cpython#127945), e.g. `TypeError: incompatible types, LP_c_float instance instead of LP_c_float instance`. The race reproduces by calling `get_ranked_memref_descriptor` from 10 threads after loading `np_to_memref.py` on its own, with no MLIR native code loaded.
- Result: without TSAN, none of the generated execution_engine tests run any more. With TSAN, they behave as before (xfail).
## Why it went unnoticed
The test only runs when the GIL is disabled (`sys._is_gil_enabled()` is False). On regular CPython builds it is a no-op that exits 0, so CI never ran it.
## Verification
- CPython 3.14t (free-threaded) on macOS arm64, MLIR Release + assertions build:
- `multithreaded_tests.py`: 5 out of 5 runs OK, 519 tests each, with no FAIL, ERROR or unexpected success.
- `check-mlir-python`: 103 passed, 8 unsupported, 0 failed.
- Before this change, on the same setup: 0 ok, 18 FAIL/ERROR, then SIGABRT (exit 134).
- On a GIL build the test is still a no-op (exit 0).
This change was prepared with AI assistance.
>From f65727e1d1a211a9331e31ccaf27b6bc79ff9ed0 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Thu, 29 Oct 2026 16:46:06 -0700
Subject: [PATCH] [MLIR][Python] Fix stale multithreaded_tests.py
multithreaded_tests.py only runs when the GIL is disabled, so it is a
no-op on regular CPython builds and has silently rotted. On a
free-threaded CPython 3.14t build it failed every test and then aborted:
- `_original_test_create_context` called `_get_live_operation_count`,
`_get_live_operation_objects` and `_clear_live_operations`, which were
removed in #155114. Drop those calls.
- Each generated test asserted `Context._get_live_count() == 0`, but
ir/array_attributes.py intentionally keeps a module-level Context alive
(testDanglingResource), so every test failed. Compare against the
number of live contexts before the test instead.
- Importing _mlirPythonTestNanobind (done by ir/diagnostic_handler.py)
registered TestIntegerRankedTensorType, whose isa check only accepts
integer element types, as the caster for the builtin RankedTensorType
TypeID. Every non-integer ranked tensor then failed to downcast, which
aborted the process through the noexcept MlirType from_cpp caster
(fixed separately in #237010). Restore the builtin caster after binding
the test type; python_test.py registers the custom caster explicitly.
- Test collection picked up imported callables whose names start with
`test`, such as the C++ helper test_diagnostics_with_errors_and_notes.
Only collect functions defined in the source test module.
- testMatchInterfaceEnumReplaceAttributeBuilder globally replaces the
MatchInterfaceEnum attribute builder, which breaks testMatchInterfaceEnum
when it runs afterwards. Skip it.
- TESTS_TO_XFAIL entries only fail reliably when TSAN reports races;
without TSAN they fail nondeterministically (unexpected successes or
ctypes TypeErrors). Skip them unless CPython is built with TSAN, and
treat all tests using the np_to_memref descriptor helpers, which race
on ctypes type creation (python/cpython#127945), the same way.
---
mlir/test/python/dialects/python_test.py | 4 +-
.../python/lib/PythonTestModuleNanobind.cpp | 10 +++
mlir/test/python/multithreaded_tests.py | 67 +++++++++++++++++--
3 files changed, 73 insertions(+), 8 deletions(-)
diff --git a/mlir/test/python/dialects/python_test.py b/mlir/test/python/dialects/python_test.py
index 917d1eda747d09..1b4ea79456c303 100644
--- a/mlir/test/python/dialects/python_test.py
+++ b/mlir/test/python/dialects/python_test.py
@@ -787,8 +787,8 @@ def type_caster(pytype):
except RuntimeError as e:
print(e)
- # python_test dialect registers a caster for RankedTensorType in its extension (pybind) module.
- # So this one replaces that one (successfully). And then just to be sure we restore the original caster below.
+ # A caster for RankedTensorType is already registered (the builtin one).
+ # So this one replaces that one (successfully). And then we register the custom caster below.
@register_type_caster(c.typeid, replace=True)
def type_caster(pytype):
return RankedTensorType(pytype)
diff --git a/mlir/test/python/lib/PythonTestModuleNanobind.cpp b/mlir/test/python/lib/PythonTestModuleNanobind.cpp
index 36fc7b433752d1..f7003946d88161 100644
--- a/mlir/test/python/lib/PythonTestModuleNanobind.cpp
+++ b/mlir/test/python/lib/PythonTestModuleNanobind.cpp
@@ -178,5 +178,15 @@ NB_MODULE(_mlirPythonTestNanobind, m) {
PyTestAttr::bind(m);
PyTestType::bind(m);
PyTestIntegerRankedTensorType::bind(m);
+ // `bind` registers `TestIntegerRankedTensorType` as the type caster for the
+ // builtin `RankedTensorType` TypeID. Its isa check only accepts integer
+ // element types, so every other ranked tensor type would then fail to
+ // downcast as soon as this module is imported. Restore the builtin caster;
+ // tests opt into the custom caster explicitly via `register_type_caster`.
+ PyGlobals::get().registerTypeCaster(
+ mlirRankedTensorTypeGetTypeID(),
+ nb::cast<nb::callable>(nb::cpp_function(
+ [](PyType pyType) { return PyRankedTensorType(pyType); })),
+ /*replace=*/true);
PyTestTensorValue::bind(m);
}
diff --git a/mlir/test/python/multithreaded_tests.py b/mlir/test/python/multithreaded_tests.py
index 6e1a6683468729..69fbc4bdcc8e17 100644
--- a/mlir/test/python/multithreaded_tests.py
+++ b/mlir/test/python/multithreaded_tests.py
@@ -40,8 +40,10 @@
import importlib.util
import os
import sys
+import sysconfig
import threading
import tempfile
+import types
import unittest
from contextlib import contextmanager
@@ -284,6 +286,7 @@ def run_construct_and_print_in_module(f):
"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", # RuntimeError: Value caster is already registered: <function testValueCasters.<locals>.dont_cast_int, even with GIL
+ "test_dialects_transform_structured_ext__testMatchInterfaceEnumReplaceAttributeBuilder_multi_threaded", # Globally replaces the MatchInterfaceEnum attribute builder, which breaks testMatchInterfaceEnum, even with GIL
# tests indirectly calling thread-unsafe llvm::raw_ostream
"test_execution_engine__testInvalidModule_multi_threaded", # mlirExecutionEngineCreate calls thread-unsafe llvm::raw_ostream
"test_pass_manager__testPrintIrAfterAll_multi_threaded", # IRPrinterInstrumentation::runAfterPass calls thread-unsafe llvm::raw_ostream
@@ -311,6 +314,8 @@ def run_construct_and_print_in_module(f):
"test_execution_engine__testComplexUnrankedMemrefAdd_multi_threaded",
"test_execution_engine__testDynamicMemrefAdd2D_multi_threaded",
"test_execution_engine__testF16MemrefAdd_multi_threaded",
+ "test_execution_engine__testF8E3M4Memref_multi_threaded",
+ "test_execution_engine__testF8E4M3Memref_multi_threaded",
"test_execution_engine__testF8E5M2Memref_multi_threaded",
"test_execution_engine__testInvokeFloatAdd_multi_threaded",
"test_execution_engine__testInvokeVoid_multi_threaded", # a ctypes race
@@ -340,6 +345,40 @@ def run_construct_and_print_in_module(f):
]
+# The tests above are expected to fail only when ThreadSanitizer reports
+# warnings. Without TSAN the underlying races (e.g. in ctypes) make them fail
+# nondeterministically, so skip them instead. The check mirrors CPython's
+# test.support.check_sanitizer(thread=True).
+def is_tsan_enabled() -> bool:
+ cflags = sysconfig.get_config_var("CFLAGS") or ""
+ config_args = sysconfig.get_config_var("CONFIG_ARGS") or ""
+ return "-fsanitize=thread" in cflags or "--with-thread-sanitizer" in config_args
+
+
+if not is_tsan_enabled():
+ TESTS_TO_SKIP += TESTS_TO_XFAIL
+ TESTS_TO_XFAIL = []
+
+# Tests converting numpy arrays to memref descriptors with these helpers from
+# mlir.runtime.np_to_memref race on concurrent ctypes type creation
+# (https://github.com/python/cpython/issues/127945), e.g. "TypeError:
+# incompatible types, LP_c_float instance instead of LP_c_float instance".
+# The race reproduces with np_to_memref alone, without any MLIR native code.
+# Like TESTS_TO_XFAIL, they are expected to fail with TSAN and skipped otherwise.
+NP_TO_MEMREF_DESCRIPTOR_FNS = {
+ "get_ranked_memref_descriptor",
+ "get_unranked_memref_descriptor",
+ "make_nd_memref_descriptor",
+}
+
+
+def uses_np_to_memref_descriptors(code: types.CodeType) -> bool:
+ return any(name in NP_TO_MEMREF_DESCRIPTOR_FNS for name in code.co_names) or any(
+ isinstance(const, types.CodeType) and uses_np_to_memref_descriptors(const)
+ for const in code.co_consts
+ )
+
+
def add_existing_tests(test_modules, test_prefix: str = "_original_test"):
def decorator(test_cls):
this_folder = Path(__file__).parent.absolute()
@@ -367,7 +406,12 @@ def decorator(test_cls):
)
if is_test_fn:
obj = getattr(test_mod, attr_name)
- if callable(obj):
+ # Only collect functions defined in the test module itself,
+ # not imported helpers (e.g. test_diagnostics_with_errors_and_notes).
+ if (
+ callable(obj)
+ and getattr(obj, "__module__", None) == test_mod.__name__
+ ):
test_name = f"{test_prefix}_{test_module_name.replace('/', '_')}__{attr_name}"
def wrapped_test_fn(
@@ -375,6 +419,9 @@ def wrapped_test_fn(
):
__exec_fn__(__test_fn__)
+ wrapped_test_fn.uses_np_to_memref_descriptors = hasattr(
+ obj, "__code__"
+ ) and uses_np_to_memref_descriptors(obj.__code__)
setattr(test_cls, test_name, wrapped_test_fn)
return test_cls
@@ -434,6 +481,11 @@ def decorator(test_cls):
continue
def multi_threaded_test_fn(self, *args, __test_fn__=test_fn, **kwargs):
+ # Some source test modules intentionally keep a Context alive at
+ # module scope (e.g. ir/array_attributes.py), so compare against
+ # the number of live contexts before the test instead of zero.
+ gc.collect()
+ num_live_contexts = Context._get_live_count()
with capture_stdout(), capture_stderr() as get_output:
barrier = threading.Barrier(num_workers)
@@ -453,7 +505,7 @@ def closure():
assert len(list(f.result() for f in futures)) == num_workers
gc.collect()
- assert Context._get_live_count() == 0
+ assert Context._get_live_count() == num_live_contexts
captured = get_output()
if len(captured) > 0 and "ThreadSanitizer" in captured:
@@ -462,7 +514,13 @@ def closure():
)
test_new_name = f"{name}{multithreaded_test_postfix}"
- if xfail_tests is not None and test_new_name in xfail_tests:
+ is_xfail = xfail_tests is not None and test_new_name in xfail_tests
+ if getattr(test_fn, "uses_np_to_memref_descriptors", False):
+ # See NP_TO_MEMREF_DESCRIPTOR_FNS.
+ if not is_tsan_enabled():
+ continue
+ is_xfail = True
+ if is_xfail:
multi_threaded_test_fn = unittest.expectedFailure(
multi_threaded_test_fn
)
@@ -491,10 +549,7 @@ 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]
More information about the Mlir-commits
mailing list