Much slower than `torch.linalg.solve`

[uranium11010]

I wrote the following code to benchmark against torch.linalg.solve for large sparse matrices. Here, A is a 5000x5000 coalesced COO matrix with 10 nonzero elements in every row. b is a length-5000 vector.

import torch
from torch_sparse_solve import solve

import time

torch.manual_seed(0)

N = 5000
k = 10
row_inds = torch.arange(N).reshape(-1, 1).expand(-1, k)
col_inds = torch.multinomial(torch.ones(N, N, dtype=torch.float64), num_samples=k).sort(axis=1).values
indices = torch.cat([torch.zeros(1, N * k, dtype=torch.int64), row_inds.reshape(1, -1), col_inds.reshape(1, -1)], dim=0)
values = torch.randn(N * k, dtype=torch.float64)
A = torch.sparse_coo_tensor(indices, values, size=(1, N, N))
b = torch.randn(1, N, 1, dtype=torch.float64)

start_time = time.time()
x = solve(A, b)
print(f"sparse solve: {time.time() - start_time: .7f} seconds")

# compare to torch.solve:
A = A.to_dense()
start_time = time.time()
x_true = torch.linalg.solve(A, b)
print(f"dense solve: {time.time() - start_time: .7f} seconds")
print("error: {}".format(torch.max(torch.abs(x - x_true))))

The output:

sparse solve:  45.2241280 seconds
dense solve:  0.9555798 seconds
error: 1.7035422615663265e-05

I was hoping to use torch_sparse_solve to solve some large systems quickly and in a memory-efficient manner, so it would be great if this performance issue can be looked into. Thanks!

[flaport]

Hi @uranium11010 ,

Sparse solvers are tricky. There are lots of sparse solvers out there (see for example SuiteSparse). Each of them have different use cases and are ideal for different purposes.

In this case, torch_sparse_solve wraps the SuiteSparse KLU algorithm, which is best for sparse matrices arising from circuit architectures. In practice this means that KLU is ideal for non-singular sparse adjacency matrices with around 1 or 2 elements per row/column. If I adjust your example to work for matrices like that you can check that it's indeed faster to use a sparse solve:

row_inds = np.arange(N, dtype=np.int64)
np.random.shuffle(row_inds)
col_inds = np.arange(N, dtype=np.int64)
np.random.shuffle(row_inds)
values = np.asarray(np.random.randn(N), dtype=np.float64)
A = torch.sparse_coo_tensor(
    indices=np.stack([np.zeros_like(row_inds), row_inds, col_inds], 0),
    values=values,
).coalesce()
b = torch.randn(1, N, 1, dtype=torch.float64)

in any case choosing the right algorithm for a sparse solve is important. I've been thinking of adding additional backends (LU, UMFPACK) to this package but haven't gotten to it yet. Maybe in the future.

[uranium11010]

Ah I see--that makes sense to me. Thanks a lot and looking forward to seeing the additional backends :)