[RFC] Enhancement on support for linear equation solvers

[houkensjtu]

Purpose

The purpose of this RFC is to discuss the plan on refactoring the linear solvers available under ti.linalg. Generally speaking, all solvers target at solving

$$ Ax = b $$

where $A$ is a $n \times n$ matrix and $n$ is the number of unknowns.

Current status

Currently, users have four options if they need to solve a set of large-scale linear equations.

1. Use the SparseMatrix API to build a sparse matrix, and then solve it using ti.linalg.SparseSolver. [example]
2. Use the SparseMatrix API to build a sparse matrix, and then solve it using ti.linalg.CG. [example]
3. Use the matrix-free CG solver ti.linalg.taichi_cg_solver to which you pass a LinearOperator to represent $A$. [example]
4. Implement their own linear solvers in Taichi language.

Note 1: We limit the scope of this topic to "large-scale" problems. Inverse of relative small matrices can be solved using the inverse() method provided in the taichi.math module.

Note 2: The development of the SparseMatrix solvers is well-documented in Issue #2906

A brief summary of the supported solvers:

matrix type	precision	backend	support
SparseMatrix	f32	CPU	✅
SparseMatrix	f64	CPU	✅
SparseMatrix	f32	CUDA	✅
SparseMatrix	f64	CUDA	❎
matrix-free	f32	CPU	✅
matrix-free	f64	CPU	✅

We notice that the naming of the provided solvers are not consistent and intuitive. For example,

• ti.linalg.taichi_cg_solver and ti.linalg.CG exist simultaneously, and it's unclear which one should be used for a certain circumstance
• ti.linalg.CG is intended (only) for SparseMatrix, but a user might expect it to be able to handle more data types

Proposed plan

I proposed the following plans to resolve the naming issue, as well as promote a better separation of concerns:

Plan 1. Keep the file organization untouched; use a more intuitive naming to suggest their usage. Namely,

ti.linalg.SparseSolver => no change 
ti.linalg.CG => ti.linalg.SparseCG
ti.linalg.taichi_cg_solver => ti.linalg.MatrixFreeCG

Plan 2. Construct an extensible module structure for each matrix type. Specifically, the solvers will be reorganized into subdirectories suggesting its target matrix type:

linalg/
│
├── sparse/
│   ├── __init__.py
│   ├── direct_solver.py
│   └── cg.py
│
└── matrix_free/
    ├── __init__.py
    └── cg.py

This way, user can use the solvers by the following imports:

from ti.linalg.sparse import DirectSolver, CG
from ti.linalg.matrix_free import CG as MatrixFreeCG

Plan 3. Unify the solver API, and let the solver to determine which version it should use. In this case, ti.linalg.CG and ti.linalg.taichi_cg_solver will be combined into one ti.linalg.CG. From a user's perspective, calling the CG solver will simply be

cg = ti.linalg.CG(A, b, x0, max_iter=50, atol=1e-6)
x = cg.solve()

no matter A is a LinearOperator or a SparseMatrix.

Additional comments

Which plan do you prefer? To be sure, each plan has its own pros and cons, and we welcome all users and developers' opinion on this issue!