Repository: pytorch
Issue: Slow backward for matrix multiplication of two sparse COO tensors on CPU
Labels: module: sparse,triaged

### 🐛 Describe the bug

It was originally reported in https://github.com/pytorch/pytorch/issues/2389#issuecomment-1009356963 that the performance of backward computation is substantially slower than forward.

Copying the demonstration snippet here:
```python
import torch
import numpy as np
from time import time

def test_sparse_speed(a1, a2, b2, density_A, density_B):
    b1 = a2
    nnz_A = int(density_A * a1 * a2)
    nnz_B = int(density_B * b1 * b2)
    
    print(f'Matrix A shape: [{a1}, {a2}], nonzero entries: {nnz_A}')
    print(f'Matrix B shape: [{b1}, {b2}], nonzero entries: {nnz_B}\n')
    
    # create first matrix
    indexA = np.array([np.random.randint(0, a1, (nnz_A, )), np.random.randint(0, a2, (nnz_A, ))])
    valueA = torch.ones(nnz_A, requires_grad=True)
    A = torch.sparse_coo_tensor(
        indexA,
        valueA,
        size=[a1, a2],
    )

    # create second matrix
    indexB = np.array([np.random.randint(0, b1, (nnz_B, )), np.random.randint(0, b2, (nnz_B, ))])
    valueB = torch.ones(nnz_B, requires_grad=True)
    B = torch.sparse_coo_tensor(
        indexB,
        valueB,
        size=[b1, b2],
    )
    
    start = time()
    y = torch.sparse.mm(A, B)
    print(f'Multiplication time: {time()-start} seconds')
    
    loss = y.coalesce().values().sum()
    
    start = time()
    loss.backward()
    print(f'Backward time: {time()-start} seconds')
    return A, B

A, B = test_sparse_speed(a1=100, a2=10**5, b2=10**4, density_A=0.005, density_B=0.0002)
```
```
Matrix A shape: [100, 100000], nonzero entries: 50000
Matrix B shape: [100000, 10000], nonzero entries: 200000

Multiplication time: 0.702399730682373 seconds

Backward time: 218.6345179080963 seconds
```

Profiling [`SparseSparseMatmulBackward0`](https://github.com/pytorch/pytorch/blob/40b80aa49011849d916f71d9ba3ad6026cbfb662/torch/csrc/autograd/FunctionsManual.cpp#L861-L891) for the given shapes and density gives:
```python
from torch.profiler import profile, record_function, ProfilerActivity
with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
    with record_function("SparseSparseMatmulBackward0"):
        A, B = test_sparse_speed(a1=100, a2=10**5, b2=10**4, density_A=0.005, density_B=0.0002)
```
```
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                            SparseSparseMatmulBackward0         1.73%        3.801s       199.26%      437.498s      218.749s             2
autograd::engine::evaluate_function: SparseSparseMat...         0.00%      47.000us        99.26%      217.941s      217.941s             1
                                         aten::coalesce         0.00%     206.000us        83.35%      183.000s       14.077s            13
                                        aten::_coalesce         2.58%        5.666s        83.35%      183.000s       16.636s            11
                                             aten::sort        80.61%      176.975s        80.66%      177.092s       16.099s            11
                              aten::_sparse_mask_helper        14.08%       30.903s        14.19%       31.145s        7.786s             4
                            aten::_sparse_sparse_matmul         0.72%        1.583s         1.32%        2.900s     966.551ms             3
autograd::engine::evaluate_function: SparseCooTensor...         0.00%      94.000us         0.31%     691.524ms     345.762ms             2
             SparseCooTensorWithDimsAndTensorsBackward0         0.00%      72.000us         0.31%     691.430ms     345.715ms             2
                                              aten::mul         0.15%     318.643ms         0.15%     318.643ms      16.771ms            19
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 219.557s
```
`aten::sort` (likely called inside `.coalesce()`) and `aten::_sparse_mask_helper` are very slow taking 80% and 14% of the total time.



### Versions

```
Collecting environment information...
PyTorch version: 1.11.0a0+git5989b9d
Is debug build: True
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (crosstool-NG 1.24.0.131_87df0e6_dirty) 7.5.0
Clang version: 13.0.0 (https://github.com/conda-forge/clangdev-feedstock f94b4cb272bf2163513b577cc343fc2867bbf229)
CMake version: version 3.19.2
Libc version: glibc-2.31

Python version: 3.9.1 | packaged by conda-forge | (default, Dec 21 2020, 22:08:58)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-89-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.3.109
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 2060
GPU 1: NVIDIA GeForce RTX 2060

Nvidia driver version: 465.19.01
cuDNN version: Probably one of the following:
/usr/local/cuda-10.1.243/targets/x86_64-linux/lib/libcudnn.so.8
/usr/local/cuda-10.1.243/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8
/usr/local/cuda-10.1.243/targets/x86_64-linux/lib/libcudnn_adv_train.so.8
/usr/local/cuda-10.1.243/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8
/usr/local/cuda-10.1.243/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8
/usr/local/cuda-10.1.243/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8
/usr/local/cuda-10.1.243/targets/x86_64-linux/lib/libcudnn_ops_train.so.8
/usr/local/cuda-10.2.89/targets/x86_64-linux/lib/libcudnn.so.8
/usr/local/cuda-10.2.89/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8
/usr/local/cuda-10.2.89/targets/x86_64-linux/lib/libcudnn_adv_train.so.8
/usr/local/cuda-10.2.89/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8
/usr/local/cuda-10.2.89/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8
/usr/local/cuda-10.2.89/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8
/usr/local/cuda-10.2.89/targets/x86_64-linux/lib/libcudnn_ops_train.so.8
/usr/local/cuda-11.4.2/targets/x86_64-linux/lib/libcudnn.so.8
/usr/local/cuda-11.4.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8
/usr/local/cuda-11.4.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8
/usr/local/cuda-11.4.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8
/usr/local/cuda-11.4.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8
/usr/local/cuda-11.4.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8
/usr/local/cuda-11.4.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.812
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.4
[pip3] pytorch-sphinx-theme==0.0.24
[pip3] torch==1.9.0a0+giteb6445a
[conda] magma-cuda112             2.5.2                         1    pytorch
[conda] mkl                       2021.4.0           h8d4b97c_729    conda-forge
[conda] mkl-devel                 2021.4.0           ha770c72_730    conda-forge
[conda] mkl-include               2021.4.0           h8d4b97c_729    conda-forge
[conda] mypy                      0.812            py39h3811e60_2    conda-forge
[conda] mypy_extensions           0.4.3            py39hf3d152e_2    conda-forge
[conda] numpy                     1.21.4           py39hdbf815f_0    conda-forge
[conda] pytorch-sphinx-theme      0.0.24                    dev_0    <develop>
[conda] sparse-dot-mkl            0.7.2                    pypi_0    pypi
[conda] torch                     1.11.0a0+gite7b8b62           dev_0    <develop>
```

cc @nikitaved @pearu @cpuhrsch