[Mlir-commits] [mlir] [mlir][python] Fix segfault at interpreter shutdown with entered contexts (PR #203826)

[Maksim Levental]

https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/203826

>From 97fec7cf6c2288fe77ee692ad13130b9595c3e92 Mon Sep 17 00:00:00 2001
From: makslevental <m_levental at apple.com>
Date: Sun, 14 Jun 2026 22:00:30 -0700
Subject: [PATCH] [mlir][python] Fix segfault at interpreter shutdown with
 entered contexts
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The thread-local context stack (`PyThreadContextEntry::getStack()`)
holds `nb::object` references to Python Context, Location, and
InsertionPoint objects. When a Context is entered via `__enter__` but
never exited before the interpreter shuts down, these references
cause a segfault during process teardown.

The crash sequence:
1. User calls `ctx.__enter__()`, pushing a frame onto the
   `static thread_local vector<PyThreadContextEntry>`.
2. The script ends; CPython runs `Py_FinalizeEx()` which tears down
   the interpreter (clears modules, destroys remaining objects).
3. `main()` returns.
4. The C runtime destroys static/thread_local storage. On the main
   thread, thread_local variables have the same destruction timing
   as static storage — they are destroyed *after* main() returns.
5. The vector destructor runs, and each `PyThreadContextEntry`'s
   `nb::object` members call `Py_DECREF` — but the interpreter is
   already dead. This dereferences freed memory → SIGSEGV.

The fix registers a Python `atexit` handler that clears the
thread-local stack. Python's atexit handlers run *before*
`Py_FinalizeEx()`, while the interpreter is still fully alive,
so the `nb::object` decrefs execute safely.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply at anthropic.com>
---
 mlir/lib/Bindings/Python/IRCore.cpp  | 11 +++++++++++
 mlir/test/python/context_shutdown.py | 26 ++++++++++++++++++++++++++
 2 files changed, 37 insertions(+)
 create mode 100644 mlir/test/python/context_shutdown.py

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 92e9ecf3f2c20..e0a141d09339d 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -5292,6 +5292,17 @@ void populateIRCore(nb::module_ &m) {
 
   // MLIRError exception.
   MLIRError::bind(m);
+
+  // Register an atexit handler to clear the thread-local context stack.
+  // The stack holds nb::object references that prevent Python GC of Contexts.
+  // At interpreter shutdown, thread_local storage outlives Py_Finalize() on
+  // the main thread. When the thread_local vector destructs, its nb::object
+  // members call Py_DECREF through the dead runtime, causing a segfault.
+  // Clearing the stack in atexit releases references while alive.
+  nb::module_::import_("atexit")
+      .attr("register")(nb::cpp_function([]() {
+        PyThreadContextEntry::getStack().clear();
+      }));
 }
 } // namespace MLIR_BINDINGS_PYTHON_DOMAIN
 } // namespace python
diff --git a/mlir/test/python/context_shutdown.py b/mlir/test/python/context_shutdown.py
new file mode 100644
index 0000000000000..cb7b12cd5a4d6
--- /dev/null
+++ b/mlir/test/python/context_shutdown.py
@@ -0,0 +1,26 @@
+# RUN: %PYTHON %s
+# Regression test: entering a Context (or Location/InsertionPoint) without
+# exiting before interpreter shutdown used to segfault. The thread-local
+# context stack holds nb::object references; if not cleared before
+# Py_Finalize(), the thread_local destructor calls Py_DECREF through the
+# dead runtime.
+
+from mlir.ir import *
+
+# Case 1: Single context entered, not exited.
+ctx = Context()
+ctx.__enter__()
+ctx.enable_multithreading(False)
+with Location.unknown():
+    m = Module.parse("func.func @f() { return }")
+
+# Case 2: Multiple contexts entered, not exited.
+ctx2 = Context()
+ctx2.__enter__()
+
+# Case 3: Location entered, not exited (also uses the same thread-local stack).
+loc = Location.unknown()
+loc.__enter__()
+
+# Interpreter shutdown proceeds with contexts/locations still on the
+# thread-local stack. Before the fix, this would segfault (exit code 139).