Implement hybrid sparse to/from dense conversions.

[alexsamardzic]

Stack from ghstack (oldest at bottom):

• -> Implement hybrid sparse to/from dense conversions. #90177

cc @nikitaved @pearu @cpuhrsch @amjames @bhosmer @gujinghui @PenghuiCheng @XiaobingSuper @jianyuh @jgong5 @mingfeima @sanchitintel @ashokei @jingxu10 @min-jean-cho @yanbing-j @Guobing-Chen @Xia-Weiwen

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 0498f42:
💚 Looks good so far! There are no failures yet. 💚

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

[alexsamardzic]

Just the first version of hybrid CSR/CSC to dense conversions... Mimics BSR/BSC to dense conversions, so it remains to:

1. unify those
2. try to avoid looping over batches

Testing script (create a CSR matrix, then create CSC with transposed rows/columns, convert both to dense, and check are these transposes, along the sparse dimensions, of each other)

import torch

BATCH_DIMS_MIN = 1
BATCH_DIMS_MAX = 3
BATCH_SIZE_MIN = 1
BATCH_SIZE_MAX = 5
SPARSE_SIZE_MIN = 1
SPARSE_SIZE_MAX = 10
HYBRID_DIMS_MIN = 1
HYBRID_DIMS_MAX = 5
HYBRID_SIZE_MIN = 1
HYBRID_SIZE_MAX = 10
NTESTS = 10000

for i in range(NTESTS):
    batch_dims = torch.randint(BATCH_DIMS_MIN, BATCH_DIMS_MAX + 1, (1,)).item()
    batch_size = tuple(torch.randint(BATCH_SIZE_MIN, BATCH_SIZE_MAX + 1, (batch_dims,)).tolist())
    (nrows, ncols) = sparse_size = tuple(torch.randint(SPARSE_SIZE_MIN, SPARSE_SIZE_MAX + 1, (2,)).tolist())
    hybrid_dims = torch.randint(HYBRID_DIMS_MIN, HYBRID_DIMS_MAX + 1, (1,)).item()
    hybrid_size = tuple(torch.randint(HYBRID_SIZE_MIN, HYBRID_SIZE_MAX + 1, (hybrid_dims,)).tolist())

    dense = torch.randn(sparse_size, dtype=torch.float32)
    dense = dense * dense.relu().bool()
    dense = dense.repeat(batch_size + (1, 1))
    csr = dense.to_sparse_csr()
    hybrid_values = torch.randn(csr.values().shape + hybrid_size, dtype=torch.float32)
    hybrid_csr = torch.sparse_compressed_tensor(
        csr.crow_indices(),
        csr.col_indices(),
        hybrid_values,
        batch_size + (nrows, ncols) + hybrid_size,
        layout=torch.sparse_csr,
        dtype=torch.float32)
    #print("hybrid_csr:")
    #print(hybrid_csr)
    dense_hybrid_csr = hybrid_csr.to_dense()
    #print(dense_hybrid_csr)

    hybrid_csc = torch.sparse_compressed_tensor(
        csr.crow_indices(),import torch
        csr.col_indices(),
        hybrid_values,
        batch_size + (ncols, nrows) + hybrid_size,
        layout=torch.sparse_csc,
        dtype=torch.float32)
    #print("hybrid_csc:")
    #print(hybrid_csc)
    dense_hybrid_csc = hybrid_csc.to_dense()
    #print(dense_hybrid_csc)

    assert(torch.all(torch.transpose(dense_hybrid_csr, batch_dims, batch_dims + 1) == dense_hybrid_csc))

[nikitaved]

nit: mixed-style in-place unsqueeze. It is probably better to do values.unsqueeze_(...); dense.unsqueeze_(...) for more clarity.

[nikitaved]

It is possible to remove the for-loop and use compressed_indices_to_plain_indices from https://github.com/pytorch/pytorch/pull/88078/files#r1013344803 reimplemented in CPP as a utility function so that a single call to _convert_indices_from_csr_to_coo is sufficient. Similar pattern but reversed was used in #82122.

[alexsamardzic]

Thanks for the pointer - as mentioned in my comment, that would be good optimization to have.

[alexsamardzic]

Added support for hybrid BSR/BSC to dense conversions. The unification with CSR/CSC to dense conversions, as well as optimization on eliminating looping for batches handling, are next to come.

[alexsamardzic]

Unified CSR/CSC and BSR/BSC conversion to dense.

[amjames]

Great start, I found one error, and I have a cleaner-code note. See inline comments.

[amjames]

This condition should either be batch_ndim > 1 or there should be a notee that when there is already a single batch dim (batch_ndim = 1) the flatten ops below result in flatten(0, 0) and return the input tensors unmodified hence we are not doing anything incorrect or unnecessary.

[amjames]

values was reshaped above so it contains exactly one batch dimension, but batch_ndim tracks the number of batch dims in the original shape.

After the flatten/unsqueeze of is applied to values the generalized shape would be (n_batch, nnz) + blocksize + dense_dims where n_batch, and nnz are scalars, blocksize is a length 2 tuple and dense_dims is a tuple with as many members as there are dense dimensions in the tensor, possibly zero.

So at this point the blocksizes can always be computed as values.size(2) and values.size(3) for any input which enters this region.

[alexsamardzic]

Fixed - thanks! Also updated, in the next commit below, according to other comments of yours.

[amjames]

It looks like blocksize is only used to compute the dense shape in the else block of the conditional below, so it can be moved into that scope and need not exist at this level.

[alexsamardzic]

Removed looping over batches. Please do not review yet, until comments added, and tests eventually enabled.

[pearu]

Thanks, @alexsamardzic, this looks good!

I have a suggestion to add support for

to_sparse(layout=torch.tensor_coo, dense_dim=...)

as a low handing fruit.

Also, I have a question about the non-sparse-non-dense dimensions in to_sparse_csr/csc/bsr/bsc.

[cpuhrsch]

Is this additional whitespace really how we print BSR Tensors?

@alexsamardzic - you can see the preview of the Python documentation as part of the Helpful Links. You'll just need to search for the function name that you changed to see how it renders on the web.

[pearu]

@cpuhrsch printing BSR tensors are defined as how we print strided tensors. For example:

>>> torch.ones(4, 4).to_sparse_bsr((2, 2)).values()
tensor([[[1., 1.],
         [1., 1.]],

        [[1., 1.],
         [1., 1.]],

        [[1., 1.],
         [1., 1.]],

        [[1., 1.],
         [1., 1.]]])

Personally, I also find the whitespace annoying but I guess this is supposed to help to view high-dimensional tensors.

[alexsamardzic]

Indeed, I've copied output from the interpreter for these examples.

[cpuhrsch]

Should this be = 1?

[alexsamardzic]

It should, strange that output is correct i.e. as that = 1 was there. In any case, an update pushed.

[pearu]

I have a couple of nits and questions. Thanks, @alexsamardzic!

[pearu]

I have a couple of more clean-up nits but otherwise looks good to me!

[pearu]

LGTM! Thanks, @alexsamardzic!

[cpuhrsch]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your 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

Advanced Debugging

Check the merge workflow status
here