[MLIR] Fix thread safety of the deleter in PyDenseResourceElementsAttribute

[naibaf7]

In general, PyDenseResourceElementsAttribute can get deleted at any time and any thread, where unlike the getFromBuffer call, the Python interpreter may not be initialized and the GIL may not be held.

This PR fixes segfaults caused by PyBuffer_Release when the GIL is not being held by the thread calling the deleter.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Fabian Tschopp (naibaf7)

Changes

In general, PyDenseResourceElementsAttribute can get deleted at any time and any thread, where unlike the getFromBuffer call, the Python interpreter may not be initialized and the GIL may not be held.

This PR fixes segfaults caused by PyBuffer_Release when the GIL is not being held by the thread calling the deleter.

---

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

1 Files Affected:

• (modified) mlir/lib/Bindings/Python/IRAttributes.cpp (+3)

diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index 7bc21a31c3c84c..142b6eca11c438 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -1468,7 +1468,10 @@ class PyDenseResourceElementsAttribute
     // The userData is a Py_buffer* that the deleter owns.
     auto deleter = [](void *userData, const void *data, size_t size,
                       size_t align) {
+      if (!Py_IsInitialized())
+        Py_Initialize();
       Py_buffer *ownedView = static_cast<Py_buffer *>(userData);
+      nb::gil_scoped_acquire gil;
       PyBuffer_Release(ownedView);
       delete ownedView;
     };

[chrsmcgrr]

@naibaf7 can you give a bit more context when this happens? Do you have a reproducer for this? Py_Initialize() from the documentation initializes the interpreter, but at this point in the bindings there should already be an interpreter active. If not, then something else seems wrong.

I'm looking at our end why this change causes problems for us. I'll post with more info later.

[chrsmcgrr]

@naibaf7 @lattner @weiweichen I now see this PR referenced in another project. Could you share how to reproduce the error you saw?

[naibaf7]

@naibaf7 @lattner @weiweichen I now see this PR referenced in another project. Could you share how to reproduce the error you saw?

@chrsmcgrr
The issue will be if the thread that releases is not the Python thread that called into whatever API is exposed to Python. This change ensured that if the thread calling the callback isn't that same thread, it still works.
It's quite easy to reproduce: Make sure a different thread causes the release of that resource, e.g. by a DCE pass.

This change can cause issues if the GIL, for example, is held by another thread than the releasing one. In that case it may deadlock, but that would be correct behaviour, and the other thread should release the GIL. Though without this change it's not thread safe and may fail with segfaults.

[asl]

@naibaf7 This causes deadlocks when deleter is called during context destruction as the latter already holds a GIL:

PyMlirContext::~PyMlirContext() {
  // Note that the only public way to construct an instance is via the
  // forContext method, which always puts the associated handle into
  // liveContexts.
  nb::gil_scoped_acquire acquire;
  {
    nb::ft_lock_guard lock(live_contexts_mutex);
    getLiveContexts().erase(context.ptr);
  }
  mlirContextDestroy(context);
}

This leads to ResourceBlobManagerDialectInterfaceBase::~ResourceBlobManagerDialectInterfaceBase call which in turn calls deleter with GIL already held. This is the reason of llvm/torch-mlir#4157

[asl]

@naibaf7 @weiweichen a bit more context. Here is the finalizer for the dense resource:

    // The userData is a Py_buffer* that the deleter owns.
    auto deleter = [](void *userData, const void *data, size_t size,
                      size_t align) {
      if (!Py_IsInitialized())
        Py_Initialize();
      Py_buffer *ownedView = static_cast<Py_buffer *>(userData);
      nb::gil_scoped_acquire gil;
      PyBuffer_Release(ownedView);
      delete ownedView;
    };

when invoked during the finalization Py_IsInitialized() already returns false. As a result, Py_Initialize() is called again causing all kinds of problems including broken thread state. It seems it is missing Py_IsFinalizing() check here. But the problem is that it is a part of stable ABI only from 3.13.

[makslevental]

@naibaf7 I don't understand why the solution here wasn't to just leak the buffer? The interpreter is shutting down right? Why restart the interpreter just to release an object?

Note, there's a very similar kind of thing in lldb.