[mlir][python] fix PyDenseResourceElementsAttribute finalizer

[makslevental]

This PR melds #150137 and #149414 and partially reverts #124832.

The summary is the PyDenseResourceElementsAttribute finalizer/deleter has/had two problems

1. wasn't threadsafe (can be called from a different thread than that which currently holds the GIL)
2. can be called while the interpreter is "not initialized"

#124832 for some reason decides to re-initialize the interpreter to avoid case 2 and runs afoul of the fact that Py_IsInitialized can be false during the finalization of the interpreter itself (e.g., at the end of a script).

I don't know why this decision was made (I missed the PR) but I believe we should never be calling Py_Initialize:

In an application ****embedding Python****, this should be called before using any other Python/C API functions

but we aren't embedding Python!

So therefore we will only be in case 2 when the interpreter is being finalized and in that case we should just leak the buffer.

Note, lldb does a similar sort of thing for its finalizers.

Co-authored-by: Anton Korobeynikov [email protected]
Co-authored-by: Max Manainen [email protected]

[makslevental]

cc @chrsmcgrr @mamanain

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

This PR melds #150137 and #149414 and partially reverts #124832.

The summary is the PyDenseResourceElementsAttribute finalizer/deleter has/had two problems

1. wasn't threadsafe (can be called from a different thread than that which currently holds the GIL)
2. can be called while the interpreter is not initialized

#124832 for some reason decides to re-initialize the interpreter to avoid case 2 and runs afoul of the fact that Py_IsInitialized can be false during the finalization of the interpreter itself (e.g., at the end of a script).

I don't know why this decision was made (I missed the PR) but I believe we should never be calling Py_Initialize:

> In an application embedding Python, this should be called before using any other Python/C API functions

but we aren't embedding Python!

So therefore we will only be in case when the interpreter is being finalized and in that case we should just speak the buffer.

Note, lldb does a similar sort of thing for its finalizers.

---

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

2 Files Affected:

• (modified) mlir/lib/Bindings/Python/IRAttributes.cpp (+9-2)
• (modified) mlir/test/python/ir/array_attributes.py (+15)

diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index 8f79caf08a6d0..7ddae4b5fd358 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -1428,6 +1428,12 @@ class PyDenseIntElementsAttribute
   }
 };
 
+// Check if the python version is less than 3.13. Py_IsFinalizing is a part
+// of stable ABI since 3.13 and before it was available as _Py_IsFinalizing.
+#if PY_VERSION_HEX < 0x030d0000
+#define Py_IsFinalizing _Py_IsFinalizing
+#endif
+
 class PyDenseResourceElementsAttribute
     : public PyConcreteAttribute<PyDenseResourceElementsAttribute> {
 public:
@@ -1474,8 +1480,9 @@ 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();
+      if (Py_IsFinalizing())
+        return;
+      assert(Py_IsInitialized() && "expected interpreter to be initialized");
       Py_buffer *ownedView = static_cast<Py_buffer *>(userData);
       nb::gil_scoped_acquire gil;
       PyBuffer_Release(ownedView);
diff --git a/mlir/test/python/ir/array_attributes.py b/mlir/test/python/ir/array_attributes.py
index ef1d835fc6401..67b934dd19bdb 100644
--- a/mlir/test/python/ir/array_attributes.py
+++ b/mlir/test/python/ir/array_attributes.py
@@ -617,3 +617,18 @@ def test_attribute(context, mview):
     # CHECK: BACKING MEMORY DELETED
     # CHECK: EXIT FUNCTION
     print("EXIT FUNCTION")
+
+
+# CHECK-LABEL: TEST: testDanglingResource
+print("TEST: testDanglingResource")
+# see https://github.com/llvm/llvm-project/pull/149414, https://github.com/llvm/llvm-project/pull/150137
+# This error occurs only when there is an alive context with a DenseResourceElementsAttr
+# in the end of the program, so we put it here without an encapsulating function.
+ctx = Context()
+
+with ctx, Location.unknown():
+    DenseResourceElementsAttr.get_from_buffer(
+        memoryview(np.array([1, 2, 3])),
+        "some_resource",
+        RankedTensorType.get((3,), IntegerType.get_signed(32)),
+    )

[makslevental]

@mamanain if you tell me your email I can update the Co-authored to cite you.

[mamanain]

[email protected]

[chrsmcgrr]

Thanks for this! This seems to be what was causing issues on our end. It now works in python 3.11.

[asl]

@makslevental Likely we'd not need the check at all. See #152226 for similar LLDB change (and notes around).

[makslevental]

Likely we'd not need the check at all.

Isn't that the original version prior to #124832? Ie isn't that bug reporting that we do indeed sometimes release when (Py_IsFinalizing()) is true?

[asl]

Isn't that the original version prior to #124832? Ie isn't that bug reporting that we do indeed sometimes release when (Py_IsFinalizing()) is true?

It seems, there are multiple cases here:

• We're deleting an attribute from the thread with active interpreter
• We're deleting an attribute from the thread without active interpreter
• Dense resource is deleted when holding context is deleted during overall shutdown

If I'm reading the pytorch version and pylifecycle.c correctly, then we'd need to hold the gil if the interpreter is initialized. And simply destroy the buffer otherwise. Currently we're leaking ownedView but this seems to be fine only if the whole process is terminating. Which might not be ok in the shared library context... So, I'm thinking if something like this would be more correct?

      Py_buffer *ownedView = static_cast<Py_buffer *>(userData);
      if (Py_IsInitialized()) {
        nb::gil_scoped_acquire gil;
        PyBuffer_Release(ownedView);
      }
      delete ownedView;

[makslevental]

yea sure this seems reasonable - do you want to send the PR that removes Py_IsFinalizing and etc?