Make from_dlpack handle cuda synchronization implicitly for input tensors that have __dlpack__ and __dlpack_device__ attributes.#7125
Merged
vanbasten23 merged 6 commits intomasterfrom May 30, 2024
Conversation
Collaborator
|
Let me know when this is ready for review |
Collaborator
Author
Thanks. It's ready for review. |
JackCaoG
reviewed
May 28, 2024
JackCaoG
reviewed
May 28, 2024
…onversion" This reverts commit 06b796e.
d42d453 to
c489851
Compare
JackCaoG
reviewed
May 30, 2024
JackCaoG
reviewed
May 30, 2024
Collaborator
Author
|
Thanks Jack for the review! |
Merged
yitongh
pushed a commit
to AlibabaPAI/xla
that referenced
this pull request
Oct 11, 2024
…sors that have __dlpack__ and __dlpack_device__ attributes. (pytorch#7125)
yitongh
pushed a commit
to AlibabaPAI/xla
that referenced
this pull request
Dec 11, 2024
…sors that have __dlpack__ and __dlpack_device__ attributes. (pytorch#7125)
yitongh
pushed a commit
to AlibabaPAI/xla
that referenced
this pull request
Dec 11, 2024
…sors that have __dlpack__ and __dlpack_device__ attributes. (pytorch#7125)
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 join this conversation on GitHub.
Already have an account?
Sign in to comment
2 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.
from_dlpack should leverage
__dlpack__and__dlpack_device__attributes of a tensor. That way, from_dlpack will handle the cuda synchronization implicity. Similar approaches can be found in pytorch and jax.This PR fixes the direction when we convert an external tensor such as cuda tensor to an XLA tensor via from_dlpack. For the other direction (convert an XLA tensor to CUDA tensor) requires a change in upstream pytorch.