Repository: pytorch
Issue: Sparse CSR layout GPU backend tracking issue
Labels: module: sparse,module: cuda,triaged,tracker

## PyTorch 1.11

Currently, cuSPARSE is already used in PyTorch for some operations with COO sparse matrix format. Internally COO indices are converted to a low-level CSR representation that is used to call cuSPARSE routines and reconstruct the result back to COO.

For PyTorch 1.11 we're focusing on improving sparse CSR support and performance on GPUs.
This work aims to utilize all available functionality of cuSPARSE to expand operator support for PyTorch CSR matrices.

Go to [the list of tasks](https://github.com/pytorch/pytorch/issues/60854#pytorch111-tasks).

## cuSPARSE Generic API
cuSPARSE started to deprecate low-level functions that are similar to BLAS to more high level "generic" interface:
> The cuSPARSE Generic APIs allow computing the most common sparse linear algebra operations, such as sparse matrix-vector (SpMV) and sparse matrix-matrix multiplication (SpMM), in a flexible way. The new APIs have the following capabilities and features:
> Set matrix data layouts, number of batches, and storage formats (for example, CSR, COO, and so on)
> Set input/output/compute data types. This also allows mixed data-type computation
> Set types of sparse matrix indices
> Choose the algorithm for the computation
> Provide external device memory for internal operations
> Provide extensive consistency checks across input matrices and vectors for a given routine. This includes the validation of matrix sizes, data types, layout, allowed operations, etc.

Hopefully, the price for the generic vs low-level API is not large, and we have no choice here as the low-level BLAS-like functions are deprecated and being removed.

## CUDA version requirement for cuSPARSE
cuSPARSE Generic API was introduced in CUDA 10.1, Windows distribution was missing until CUDA 11.
For now, the requirement will be CUDA 11+, because it has more functionality. This requirement can be revisited later on a function-by-function basis.
PyTorch will raise a runtime error if an older CUDA version is used while calling sparse operators on GPU.

Not all architectures are supported. If a specific GPU model does not provide native support for a given data type, the routine returns the `CUSPARSE_STATUS_ARCH_MISMATCH` error.
Unsupported data types and Compute Capability (CC):
`__half` on GPUs with CC < 53 (e.g. Kepler)
`__nv_bfloat16` on GPUs with CC < 80 (e.g. Kepler, Maxwell, Pascal, Volta, Turing).

## Sparse/Dense Matrix/Vector descriptors
cuSPARSE Generic API uses descriptor structs to collect the information about the tensor (size, indices, values, etc.)
There are 4 descriptors:
* cusparseSpVecDescr - sparse vector
* cusparseSpMatDescr - sparse matrix (COO or CSR format)
* cusparseDnVecDescr - dense vector
* cusparseDnMatDescr - dense matrix

Each descriptor struct has constructors, destructors and a few attribute getters/setters defined in cuSPARSE.
cuDNN follows a similar model and in PyTorch, the descriptors are wrapped in `aten/src/ATen/cudnn/Descriptors.h`.

## Indices dtype

PyTorch uses `int64` indices for COO format. The current implementation of CSR supports both `int32` and `int64` dtypes for indices.

cuSPARSE supports choosing 32-bit or 64-bit indexing in runtime. The generic API only SpGEMM function doesn't support 64-bit indexing, and `cusparse<t>csrgeam2` doesn't support it.

<h2 id="batched-CSR-computation">Batched computation support</h2>
cuSPARSE descriptors for dense and sparse matrices have the functionality to set batch count and batch stride.
Currently, only the cusparseSpMM function can use this information.
This poses the requirement that indices and values are batch strided (contiguous) tensors, that is


| attribute | batched shape    | single matrix shape
| -------- | -------- | ---- |
| crow_indices     | (batch_shape, num_rows+1) | (num_rows+1,) |
| col_indices     | (batch_shape, nnz)     | (nnz,) |
| values     | (batch_shape, nnz)     | (nnz,)

For functions that do not have batched support, we can have a for loop on the batch dimension.

<h2 id="pytorch111-tasks">PyTorch 1.11 Tasks</h2>

If you are planning to work on a specific task please put your name next to the task or leave a comment below. For any related pull requests use the `module: sparse` label. Feel free to tag @ivanyashchuk as a reviewer.

A list of tasks enumerated by cuSPARSE routine (ordered by priority):

- [ ] [cusparseSpMM](https://docs.nvidia.com/cuda/cusparse/index.html#cusparse-generic-function-spmm)  (sparse matrix - dense matrix multiplication, dense result) (@IvanYashchuk)
    cuSPARSE supports batching Cᵢ = A ⋅ Bᵢ, Cᵢ = Aᵢ ⋅ B, Cᵢ = Aᵢ ⋅ Bᵢ.
    The deterministic algorithm is available (only column-major dense matrices, no batching).
    Both 32-bit and 64-bit indices are supported.
    Uniform precision and mixed-precision computations are supported.
    Dtypes: float32, float64, complex64, complex128, float16, bfloat16 (`AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND2(kHalf, kBFloat16`)
    - [x] Implement descriptor wrappers for CSR and dense matrices.
    - [x] https://github.com/pytorch/pytorch/issues/58899
    - [ ] https://github.com/pytorch/pytorch/issues/59099
    
    Backend for:
    - [x] torch.mm
    - [x] torch.addmm
    - [x] torch.baddmm (if batched CSR is enabled in PyTorch)
       PR: https://github.com/pytorch/pytorch/pull/68711
       done via a hack https://github.com/pytorch/pytorch/blob/bd3db019a039d101124d35c9feba16118c96e458/test/test_sparse_csr.py#L649-L651
    - [x] torch._sparse_mm
        PR in progress: https://github.com/pytorch/pytorch/pull/60838

- [ ] [cusparseSpGEMM](https://docs.nvidia.com/cuda/cusparse/index.html#cusparse-generic-function-spgemm) (sparse matrix - sparse matrix multiplication, sparse result) (@IvanYashchuk)
    Only deterministic algorithm is available.
    Only 32-bit indices are supported.
    Only opA, opB equal to `CUSPARSE_OPERATION_NON_TRANSPOSE` are supported.
    Only uniform-precision computation
    Dtypes: float32, float64, complex64, complex128, float16, bfloat16 (`AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND2(kHalf, kBFloat16`)
    - [x] torch.mm
    - [x] torch.addmm
    - [ ] torch.baddmm (if batched CSR is enabled in PyTorch)
    - [x] torch._sparse_sparse_matmul
    PR in progress: https://github.com/pytorch/pytorch/pull/63511

- [x] [cusparseSpMV](https://docs.nvidia.com/cuda/cusparse/index.html#cusparse-generic-function-spmv) (sparse matrix - dense vector multiplication, dense result)  (@IvanYashchuk)

    Deterministic algorithm is available (only `opA == CUSPARSE_OPERATION_NON_TRANSPOSE`).
    Both 32-bit and 64-bit indices are supported.
    Uniform precision and mixed-precision computations are supported.
    Dtypes: float32, float64, complex64, complex128, float16, bfloat16 (`AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND2(kHalf, kBFloat16`)
    - [x] Implement descriptor wrappers for dense vectors.

    Backend for:
    - [x] torch.mv
    - [x] torch.addmv
        PR in progress: https://github.com/pytorch/pytorch/pull/61407

- [x] [cusparseSpSM](https://docs.nvidia.com/cuda/cusparse/index.html#cusparse-generic-function-spsm) (sparse matrix - dense matrix triangular solver, dense result) (@IvanYashchuk)
    [cusparseSpSV](https://docs.nvidia.com/cuda/cusparse/index.html#cusparse-generic-function-spsv) (sparse matrix - dense vector triangular solver, dense result)
    Only deterministic algorithm is available.
    Both 32-bit and 64-bit indices are supported.
    Dtypes: float32, float64, complex64, complex128 (`AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES`)

    Backend for:
    - [x] torch.triangular_solve
    PR in progress: https://github.com/pytorch/pytorch/pull/61858

- [x] [cusparse\<t\>csrgeam2](https://docs.nvidia.com/cuda/cusparse/index.html#csrgeam2) (sparse matrix - sparse matrix addition, sparse result) (@IvanYashchuk)

    Only 32-bit indices are supported.
    Dtypes: float32, float64, complex64, complex128 (`AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES`)

    Backend for:
    - [x] torch.Tensor.add https://github.com/pytorch/pytorch/issues/59060
    PR is progress: https://github.com/pytorch/pytorch/pull/63948

- [x] [cusparseSDDMM](https://docs.nvidia.com/cuda/cusparse/index.html#cusparse-generic-function-sddmm) (sparse mask, dense matrix - dense matrix multiplication, sparse result)

    Only a deterministic algorithm is available.
    Both 32-bit and 64-bit indices are supported.
    Dtypes: float32, float64, complex64, complex128 (`AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES`)

    Backend for:
    - [x] [new] torch.sparse.sampled_addmm 
    PR: https://github.com/pytorch/pytorch/pull/68007

- [ ] all above should support batched CSR input

Note that tensors with the CSR layout on CUDA device use `SparseCsrCUDA` dispatch key in `native_functions.yaml` file.

## Possible future work
There is a new library for structured sparse-dense matrix multiplication [cuSPARSELt](https://docs.nvidia.com/cuda/cusparselt/index.html), it might have better performance than cuSPARSE for some sparsity patterns https://github.com/pytorch/pytorch/issues/62153.



cc @ngimel @nikitaved @pearu @cpuhrsch @IvanYashchuk