[mlir][python] fix PyDenseResourceElementsAttribute finalizer (#150561)

This PR melds https://github.com/llvm/llvm-project/pull/150137 and
https://github.com/llvm/llvm-project/pull/149414 *and* partially reverts
https://github.com/llvm/llvm-project/pull/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"

https://github.com/llvm/llvm-project/pull/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](https://docs.python.org/3/c-api/init.html#c.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](548ca9e976/lldb/source/Plugins/ScriptInterpreter/Python/PythonDataObjects.cpp (L81-L93))
does a similar sort of thing for its finalizers.

Co-authored-by: Anton Korobeynikov <anton@korobeynikov.info>
Co-authored-by: Max Manainen <maximmanainen@gmail.com>

Co-authored-by: Anton Korobeynikov <anton@korobeynikov.info>
Co-authored-by: Max Manainen <maximmanainen@gmail.com>
This commit is contained in:
Maksim Levental 2025-07-25 07:05:30 -05:00 committed by GitHub
parent 8005c6a108
commit 21774489f0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 30 additions and 5 deletions

View File

@ -16,8 +16,8 @@
#include "NanobindUtils.h"
#include "mlir-c/BuiltinAttributes.h"
#include "mlir-c/BuiltinTypes.h"
#include "mlir/Bindings/Python/NanobindAdaptors.h"
#include "mlir/Bindings/Python/Nanobind.h"
#include "mlir/Bindings/Python/NanobindAdaptors.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/Support/raw_ostream.h"
@ -1428,6 +1428,12 @@ public:
}
};
// 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 @@ public:
// 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);

View File

@ -31,6 +31,7 @@ def testGetDenseElementsUnsupported():
# CHECK: unimplemented array format conversion from format:
print(e)
# CHECK-LABEL: TEST: testGetDenseElementsUnSupportedTypeOkIfExplicitTypeProvided
@run
def testGetDenseElementsUnSupportedTypeOkIfExplicitTypeProvided():
@ -41,8 +42,9 @@ def testGetDenseElementsUnSupportedTypeOkIfExplicitTypeProvided():
# realistic example would be a NumPy extension type like the bfloat16
# type from the ml_dtypes package, which isn't a dependency of this
# test.
attr = DenseElementsAttr.get(array.view(np.datetime64),
type=IntegerType.get_signless(64))
attr = DenseElementsAttr.get(
array.view(np.datetime64), type=IntegerType.get_signless(64)
)
# CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi64>
print(attr)
# CHECK: {{\[}}[1 2 3]
@ -135,6 +137,7 @@ def testGetDenseElementsFromListMixedTypes():
# Splats.
################################################################################
# CHECK-LABEL: TEST: testGetDenseElementsSplatInt
@run
def testGetDenseElementsSplatInt():
@ -617,3 +620,18 @@ def testGetDenseResourceElementsAttr():
# 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, https://github.com/llvm/llvm-project/pull/150561
# 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)),
)