[Mlir-commits] [mlir] f2f65ac - [MLIR][Python] Release GIL during long-running operations. (#236861)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Oct 31 10:37:18 PDT 2026
Author: Esteban Pauli
Date: 2026-10-31T10:37:14-07:00
New Revision: f2f65acafb421e651d4d1d679bee7a997e4d305e
URL: https://github.com/llvm/llvm-project/commit/f2f65acafb421e651d4d1d679bee7a997e4d305e
DIFF: https://github.com/llvm/llvm-project/commit/f2f65acafb421e651d4d1d679bee7a997e4d305e.diff
LOG: [MLIR][Python] Release GIL during long-running operations. (#236861)
Made changes to release the GIL during some potentially long-running
operations. These are running a pass manager and creating an execution
engine. Because these may call back to Python, several changes had to be
made to re-acquire the GIL at certain points.
Fixes https://github.com/llvm/llvm-project/issues/236564.
These changes were made with the help of AI. I used it to help analyze
the code for thread safety issues and to help find where the GIL would
need to be re-acquired. I also used it to help write the code.
Co-authored-by: Maksim Levental <maksim.levental at gmail.com>
Added:
mlir/test/python/lib/TestConcurrencyPasses.cpp
mlir/test/python/pass_manager_concurrency.py
Modified:
mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
mlir/lib/Bindings/Python/IRCore.cpp
mlir/lib/Bindings/Python/IRInterfaces.cpp
mlir/lib/Bindings/Python/Pass.cpp
mlir/python/mlir/dialects/python_test.py
mlir/test/python/lib/CMakeLists.txt
mlir/test/python/lib/PythonTestCAPI.h
mlir/test/python/lib/PythonTestModuleNanobind.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
index 76dbd54ef04670..4d08c4da5d0da5 100644
--- a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
+++ b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
@@ -89,9 +89,17 @@ NB_MODULE(_mlirExecutionEngine, m) {
libPaths.reserve(sharedLibPaths.size());
for (const std::string &path : sharedLibPaths)
libPaths.push_back({path.c_str(), path.length()});
- MlirExecutionEngine executionEngine = mlirExecutionEngineCreate(
- module.get(), optLevel, libPaths.size(), libPaths.data(),
- enableObjectDump, enablePIC);
+
+ // Access any Python objects before releasing the GIL.
+ MlirModule mlirModule = module.get();
+
+ MlirExecutionEngine executionEngine;
+ {
+ nb::gil_scoped_release release;
+ executionEngine = mlirExecutionEngineCreate(
+ mlirModule, optLevel, libPaths.size(), libPaths.data(),
+ enableObjectDump, enablePIC);
+ }
if (mlirExecutionEngineIsNull(executionEngine))
throw std::runtime_error(
"Failure while creating the ExecutionEngine.");
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index bb0a49ba93ce23..ccf897049d6524 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -516,29 +516,28 @@ 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.
+ 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();
return result ? mlirLogicalResultSuccess() : mlirLogicalResultFailure();
};
auto deleteCallback = +[](void *userData) {
+ // May run during context destruction, so always get the GIL.
+ nb::gil_scoped_acquire gil;
auto *pyHandler = static_cast<PyDiagnosticHandler *>(userData);
assert(pyHandler->registeredID && "handler is not registered");
pyHandler->registeredID.reset();
@@ -564,6 +563,10 @@ MlirLogicalResult PyMlirContext::ErrorCapture::handler(MlirDiagnostic diag,
MlirDiagnosticSeverity::MlirDiagnosticError)
return mlirLogicalResultFailure();
+ // This may be invoked from a thread that is not attached to the
+ // interpreter (e.g. while PassManager.run has released the GIL),
+ // so reacquire the GIL.
+ nb::gil_scoped_acquire acquire;
self->errors.emplace_back(PyDiagnostic(diag).getInfo());
return mlirLogicalResultSuccess();
}
diff --git a/mlir/lib/Bindings/Python/IRInterfaces.cpp b/mlir/lib/Bindings/Python/IRInterfaces.cpp
index cfe9bc5eefa5b0..e6474367032713 100644
--- a/mlir/lib/Bindings/Python/IRInterfaces.cpp
+++ b/mlir/lib/Bindings/Python/IRInterfaces.cpp
@@ -473,9 +473,14 @@ class PyConditionallySpeculatableOpInterface
nb::handle(static_cast<PyObject *>(callbacks.userData)).inc_ref();
callbacks.construct = nullptr;
callbacks.destruct = [](void *userData) {
+ nb::gil_scoped_acquire acquire;
nb::handle(static_cast<PyObject *>(userData)).dec_ref();
};
callbacks.getSpeculatability = [](MlirOperation op, void *userData) {
+ // Interface methods can be queried from native passes while
+ // PassManager.run has detached from the interpreter; re-attach before
+ // calling back into Python.
+ nb::gil_scoped_acquire acquire;
nb::handle pyClass(static_cast<PyObject *>(userData));
auto pyGetSpeculatability =
@@ -541,11 +546,16 @@ class PyMemoryEffectsOpInterface
nb::handle(static_cast<PyObject *>(callbacks.userData)).inc_ref();
callbacks.construct = nullptr;
callbacks.destruct = [](void *userData) {
+ nb::gil_scoped_acquire acquire;
nb::handle(static_cast<PyObject *>(userData)).dec_ref();
};
callbacks.getEffects = [](MlirOperation op,
MlirMemoryEffectInstancesCallback callback,
void *callbackUserData, void *userData) {
+ // Interface methods can be queried from native passes while
+ // PassManager.run has detached from the interpreter; re-attach before
+ // calling back into Python.
+ nb::gil_scoped_acquire acquire;
nb::handle pyClass(static_cast<PyObject *>(userData));
// Get the 'get_effects' method from the Python class.
diff --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp
index 6cf5066b244ebc..9da1dbc92ab7bf 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -209,9 +209,11 @@ void populatePassManagerSubmodule(nb::module_ &m) {
MlirTypeID passID = PyGlobals::get().allocateTypeID();
MlirExternalPassCallbacks callbacks;
callbacks.construct = [](void *obj) {
+ nb::gil_scoped_acquire acquire;
(void)nb::handle(static_cast<PyObject *>(obj)).inc_ref();
};
callbacks.destruct = [](void *obj) {
+ nb::gil_scoped_acquire acquire;
(void)nb::handle(static_cast<PyObject *>(obj)).dec_ref();
};
callbacks.initialize = nullptr;
@@ -220,6 +222,7 @@ void populatePassManagerSubmodule(nb::module_ &m) {
};
callbacks.run = [](MlirOperation op, MlirExternalPass pass,
void *userData) {
+ nb::gil_scoped_acquire acquire;
nb::handle(static_cast<PyObject *>(userData))(
op, PyMlirExternalPass{pass.ptr});
};
@@ -252,8 +255,13 @@ void populatePassManagerSubmodule(nb::module_ &m) {
[](PyPassManager &passManager, PyOperationBase &op) {
// Actually run the pass manager.
PyMlirContext::ErrorCapture errors(op.getOperation().getContext());
- MlirLogicalResult status = mlirPassManagerRunOnOp(
- passManager.get(), op.getOperation().get());
+ MlirPassManager pm = passManager.get();
+ MlirOperation operation = op.getOperation().get();
+ MlirLogicalResult status;
+ {
+ nb::gil_scoped_release release;
+ status = mlirPassManagerRunOnOp(pm, operation);
+ }
if (mlirLogicalResultIsFailure(status))
throw MLIRError("Failure while executing pass pipeline",
errors.take());
diff --git a/mlir/python/mlir/dialects/python_test.py b/mlir/python/mlir/dialects/python_test.py
index f0bfcd8b68b35d..f88c9724ef5fdd 100644
--- a/mlir/python/mlir/dialects/python_test.py
+++ b/mlir/python/mlir/dialects/python_test.py
@@ -10,3 +10,9 @@ def register_python_test_dialect(registry):
from .._mlir_libs import _mlirPythonTestNanobind
_mlirPythonTestNanobind.register_dialect(registry)
+
+
+def register_rendezvous_pass():
+ from .._mlir_libs import _mlirPythonTestNanobind
+
+ _mlirPythonTestNanobind.register_rendezvous_pass()
diff --git a/mlir/test/python/lib/CMakeLists.txt b/mlir/test/python/lib/CMakeLists.txt
index f51a7b44e64dd1..0b84d456b3c597 100644
--- a/mlir/test/python/lib/CMakeLists.txt
+++ b/mlir/test/python/lib/CMakeLists.txt
@@ -2,6 +2,7 @@ set(LLVM_OPTIONAL_SOURCES
PythonTestCAPI.cpp
PythonTestDialect.cpp
PythonTestModuleNanobind.cpp
+ TestConcurrencyPasses.cpp
)
add_mlir_library(MLIRPythonTestDialect
@@ -20,6 +21,7 @@ mlir_target_link_libraries(MLIRPythonTestDialect PUBLIC
add_mlir_public_c_api_library(MLIRCAPIPythonTestDialect
PythonTestCAPI.cpp
+ TestConcurrencyPasses.cpp
DEPENDS
MLIRPythonTestIncGen
@@ -27,6 +29,7 @@ add_mlir_public_c_api_library(MLIRCAPIPythonTestDialect
LINK_LIBS PUBLIC
MLIRCAPIInterfaces
MLIRCAPIIR
+ MLIRPass
MLIRPythonTestDialect
)
diff --git a/mlir/test/python/lib/PythonTestCAPI.h b/mlir/test/python/lib/PythonTestCAPI.h
index 5ffd0cbf6201eb..c4e6ca349211d2 100644
--- a/mlir/test/python/lib/PythonTestCAPI.h
+++ b/mlir/test/python/lib/PythonTestCAPI.h
@@ -36,6 +36,10 @@ MLIR_CAPI_EXPORTED bool mlirTypeIsAPythonTestTestTensorValue(MlirValue value);
MLIR_CAPI_EXPORTED void mlirPythonTestEmitDiagnosticWithNote(MlirContext ctx);
+/// Registers the `test-python-rendezvous` pass (see
+/// TestConcurrencyPasses.cpp). Idempotent.
+MLIR_CAPI_EXPORTED void mlirPythonTestRegisterRendezvousPass(void);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/test/python/lib/PythonTestModuleNanobind.cpp b/mlir/test/python/lib/PythonTestModuleNanobind.cpp
index f7003946d88161..82874159fa6e79 100644
--- a/mlir/test/python/lib/PythonTestModuleNanobind.cpp
+++ b/mlir/test/python/lib/PythonTestModuleNanobind.cpp
@@ -138,6 +138,11 @@ NB_MODULE(_mlirPythonTestNanobind, m) {
nb::sig("def register_dialect(registry: " MAKE_MLIR_PYTHON_QUALNAME("ir.DialectRegistry") ") -> None"));
// clang-format on
+ m.def(
+ "register_rendezvous_pass",
+ []() { mlirPythonTestRegisterRendezvousPass(); },
+ "Registers the test-python-rendezvous pass used by concurrency tests.");
+
m.def(
"test_diagnostics_with_errors_and_notes",
[](DefaultingPyMlirContext ctx) {
diff --git a/mlir/test/python/lib/TestConcurrencyPasses.cpp b/mlir/test/python/lib/TestConcurrencyPasses.cpp
new file mode 100644
index 00000000000000..7a313a0c21b85f
--- /dev/null
+++ b/mlir/test/python/lib/TestConcurrencyPasses.cpp
@@ -0,0 +1,124 @@
+//===- TestConcurrencyPasses.cpp - Passes for binding concurrency tests ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// A test pass that meets other threads at a named native barrier.
+// The Python bindings tests use it to prove that several
+// PassManagers, on independent contexts, execute native pass code
+// concurrently: the barrier can only be satisfied while all participating
+// threads are inside Pass::runOnOperation at the same time, which requires
+// PassManager.run to detach from the Python interpreter.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PythonTestCAPI.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+
+#include <chrono>
+#include <condition_variable>
+#include <map>
+#include <mutex>
+#include <optional>
+#include <string>
+
+namespace {
+/// A single-use counting barrier.
+class RendezvousBarrier {
+public:
+ /// Returns the process-global barrier named `id`, creating it on first use.
+ static RendezvousBarrier &get(const std::string &id) {
+ static std::mutex mapMutex;
+ // std::map rather than an LLVM container: barriers must never move, and
+ // the returned reference must stay valid across later insertions.
+ static std::map<std::string, RendezvousBarrier> barriers;
+ std::lock_guard<std::mutex> lock(mapMutex);
+ return barriers[id];
+ }
+
+ /// Arrives at the barrier and blocks until `participants` threads have
+ /// arrived. Returns std::nullopt on success. On timeout, or if the barrier
+ /// was already poisoned, poisons the barrier and returns the number of
+ /// arrivals observed.
+ std::optional<unsigned> arriveAndWait(unsigned participants,
+ std::chrono::seconds timeout) {
+ std::unique_lock<std::mutex> lock(mutex);
+ if (!poisoned) {
+ if (++arrived >= participants) {
+ cv.notify_all();
+ return std::nullopt;
+ }
+ bool satisfied = cv.wait_for(
+ lock, timeout, [&] { return arrived >= participants || poisoned; });
+ if (satisfied && !poisoned)
+ return std::nullopt;
+ }
+ poisoned = true;
+ cv.notify_all();
+ return arrived;
+ }
+
+private:
+ std::mutex mutex;
+ std::condition_variable cv;
+
+ /// Number of threads that have reached the barrier.
+ unsigned arrived = 0;
+
+ /// Set when a participant times out, so that latecomers fail fast instead
+ /// of waiting out their own timeout.
+ bool poisoned = false;
+};
+
+struct TestPythonRendezvousPass
+ : public mlir::PassWrapper<TestPythonRendezvousPass,
+ mlir::OperationPass<mlir::ModuleOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPythonRendezvousPass)
+
+ TestPythonRendezvousPass() = default;
+ TestPythonRendezvousPass(const TestPythonRendezvousPass &other)
+ : PassWrapper(other) {}
+
+ llvm::StringRef getArgument() const final { return "test-python-rendezvous"; }
+ llvm::StringRef getDescription() const final {
+ return "Blocks natively until `participants` threads have entered this "
+ "pass with the same barrier-id; fails after timeout-seconds.";
+ }
+
+ Option<std::string> barrierId{
+ *this, "barrier-id",
+ llvm::cl::desc("Name of the process-global barrier to meet at"),
+ llvm::cl::init("default")};
+ Option<unsigned> participants{
+ *this, "participants",
+ llvm::cl::desc("Number of threads that must arrive at the barrier"),
+ llvm::cl::init(2)};
+ Option<unsigned> timeoutSeconds{
+ *this, "timeout-seconds",
+ llvm::cl::desc("Seconds to wait for the other participants"),
+ llvm::cl::init(20)};
+
+ void runOnOperation() final {
+ std::optional<unsigned> arrivedAtFailure =
+ RendezvousBarrier::get(barrierId).arriveAndWait(
+ participants, std::chrono::seconds(timeoutSeconds.getValue()));
+ if (!arrivedAtFailure)
+ return;
+ getOperation().emitError()
+ << "rendezvous '" << barrierId.getValue() << "' timed out with "
+ << *arrivedAtFailure << " of " << participants.getValue()
+ << " participants";
+ signalPassFailure();
+ }
+};
+} // namespace
+
+extern "C" void mlirPythonTestRegisterRendezvousPass(void) {
+ // Thread-safe and idempotent: registration happens once per process.
+ static mlir::PassRegistration<TestPythonRendezvousPass> registration;
+ (void)registration;
+}
diff --git a/mlir/test/python/pass_manager_concurrency.py b/mlir/test/python/pass_manager_concurrency.py
new file mode 100644
index 00000000000000..df38db27e90a10
--- /dev/null
+++ b/mlir/test/python/pass_manager_concurrency.py
@@ -0,0 +1,126 @@
+# RUN: %PYTHON %s
+"""Tests that PassManager.run detaches from the Python interpreter.
+
+The core test runs two PassManagers, on independent Contexts, in two threads.
+Each pipeline contains the native `test-python-rendezvous` pass (defined in
+lib/TestConcurrencyPasses.cpp), which blocks inside runOnOperation until both
+threads have arrived at a shared native barrier. The barrier can only be
+satisfied while both threads execute native pass code at the same time:
+
+* If PassManager.run releases the GIL (or the interpreter is free-threaded),
+ both threads enter their pipelines, meet at the barrier, and both runs
+ succeed.
+* If PassManager.run holds the GIL, whichever thread enters run() first blocks
+ natively while still attached to the interpreter, the second thread can
+ never start its own run(), and the first thread fails deterministically
+ once the barrier times out.
+
+No part of this relies on timing margins: the failure mode is a guaranteed
+timeout, and the success mode only requires that a runnable thread is
+eventually scheduled, which the barrier waits for.
+"""
+
+import threading
+import uuid
+
+from mlir.ir import Context, MLIRError, Module
+from mlir.passmanager import PassManager
+import mlir.dialects.python_test as python_test
+
+python_test.register_rendezvous_pass()
+
+
+def unique_barrier_id(prefix):
+ """Barriers are process-global and one-shot, so give each use a fresh id."""
+ return f"{prefix}-{uuid.uuid4().hex}"
+
+
+def run_rendezvous(barrier_id, timeout_seconds, participants=2, ctx=None):
+ """Runs one rendezvous pipeline on the given (default: fresh) Context."""
+ with Context() if ctx is None else ctx:
+ module = Module.parse("module {}")
+ pm = PassManager.parse(
+ "builtin.module(test-python-rendezvous{"
+ f"barrier-id={barrier_id} "
+ f"participants={participants} "
+ f"timeout-seconds={timeout_seconds}}})"
+ )
+ pm.run(module.operation)
+
+
+# Sanity-check the failure mode first: a lone participant must time out, the
+# pipeline must fail, and the pass's diagnostic must be captured and attached
+# to the MLIRError. This also exercises the ErrorCapture diagnostic path.
+def testSoloRendezvousTimesOut():
+ barrier_id = unique_barrier_id("solo")
+ try:
+ run_rendezvous(barrier_id, timeout_seconds=1)
+ raise AssertionError("expected the pipeline to fail")
+ except MLIRError as e:
+ messages = [d.message for d in e.error_diagnostics]
+ assert any(
+ f"rendezvous '{barrier_id}' timed out" in m for m in messages
+ ), messages
+
+
+# The core test: two PassManagers on independent Contexts must be able to
+# execute native pass code simultaneously.
+def testConcurrentPassManagersRendezvous():
+ barrier_id = unique_barrier_id("pair")
+ errors = [None, None]
+
+ def worker(idx):
+ try:
+ run_rendezvous(barrier_id, timeout_seconds=20)
+ except Exception as e: # pylint: disable=broad-except
+ errors[idx] = e
+
+ threads = [threading.Thread(target=worker, args=(i,)) for i in range(2)]
+ for t in threads:
+ t.start()
+ for t in threads:
+ t.join(timeout=60)
+ assert not any(t.is_alive() for t in threads), "worker thread hung"
+ if any(errors):
+ raise AssertionError(
+ "two PassManagers could not execute concurrently; PassManager.run "
+ f"is not releasing the GIL: {errors}"
+ )
+
+
+# A user-attached diagnostic handler must work while the calling thread is
+# detached from the interpreter inside PassManager.run: the handler trampoline
+# re-attaches around everything it does (creating the Diagnostic object,
+# invoking the callback, and tearing the object down afterwards). With
+# emit_error_diagnostics set, the pipeline failure diagnostic bypasses the
+# internal error capture and reaches the handler instead.
+def testDiagnosticHandlerDuringDetachedRun():
+ seen = []
+
+ def handler(d):
+ # Materialize everything, including the notes tuple, so that the
+ # diagnostic's teardown after the callback has Python objects to drop.
+ seen.append((d.severity, str(d.message), [str(n) for n in d.notes]))
+ return True
+
+ with Context() as ctx:
+ ctx.emit_error_diagnostics = True
+ handle = ctx.attach_diagnostic_handler(handler)
+ barrier_id = unique_barrier_id("diag")
+ try:
+ run_rendezvous(barrier_id, timeout_seconds=1, ctx=ctx)
+ raise AssertionError("expected the pipeline to fail")
+ except MLIRError as e:
+ # The handler consumed the diagnostic, so nothing was captured.
+ assert not e.error_diagnostics, e.error_diagnostics
+ finally:
+ handle.detach()
+ assert len(seen) == 1, seen
+ assert f"rendezvous '{barrier_id}' timed out" in seen[0][1], seen
+ assert not handle.had_error
+
+
+if __name__ == "__main__":
+ testSoloRendezvousTimesOut()
+ testConcurrentPassManagersRendezvous()
+ testDiagnosticHandlerDuringDetachedRun()
More information about the Mlir-commits
mailing list