Repository: pytorch
Issue: Unify matrix multiplications operations
Labels: module: sparse,triaged,open source,module: linear algebra

## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

Unify matrix multiplication operations

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

Currently, there exists a number of different ways to perform matrix multiplications depending on the layouts of the inputs and the desired output layout as summarised in the following table:
| Function signature | Operation layout signature |
| :--- | :--- |
| `torch.matmul(input, other, *, out=None)` | `M[sparse_coo/strided] @ M[strided] -> M[strided]` |
| `torch.mm(input, mat2, *, out=None)` | `M[sparse_coo/strided] @ M[strided] -> M[strided]` |
| `torch.spmm(input, mat2, *, out=None)` | `M[sparse_coo/strided] @ M[strided] -> M[strided]` |
| `torch.dsmm(input, mat2, *, out=None)` | `M[sparse_coo/strided] @ M[strided] -> M[strided]` |
| `torch.sparse.mm(mat1, mat2)` | `M[sparse_coo] @ M[strided] -> M[strided]` |
| `torch.smm(input, mat)` | `M[sparse_coo] @ M[strided] -> M[sparse_coo]` |
| `torch.hspmm(mat1, mat2, *, out=None)` | `M[sparse_coo] @ M[strided] -> M[1+1 hybrid sparse_coo]` |
| `torch.bmm(input, mat2, *, deterministic=False, out=None)` | `T[sparse_coo/strided] @ T[strided] -> T[strided]` |
| `torch.Tensor.__matmul__(self, other)` | same as `torch.matmul()` |
| `torch.chain_matmul(*matrices)`| chained `torch.matmul()` |

where `T` denotes a tensor and `M` denotes a 2-D tensor.

While the different behavior of these functions is needed, the users have difficulties in choosing the matmul function they need. 

Related discussions:
- https://github.com/pytorch/pytorch/issues/44635#issuecomment-692159757
- https://github.com/pytorch/pytorch/pull/45400#discussion_r500041980
- https://github.com/pytorch/pytorch/pull/45400#discussion_r500042358
- https://github.com/pytorch/pytorch/pull/45400#discussion_r500105693
- https://github.com/pytorch/pytorch/pull/45400#discussion_r500047617

## Pitch

<!-- A clear and concise description of what you want to happen. -->

The number of possible functions to perform matrix multiplications needs to be reduced. The matmul table from above should transform to:

| Function signature | Operation layout signature | Notes |
| :--- | :--- | :--- |
| `torch.matmul(input, other, *, out=None)` | `T[sparse_coo/strided] @ T[sparse_coo/strided] -> T[sparse_coo/strided]` | supports broadcasting |
| `torch.mm(input, other, *, out=None)` | `M[sparse_coo/strided] @ M[sparse_coo/strided] -> M[sparse_coo/strided]` | operates on 2-D tensors only |
| `torch.bmm(input, other, *, deterministic=False, out=None)` | `T[sparse_coo/strided] @ T[sparse_coo/strided] -> T[sparse_coo/strided]` | operates on 3-D tensors only |
| `torch.spmm(input, mat2, *, out=None)` | `M[sparse_coo] @ M[strided] -> M[strided]` | deprecate[1] |
| `torch.dsmm(input, mat2, *, out=None)` | `M[sparse_coo] @ M[strided] -> M[strided]` | deprecate[2] |
| `torch.sparse.mm(mat1, mat2)` | `M[sparse_coo] @ M[strided] -> M[strided]` | deprecate[3] |
| `torch.smm(input, mat)` | `M[sparse_coo] @ M[strided] -> M[sparse_coo]` | deprecate[4] |
| `torch.hspmm(mat1, mat2, *, out=None)` | `M[sparse_coo] @ M[strided] -> M[(1+1)-D hybrid sparse_coo]` | deprecate[5] |
| `torch.Tensor.__matmul__(self, other)` | same as `torch.matmul()` | |
| `torch.chain_matmul(*tensors)`| chained `torch.matmul()` | |

where
- `deprecate[1, 2]` is straightforward as `spmm` and `dsmm` are aliases to `torch.mm`.
- `deprecate[3]` is possible via supporting sparse tensor inputs to `torch.mm`
- `deprecate[4, 5]` is possible via supporting the specification of output layout in the `out` keyword argument of `torch.mm`
- `M[sparse_coo] @ M[sparse_coo]` is possible via https://github.com/pytorch/pytorch/pull/39526

Notes:
- when the `out` keyword argument is not specified, the output layout will be the same as the layout of the `input` argument.
- do we need `torch.chain_mm(*matrices)` and `torch.chain_bmm`?
- this issue is analogous for `addmm` operations (`torch.addmm`, `torch.sspaddmm`, `torch.sparse.addmm`)

cc @vincentqb @aocsa @nikitaved @pearu @mruberry @vishwakftw @jianyuh