Add pinned memory support to sparse COO/CSR/CSC/BSR/BSC tensors #129645
Conversation
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
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/129645
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 58348cb with merge base bdd83c4 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
…nsors" As in the title. Fixes #115330 [ghstack-poisoned]
…nsors" As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 [ghstack-poisoned]
…nsors" As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 [ghstack-poisoned]
…nsors" As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 [ghstack-poisoned]
…nsors" As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 [ghstack-poisoned]
…nsors" As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 [ghstack-poisoned]
…nsors" As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 [ghstack-poisoned]
…nsors" As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 [ghstack-poisoned]
…nsors" As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 [ghstack-poisoned]
amjames
requested changes
Jun 28, 2024
There was a problem hiding this comment.
I see one major issue that should be addressed.
What should happen when one component tensor passed to _with_tensorsconstructors, but not all. In particular consider when thevalues` are pinned, but indices tensor(s) are not.
a = torch.randn(10, 10)
a[::2, ::2] = 0
sp1 = a.to_sparse_coo()
values = a._values().clone().pin_memory()
indices = a._indices()
sp2 = torch.sparse_coo_tensor(indices, values, size=(10, 10))Here sp2.is_pinned() would return true, but sp2._indices().is_pinned() is false. This is inconsistent with
sp3 = sp1.pin_memory()which also has sp3.is_pinned() True, but sp3._indices().is_pinned() is True. Inconsistency like this will be a pain point for some users.
I see two options:
- Expect this meta data field to match across members -> We should add this to the invariant checks, or enforce this by pinning out-of-sync members if any member is pinned in the constructors.
- Allow for these to be out-of-sync
sparse.is_pinned()should be a short hand forvalues.is_pinned() && indices.is_pinned()(and the plain/compressed indices for sparse compressed layouts) the same way thatsparse.pin_memory()is short hand for pinning all member tensors.
Personally I think option 1) is better for the simplicity and I can't see a reason why one would want some members to be pinned while others are not.
Registering a CPU tensor with CUDA is essentially a copy operation to another device (although the pinned tensor can still be accessed by CPU processes). In the case of sparse tensors, we don't allow indices and values to be on different devices: an error is raised when trying to create a sparse tensor from indices and values that have different devices. I suggest that we treat pinned and non-pinned tensors as having different devices. So, when one tries to create a sparse tensor from indices and values that have different pinning state then an exception will be raised. This corresponds to the first part of option 1). |
…nsors" As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 cc alexsamardzic nikitaved cpuhrsch amjames bhosmer jcaip [ghstack-poisoned]
amjames
approved these changes
Jul 1, 2024
cpuhrsch
approved these changes
Jul 30, 2024
eqy
approved these changes
Jul 30, 2024
…nsors" As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 cc alexsamardzic nikitaved cpuhrsch amjames bhosmer jcaip [ghstack-poisoned]
…nsors" As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 cc alexsamardzic nikitaved cpuhrsch amjames bhosmer jcaip [ghstack-poisoned]
…nsors" As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 cc alexsamardzic nikitaved cpuhrsch amjames bhosmer jcaip [ghstack-poisoned]
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
aartbik
pushed a commit
that referenced
this pull request
Aug 3, 2024
) As in the title: To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported - `sparse_xyz_tensor(indices, values, pin_memory=True)` - `sparse_xyz_tensor(indices, values).pin_memory()` - `sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())` Fixes #115330 Pull Request resolved: #129645 Approved by: https://github.com/amjames, https://github.com/cpuhrsch, https://github.com/eqy
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
As in the title:
To register indices/values of a sparse XYZ tensor with CUDA, the following methods are supported
sparse_xyz_tensor(indices, values, pin_memory=True)sparse_xyz_tensor(indices, values).pin_memory()sparse_xyz_tensor(indices.pin_memory(), values.pin_memory())Fixes #115330
Stack from ghstack (oldest at bottom):
cc @alexsamardzic @nikitaved @cpuhrsch @amjames @bhosmer @jcaip