[RFC] Provide built-in matrix-free iterative solvers

[houkensjtu]

Proposed feature

The purpose of this RFC is to discuss the possibility of adding a set of basic matrix-free iterative solvers, including conjugate-gradient (CG), BiCG, and BiCGSTAB to ti.linalg. The 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 two 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.
2. 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

Solution 1 is reasonably easy to use except that it does require the user to explicitly build-up the SparseMatrix themselves before solving (see documentation here).

As for solution 2, while writing an iterative solver on one's own provide great flexibility to users and does keep the language itself simple, we found the needs for a set of basic iterative solvers are quite common in entry-level users.

Overview of the plan

We can implement a set of basic matrix-free iterative solvers in native Taichi language. To represent the coefficient matrix A, we can implement a LinearOperator class for a user to transform their compute kernels to a "matrix-like" object.

# Implementation code not visible to users
@ti.data_oriented
class LinearOperator:
    def __init__(self, matvec):
        self.matvec = matvec

    def _matvec(self, x):
        # Default matvec behavior here
        
    def matvec(self, x):
        # Extra shape check here
        self._matvec(x, Ax)

The class will be implemented in a separate module iterative_solver.py under python/taichi/linalg so that it can be distinguished from the current SparseMatrixSolver family.

Use case

The user can pass a Taichi kernel to LinearOperator to create their "matrix-free" matrix. Here, the input of the kernel will be a ti.field of arbitrary shape, as long as the dimension of v and mv matches:

# User code
@ti.kernel
def compute_Ax(v:ti.template(), mv:ti.template()):  # v : vector, mv: matrix-vector product
    for i, j in v:
        l = v[i - 1, j] if i - 1 >= 0 else 0.0
        r = v[i + 1, j] if i + 1 <= GRID - 1 else 0.0
        t = v[i, j + 1] if j + 1 <= GRID - 1 else 0.0
        b = v[i, j - 1] if j - 1 >= 0 else 0.0
        mv[i, j] = 4 * v[i, j] - l - r - t - b

A = LinearOperator(compute_Ax)

This way, there will be no explicit stored matrix (thus the name matrix-free) and the A can be used for matrix-vector multiplication (which is required in most iterative methods). Notice that in this case, there will be no shape checking because the matvec() method is overwritten by the user.

Alternatively, user can create subclass of the LinearOperator class if they wish to implement their own:

class MyLinearOperator(LinearOperator):
    def __init__(self):
        ...
    def _matvec(self, x):
        ...

Here, when a user call matvec() method, extra shape checking will happen as it's inherited from the LinearOperator parent class.
Finally, we will provide a set of iterative solvers as Python functions, with similar API design to Scipy:

# Implementation code not visible to users
def cg(A, b, x, tol=1e-6, maxiter=5000):
    ...

def bicg(A, b, x, tol=1e-6, maxiter=5000):
    ...

def bicgstab(A, b, x, tol=1e-6, maxiter=5000):
    ...

and the user can call our pre-coded solvers simply by:

cg(A, b, x)

Additional comments

In Scipy, the LinearOperator class encourages the user to create their own subclass of it instead of creating an instance of it directly. Shall we follow the same path here in Taichi? Please let me know!

[houkensjtu]

The experimental matrix-free CG solver is now available under ti.linalg per PR#7690. This thread will be closed for now.

[liuyunpu]

I have a question. If the Taichi kernel function "compute_Ax" need more input to compute Ax, for example i need variables from a class, so my compute_Ax should look like this:

@ti.kernel
def compute_Ax(self, v:ti.template(), mv:ti.template()):
    #compute Ax

how can i pass a kernel like this to LinearOperator? (writting all things needed to compute Ax in one kernel seems impossible in most physical simulations)
Another question, will there be a bicgstab solver for sparse matrix in cuda backend? the current cg solver in cuda can only solve positive-definite matrix. And more iterative solver would be nice.