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

[Maksim Levental]

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

# TL;DR:

The `static thread_local std::vector<PyThreadContextEntry>` holds `nb::object` references to Python Context/Location/InsertionPoint objects. When a Context is entered (pushed onto the stack) but never exited before interpreter shutdown, the thread-local storage destructor runs after `Py_Finalize()` on the main thread, attempting `Py_DECREF` through the dead runtime → SIGSEGV.

Fix: Register an `atexit` handler that clears the stack while the interpreter is still alive, releasing all held Python references before finalization.

The checked in tests will SIGSEGV on a commit that doesn't have this fix (try it!).

# Full explanation

Per

> Destructors for initialized objects with thread storage duration within a given thread are called as a result of returning from the initial function of that thread and as a result of that thread calling `std::exit`.[^1]

For the main thread, the "initial function" is CPython's `main()` of course. So thread_local destructors fire after `main()` returns. But CPython calls `Py_FinalizeEx()` explicitly before returning:

```
// CPython's Modules/main.c (simplified)
int main() {
    Py_Initialize();
    run_script();
    Py_FinalizeEx();   // ← tears down interpreter (calls Python atexit, GC, etc.)
    return 0;          // ← thread_local destructors fire AFTER this
}
```

So the sequence is:
1. `Py_FinalizeEx()` runs — calls Python `atexit` callbacks (interpreter alive)
2. `Py_FinalizeEx()` continues — GC collects, clears modules, tears down interpreter
3. `Py_FinalizeEx()` returns; `main()` returns
4. Thread_local destructors fire (per [basic.start.term]) ← **the `vector<PyThreadContextEntry>` dies here**
5. `nb::object` destructors call `Py_DECREF` → **interpreter is already dead → segfault**

The `atexit.register()` fix works because Python's `atexit` callbacks execute at step 1 — the very first thing inside `Py_FinalizeEx()` — while the interpreter is still fully functional.

Note: Python's `atexit` module is distinct from C's `atexit()`. Python's runs inside `Py_FinalizeEx()` (before main returns). C's runs after main returns (interleaved with static destructors).

Assisted by Claude

[^1]: https://timsong-cpp.github.io/cppwp/n3337/basic.start.term



>From 3f55f00db6ebcb886c58881a97dd759a9a57f28a 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 | 27 +++++++++++++++++++++++++++
 2 files changed, 38 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..83bb7cc7b3d63
--- /dev/null
+++ b/mlir/test/python/context_shutdown.py
@@ -0,0 +1,27 @@
+# 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).