[pull] master from tensorflow:master#1626
Merged
Merged
Conversation
PiperOrigin-RevId: 912256028
PiperOrigin-RevId: 912272639
… single large tensor and slice them out for later usage. PiperOrigin-RevId: 912273572
PiperOrigin-RevId: 912275022
PiperOrigin-RevId: 912279501
…s to call into CrossHostTransferBuffers Imported from GitHub PR openxla/xla#41324 📝 Summary of Changes This PR builds on #41055, and is the next in a sequence of PRs that will refactor cross-host data transfer implementations to rely on a shared helper function `CrossHostTransferBuffers`. `CrossHostTransferBuffers` is planned to eventually be integrated into the PJRT APIs to enable receiving data into preallocated receive buffers (this feature is being planned in collaboration with @gspschmid, @emilyfertig, and @pschuh). This PR refactors `CrossHostReceiveBuffers` to call into `CrossHostTransferBuffers`. 🎯 Justification It is difficult to achieve good comm/compute overlap with cross-host data transfers as the current implementation always allocates receive-buffers 'just-in-time', and because the GPU memory allocator blocks on the compute stream. CrossHostTransferBuffers will enable users to receive into preallocated receive buffers, making it easier to avoid the allocator blocking issue. This PR refactors CrossHostReceiveBuffers to call into CrossHostTransferBuffers. Later PRs will expose CrossHostTransferBuffers through the PJRT C/C++ APIs. 🚀 Kind of Contribution ♻️ Cleanup (eventually ✨ New Feature) 🧪 Unit Tests: This PR only refactors the implementation of CrossHost{Send/Receive}Buffers, so the pre-existing unit tests for those methods already test this PR. 🧪 Execution Tests: Verified that [these 4 correctness tests](https://gist.github.com/rao-ashish/24ac0df0cb18243c649ac535964b31b8) continue to pass. Copybara import of the project: -- 0260d4efb3924afe56aac024a5367796e3553b6c by Ashish Rao <asrao@nvidia.com>: Refactor CrossHostReceiveBuffers to call into CrossHostTransferBuffers -- 254651affff4f6a838c01ac87c32cbb0caea99d2 by Ashish Rao <asrao@nvidia.com>: In CrossHostReceiveBuffers, retain a reference to the raw buffers until the transfer is done (also revert indentation change) Merging this change closes #41324 PiperOrigin-RevId: 912279538
PiperOrigin-RevId: 912290144
…when the memory pressure increase equals to zero (vs a node with positive increase). A zero-memory-increase node is obviously better than a positive-increase node when negative-increase nodes do not exist. PiperOrigin-RevId: 912294734
… utilities. `pjrt_c_api_status_utils.h` was split out of `pjrt_c_api_helpers.h` to avoid a circular dependency. PiperOrigin-RevId: 912295423
Also fix deprecated uses of `dimensions_size` PiperOrigin-RevId: 912305567
PiperOrigin-RevId: 912312358
PiperOrigin-RevId: 912330226
PiperOrigin-RevId: 912332394
PiperOrigin-RevId: 912332539
PiperOrigin-RevId: 912333107
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
11 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )