Use correct pyobject for ref counting in `VmModule` pybindings (#13759)
Addresses memory leak by ref counting in the pyobject instead of the
raw_ptr of the buffer.
Fixes #13742
diff --git a/runtime/bindings/python/vm.cc b/runtime/bindings/python/vm.cc
index 5f9acfa..455bd3f 100644
--- a/runtime/bindings/python/vm.cc
+++ b/runtime/bindings/python/vm.cc
@@ -159,16 +159,16 @@
iree_vm_module_t* module = nullptr;
// Bridge to the C-based deallocator API.
- auto* raw_ptr = flatbuffer_blob.ptr();
+ PyObject* pyobject_ptr = flatbuffer_blob_object.ptr();
auto ctl_fn = +([](void* self, iree_allocator_command_t command,
const void* params, void** inout_ptr) {
assert(command == IREE_ALLOCATOR_COMMAND_FREE);
- PyObject* object_ptr = static_cast<PyObject*>(*inout_ptr);
- Py_XDECREF(object_ptr);
+ PyObject* pyobject_ptr = static_cast<PyObject*>(self);
+ Py_XDECREF(pyobject_ptr);
return iree_ok_status();
});
- flatbuffer_blob.inc_ref();
- iree_allocator_t deallocator{/*self=*/NULL, /*ctl=*/ctl_fn};
+ Py_XINCREF(pyobject_ptr);
+ iree_allocator_t deallocator{/*self=*/pyobject_ptr, /*ctl=*/ctl_fn};
auto status = iree_vm_bytecode_module_create(
instance->raw_ptr(),
@@ -176,7 +176,7 @@
static_cast<iree_host_size_t>(buffer_info.size)},
deallocator, iree_allocator_system(), &module);
if (!iree_status_is_ok(status)) {
- iree_allocator_free(deallocator, raw_ptr);
+ Py_XDECREF(pyobject_ptr);
}
CheckApiStatus(status, "Error creating vm module from FlatBuffer");
