[Mlir-commits] [mlir] [MLIR][Python] Fix GIL handling in verification and printing (PR #215848)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Aug 12 10:04:00 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Lei Fengxiang (hahalfx)
<details>
<summary>Changes</summary>
## Summary
I encountered this issue while running an MLIR-based outlining workflow and reduced it to the standalone reproducer reported in #<!-- -->215781.
When verification runs in a multithreaded `MLIRContext`, MLIR may verify nested `IsolatedFromAbove` operations on worker threads. If that verification reaches a Python-defined `DynamicOpTrait` verifier, the worker can enter CPython without holding the GIL, causing a native segmentation fault.
Acquiring the GIL only in the callback is insufficient: if the originating Python thread still holds the GIL while waiting for the MLIR worker, the callback can deadlock while trying to acquire it. This change therefore handles both sides of the affected call paths:
- release the GIL around Python `Operation.verify()` and printing/`AsmState` entry points that may invoke verification internally;
- ensure the GIL is held before accessing Python objects in callbacks reachable from those paths, including dynamic operation verifiers, diagnostic callbacks and capture paths, and print accumulators;
- acquire the GIL in the Python-backed `MemoryEffectsOpInterface` `getEffects` callback because parent Transform dialect operation verifiers can query the effects of nested transform operations.
Regression tests added by this change cover successful and failing multithreaded Python verification, diagnostics emitted by native verifiers, printing, dumping, and `AsmState` construction. The Python-backed `getEffects` path is also exercised by existing calls to `named_seq.verify()` in `transform_op_interface.py`; without the callback-side GIL acquisition, that test segfaults after verification releases the GIL.
Fixes #<!-- -->215781.
## Scope
This change is intentionally limited to Python `Operation.verify()` and printing/`AsmState` entry points that may invoke verification internally. It does not attempt to establish GIL safety for every C++-to-Python callback in the MLIR Python bindings.
`PassManager.run` is left unchanged in this PR. Releasing the GIL around pass execution would require auditing the Python callbacks reachable from that path, including external Python passes; consequently, multithreaded post-pass verification that reaches a Python callback may still deadlock.
## Testing
- `MLIRPythonModules`
- `check-mlir-python`
- Existing `transform_op_interface.py` coverage of Python-backed `getEffects` callbacks reached from `named_seq.verify()`
- Repeated stress testing of multithreaded verification, diagnostics, operation/value/block printing, dumping, and `AsmState` construction
---
Full diff: https://github.com/llvm/llvm-project/pull/215848.diff
6 Files Affected:
- (modified) mlir/include/mlir/Bindings/Python/IRCore.h (+5-2)
- (modified) mlir/include/mlir/Bindings/Python/NanobindUtils.h (+1)
- (modified) mlir/lib/Bindings/Python/IRCore.cpp (+59-26)
- (modified) mlir/lib/Bindings/Python/IRInterfaces.cpp (+2)
- (modified) mlir/test/python/dialects/ext.py (+69-1)
- (modified) mlir/test/python/ir/exception.py (+21)
``````````diff
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 3314e0b2a8fcf..47d579d586e30 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1690,8 +1690,11 @@ class MLIR_PYTHON_API_EXPORTED PyConcreteValue : public PyValue {
cls.def("__str__", [](PyValue &self) {
PyPrintAccumulator printAccum;
printAccum.parts.append(std::string(DerivedTy::pyClassName) + "(");
- mlirValuePrint(self.get(), printAccum.getCallback(),
- printAccum.getUserData());
+ {
+ nanobind::gil_scoped_release gil;
+ mlirValuePrint(self.get(), printAccum.getCallback(),
+ printAccum.getUserData());
+ }
printAccum.parts.append(")");
return printAccum.join();
});
diff --git a/mlir/include/mlir/Bindings/Python/NanobindUtils.h b/mlir/include/mlir/Bindings/Python/NanobindUtils.h
index ea43356d2cf54..594fc31f9f47b 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindUtils.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindUtils.h
@@ -177,6 +177,7 @@ struct PyPrintAccumulator {
MlirStringCallback getCallback() {
return [](MlirStringRef part, void *userData) {
+ nanobind::gil_scoped_acquire gil;
PyPrintAccumulator *printAccum =
static_cast<PyPrintAccumulator *>(userData);
nanobind::str pyPart(part.data,
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 75cfd2a0a1c0b..1faff0b827048 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -516,23 +516,21 @@ nb::object PyMlirContext::attachDiagnosticHandler(nb::object callback) {
// guaranteed to be known to pybind.
auto handlerCallback =
+[](MlirDiagnostic diagnostic, void *userData) -> MlirLogicalResult {
+ // Since this can be called from arbitrary C++ contexts, always get the
+ // GIL before creating any Python objects.
+ nb::gil_scoped_acquire gil;
PyDiagnostic *pyDiagnostic = new PyDiagnostic(diagnostic);
nb::object pyDiagnosticObject =
nb::cast(pyDiagnostic, nb::rv_policy::take_ownership);
auto *pyHandler = static_cast<PyDiagnosticHandler *>(userData);
bool result = false;
- {
- // Since this can be called from arbitrary C++ contexts, always get the
- // gil.
- nb::gil_scoped_acquire gil;
- try {
- result = nb::cast<bool>(pyHandler->callback(pyDiagnostic));
- } catch (std::exception &e) {
- fprintf(stderr, "MLIR Python Diagnostic handler raised exception: %s\n",
- e.what());
- pyHandler->hadError = true;
- }
+ try {
+ result = nb::cast<bool>(pyHandler->callback(pyDiagnostic));
+ } catch (std::exception &e) {
+ fprintf(stderr, "MLIR Python Diagnostic handler raised exception: %s\n",
+ e.what());
+ pyHandler->hadError = true;
}
pyDiagnostic->invalidate();
@@ -564,6 +562,7 @@ MlirLogicalResult PyMlirContext::ErrorCapture::handler(MlirDiagnostic diag,
MlirDiagnosticSeverity::MlirDiagnosticError)
return mlirLogicalResultFailure();
+ nb::gil_scoped_acquire gil;
self->errors.emplace_back(PyDiagnostic(diag).getInfo());
return mlirLogicalResultSuccess();
}
@@ -1051,8 +1050,12 @@ void PyOperationBase::print(std::optional<int64_t> largeElementsLimit,
mlirOpPrintingFlagsPrintNameLocAsPrefix(flags);
PyFileAccumulator accum(fileObject, binary);
- mlirOperationPrintWithFlags(operation, flags, accum.getCallback(),
- accum.getUserData());
+ // Printing creates an AsmState that may recursively invoke Python verifiers.
+ {
+ nb::gil_scoped_release gil;
+ mlirOperationPrintWithFlags(operation, flags, accum.getCallback(),
+ accum.getUserData());
+ }
mlirOpPrintingFlagsDestroy(flags);
}
@@ -1177,7 +1180,15 @@ bool PyOperationBase::isBeforeInBlock(PyOperationBase &other) {
bool PyOperationBase::verify() {
PyOperation &op = getOperation();
PyMlirContext::ErrorCapture errors(op.getContext());
- if (!mlirOperationVerify(op.get()))
+ bool verified;
+ {
+ // Recursive verification may invoke Python callbacks on MLIR worker
+ // threads. Release the GIL here; Python callbacks reachable during
+ // verification reacquire it.
+ nb::gil_scoped_release gil;
+ verified = mlirOperationVerify(op.get());
+ }
+ if (!verified)
throw MLIRError("Verification failed", errors.take());
return true;
}
@@ -1758,7 +1769,10 @@ PyAsmState::PyAsmState(MlirValue value, bool useLocalScope) {
// associate lifetime with the state.
if (useLocalScope)
mlirOpPrintingFlagsUseLocalScope(flags);
- state = mlirAsmStateCreateForValue(value, flags);
+ {
+ nb::gil_scoped_release gil;
+ state = mlirAsmStateCreateForValue(value, flags);
+ }
}
PyAsmState::PyAsmState(PyOperationBase &operation, bool useLocalScope) {
@@ -1767,7 +1781,11 @@ PyAsmState::PyAsmState(PyOperationBase &operation, bool useLocalScope) {
// associate lifetime with the state.
if (useLocalScope)
mlirOpPrintingFlagsUseLocalScope(flags);
- state = mlirAsmStateCreateForOperation(operation.getOperation().get(), flags);
+ {
+ nb::gil_scoped_release gil;
+ state =
+ mlirAsmStateCreateForOperation(operation.getOperation().get(), flags);
+ }
}
//------------------------------------------------------------------------------
@@ -2549,6 +2567,7 @@ void PyOpAdaptor::bind(nb::module_ &m) {
static MlirLogicalResult verifyTraitByMethod(MlirOperation op, void *userData,
const char *methodName) {
+ nb::gil_scoped_acquire gil;
nb::handle targetObj(static_cast<PyObject *>(userData));
if (!nb::hasattr(targetObj, methodName))
return mlirLogicalResultSuccess();
@@ -3890,6 +3909,7 @@ void populateIRCore(nb::module_ &m) {
.def(
"dump",
[](PyModule &self) {
+ nb::gil_scoped_release gil;
mlirOperationDump(mlirModuleGetOperation(self.get()));
},
kDumpDocstring)
@@ -4637,8 +4657,11 @@ void populateIRCore(nb::module_ &m) {
[](PyBlock &self) {
self.checkValid();
PyPrintAccumulator printAccum;
- mlirBlockPrint(self.get(), printAccum.getCallback(),
- printAccum.getUserData());
+ {
+ nb::gil_scoped_release gil;
+ mlirBlockPrint(self.get(), printAccum.getCallback(),
+ printAccum.getUserData());
+ }
return printAccum.join();
},
"Returns the assembly form of the block.")
@@ -5034,7 +5057,11 @@ void populateIRCore(nb::module_ &m) {
},
"Context in which the value lives.")
.def(
- "dump", [](PyValue &self) { mlirValueDump(self.get()); },
+ "dump",
+ [](PyValue &self) {
+ nb::gil_scoped_release gil;
+ mlirValueDump(self.get());
+ },
kDumpDocstring)
.def_prop_ro(
"owner",
@@ -5084,8 +5111,11 @@ void populateIRCore(nb::module_ &m) {
[](PyValue &self) {
PyPrintAccumulator printAccum;
printAccum.parts.append("Value(");
- mlirValuePrint(self.get(), printAccum.getCallback(),
- printAccum.getUserData());
+ {
+ nb::gil_scoped_release gil;
+ mlirValuePrint(self.get(), printAccum.getCallback(),
+ printAccum.getUserData());
+ }
printAccum.parts.append(")");
return printAccum.join();
},
@@ -5105,11 +5135,14 @@ void populateIRCore(nb::module_ &m) {
mlirOpPrintingFlagsUseLocalScope(flags);
if (useNameLocAsPrefix)
mlirOpPrintingFlagsPrintNameLocAsPrefix(flags);
- MlirAsmState valueState =
- mlirAsmStateCreateForValue(self.get(), flags);
- mlirValuePrintAsOperand(self.get(), valueState,
- printAccum.getCallback(),
- printAccum.getUserData());
+ MlirAsmState valueState;
+ {
+ nb::gil_scoped_release gil;
+ valueState = mlirAsmStateCreateForValue(self.get(), flags);
+ mlirValuePrintAsOperand(self.get(), valueState,
+ printAccum.getCallback(),
+ printAccum.getUserData());
+ }
mlirOpPrintingFlagsDestroy(flags);
mlirAsmStateDestroy(valueState);
return printAccum.join();
diff --git a/mlir/lib/Bindings/Python/IRInterfaces.cpp b/mlir/lib/Bindings/Python/IRInterfaces.cpp
index d16015acaf3dc..57695aaf15643 100644
--- a/mlir/lib/Bindings/Python/IRInterfaces.cpp
+++ b/mlir/lib/Bindings/Python/IRInterfaces.cpp
@@ -496,6 +496,8 @@ class PyMemoryEffectsOpInterface
callbacks.getEffects = [](MlirOperation op,
MlirMemoryEffectInstancesList effects,
void *userData) {
+ // Parent transform op verifiers query effects of nested transform ops.
+ nb::gil_scoped_acquire gil;
nb::handle pyClass(static_cast<PyObject *>(userData));
// Get the 'get_effects' method from the Python class.
diff --git a/mlir/test/python/dialects/ext.py b/mlir/test/python/dialects/ext.py
index 6e83da4a4a78a..3be17794c25fb 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -1,7 +1,7 @@
# RUN: %PYTHON %s 2>&1 | FileCheck %s
from mlir.ir import *
-from mlir.dialects import arith
+from mlir.dialects import arith, func
from mlir.dialects.ext import *
from mlir.rewrite import *
from mlir import ir
@@ -637,6 +637,74 @@ class NoTermOp(TestRegion.Operation, name="no_term", traits=[NoTerminatorTrait])
print(e)
+# CHECK: TEST: testDynamicOpTraitMultithreadedVerification
+ at run
+def testDynamicOpTraitMultithreadedVerification():
+ class TestParallelVerify(Dialect, name="ext_parallel_verify"):
+ pass
+
+ class ValidOp(TestParallelVerify.Operation, name="valid"):
+ def verify_invariants(self) -> bool:
+ return True
+
+ class InvalidOp(TestParallelVerify.Operation, name="invalid"):
+ def verify_invariants(self) -> bool:
+ self.location.emit_error("parallel Python verifier failed")
+ return False
+
+ def add_function(module, name, op_type, add_result=False):
+ i32 = IntegerType.get_signless(32)
+ with InsertionPoint(module.body):
+ result_types = [i32] if add_result else []
+ function = func.FuncOp(name, ([], result_types))
+ block = function.add_entry_block()
+ with InsertionPoint(block):
+ value = None
+ return_values = []
+ if add_result:
+ value = arith.constant(i32, 0)
+ return_values = [value]
+ op_type()
+ func.ReturnOp(return_values)
+ return function, value
+
+ context = Context()
+ context.enable_multithreading(True)
+ with context, Location.unknown():
+ TestParallelVerify.load()
+
+ module = Module.create()
+ first_function, value = add_function(
+ module, "first_valid", ValidOp, add_result=True
+ )
+ add_function(module, "second_valid", ValidOp, add_result=True)
+ assert module.operation.verify()
+ assert "ext_parallel_verify.valid" in str(module)
+ # AsmState construction implicitly verifies the parent operation.
+ AsmState(module.operation)
+ AsmState(value)
+ assert value.get_name()
+ assert "ext_parallel_verify.valid" in str(first_function.body.blocks[0])
+ assert "arith.constant" in str(value)
+ assert "arith.constant" in str(Value(value))
+ module.dump()
+ value.dump()
+ # CHECK: parallel verification succeeded
+ print("parallel verification succeeded")
+
+ module = Module.create()
+ add_function(module, "first_invalid", InvalidOp)
+ add_function(module, "second_invalid", InvalidOp)
+ try:
+ module.operation.verify()
+ except MLIRError as e:
+ assert "parallel Python verifier failed" in str(e)
+ # CHECK: parallel verification failure captured
+ print("parallel verification failure captured")
+ else:
+ raise AssertionError("expected parallel verification to fail")
+
+
# CHECK: TEST: testIsIsolatedFromAboveTrait
@run
def testIsIsolatedFromAboveTrait():
diff --git a/mlir/test/python/ir/exception.py b/mlir/test/python/ir/exception.py
index 74085cd349643..7d323282c25ea 100644
--- a/mlir/test/python/ir/exception.py
+++ b/mlir/test/python/ir/exception.py
@@ -2,6 +2,7 @@
import gc
from mlir.ir import *
+from mlir.dialects import func
def run(f):
@@ -93,3 +94,23 @@ def handler(d):
print(f"emit_error_diagnostics=True:")
print(f"e.error_diagnostics: {[str(diag) for diag in e.error_diagnostics]}")
print(f"handler_diags: {handler_diags}")
+
+
+# CHECK-LABEL: TEST: test_multithreaded_native_verifier_diagnostics
+ at run
+def test_multithreaded_native_verifier_diagnostics():
+ context = Context()
+ context.enable_multithreading(True)
+ with context, Location.unknown():
+ module = Module.create()
+ with InsertionPoint(module.body):
+ func.FuncOp("first_native_invalid", ([], [])).add_entry_block()
+ func.FuncOp("second_native_invalid", ([], [])).add_entry_block()
+ try:
+ module.operation.verify()
+ except MLIRError as e:
+ assert str(e).count("empty block: expect at least a terminator") == 2
+ # CHECK: parallel native diagnostics captured
+ print("parallel native diagnostics captured")
+ else:
+ raise AssertionError("expected native verification to fail")
``````````
</details>
https://github.com/llvm/llvm-project/pull/215848
More information about the Mlir-commits
mailing list