Matrix-free operators

In addition to supporting computation with the workhorse of sparse linear algebra, an assembled sparse matrix, Firedrake also supports computing “matrix-free”. In this case, the matrix returned from assemble() implements matrix-vector multiplication by the assembly of a 1-form subject to boundary conditions rather than direct construction of a sparse matrix (“aij” format) followed by traditional CSR algorithms. This functionality is documented in more detail in [KM18].

There are two ways of accessing this functionality. One can either request a matrix-free operator by passing mat_type="matfree" to assemble(). In this case, the returned object is an ImplicitMatrix. This object can be used in the normal way with a LinearSolver. Alternately, when solving a variational problem, an ImplicitMatrix is requested through the solver_parameters dict, by setting the option mat_type to matfree. The type of the preconditioning matrix can be controlled separately by setting pmat_type.

Generically, one can expect such a matrix to be cheaper to “assemble” and to use less memory, especially for high-order discretizations or complicated systems. The downside is that traditional algebraic preconditioners will not work with such unassembled matrices. To take advantage of these features, we need to configure our solvers correctly. To expedite this, the matrix-free operator, implemented as a PETSc shell matrix, contains an application context of type ImplicitMatrixContext. This context provides some important features to enabled advanced solver configuration.

Splitting unassembled matrices

For the purposes of fieldsplit preconditioners, the PETSc matrix object must be able to extract submatrices. For unassembled matrices, this is achieved through a specialized ImplicitMatrixContext.getSubMatrix() method that partitions the UFL form defining the operator into pieces corresponding to the integer labels of the unknown fields. This is in contrast to the normal splitting of assembled matrices which operates at a purely algebraic level. With unassembled operators, the PDE description is available in the matrix, and is therefore propagated down to the split operators.

Preconditioning unassembled matrices

As well as providing symbolic field splitting, the ImplicitMatrixContext object is available to preconditioners. Since it contains a complete UFL description of the bilinear form, preconditioners can query or manipulate it as desired. As a particularly simple example, the class AssembledPC simply passes the UFL into assemble() to produce an explicit matrix during set up. It also sets up a new PETSc PC context acting on this assembled matrix so that the user can configure it at run-time via the options database. This allows the use of matrix-free actions in the Krylov solve, preconditioned using an assembled matrix.

Providing application context to preconditioners

Frequently, such custom preconditioners require some additional information that will not be fully available from the UFL description in ImplicitMatrixContext. For example, it is not possible to extract physical parameters such as the Reynolds number from a UFL bilinear form. In this case, the solver accepts a dictionary "appctx" as an optional keyword argument, the same argument may also be passed to assemble() in the case of preassembled solves. Firedrake passes that down into the ImplicitMatrixContext so that it is accessible to preconditioners.

Example usage

To demonstrate basic usage of matrix-free operators and preconditioners, we show a simple Poisson problem, introducing some of the additional solver options.

Poisson equation

It is what it is, a conforming discretization on a regular mesh using piecewise quadratic elements.

As usual we start by importing firedrake and setting up the problem.:

from firedrake import *

N = 128

mesh = UnitSquareMesh(N, N)

V = FunctionSpace(mesh, "CG", 2)

u = TrialFunction(V)
v = TestFunction(V)

a = inner(grad(u), grad(v)) * dx

x = SpatialCoordinate(mesh)
F = Function(V)
L = F*v*dx

bcs = [DirichletBC(V, Constant(2.0), (1,))]

uu = Function(V)

With the setup out of the way, we now demonstrate various ways of configuring the solver. First, a direct solve with an assembled operator.:

solve(a == L, uu, bcs=bcs, solver_parameters={"ksp_type": "preonly",
                                              "pc_type": "lu"})

Next, we use unpreconditioned conjugate gradients using matrix-free actions. This is not very efficient due to the \(h^{-2}\) conditioning of the Laplacian, but demonstrates how to request an unassembled operator using the "mat_type" solver parameter.:

solve(a == L, uu, bcs=bcs, solver_parameters={"mat_type": "matfree",
                                              "ksp_type": "cg",
                                              "pc_type": "none",
                                              "ksp_monitor": None})

Finally, we demonstrate the use of a AssembledPC preconditioner. This uses matrix-free actions but preconditions the Krylov iterations with an incomplete LU factorisation of the assembled operator.:

solve(a == L, uu, bcs=bcs, solver_parameters={"mat_type": "matfree",
                                              "ksp_type": "cg",
                                              "ksp_monitor": None,

To use the assembled matrix for the preconditioner we select a "python" type:

"pc_type": "python",

and set its type, by providing the name of the class constructor to PETSc.:

"pc_python_type": "firedrake.AssembledPC",

Finally, we set the preconditioner type for the assembled operator:

"assembled_pc_type": "ilu"})

This demo is available as a runnable python file here.