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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

<details>
<summary>Changes</summary>

# 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



---
Full diff: https://github.com/llvm/llvm-project/pull/203826.diff


2 Files Affected:

- (modified) mlir/lib/Bindings/Python/IRCore.cpp (+9) 
- (added) mlir/test/python/context_shutdown.py (+26) 


``````````diff
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 92e9ecf3f2c20..fc69ef8884466 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -5292,6 +5292,15 @@ 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).

``````````

</details>


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