Support different NSE in batches of CSR and CSC tensors #84843
Conversation
This file contains 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/84843
Note: Links to docs will display an error until the docs builds have been completed. ❌ 4 Failures, 1 PendingAs of commit 18f4bfc: The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
[ghstack-poisoned]
This PR enables batched CSR/CSC tensors that batches may have different NSE counts.
For instance, with the current master we have
```python
>>> a = torch.tensor([[[1, 2], [3, 4]], [[0, 12], [21, 0]]])
>>> a.to_sparse_csr()
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
RuntimeError: Expect the same number of specified elements per batch.
```
because the NSE of the first and second batches are different, 4 and 2, respectively.
This PR implements a strided-to-sparse-CSR/CSC conversion algorithm that supports CSR/CSC batches with different NSE counts. For instance:
```python
>>> a = torch.tensor([[[1, 2], [3, 4]], [[0, 12], [21, 0]]])
>>> b = a.to_sparse_csr()
>>> b
tensor(crow_indices=tensor([[0, 2, 4],
[0, 1, 2]]),
col_indices=tensor([[0, 1, 0, 1],
[1, 0, 0, 0]]),
values=tensor([[ 1, 2, 3, 4],
[12, 21, 0, 0]]), size=(2, 2, 2), nnz=4,
layout=torch.sparse_csr)
>>> b[0]
tensor(crow_indices=tensor([0, 2, 4]),
col_indices=tensor([0, 1, 0, 1]),
values=tensor([1, 2, 3, 4]), size=(2, 2), nnz=4,
layout=torch.sparse_csr)
>>> b[1]
tensor(crow_indices=tensor([0, 1, 2]),
col_indices=tensor([1, 0]),
values=tensor([12, 21]), size=(2, 2), nnz=2, layout=torch.sparse_csr)
```
that is, if the NSE of a batch is smaller than the maximum NSE over all batches, the corresponding rows in `col_indices`/`values` are padded with zeros as placeholders. Algorithms on batched CSR/CSC tensors must not access the padded parts of these tensors, that is, the algorithms should use the last element of the corresponding `crow_indices` row as the NSE value rather than the value of `.values().shape[0]` that holds the maximum NSE over all batches.
Performance-wise, the strided-to-sparse-CSR/CSC conversion algorithms in master and in this PR, are roughly equivalent:
```python
# master branch:
n [2]: a = torch.rand(10, 10, 1000, 1000)
In [3]: a = torch.where(a==0, 0.1, a) # required for master, optional for the PR
In [4]: %timeit a.to_sparse_csr()
2.25 s ± 9.84 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [5]: a_cuda = a.cuda()
In [6]: %timeit a_cuda.to_sparse_csr()
55.2 ms ± 6.95 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```
```python
# this PR
In [2]: a = torch.rand(10, 10, 1000, 1000)
In [3]: a = torch.where(a==0, 0.1, a) # required for master, optional for the PR
In [4]: %timeit a.to_sparse_csr()
2.13 s ± 7.73 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [5]: a_cuda = a.cuda()
In [6]: %timeit a_cuda.to_sparse_csr()
54.3 ms ± 20.5 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```
The performance of the PR is only slightly better than the master branch.
A strided-to-sparse-BSR/BSC conversion with variable NSE support will be implemented as a follow-up.
[ghstack-poisoned]
This PR enables batched CSR/CSC tensors that batches may have different NSE counts.
For instance, with the current master we have
```python
>>> a = torch.tensor([[[1, 2], [3, 4]], [[0, 12], [21, 0]]])
>>> a.to_sparse_csr()
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
RuntimeError: Expect the same number of specified elements per batch.
```
because the NSE of the first and second batches are different, 4 and 2, respectively.
This PR implements a strided-to-sparse-CSR/CSC conversion algorithm that supports CSR/CSC batches with different NSE counts. For instance:
```python
>>> a = torch.tensor([[[1, 2], [3, 4]], [[0, 12], [21, 0]]])
>>> b = a.to_sparse_csr()
>>> b
tensor(crow_indices=tensor([[0, 2, 4],
[0, 1, 2]]),
col_indices=tensor([[0, 1, 0, 1],
[1, 0, 0, 0]]),
values=tensor([[ 1, 2, 3, 4],
[12, 21, 0, 0]]), size=(2, 2, 2), nnz=4,
layout=torch.sparse_csr)
>>> b[0]
tensor(crow_indices=tensor([0, 2, 4]),
col_indices=tensor([0, 1, 0, 1]),
values=tensor([1, 2, 3, 4]), size=(2, 2), nnz=4,
layout=torch.sparse_csr)
>>> b[1]
tensor(crow_indices=tensor([0, 1, 2]),
col_indices=tensor([1, 0]),
values=tensor([12, 21]), size=(2, 2), nnz=2, layout=torch.sparse_csr)
```
that is, if the NSE of a batch is smaller than the maximum NSE over all batches, the corresponding rows in `col_indices`/`values` are padded with zeros as placeholders. Algorithms on batched CSR/CSC tensors must not access the padded parts of these tensors, that is, the algorithms should use the last element of the corresponding `crow_indices` row as the NSE value rather than the value of `.values().shape[0]` that holds the maximum NSE over all batches.
Performance-wise, the strided-to-sparse-CSR/CSC conversion algorithms in master and in this PR, are roughly equivalent:
```python
# master branch:
n [2]: a = torch.rand(10, 10, 1000, 1000)
In [3]: a = torch.where(a==0, 0.1, a) # required for master, optional for the PR
In [4]: %timeit a.to_sparse_csr()
2.25 s ± 9.84 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [5]: a_cuda = a.cuda()
In [6]: %timeit a_cuda.to_sparse_csr()
55.2 ms ± 6.95 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```
```python
# this PR
In [2]: a = torch.rand(10, 10, 1000, 1000)
In [3]: a = torch.where(a==0, 0.1, a) # required for master, optional for the PR
In [4]: %timeit a.to_sparse_csr()
2.12 s ± 2.13 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [5]: a_cuda = a.cuda()
In [6]: %timeit a_cuda.to_sparse_csr(); torch.cuda.synchronize()
47.2 ms ± 10.4 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```
The performance of `to_sparse_csr()` on CUDA tensors increased by 15% with this PR.
A strided-to-sparse-BSR/BSC conversion with variable NSE support will be implemented as a follow-up.
[ghstack-poisoned]
This PR enables batched CSR/CSC tensors that batches may have different NSE counts.
For instance, with the current master we have
```python
>>> a = torch.tensor([[[1, 2], [3, 4]], [[0, 12], [21, 0]]])
>>> a.to_sparse_csr()
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
RuntimeError: Expect the same number of specified elements per batch.
```
because the NSE of the first and second batches are different, 4 and 2, respectively.
This PR implements a strided-to-sparse-CSR/CSC conversion algorithm that supports CSR/CSC batches with different NSE counts. For instance:
```python
>>> a = torch.tensor([[[1, 2], [3, 4]], [[0, 12], [21, 0]]])
>>> b = a.to_sparse_csr()
>>> b
tensor(crow_indices=tensor([[0, 2, 4],
[0, 1, 2]]),
col_indices=tensor([[0, 1, 0, 1],
[1, 0, 0, 0]]),
values=tensor([[ 1, 2, 3, 4],
[12, 21, 0, 0]]), size=(2, 2, 2), nnz=4,
layout=torch.sparse_csr)
>>> b[0]
tensor(crow_indices=tensor([0, 2, 4]),
col_indices=tensor([0, 1, 0, 1]),
values=tensor([1, 2, 3, 4]), size=(2, 2), nnz=4,
layout=torch.sparse_csr)
>>> b[1]
tensor(crow_indices=tensor([0, 1, 2]),
col_indices=tensor([1, 0]),
values=tensor([12, 21]), size=(2, 2), nnz=2, layout=torch.sparse_csr)
```
that is, if the NSE of a batch is smaller than the maximum NSE over all batches, the corresponding rows in `col_indices`/`values` are padded with zeros as placeholders. Algorithms on batched CSR/CSC tensors must not access the padded parts of these tensors, that is, the algorithms should use the last element of the corresponding `crow_indices` row as the NSE value rather than the value of `.values().shape[0]` that holds the maximum NSE over all batches.
Performance-wise, the strided-to-sparse-CSR/CSC conversion algorithms in master and in this PR, are roughly equivalent:
```python
# master branch:
n [2]: a = torch.rand(10, 10, 1000, 1000)
In [3]: a = torch.where(a==0, 0.1, a) # required for master, optional for the PR
In [4]: %timeit a.to_sparse_csr()
2.25 s ± 9.84 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [5]: a_cuda = a.cuda()
In [6]: %timeit a_cuda.to_sparse_csr()
55.2 ms ± 6.95 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```
```python
# this PR
In [2]: a = torch.rand(10, 10, 1000, 1000)
In [3]: a = torch.where(a==0, 0.1, a) # required for master, optional for the PR
In [4]: %timeit a.to_sparse_csr()
2.12 s ± 2.13 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [5]: a_cuda = a.cuda()
In [6]: %timeit a_cuda.to_sparse_csr(); torch.cuda.synchronize()
47.2 ms ± 10.4 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```
The performance of `to_sparse_csr()` on CUDA tensors increased by 15% with this PR.
A strided-to-sparse-BSR/BSC conversion with variable NSE support will be implemented as a follow-up.
[ghstack-poisoned]
This was referenced Sep 19, 2022
pearu
added a commit
that referenced
this pull request
Sep 19, 2022
ghstack-source-id: 19d3b9616b62846c14ef45a85ea4468db42e0836 Pull Request resolved: #84843
|
/easycla As part of the transition to the PyTorch Foundation, this project now requires contributions be covered under the new CLA. See #85559 for additional details. This comment will trigger a new check of this PR. If you are already covered, you will simply see a new "EasyCLA" check that passes. If you are not covered, a bot will leave a new comment with a link to sign. |
|
|
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
|
Converting to draft as the used approach in this PR requires further discussion and there exists other alternatives, see #104193 |
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.
This PR enables batched CSR/CSC tensors that batches may have different NSE counts.
For instance, with the current master we have
because the NSE of the first and second batches are different, 4 and 2, respectively.
This PR implements a strided-to-sparse-CSR/CSC conversion algorithm that supports CSR/CSC batches with different NSE counts. For instance:
that is, if the NSE of a batch is smaller than the maximum NSE over all batches, the corresponding rows in
col_indices/valuesare padded with zeros as placeholders. Algorithms on batched CSR/CSC tensors must not access the padded parts of these tensors, that is, the algorithms should use the last element of the correspondingcrow_indicesrow as the NSE value rather than the value of.values().shape[0]that holds the maximum NSE over all batches.Performance-wise, the strided-to-sparse-CSR/CSC conversion algorithms in master and in this PR, are roughly equivalent:
The performance of
to_sparse_csr()on CUDA tensors increased by 15% with this PR.A strided-to-sparse-BSR/BSC conversion with variable NSE support will be implemented as a follow-up.
Stack from ghstack (oldest at bottom):
cc @nikitaved @cpuhrsch @amjames @bhosmer