[Runtime][Bindings] Fix several memory leaks in python bindings (#25014) Fixed several cases in hal.cc that could trigger memory to be leaked in python in case of error Signed-off-by: Florian Walbroel <walbroel@roofline.ai>
diff --git a/runtime/bindings/python/hal.cc b/runtime/bindings/python/hal.cc index d6d7a91..b07e286 100644 --- a/runtime/bindings/python/hal.cc +++ b/runtime/bindings/python/hal.cc
@@ -185,6 +185,8 @@ iree_hal_allocator_query_statistics(raw_ptr(), &stats); auto status = iree_hal_allocator_statistics_format(&stats, &builder); iree_string_view_t view = iree_string_builder_view(&builder); + // TODO(#25083): Guard the builder and status if Python string construction + // throws. py::str result = py::str(view.data, view.size); iree_string_builder_deinitialize(&builder); @@ -229,11 +231,11 @@ IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout()); } } + HalBuffer owned_buffer = HalBuffer::StealFromRawPtr(hal_buffer); CheckApiStatus(status, "Failed to allocate buffer copy"); if (!raw_element_type) { - return py::cast(HalBuffer::StealFromRawPtr(hal_buffer), - py::rv_policy::move); + return py::cast(std::move(owned_buffer), py::rv_policy::move); } // Create the buffer_view. (note that numpy shape is ssize_t, so we need to @@ -245,12 +247,10 @@ std::vector<iree_hal_dim_t> dims(py_view.ndim); std::copy(py_view.shape, py_view.shape + py_view.ndim, dims.begin()); iree_hal_buffer_view_t* hal_buffer_view; - CheckApiStatus( - iree_hal_buffer_view_create( - hal_buffer, dims.size(), dims.data(), element_type, encoding_type, - iree_hal_allocator_host_allocator(raw_ptr()), &hal_buffer_view), - "Error allocating buffer_view"); - iree_hal_buffer_release(hal_buffer); + status = iree_hal_buffer_view_create( + hal_buffer, dims.size(), dims.data(), element_type, encoding_type, + iree_hal_allocator_host_allocator(raw_ptr()), &hal_buffer_view); + CheckApiStatus(status, "Error allocating buffer_view"); return py::cast(HalBufferView::StealFromRawPtr(hal_buffer_view), py::rv_policy::move); @@ -292,9 +292,10 @@ IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout()); } } + HalBuffer owned_buffer = HalBuffer::StealFromRawPtr(hal_buffer); CheckApiStatus(status, "Failed to allocate device visible buffer"); - return HalBuffer::StealFromRawPtr(hal_buffer); + return owned_buffer; } //------------------------------------------------------------------------------ @@ -949,6 +950,7 @@ iree_allocator_system(), &driver_info_count, &driver_infos), "Error enumerating HAL drivers"); + // TODO(#25083): Guard driver_infos if constructing the result throws. std::vector<std::string> driver_names(driver_info_count); for (iree_host_size_t i = 0; i < driver_info_count; ++i) { driver_names[i] = std::string(driver_infos[i].driver_name.data, @@ -1002,6 +1004,7 @@ CheckApiStatus(iree_hal_driver_query_available_devices( raw_ptr(), iree_allocator_system(), &count, &device_infos), "Error querying devices"); + // TODO(#25083): Guard device_infos if constructing the Python result throws. py::list results; for (iree_host_size_t i = 0; i < count; ++i) { py::dict device_data; @@ -1068,6 +1071,7 @@ static HalDevice StealConfiguredDevice(iree_hal_device_t* device, std::optional<py::list> allocators) { + HalDevice owned_device = HalDevice::StealFromRawPtr(device); iree_hal_device_group_t* device_group = nullptr; iree_status_t status = ConfigureDevice(device, allocators); if (iree_status_is_ok(status)) { @@ -1076,12 +1080,12 @@ } if (!iree_status_is_ok(status)) { iree_hal_device_group_release(device_group); - iree_hal_device_release(device); CheckApiStatus( status, "Error configuring the device and assigning its frontier tracker"); } - return HalDevice::StealFromRawPtrAndDeviceGroup(device, device_group); + return HalDevice::StealFromRawPtrAndDeviceGroup(owned_device.steal_raw_ptr(), + device_group); } HalDevice HalDriver::CreateDefaultDevice(std::optional<py::list> allocators) { @@ -1624,10 +1628,15 @@ [](HalDriver& self, iree_hal_device_id_t device_id) { iree_string_builder_t builder; iree_string_builder_initialize(iree_allocator_system(), &builder); - CheckApiStatus(iree_hal_driver_dump_device_info( - self.raw_ptr(), device_id, &builder), - "Querying device info"); + iree_status_t status = iree_hal_driver_dump_device_info( + self.raw_ptr(), device_id, &builder); + if (!iree_status_is_ok(status)) { + iree_string_builder_deinitialize(&builder); + CheckApiStatus(status, "Querying device info"); + } iree_string_view_t view = iree_string_builder_view(&builder); + // TODO(#25083): Guard the builder if Python string construction + // throws. py::str result(view.data, view.size); iree_string_builder_deinitialize(&builder); return result;
diff --git a/runtime/bindings/python/tests/hal_test.py b/runtime/bindings/python/tests/hal_test.py index b95fb8b..c6fbd2e 100644 --- a/runtime/bindings/python/tests/hal_test.py +++ b/runtime/bindings/python/tests/hal_test.py
@@ -194,6 +194,37 @@ "<HalBufferView (3, 4), element_type=0x20000011, 48 bytes (at offset 0 into 48), memory_type=DEVICE_LOCAL|HOST_VISIBLE, allowed_access=ALL, allowed_usage=TRANSFER|DISPATCH_STORAGE|MAPPING|MAPPING_PERSISTENT>", ) + def testAllocateBufferViewCopyFailureReleasesBuffer(self): + ary = np.zeros(16, dtype=np.uint8) + before = self.allocator.statistics + # HAL buffer view creation rejects a view larger than its backing buffer. + with self.assertRaisesRegex(IndexError, "exceeds backing buffer byte length"): + self.allocator.allocate_buffer_copy( + memory_type=iree.runtime.MemoryType.DEVICE_LOCAL, + allowed_usage=iree.runtime.BufferUsage.DEFAULT, + device=self.device, + buffer=ary, + element_type=iree.runtime.HalElementType.SINT_32, + ) + if self.allocator.has_statistics: + after = self.allocator.statistics + self.assertGreater( + sum( + after[f"{memory_type}_bytes_allocated"] + - before[f"{memory_type}_bytes_allocated"] + for memory_type in ("host", "device") + ), + 0, + ) + for memory_type in ("host", "device"): + with self.subTest(memory_type=memory_type): + allocated = f"{memory_type}_bytes_allocated" + freed = f"{memory_type}_bytes_freed" + self.assertEqual( + after[allocated] - before[allocated], + after[freed] - before[freed], + ) + def testAllocateHostStagingBufferCopy(self): buffer = self.allocator.allocate_host_staging_buffer_copy( self.device, np.int32(0)