Sparse matrix multiplication is too slow

[stefanonardo]

I'm comparing it to SciPy and it is clearly too slow. It can be 100 times slower on CPU, which makes it quite unusable. My GPU improves the times but it is still very slow. I show you my benchmarks and the code I used.

CPU: i5 4690k @ 4.4GHz
GPU: GTX 1060

SIZE: 5000 DENSITY: 0.01 DEVICE: cpu
torch: 0.0306358 seconds
np:    0.000252247 seconds
torch/np: 121.452
----------------------------------------
SIZE: 5000 DENSITY: 0.01 DEVICE: cuda
torch: 0.0127137 seconds
np:    0.000259161 seconds
torch/np: 49.057
----------------------------------------

SIZE: 10000 DENSITY: 0.01 DEVICE: cpu
torch: 0.155527 seconds
np:    0.00106144 seconds
torch/np: 146.524
----------------------------------------
SIZE: 10000 DENSITY: 0.01 DEVICE: cuda
torch: 0.0476248 seconds
np:    0.000991583 seconds
torch/np: 48.0291
----------------------------------------

SIZE: 50000 DENSITY: 0.01 DEVICE: cpu
torch: 5.94856 seconds
np:    0.0456181 seconds
torch/np: 130.399
----------------------------------------
SIZE: 50000 DENSITY: 0.01 DEVICE: cuda
torch: 1.06403 seconds
np:    0.0419693 seconds
torch/np: 25.3527

SIZE: 50000 DENSITY: 0.0001 DEVICE: cpu
torch: 0.0423768 seconds
np:    0.000562191 seconds
torch/np: 75.3779
----------------------------------------
SIZE: 50000 DENSITY: 0.0001 DEVICE: cuda
torch: 0.0175352 seconds
np:    0.000589371 seconds
torch/np: 29.7524
----------------------------------------

import torch
import numpy as np
from scipy import sparse
import time

size = 50000
density = 0.0001
device = 'cuda'
print('SIZE:', size, 'DENSITY:', density, 'DEVICE:', device)

A = sparse.rand(size, size, format='coo', density=density).astype(np.float32)
b = torch.rand(size, 1, device=device)

values = A.data
indices = np.vstack((A.row, A.col))

i = torch.LongTensor(indices)
v = torch.FloatTensor(values)
shape = A.shape

A_torch = torch.sparse.FloatTensor(i, v, torch.Size(shape)).to(device)

s = time.time()
A_torch.mm(b)
t_torch = time.time() - s
print('torch: {:g} seconds'.format(t_torch))

b = b.cpu().numpy()

s = time.time()
A.dot(b)
t_np = time.time() - s
print('np:    {:g} seconds'.format(t_np))

print('torch/np: {:g}'.format(t_torch / t_np))
print('-'*40)

cc @nikitaved @pearu @cpuhrsch @IvanYashchuk

[musikisomorphie]

Hi @zou3519, I am quite interested in working on this issue.
Currently, I am a PhD candidate focusing on computer vision topics. So during the research, I use pytorch, which means I have some experiences studying the source code of pytorch. So if you could help me pinning down to the function (path, name etc., my current conclusion is the functions in SparseTensorMath.cpp) that cause this problem, I would like to give it a try. Thank you.

[fmassa]

Precisely, the functions in https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/sparse/SparseTensorMath.cpp correspond to the VPU implementation.

Matrix multiplication is implemented in the addmm function

[musikisomorphie]

Thanks @fmassa, once I make some progress, I will let you know.

[wranai]

Workaround (or solution): t.coalesce() after creating the sparse tensor fixes this problem. It seems sparse tensors, unless created by to_sparse(), are non-coalesced by default and I assume a bunch of optimizations are not applied during mm as a result.

---

I noticed something real funny. If I use t.to_sparse() then torch.mm is extremely fast. However, if I create a sparse tensor in any other way (EDIT with the exception of t.clone()), it's extremely slow. Like, 100x slower.

import time
import torch

DEVICE = torch.device('cuda')

n = 8192
c = .01

t = torch.randn((n, n), device = DEVICE)
t[torch.rand_like(t) > c] = 0
t = t.to_sparse()
torch.cuda.empty_cache()

t1 = torch.sparse_coo_tensor(t._indices(), t._values(), size = (n,n))
t2 = torch.sparse_coo_tensor(t._indices().clone(), t._values().clone(), size = (n,n))

a =torch.randn((n,1), device = DEVICE)

def bb(t,n):
    t0 = time.monotonic()
    b = torch.mm(t, a).mul_(c)
    for i in range(n-2):
        b = torch.mm(t, b)
    b = torch.mm(t, b).mul_(c)
    torch.cuda.synchronize()
    print((time.monotonic()-t0)/n)

bb(t, 1000)
# 0.00034971597511321306

bb(t1, 1000)
# 0.03134377552382648

bb(t2, 1000)
# 0.030944300027564167

The 3 tensors look identical:

print(t)
# tensor(indices=tensor([[   0,    0,    0,  ..., 8191, 8191, 8191],
#                        [ 157,  190,  285,  ..., 7629, 7852, 8121]]),
#        values=tensor([ 0.5541,  0.2092, -0.9584,  ..., -0.2249, -1.5567,
#                       -0.1699]),
#        device='cuda:0', size=(8192, 8192), nnz=671027, layout=torch.sparse_coo)

print(t1)
# tensor(indices=tensor([[   0,    0,    0,  ..., 8191, 8191, 8191],
#                        [ 157,  190,  285,  ..., 7629, 7852, 8121]]),
#        values=tensor([ 0.5541,  0.2092, -0.9584,  ..., -0.2249, -1.5567,
#                       -0.1699]),
#        device='cuda:0', size=(8192, 8192), nnz=671027, layout=torch.sparse_coo)

print(t2)
# tensor(indices=tensor([[   0,    0,    0,  ..., 8191, 8191, 8191],
#                        [ 157,  190,  285,  ..., 7629, 7852, 8121]]),
#        values=tensor([ 0.5541,  0.2092, -0.9584,  ..., -0.2249, -1.5567,
#                       -0.1699]),
#        device='cuda:0', size=(8192, 8192), nnz=671027, layout=torch.sparse_coo)

[musikisomorphie]

Hi @wranai and @fmassa, after examining the running time of various function blocks in SparseTensorMath.cpp on my Macbook CPU, indeed, the dominant time complexity comes from coalesce(). It involves two expansive operations-- indices sort and a for loop (See coalesce_sparse_cpu(...) in SparseTensor.cpp). @wranai, I guess if you also take the running time of t.coalesce() into account, the whole time complexity can be again expansive.

I suppose the reason to apply coalesce() is for more efficient cuda computation. Ironically, it significantly slows down the related CPU computation. Meanwhile, I studied the related scipy source code, it does not apply coalesce().

So one solution could be that I follow scipy implementation and replace coalesce() with it for the CPU computation. The main difference is coalesce() remove the duplicates while scipy code does not. According to my experiments, though the computation remains slower than scipy on CPU, it is relatively acceptable. What do you guys think? Or should the cpu code be consistent with gpu one and keep the coalesce() implementation?

Btw, I used the std::chrono::high_resolution_clock::time_point to monitor the c++ code block and then output the running time by AT_WARN(...). I was wondering If there exists a more natural way to output the running time of c++ code block, if so please let me know.

[yf225]

@musikisomorphie Thanks for the investigation! I think it's ok to have different implementations for CPU vs. CUDA

[wranai]

@wranai, I guess if you also take the running time of t.coalesce() into account, the whole time complexity can be again expansive.

You may be right but I did't measure that so I don't know for sure. My application (an echo state network) required it exactly once so mm did dominate.

[musikisomorphie]

Hi guys @yf225 @wranai @fmassa, this week I have to deal with my paper rebuttal and may not have enough free time . Once the rebuttal is done, I am going to keep working on our problem. You are going to be informed by that time.

[musikisomorphie]

Hi guys @yf225 @wranai @fmassa, this week I finally had some time to work on this problem.
The original strategy of the code is first convert coo to csr format of the sparse matrix then do the matrix multiplication by THBlas_axpy. coo to csr is a widely-used optimization step which supposes to speed up the computation. Unfortunately, for large framework such as Pytorch this step can be surprisingly expansive. So, my solution is to directly use the coo format to do matrix multiplication.
Here is the time difference on my linux machine:

(original code)
SIZE: 15000 DENSITY: 0.01 DEVICE: cpu
torch: 0.39403 seconds
np:    0.00496674 seconds
torch/np: 79.3338

----------------------------------------

(my unpdate)
SIZE: 15000 DENSITY: 0.01 DEVICE: cpu
torch: 0.0812583 seconds
np:    0.00501871 seconds
torch/np: 16.1911

Current, my code is not well organized. So it requires more time for me to polish it and open a PR.
If needed, I can send you my code per Email.

[zou3519]

I believe this is fixed by #16905 (thanks @musikisomorphie !)

[stefanonardo]

Honestly, @musikisomorphie's effort is really appreciated and it is for sure a great improvement, but I do not think this really "fix" the problem. PyTorch is still 16 times slower. This means that if my experiments run for 1 day on NumPy, the same experiments will need 16 days on PyTorch. I hope someone will find a solution (unfortunately I do not have the time to do it by myself).

[zou3519]

Thanks for the report, @stefanonardo , my apologies for not looking at the numbers. I agree that it can definitely be even faster than it is now given that numpy has a faster implementation