-
Notifications
You must be signed in to change notification settings - Fork 26.3k
add narrow() support for sparse tensors re: #8853 #11342
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
c5bd567 to
bce0904
Compare
|
I didn't review the code closely, but IIUC you are not sharing storage with the old tensor are you? Since this is the case, you should not call it |
aten/src/ATen/native/TensorShape.cpp
Outdated
| Tensor _narrow_sparse(const Tensor& self, int64_t dim, int64_t start, int64_t length){ | ||
| LongTensor indices = self._indices(); | ||
| Tensor values = self._values(); | ||
| int64_t numCoords = indices.size(1); |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
ezyang
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Unfortunately, this needs some work.
The big problem is that, on CUDA, this kernels is written in a hugely inefficient way, because the iterated calls to toCLong each do a CUDA synchronization. That's a lot of synchronizations. There is also a code smell, where loops are written by hand (which means there's no opportunity for vectorization or parallelism).
Assuming that you want to write a non-storage sharing narrow, I think a better strategy is to compute a boolean mask of indices to keep, and then use masked_select to grab them from indices and values. That should eliminate all of the loops.
|
To answer your questions:
skipIfROCM turns off the test for AMD GPUs. I would advise you to omit that for now, and then add it if the AMD tests fail.
Moot if you don't call this narrow. |
|
@ezyang Thanks for the comments. My intention is to mirror the functionality of dense.narrow() as closely as I can; so sharing storage is the preferred approach as much as possible. A) For B) For the If I'm correct in A & B do you still suggest a name change for the function? |
8187c8b to
0cc6561
Compare
|
I've pushed a new version w/ rename & replaced the for loops. This implementation does not share storage, but it is an improvement over to_dense().narrow(). If there is an easy way to do |
|
Yes, I think I agree that writing a storage sharing narrow seems difficult. It might be possible for coalesced tensors, if you only narrow a trailing dimension, but that seems like a limited enough case that we shouldn't bother.
I don't see how your second sentence follows from the first. There's no way to do a |
aten/src/ATen/native/TensorShape.cpp
Outdated
| Tensor newIndices = indices.masked_select(mask).view({dims, -1}); | ||
| return self.type().sparse_coo_tensor(newIndices, newValues, newSizes); | ||
| }else{ | ||
| return self.clone().narrow(dim,start,length); |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
|
CC @yf225 @weiyangfb @gchanan if you have any opinions about the naming. Maybe |
aten/src/ATen/native/TensorShape.cpp
Outdated
| newSizes[dim]=length; | ||
|
|
||
| Tensor narrowDim = at::zeros_like(indices[dim]); | ||
| narrowDim.copy_(indices[dim]); |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
aten/src/ATen/native/TensorShape.cpp
Outdated
| narrowDim.copy_(indices[dim]); | ||
| Tensor mask = (narrowDim >= start).__and__((narrowDim < end)); | ||
|
|
||
| indices[dim] = indices[dim].add(-start); |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
|
New functions would need docs. |
ezyang
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Still buggy.
Sorry, I skipped a few logical steps there. I was just referring to the fact that you can theoretically do |
|
alrighty, added docs & a simplifications re: convo above. |
torch/_tensor_docs.py
Outdated
| r""" | ||
| narrow(dimension, start, length) -> Tensor | ||
| Same functionality as :meth:`Tensor.narrow` except returning a full copy, |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
aten/src/ATen/native/TensorShape.cpp
Outdated
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
weiyangfb
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This PR still requires some changes
aten/src/ATen/native/TensorShape.cpp
Outdated
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
aten/src/ATen/native/TensorShape.cpp
Outdated
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
0d9577a to
c4f0628
Compare
test/test_sparse.py
Outdated
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
75baf75 to
126f6a2
Compare
| narrow_copy(dimension, start, length) -> Tensor | ||
| Same as :meth:`Tensor.narrow` except returning a copy rather | ||
| than shared storage. This is primarily for sparse tensors, which |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
aten/src/ATen/native/TensorShape.cpp
Outdated
| Tensor narrow_copy_sparse(const Tensor& self, int64_t dim, int64_t start, int64_t length){ | ||
| AT_CHECK(self.dim() > 0, "narrow() cannot be applied to a 0-dim tensor."); | ||
| auto cur_size = self.size(dim); | ||
| AT_CHECK(length >= 0 && start <= cur_size - length, |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
|
|
||
| with_dense, _, _ = self._gen_sparse(2, 7, shape) | ||
| for narrow_args in self._all_narrow_combs(shape): | ||
| self._test_narrow(with_dense, narrow_args) |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
test/test_sparse.py
Outdated
| self._test_narrow(with_dense, narrow_args) | ||
| self._test_narrow(with_dense, narrow_args) | ||
|
|
||
| self.assertRaises(RuntimeError, lambda: input.narrow_copy(10, 0, 3)) # dim > sparseDim + denseDim |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
test/test_sparse.py
Outdated
| input, _, _ = self._gen_sparse(4, 19, shape) | ||
| for narrow_args in self._all_narrow_combs(shape): | ||
| self._test_narrow(input, narrow_args) | ||
| self._test_narrow(input.coalesce(), narrow_args) |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
|
|
||
| Tensor narrow_copy_sparse(const Tensor& self, int64_t dim, int64_t start, int64_t length){ | ||
| int64_t allDim = self.dim(); | ||
| int64_t end = start+length; |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
facebook-github-bot
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
weiyangfb has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
Looks good! Thanks @realdoug ! |
Summary: Couple questions: 1) I used the log1p implementation in #8969 as a guide especially for testing. I'm not sure what the ```skipIfROCM``` annotation is for, so unsure if i need it for my test. 2) I implemented the branching logic in the narrow function itself; is this the right place to do so? I noticed that there a number of places where sparse-specific logic is handled with just an if statement in this file. Or should I implement a separate dispatch in native_functions.yml as in the log1p? And of course, happy to make any any other updates/changes that I may have missed as well. This is my first PR to the project. Pull Request resolved: pytorch/pytorch#11342 Differential Revision: D9978430 Pulled By: weiyangfb fbshipit-source-id: e73dc20302ab58925afb19e609e31f4a38c634ad
* upstream/master: (117 commits) Add full namespace resolution in CAFFE_DURATION (pytorch#12065) T33898723: Simple put operators for caffe2 stats (pytorch#12057) add narrow() support for sparse tensors re: pytorch#8853 (pytorch#11342) Fix ONNX bug, add symbolic for full Enable tracing of tensor factories with an out argument Fix warnings emitted when testing distributions (pytorch#12038) Unify versions across setup.py, libtorch, and libcaffe2 (pytorch#12053) add autodiff expressions for common operations (pytorch#11832) Blob doesn't allow access to destroyCall anymore (pytorch#11548) IValue can store Blob (pytorch#11414) Move Blob to ATen/core (pytorch#11924) Use tempfile during serialized test comparison (pytorch#12021) fix segfault when grad to a hook fn is None (pytorch#12028) Fallback CreateMutex/AtomicIter operators for mkl-dnn Unify all *_EXPORT and *_IMPORT macros across c++ backend (pytorch#12019) Add safety asserts for methods on TensorImpl which don't work on Variable. (pytorch#12058) Make USE_IDEEP work again (pytorch#12026) Fix "identifier following the 'template' keyword does not refer to a template" (pytorch#12037) Delete some unused variables. (pytorch#12059) Support TypeIdentifier::name() (pytorch#12036) ...
Couple questions:
I used the log1p implementation in Add log1p for sparse tensor #8969 as a guide especially for testing. I'm not sure what the
@skipIfROCMannotation is for, so unsure if i need it for my test.I implemented the branching logic in the narrow function itself; is this the right place to do so? I noticed that there a number of places where sparse-specific logic is handled with just an if statement in this file. Or should I implement a separate dispatch in native_functions.yml as in the log1p?
And of course, happy to make any any other updates/changes that I may have missed as well. This is my first PR to the project.