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

## PyTorch 1.11

Currently, MKL is used in PyTorch only for `at::sparse_csr::_sparse_mm_mkl_` (torch.addmm).

For PyTorch 1.11 we're focusing on improving sparse CSR support and performance on CPU.
This work aims to utilize all available functionality of MKL Sparse to expand operator support for PyTorch CSR matrices.

Go to [the list of tasks](https://github.com/pytorch/pytorch/issues/60858#pytorch111-tasks).

## MKL Inspector-Executor (IE) API Overview
Similar to cuSPARSE, MKL started to deprecate low-level BLAS-like functions in favor of a more high-level "inspector-executor" interface.

IE API for Sparse BLAS has two stages:
* The first stage constructs the structure of the output matrix.
* The second stage constructs other arrays and performs the desired operation.
* Two stages can be combined by performing the entire computation in a single step
* Supported data types – float32, float64, complex64, complex128
* Supported matrix formats – CSR, CSC, COO, and BSR.


## MKL version requirement
Sparse BLAS API is deprecated since MKL 2018 Update 2. The replacement for it is the sparse BLAS IE API. PyTorch will raise a runtime error if an older MKL version is used while calling sparse operators on the CPU. Use `at::get_mkl_version()` to retrieve MKL's version information.


## Sparse Matrix descriptors
MKL Sparse IE API uses descriptor structs `sparse_matrix_t` to collect the information about the tensor (size, indices, values, etc.). See [the Matrix Manipulation Routines table](https://software.intel.com/content/www/us/en/develop/documentation/onemkl-developer-reference-c/top/blas-and-sparse-blas-routines/inspector-executor-sparse-blas-routines/matrix-manipulation-routines.html#matrix-manipulation-routines_TBL1).

Each descriptor struct has constructors, destructors and a few attribute getters/setters defined in MKL. The descriptors should be wrapped similarly to cuDNN's descriptors: `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.

In MKL dtype of indices cannot be changed in runtime.
Currently, PyTorch uses a 32-bit MKL version, linking with the 64-bit (ILP64) might break the internal code that relies on MKL (see https://github.com/pytorch/pytorch/issues/51720, https://github.com/pytorch/pytorch/issues/56959), but it was not tested.

- [ ] Include the recommendation in the documentation that on CPU 32-bit indices are preferred to avoid conversions for now.

## Batched computation support
MKL doesn't have specialized batched routines, we can implement a for loop on the batch dimension.

| 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,)

<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 MKL routine (ordered by priority):

- [ ] [mkl_sparse_?_mm](https://software.intel.com/content/www/us/en/develop/documentation/onemkl-developer-reference-c/top/blas-and-sparse-blas-routines/inspector-executor-sparse-blas-routines/inspector-executor-sparse-blas-execution-routines/mkl-sparse-mm.html)  (sparse matrix - dense matrix multiplication, dense result)

    - [x] Implement descriptor wrappers for CSR matrices.
    - [ ] https://github.com/pytorch/pytorch/issues/58770

    Use as a backend for:
    - [x] torch.mm
    - [x] torch.addmm
    - [ ] torch.baddmm (if batched CSR is enabled in PyTorch)
    - [x] torch._sparse_mm
    PR in progress: https://github.com/pytorch/pytorch/pull/65606

- [x] [mkl_sparse_spmm](https://software.intel.com/content/www/us/en/develop/documentation/onemkl-developer-reference-c/top/blas-and-sparse-blas-routines/inspector-executor-sparse-blas-routines/inspector-executor-sparse-blas-execution-routines/mkl-sparse-spmm.html)  (sparse matrix - sparse matrix multiplication, sparse result)

    Use as a backend for:
    - [x] torch.mm
    - [x] torch._sparse_mm
    PR in progress: https://github.com/pytorch/pytorch/pull/65606

- [x] [mkl_sparse_?_mv](https://software.intel.com/content/www/us/en/develop/documentation/onemkl-developer-reference-c/top/blas-and-sparse-blas-routines/inspector-executor-sparse-blas-routines/inspector-executor-sparse-blas-execution-routines/mkl-sparse-mv.html) (sparse matrix - dense vector multiplication, dense result) (@IvanYashchuk)

    `y := alpha*op(A)*x + beta*y`

    Use as a backend for:
    - [x] torch.mv
    - [x] torch.addmv
    PR is progress: https://github.com/pytorch/pytorch/pull/61536

- [x] [mkl_sparse_?_trsm](https://software.intel.com/content/www/us/en/develop/documentation/onemkl-developer-reference-c/top/blas-and-sparse-blas-routines/inspector-executor-sparse-blas-routines/inspector-executor-sparse-blas-execution-routines/mkl-sparse-trsm.html)(sparse matrix - dense matrix triangular solver, dense result)
    [mkl_sparse_?_trsv](https://software.intel.com/content/www/us/en/develop/documentation/onemkl-developer-reference-c/top/blas-and-sparse-blas-routines/inspector-executor-sparse-blas-routines/inspector-executor-sparse-blas-execution-routines/mkl-sparse-trsv.html) (sparse matrix - dense vector triangular solver, dense result) (@IvanYashchuk)

    `Y := alpha*inv(op(A))*X`

    Use as a backend for:
    - [x] torch.triangular_solve
PR in progress: https://github.com/pytorch/pytorch/pull/62180

- [x] [mkl_sparse_?_add](https://software.intel.com/content/www/us/en/develop/documentation/onemkl-developer-reference-c/top/blas-and-sparse-blas-routines/inspector-executor-sparse-blas-routines/inspector-executor-sparse-blas-execution-routines/mkl-sparse-add.html) (@IvanYashchuk) 

    `C := alpha*op(A) + B`

    Use as a backend for:
    - [x] torch.Tensor.add https://github.com/pytorch/pytorch/issues/59060
PR in progress: https://github.com/pytorch/pytorch/pull/64391

- [x] [mkl_sparse_?_spmmd](https://software.intel.com/content/www/us/en/develop/documentation/onemkl-developer-reference-c/top/blas-and-sparse-blas-routines/inspector-executor-sparse-blas-routines/inspector-executor-sparse-blas-execution-routines/mkl-sparse-spmmd.html)  (sparse matrix - sparse matrix multiplication, dense result)

    `C = alpha * opA(A) * opB(B) + beta*C`

    Use as a backend for:
    - [x] torch.addmm(strided, sparse, sparse)
    PR: https://github.com/pytorch/pytorch/pull/73076

Note that tensors with the CSR layout on CPU device use `SparseCsrCPU` dispatch key in `native_functions.yaml` file.
There is `AT_MKL_ENABLED()` for checking whether MKL is available and `MKL_ILP64` for checking whether linked MKL is ILP64 (64-bit API) variant.


cc @nikitaved @pearu @cpuhrsch @IvanYashchuk