Skip to content

Migrate IFRT host-buffer cleanup callbacks to std::shared_ptr<const void>. - #130334

Draft
copybara-service[bot] wants to merge 1 commit into
masterfrom
exported_pr_1021526713
Draft

copybara-service[bot] wants to merge 1 commit into
masterfrom
exported_pr_1021526713

Conversation

@copybara-service

Copy link
Copy Markdown

Migrate IFRT host-buffer cleanup callbacks to std::shared_ptr.

xla::ifrt::Client::MakeArrayFromHostBuffer and MakeArraysFromHostBufferShards previously documented that on_done_with_host_buffer and HostBuffer::on_done were invoked only when OK was returned, while implementations wrapped or executed them mid-call before validation and multi-shard creation completed. In addition, invoking std::move(fn)() on std::function<void()> left closure-captured host buffers alive until the enclosing std::function or BasicStringArray was destroyed.

Switch MakeArrayFromHostBuffer (renaming on_done_with_host_buffer to host_buffer_hold), Client::HostBuffer::on_done, and BasicStringArray::OnDoneWithBuffer to std::shared_ptr<const void> (with MakeSharedCleanup for callable-based cleanup) so ownership transfers unconditionally to the callee, cleanup runs deterministically on both success and failure once all readers finish accessing data, and callers already holding std::shared_ptr buffers pass them directly without extra wrapper allocations.

@copybara-service
copybara-service Bot force-pushed the exported_pr_1021526713 branch from 787a955 to 3ad2132 Compare November 14, 2026 09:29
…oid>.

`xla::ifrt::Client::MakeArrayFromHostBuffer` and `MakeArraysFromHostBufferShards` previously documented that `on_done_with_host_buffer` and `HostBuffer::on_done` were invoked only when `OK` was returned, while implementations wrapped or executed them mid-call before validation and multi-shard creation completed. In addition, invoking `std::move(fn)()` on `std::function<void()>` left closure-captured host buffers alive until the enclosing `std::function` or `BasicStringArray` was destroyed.

Switch `MakeArrayFromHostBuffer` (renaming `on_done_with_host_buffer` to `host_buffer_hold`), `Client::HostBuffer::on_done`, and `BasicStringArray::OnDoneWithBuffer` to `std::shared_ptr<const void>` (with `MakeSharedCleanup` for callable-based cleanup) so ownership transfers unconditionally to the callee, cleanup runs deterministically on both success and failure once all readers finish accessing `data`, and callers already holding `std::shared_ptr` buffers pass them directly without extra wrapper allocations.

PiperOrigin-RevId: 1021526713
@copybara-service
copybara-service Bot force-pushed the exported_pr_1021526713 branch from 3ad2132 to 0e61f5d Compare November 14, 2026 10:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant