Fix DLPack stream logic. #150217

ysiraichi · 2025-03-28T19:37:14Z

Stack from ghstack (oldest at bottom):

This PR fixes the logic for dealing with CUDA and ROCm streams whenever
we are trying to create a DLPack capsule from a tensor.

In summary, this PR:

Uses the legacy default stream if tensor.__dlpack__(stream=None) is
called for a CUDA tensor.
Errors if tensor.__dlpack__(stream=2) is called for a CUDA tensor:
PyTorch doesn't support the per-thread default stream.
Errors if tensor.__dlpack__(stream=stream), where stream is 1 or
2, is called for a CUDA tensor using ROCm.

For more details, see the documentation.

[ghstack-poisoned]

pytorch-bot · 2025-03-28T19:37:18Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

✅ No Failures

As of commit 5ffb2ee with merge base 7cc1a95 ():
💚 Looks good so far! There are no failures yet. 💚

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

[ghstack-poisoned]

This PR fixes the logic for dealing with CUDA and ROCm streams whenever we are trying to create a DLPack capsule from a tensor. In summary, this PR: - Uses the legacy default stream if `tensor.__dlpack__(stream=None)` is called for a CUDA tensor. - Errors if `tensor.__dlpack__(stream=2)` is called for a CUDA tensor: PyTorch doesn't support the per-thread default stream. - Errors if `tensor.__dlpack__(stream=stream)`, where `stream` is 1 or 2, is called for a CUDA tensor using ROCm. For more details, see [the documentation][1]. [1]: https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__dlpack__.html ghstack-source-id: cc0e31c Pull Request resolved: pytorch/pytorch#150217

[ghstack-poisoned]

albanD

Sounds ok even though this doesn't fix the multi-device case.

albanD · 2025-05-14T21:15:44Z

torch/_tensor.py

-        elif stream is not None and stream != -1:
+        elif stream != -1:
            if self.device.type == "cuda":
                # NB: This logic handles the special case values for default


No update to dlpack.py ? :D

No need. If stream is None, we still need to synchronize, assuming the legacy default stream.

albanD · 2025-05-14T21:17:10Z

torch/_tensor.py

+                    if is_cuda and stream == 2:
+                        raise BufferError("per-thread default stream is not supported.")
+
+                    assert is_cuda or (is_rocm and stream not in (1, 2)), (


Shouldn't this be a BufferError like above instead of AssertionError?

I don't think so. The reason being that this assertion checks something the standard explicitly states as "unsupported" or "disallowed", i.e. something the consumer should know about. Moreover, the standard also says that:

Other errors are raised when export fails for other reasons (e.g., incorrect arguments passed or out of memory).

albanD · 2025-05-14T21:20:19Z

torch/_tensor.py

                # Only synchronize on different streams
-                sync_stream = torch.cuda.current_stream()
-                if stream != sync_stream:
+                current_stream = torch.cuda.current_stream()


Do we care if self.device.index != torch.cuda.current_device() ?

Good point. I think we should. I will add a check for that.

[ghstack-poisoned]

This PR fixes the logic for dealing with CUDA and ROCm streams whenever we are trying to create a DLPack capsule from a tensor. In summary, this PR: - Uses the legacy default stream if `tensor.__dlpack__(stream=None)` is called for a CUDA tensor. - Errors if `tensor.__dlpack__(stream=2)` is called for a CUDA tensor: PyTorch doesn't support the per-thread default stream. - Errors if `tensor.__dlpack__(stream=stream)`, where `stream` is 1 or 2, is called for a CUDA tensor using ROCm. For more details, see [the documentation][1]. [1]: https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__dlpack__.html [ghstack-poisoned]

pytorchmergebot · 2025-07-12T13:35:03Z

Starting merge as part of PR stack under #150691

This PR fixes the logic for dealing with CUDA and ROCm streams whenever we are trying to create a DLPack capsule from a tensor. In summary, this PR: - Uses the legacy default stream if `tensor.__dlpack__(stream=None)` is called for a CUDA tensor. - Errors if `tensor.__dlpack__(stream=2)` is called for a CUDA tensor: PyTorch doesn't support the per-thread default stream. - Errors if `tensor.__dlpack__(stream=stream)`, where `stream` is 1 or 2, is called for a CUDA tensor using ROCm. For more details, see [the documentation][1]. [1]: https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__dlpack__.html [ghstack-poisoned]

pytorchmergebot · 2025-07-19T20:48:30Z

Starting merge as part of PR stack under #150691

This PR introduces the rest of the keyword-arguments added in DLPack version 2023.12: `dl_device` and `copy`. In summary, we handle these arguments in the C++ implementation of `to_dlpack(...)` at _torch/csrc/Module.cpp_, by calling the `maybeCopyTensor` function at _aten/src/ATen/DLConvertor.cpp_. It also introduces the following changes: - Add a new Python API `torchDeviceToDLDevice()`, which is simply a refactoring of the `getDLDevice()` function at _aten/src/ATen/DLConvertor.cpp_. - Add both keyword-arguments to the `from_dlpack()` function at _torch/utils/dlpack.py_ and to the `Tensor.__dlpack__()` dunder method. Pull Request resolved: #150218 Approved by: https://github.com/albanD ghstack dependencies: #150216, #150217

This PR addresses the Array API documentation for [`__dlpack__`][1] and [`from_dlpack`][2] by making some buffer-related errors `BufferError` instead of `RuntimeError`, e.g. incompatible dtype, strides, or device. [1]: https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__dlpack__.html [2]: https://data-apis.org/array-api/latest/API_specification/generated/array_api.from_dlpack.html#from-dlpack Pull Request resolved: #150691 Approved by: https://github.com/Skylion007, https://github.com/albanD ghstack dependencies: #150216, #150217, #150218

…capture (#163242) Many extensions (including pybind helpers) call `Tensor.__dlpack__()` without a stream argument. Before #150217, `stream=None` behaved like “no cross-stream sync” and was safe inside CUDA Graph capture. After #150217, `stream=None` maps to the legacy default stream, adding a cross-stream wait that invalidates capture when running on a non-default stream. See this example ``` import torch s = torch.cuda.Stream() x = torch.randn(8, device="cuda") g = torch.cuda.CUDAGraph() with torch.cuda.stream(s): with torch.cuda.graph(g): _ = x + 1 cap = x.__dlpack__() _ = torch.utils.dlpack.from_dlpack(cap) ``` This PR partially reverts #150217 that stream=None defaults to no sync. Pull Request resolved: #163242 Approved by: https://github.com/ngimel

Update

1eadb0e

[ghstack-poisoned]

This was referenced Mar 28, 2025

Upgrade to DLPack 1.0. #145000

Closed

[DLPack] add NumPy exchange tests. #150216

Closed

[DLPack] Add support for missing keyword-arguments. #150218

Closed

pytorchbot added the open source label Mar 28, 2025

ysiraichi added 3 commits March 28, 2025 16:48

Update

cf810ee

[ghstack-poisoned]

Rebased.

c101dbf

[ghstack-poisoned]

Update

23422b3

[ghstack-poisoned]

ysiraichi added module: dlpack topic: not user facing topic category labels Mar 28, 2025

ysiraichi requested review from albanD and rgommers March 28, 2025 21:15

ysiraichi added 2 commits April 4, 2025 12:50

Update

56e2c32

[ghstack-poisoned]

Update

3c4f7da

[ghstack-poisoned]

ysiraichi mentioned this pull request Apr 4, 2025

Raise BufferError for DLPack buffer-related errors. #150691

Closed

Update

b2de390

[ghstack-poisoned]

ysiraichi added 2 commits April 25, 2025 20:30

Update

d32a334

[ghstack-poisoned]

Update

476b1ea

[ghstack-poisoned]

msaroufim approved these changes May 14, 2025

View reviewed changes

albanD approved these changes May 14, 2025

View reviewed changes

ysiraichi added 4 commits May 24, 2025 11:55

Update

3338ea1

[ghstack-poisoned]

Rebased.

d121ce9

[ghstack-poisoned]

Update

c2e1184

[ghstack-poisoned]

ysiraichi added 2 commits July 12, 2025 16:22

pytorchmergebot added the Merged label Jul 20, 2025

pytorchmergebot closed this in 1d526fe Jul 20, 2025

github-actions bot deleted the gh/ysiraichi/85/head branch August 19, 2025 02:16

eee4017 mentioned this pull request Sep 18, 2025

Make Tensor.__dlpack__(stream=None) capture-safe during CUDA Graph capture #163242

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fix DLPack stream logic. #150217

Fix DLPack stream logic. #150217

Uh oh!

ysiraichi commented Mar 28, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Mar 28, 2025 •

edited

Loading

Uh oh!

albanD left a comment

Uh oh!

albanD May 14, 2025

Uh oh!

ysiraichi May 24, 2025

Uh oh!

albanD May 14, 2025

Uh oh!

ysiraichi May 24, 2025

Uh oh!

albanD May 14, 2025

Uh oh!

ysiraichi May 24, 2025

Uh oh!

pytorchmergebot commented Jul 12, 2025

Uh oh!

pytorchmergebot commented Jul 19, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Fix DLPack stream logic. #150217

Fix DLPack stream logic. #150217

Uh oh!

Conversation

ysiraichi commented Mar 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Mar 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ No Failures

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

albanD May 14, 2025

Choose a reason for hiding this comment

Uh oh!

ysiraichi May 24, 2025

Choose a reason for hiding this comment

Uh oh!

albanD May 14, 2025

Choose a reason for hiding this comment

Uh oh!

ysiraichi May 24, 2025

Choose a reason for hiding this comment

Uh oh!

albanD May 14, 2025

Choose a reason for hiding this comment

Uh oh!

ysiraichi May 24, 2025

Choose a reason for hiding this comment

Uh oh!

pytorchmergebot commented Jul 12, 2025

Uh oh!

pytorchmergebot commented Jul 19, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ysiraichi commented Mar 28, 2025 •

edited

Loading

pytorch-bot bot commented Mar 28, 2025 •

edited

Loading