Skip to content

DLPack: add test using PyTorch DLPack functions.#8294

Merged
qihqi merged 4 commits intomasterfrom
ysiraichi/dlpack-tensor-test
Oct 30, 2024
Merged

DLPack: add test using PyTorch DLPack functions.#8294
qihqi merged 4 commits intomasterfrom
ysiraichi/dlpack-tensor-test

Conversation

@ysiraichi
Copy link
Copy Markdown
Collaborator

Taking over: #7213

This PR introduces a test using PyTorch's to_dlpack and from_dlpack functions. In order to use the former, PyTorch's tensor __dlpack__ and __dlpack_device__ methods had to be modified. The actual changes are in pytorch/pytorch#138470.

@qihqi qihqi merged commit dc20b2d into master Oct 30, 2024
@qihqi qihqi deleted the ysiraichi/dlpack-tensor-test branch October 30, 2024 16:46
@ysiraichi
Copy link
Copy Markdown
Collaborator Author

Not sure this was ok to merge, since it still depends on pytorch/pytorch#138470 (as noted by the .torch_pin file).

@qihqi @JackCaoG thoughts?

Comment thread test/test_operations.py
Comment on lines +2915 to +2925
@onlyIfTorchSupportsCUDA
@onlyIfPJRTDeviceIsCUDA
def test_dlpack_xla_to_pytorch_cuda_protocol_conversion(self):
xla_t1 = torch.arange(5).to(xm.xla_device())
caps_t1 = torch.utils.dlpack.to_dlpack(xla_t1)
cuda_t1 = torch.utils.dlpack.from_dlpack(caps_t1)
self.assertEqual(cuda_t1.device.type, 'cuda')
self.assertEqual(cuda_t1.device.index, xla_t1.device.index)
cuda_t1[0] = cuda_t1[0] + 20
self.assertTrue(torch.allclose(xla_t1.cpu(), cuda_t1.cpu()))

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

@vanbasten23 @JackCaoG @qihqi
I was wondering why removing .torch_pin (#8344) wasn't breaking CI. I realized that it was because I modified the test @vanbasten23 had implemented, calling to_dlpack(...), and only then from_dlpack(...). In summary: it was correct, and I messed it up. Anyway, I will open a PR fixing this, so as to have it merged after pytorch/pytorch#138470.

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.

4 participants