[MLIR][Python] Fix GIL handling in verification and printing - #215848
[MLIR][Python] Fix GIL handling in verification and printing#215848hahalfx wants to merge 2 commits into
Conversation
Recursive verification may process isolated operations on MLIR worker threads. When such a traversal invokes a Python-defined verifier, diagnostic handler, or memory-effects callback without the GIL, the process can crash. Release the GIL in Python entry points that perform verification or create and use printing state, and reacquire it in the Python callback trampolines reachable from those paths. Add regression coverage for valid and invalid multithreaded verification, diagnostic propagation, printing, and AsmState creation. PassManager execution and unrelated Python callback trampolines remain outside the scope of this change and require a separate audit.
|
Hello @hahalfx 👋 Thank you for submitting a Pull Request (PR) to the LLVM Project. Since this is your first PR, here are a few useful links covering our main contribution policies and review practices.
Please reply to this message to confirm that you have read these policies, especially the LLVM AI Tool Use Policy, and that any AI tool usage has been noted in the PR description. Frequently asked questionsHow do I add reviewers? This PR will be automatically labeled, and the relevant teams will be notified. For some parts of the project, reviewers may also be added automatically. You can also add reviewers manually using the Reviewers section on this page. If you cannot use that section, it is probably because you do not have write permissions for the repository. In that case, you can request a review by tagging reviewers in a comment using What if there are no comments? If you have not received any comments on your PR after a week, you can request a review by pinging the PR with a comment such as “Ping”. The common courtesy ping rate is once a week. Please remember that you are asking for volunteer time from other developers. Are any special GitHub settings required to contribute to LLVM? We only require contributors to have a public email address associated with their GitHub commits, see this section of LLVM Developer Policy for details. If you have questions, feel free to leave a comment on this PR, or ask on LLVM Discord or LLVM Discourse. Thank you, |
|
@llvm/pr-subscribers-mlir Author: Lei Fengxiang (hahalfx) ChangesSummaryI 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 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:
Regression tests added by this change cover successful and failing multithreaded Python verification, diagnostics emitted by native verifiers, printing, dumping, and Fixes #215781. ScopeThis change is intentionally limited to Python
Testing
Full diff: https://github.com/llvm/llvm-project/pull/215848.diff 6 Files Affected:
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
+@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
+@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")
|
makslevental
left a comment
makslevental
left a comment
There was a problem hiding this comment.
Highly suspect this much GIL manipulation is actually necessary. Not blocking indefinitely but just flagging we need to take a much closer look than just spamming nb::gil_scoped_acquire.
|
I think this problem can be modelled as: And the fix can be modelled as: So I think this fix can indeed solve the issue. Though I don't know if this is the best fix. This is a very general problem than just for operation verification. It exists for any pattern of One issue I saw for such fix is that we need to be careful that |
|
Thanks @PragmaTwice — that matches my understanding as well. And thanks @makslevental for pushing back on the amount of GIL handling here. I went back to the smallest fix first, then re-tested the remaining changes against failures I could reproduce independently. 1. Starting with the original
|
|
@hahalfx thanks for the investigation. I still haven't had time to look at this closely (though I did skim your current summary). Just gonna paste here what I told @PragmaTwice offline: The problem is the c++ back to python call isn't the common cse. In this specific case obviously most users don't have verifiers written in python. So if you put explicitly do GIL manipulation you're pessimizing the common path. My preference would be to just warn people about the sharp edge ehre: if you have Python callbacks then you need to turn off multithreading. Or at minimum make the GIL optional in those APIs (via a param/flag). But I'll come back soon. |
3 participants
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 nestedIsolatedFromAboveoperations on worker threads. If that verification reaches a Python-definedDynamicOpTraitverifier, 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:
Operation.verify()and printing/AsmStateentry points that may invoke verification internally;MemoryEffectsOpInterfacegetEffectscallback 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
AsmStateconstruction. The Python-backedgetEffectspath is also exercised by existing calls tonamed_seq.verify()intransform_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/AsmStateentry 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.runis 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
MLIRPythonModulescheck-mlir-pythontransform_op_interface.pycoverage of Python-backedgetEffectscallbacks reached fromnamed_seq.verify()AsmStateconstruction