Can pytorch add sparse linear solvers like scipy.sparse.linalg.gmres, scipy.sparse.linalg.bicg etc.

[liuqi45]

🚀 The feature, motivation and pitch

I am a compute fluid dynamic software developer. Pytorch really helped me for using GPU to accelerate computing. However, in compute fluid dynamic or other mechanics domain, we usually store coefficients matrix in sparse format like CSR and solve the linear equation systems. I wondering if pytorch can develop a sparse solver which input the sparse coefficients matrix with a dense vector for the right hand of equations and output a dense vector represent the solve. In scipy there are already such solvers, but it does not support gpu acceleration. I hope pytorch can add such solvers.

Alternatives

No response

Additional context

No response

cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer @jcaip @jianyuh @mruberry @walterddr @xwang233 @lezcano

### Tasks

[vadimkantorov]

Related:

• https://github.com/facebookresearch/theseus
• [feature request] Provide some sparse eigen solver(s) for PyTorch (maybe via ARPACK as in scipy) + SPD sparse / laplace linear system solver - maybe NVidia AMGx library? #108977
• [Feature Pitch] Full-batch optimization toolkit #55279
• torch.sparse_coo_tensor argname quirks + [feature request] .numpy()/from_numpy method for sparse_coo_tensor/sparse_csr_tensor (or maybe name them as .scipy()/.from_scipy() or at least under some torch.utils.* namespace for zero-copy conversions to<>from scipy sparse arrays #109017 (for easy, zero-copy conversion between PyTorch's COO/CSR and SciPy's COO/CSR)

[alexsamardzic]

The main requirement for implementing a sparse iterative solver is to have SpMV routine available, so what you're asking for is in principle doable within PyTorch. I'll defer to other developers to eventually provide you with an answer would that be useful at all for PyTorch, but IMO you would be much better server with specialized software. So I'd suggest you take a look for example into NVIDIA AMGX, in particular you may want to check are its Python bindings still usable.

[lezcano]

For CSR sparsisty, #129856 is being merged soon. Would this work for your use case @liuqi45?

[liuqi45]

My request is not only for the sparse input matrix, but also the iterative solve algorithm like gmres, bicg etc. I think the algorithm behind torch.linalg.solve() is a direct solver like LU decomposition, this is not efficient for large linear system. We often use iterative solvers with preconditioner to solve the large linear equations. @lezcano

[zitongzhan]

My request is not only for the sparse input matrix, but also the iterative solve algorithm like gmres, bicg etc. I think the algorithm behind torch.linalg.solve() is a direct solver like LU decomposition, this is not efficient for large linear system. We often use iterative solvers with preconditioner to solve the large linear equations. @lezcano

Does the conjugate gradient method satisfy your needs?

[liuqi45]

Conjugate gradient method doesn't satisfy my needs. Because CG method is used for solving linear equations with symmetric coefficient matrix. But advection-diffusion function isn't symmetric. Biconjugate gradient, Biconjugate gradient stabilized, GMRES algorithms are useful. @zitongzhan

[lezcano]

Do cusolver / magma / blas have implementations of these algorithms? If they don't, you should probably first pitch this feature to those libraries first, as we don't implement linalg algorithms ourselves.

Also, @liuqi45, could you try with the newly added support in #129856 and provide some speed / memory numbers (does it OOM? does it work but it's too slow?) related to your problem at hand, just to understand a bit better your use cases. A self-contained minimal repro of the problems that you want to solve would be best

[liuqi45]

magma has implementations of GMRES,bicgstab algorithm. Have a look at this link https://icl.utk.edu/~luszczek/teaching/courses/2017/mhpc/numlinalg_sparse.pdf
Also I will try #129856 . @lezcano

[tvercaut]

If helpful, we have been putting together a few iterative solvers here:
https://github.com/cai4cai/torchsparsegradutils
Some one them are implemented in pytorch directly, others are just wrappers around cupy or jax.

[Litianyu141]

I also have some mirroring JAX repository implementation of CG, BiCG, GMRES algorithms::
https://github.com/Litianyu141/Pytorch_Sparse_Linalg-torch.cg-bicg-gmres-
Hope it is helpful.