Sparse CSR layout GPU backend tracking issue

[IvanYashchuk]

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.

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.

Batched computation support

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.

PyTorch 1.11 Tasks

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 (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)

  • Implement descriptor wrappers for CSR and dense matrices.
  • Add 64-bit indices support to csrmm2 #58899
  • CSR: Relaxing constraints to s_addmm_out_sparse_dense_cuda_worker #59099

  Backend for:

  • torch.mm
  • torch.addmm
  • torch.baddmm (if batched CSR is enabled in PyTorch)
    PR: Sparse CSR CUDA: Add torch.baddbmm and torch.bmm #68711
    done via a hack

    pytorch/test/test_sparse_csr.py

    Lines 649 to 651 in bd3db01

    	# a_batched is a regular CSR tensor but with a batch dimension in the shape
    	a_batched = torch._sparse_csr_tensor_unsafe(
    	a.crow_indices(), a.col_indices(), a.values(), (batch_size, m, k))

  • torch._sparse_mm
    PR in progress: Add cuSPARSE descriptors and update CSR addmm #60838

• cusparseSpGEMM (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)

  • torch.mm
  • torch.addmm
  • torch.baddmm (if batched CSR is enabled in PyTorch)
  • torch._sparse_sparse_matmul
    PR in progress: Sparse CSR CUDA: add torch.addmm with all inputs sparse #63511

• cusparseSpMV (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)

  • Implement descriptor wrappers for dense vectors.

  Backend for:

  • torch.mv
  • torch.addmv
    PR in progress: Sparse CSR CUDA: add addmv_out #61407

• cusparseSpSM (sparse matrix - dense matrix triangular solver, dense result) (@IvanYashchuk)
  cusparseSpSV (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:

  • torch.triangular_solve
    PR in progress: Sparse CSR CUDA: add triangular_solve_out #61858

• cusparse<t>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:

  • torch.Tensor.add Addition of sparse CSR tensors is not yet implemented #59060
    PR is progress: Sparse CSR CUDA: add torch.add with all inputs sparse #63948

• cusparseSDDMM (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:

  • [new] torch.sparse.sampled_addmm
    PR: Sparse CSR CUDA: Add torch.sparse.sampled_addmm #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, it might have better performance than cuSPARSE for some sparsity patterns #62153.

cc @ngimel @nikitaved @pearu @cpuhrsch @IvanYashchuk

[cpuhrsch]

Looks great! Does the listing of tasks represent the preferred order of execution / priority?

[IvanYashchuk]

Yes, it's ordered by priority.

[fmassa]

Hi,

This is great!

One thing I would consider is if we should use torch.sparse.masked_matmul as a name, or instead name it torch.sparse.sddmm (as it is known in the literature). I'm not sure how common / clear masked_matmul is, and maybe sampled_matmul or sampled_dense_dense_matmul would be more appropriate as names

[pearu]

I would wait for a little with naming functions with masked_ prefix because according to the Masked reductions proposal (TBA, @cpuhrsch), masked_op takes a pair(s) (tensor, mask) as input(s) and returns similar pair(s) as outputs, and it would be good to settle down this convention before implementing other possible masked-tensor-input-output conventions.

[IvanYashchuk]

All three options (sddmm, sampled_matmul, sampled_dense_dense_matmul) sound good to me. Maybe the latter two are better as there is a tendency in PyTorch to use more descriptive names than accronyms.

[cpuhrsch]

I am in favor of adding functions named based on consensus across the vendor libraries we're doing to bind or as specified by already existing popular libraries. Consistency across frameworks I think makes it easier for users to work on them and find more information in other contexts than a more verbose or "clear" name within our framework.

[IvanYashchuk]

The last part that needs to be done from the original list of tasks is batched CSR support.