Skip to content

Make Tensor's __dlpack__ and __dlpack_device__ account for XLA.#128176

Closed
vanbasten23 wants to merge 2 commits intopytorch:mainfrom
vanbasten23:xiowei/make_dlpack_account_for_xla
Closed

Make Tensor's __dlpack__ and __dlpack_device__ account for XLA.#128176
vanbasten23 wants to merge 2 commits intopytorch:mainfrom
vanbasten23:xiowei/make_dlpack_account_for_xla

Conversation

@vanbasten23
Copy link
Copy Markdown
Contributor

@vanbasten23 vanbasten23 commented Jun 6, 2024 •
edited
Loading

Fixes #ISSUE_NUMBER

XLA tensor is a Tensor. This PR

  • makes __dlpack_device__ account for XLA device when XLA GPU is used.
  • makes __dlpack__ call the torch_xla version of to_dlpack.

A corresponding test is added in pytorch/xla#7213

@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented Jun 6, 2024 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/128176

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 01c4a1f with merge base fa8c343 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@vanbasten23
Copy link
Copy Markdown
Contributor Author

vanbasten23 commented Jun 6, 2024

cc @ysiraichi @JackCaoG

@vanbasten23 vanbasten23 mentioned this pull request Jun 6, 2024
Closed
@vanbasten23 vanbasten23 mentioned this pull request Jun 10, 2024
Merged
@miladm
Copy link
Copy Markdown
Collaborator

miladm commented Jun 11, 2024

@emcastillo, @rgommers, @Mulberry are you the right person to review this PR? Appreciate the suggestion.

@colesbury colesbury requested review from Skylion007 and albanD June 12, 2024 01:37
@colesbury colesbury added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Jun 12, 2024
@vanbasten23
Copy link
Copy Markdown
Contributor Author

vanbasten23 commented Jun 14, 2024

Hi @albanD @Skylion007 , could you help review this PR? Thanks!

@vanbasten23
Copy link
Copy Markdown
Contributor Author

vanbasten23 commented Jun 21, 2024

Gentle ping @albanD @Skylion007

@albanD
Copy link
Copy Markdown
Collaborator

albanD commented Aug 6, 2024

I'm not sure what this environment variable is expected to do cc @JackCaoG

JackCaoG
JackCaoG reviewed Aug 6, 2024
View reviewed changes
Comment thread torch/_tensor.py Outdated
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG Aug 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if torch_device_type is already xla, I think it is better to import torch_xla and check runtime time instead of checking the env var.

Copy link
Copy Markdown
Contributor Author

@vanbasten23 vanbasten23 Aug 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed it.

@vanbasten23 vanbasten23 force-pushed the xiowei/make_dlpack_account_for_xla branch from 7c405b1 to 01c4a1f Compare August 8, 2024 23:39
@vanbasten23 vanbasten23 requested a review from JackCaoG August 8, 2024 23:40
@github-actions
Copy link
Copy Markdown
Contributor

github-actions Bot commented Oct 8, 2024

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

@github-actions github-actions Bot added the Stale label Oct 8, 2024
@vanbasten23 vanbasten23 closed this Oct 8, 2024
@ysiraichi ysiraichi mentioned this pull request Oct 21, 2024
Closed
pytorchmergebot pushed a commit that referenced this pull request Dec 5, 2024
@ysiraichi @vanbasten23 @pytorchmergebot
DLPack: add support to PyTorch/XLA (#138470)
7e49da6 
Taking over: #128176.

In summary, this PR:

- `__dlpack__`: Calls PyTorch/XLA `to_dlpack` function, if the tensor lives in an XLA:CUDA device
- `__dlpack_device__`: Correctly maps PyTorch/XLA tensors to `kDLGPU`, if XLA:CUDA is being used

The tests are introduced in pytorch/xla#7213.
Pull Request resolved: #138470
Approved by: https://github.com/albanD

Co-authored-by: iefgnoix <isaacwxf23@gmail.com>
pytorch-bot Bot pushed a commit that referenced this pull request Dec 9, 2024
@ysiraichi @vanbasten23 @AmdSampsa
DLPack: add support to PyTorch/XLA (#138470)
a5d5836 
Taking over: #128176.

In summary, this PR:

- `__dlpack__`: Calls PyTorch/XLA `to_dlpack` function, if the tensor lives in an XLA:CUDA device
- `__dlpack_device__`: Correctly maps PyTorch/XLA tensors to `kDLGPU`, if XLA:CUDA is being used

The tests are introduced in pytorch/xla#7213.
Pull Request resolved: #138470
Approved by: https://github.com/albanD

Co-authored-by: iefgnoix <isaacwxf23@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@albanD albanD Awaiting requested review from albanD

@Skylion007 Skylion007 Awaiting requested review from Skylion007

@JackCaoG JackCaoG Awaiting requested review from JackCaoG

Assignees

@vanbasten23 vanbasten23

Labels

open source Stale triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@vanbasten23 @miladm @albanD @JackCaoG @colesbury @pytorchbot