Solver Benchmarks

A fixed linear-elasticity cantilever (HEX8, float64) refined from 48k to 1M degrees of freedom, solved three ways and timed on several NVIDIA GPUs:

cuDSS direct — sparse Cholesky/LDLᵀ factorization on GPU.
AMG-CG — algebraic-multigrid-preconditioned conjugate gradient (rigid-body near-null-space).
Krylov CG — fully matrix-free conjugate gradient (Jacobi-preconditioned).

Each solver is measured in eager, JIT, grad (forward + adjoint solve), and batched (vmap) execution. Use the controls to compare by solver on one device, or by device for a fixed solver, across solve time, throughput, and one-time compile cost.

Loading benchmark…

How to read it

Time per solve is batch-amortized (a vmap batch of 10 is divided by 10); throughput counts the whole batch; compile is the one-time trace + XLA compile, measured ahead of time with jit(f).lower(...).compile() so it contains no execution (eager has none). Rows tagged legacy-dashboard predate that change and include one execution in compile.
The grad mode times jax.jit(jax.value_and_grad(objective)) with objective ½‖u‖² and the nodal Young's modulus as design variable: one forward solve plus the adjoint solve through the solver's custom_vjp. Read grad / jit at the same DOF as the adjoint overhead. The direct solver reuses its forward factorization for the adjoint (reuse_factorization=True, SPD matrix); AMG reuses its hierarchy; Krylov runs a second CG.
Colour encodes the compared dimension (solver or device); line dash encodes the execution mode. A table view (▤) and a light/dark toggle sit in the top-right.
The vmap·rhs mode batches the load — the right-hand side b varies with the operator fixed (factor-once / solve-many) — while vmap·lhs batches the material, so the left-hand-side matrix A (and its factorization) changes per case. See Periodic Boundary Conditions and the vmap transform guide for where this matters.
All five modes apply to all three solvers. AMG-CG is a matrix-free CG solve whose preconditioner, fe.amg(...), is a fixed V-cycle built once from a sample Jacobian; it is a pytree passed to the solve as a jit argument (preconditioner=P), so the solve itself jits, vmaps, and differentiates like the others.

Benchmark setup

Physical problem

A three-dimensional linear-elastic cantilever: one physical box, clamped at one face and pulled by a uniform traction on the opposite face. Only the mesh resolution changes across the sweep, so every DOF point is the same problem refined — the per-DOF numbers are directly comparable.

Static equilibrium of a linear-elastic solid on a domain $\Omega$ :

\nabla\cdot\boldsymbol{\sigma} = \mathbf{0},\qquad \boldsymbol{\sigma} = \lambda\,\mathrm{tr}(\boldsymbol{\varepsilon})\,\mathbf{I} + 2\mu\,\boldsymbol{\varepsilon},\qquad \boldsymbol{\varepsilon} = \tfrac{1}{2}\left(\nabla\mathbf{u} + \nabla\mathbf{u}^{\mathsf{T}}\right),

with Lamé parameters $\mu = E/[2(1+\nu)]$ and $\lambda = E\nu/[(1+\nu)(1-2\nu)]$ . The boundary conditions are a clamp $\mathbf{u}=\mathbf{0}$ at $x=0$ and a uniform tip traction $\mathbf{t}=(0,0,t_z)$ at $x=L_x$ , traction-free elsewhere.

Quantity	Value
Domain	$8 \times 2 \times 2$ box (fixed, scale-invariant)
Element	HEX8, `vec=3` (3-D displacement)
Young's modulus $E$	$70{,}000$
Poisson's ratio $\nu$	$0.3$
Clamp	$\mathbf{u} = \mathbf{0}$ on $x = 0$
Tip traction $\mathbf{t}$	$(0,\,0,\,10^{-3})$ on $x = L_x$
Precision	`float64`
DOF sweep	$\approx 48\text{k} \to 1\text{M}$ , with $3(n_x{+}1)(n_y{+}1)(n_z{+}1)$

Solvers

Label	Method	Key settings
cuDSS direct	GPU sparse Cholesky / LDLᵀ factorization	`reuse_factorization=True`
AMG-CG	matrix-free CG preconditioned by `fe.amg(...)`	`AMGOptions(near_nullspace="rigid_body")`, `tol=1e-8`, `maxiter=500`
Krylov CG	matrix-free CG preconditioned by `fe.jacobi()`	`tol=1e-8`, `maxiter=2000`

Reading the comparison fairly

All three solve the same assembled system $A\mathbf{u}=\mathbf{b}$ , but the timed region does not measure identical numerical work — for two reasons worth keeping in mind when reading the plots.

Setup is hoisted out of the timing — symmetrically. Whatever each solver can precompute before seeing the matrix values is built once at setup (create_solver, or fe.amg for the hierarchy — outside the timed loop), and the rest is timed per solve. The catch is that the two methods have a different "precomputable" boundary:

	Built once at setup (untimed)	Per solve (timed)
AMG-CG	the whole multigrid hierarchy — aggregation, Galerkin coarse operators $A_c = RAP$ , rigid-body near-null-space (its main cost)	matrix-free $A\mathbf{v}$ + apply the fixed V-cycle `M`
cuDSS direct	symbolic analysis / fill-reducing ordering only	numeric factorization (its main cost) + triangular solves

This is not a thumb on the scale for AMG: each solver hoists everything it can. AMG's dominant cost (the hierarchy) is value-independent enough to build from one sample matrix and reuse as values change — so it belongs at setup. A direct factorization is inseparable from the specific matrix values, so it cannot be hoisted and recurs every solve (reuse_factorization shares factors between the forward and adjoint solve and powers vmap·rhs's factor-once / solve-many — it does not cache across independent solves). The comparison therefore models "solve many systems of the same sparsity pattern but changing values" (Newton steps, parameter sweeps, topology optimization), where that asymmetry is exactly the real-world trade-off. For "one matrix, many right-hand sides," the direct factorization amortizes too — that is what the vmap·rhs column shows, and there the direct solver wins.

Accuracy is not equalized. The direct solve is exact to round-off ( $\sim$ machine precision); AMG-CG and Krylov CG stop at a relative residual (tol=1e-8). The iterative throughput therefore reflects a looser-but-sufficient answer; tightening the tolerance would trade some of it back.

Code

The problem definition and the three interchangeable solver backends, condensed from examples/benchmark/bench_linear_elasticity.py:

import feax as fe
import jax.numpy as jnp

nu = 0.3

class LinearElasticity(fe.Problem):
    def get_tensor_map(self):
        def stress(u_grad, E):                       # Cauchy stress σ(∇u, E)
            mu  = E / (2. * (1. + nu))
            lam = E * nu / ((1. + nu) * (1. - 2. * nu))
            eps = 0.5 * (u_grad + u_grad.T)
            return lam * jnp.trace(eps) * jnp.eye(self.dim) + 2. * mu * eps
        return stress

    def get_surface_maps(self):
        def traction(u, x, mag):                     # uniform tip load, +z
            return jnp.array([0., 0., mag])
        return [traction]

# fixed 8 x 2 x 2 box, refined to hit each DOF target
lx, ly, lz = 8.0, 2.0, 2.0
mesh = fe.StructuredGrid((nx, ny, nz),
                         spacing=(lx/nx, ly/ny, lz/nz)).to_mesh()

clamp = lambda p: jnp.isclose(p[0], 0.0, 1e-5)                 # x = 0
tip   = lambda p: jnp.isclose(p[0], lx, 1e-5 + 0.5 * lx / nx)  # x = Lx

problem = LinearElasticity(mesh, vec=3, dim=3, ele_type="HEX8",
                           location_fns=[tip])
bc = fe.DirichletBCConfig([
    fe.DirichletBCSpec(location=clamp, component="all", value=0.),
]).create_bc(problem)

E        = fe.TracedParams.create_node_var(problem, 70e3)      # Young's modulus
traction = fe.TracedParams.create_uniform_surface_var(problem, 1e-3)
tp = fe.TracedParams(volume_vars=(E,), surface_vars=[(traction,)])
ts = fe.TracedStructure.from_problem(problem)

# three interchangeable solver setups: options (+ preconditioner for Krylov)
direct = fe.DirectSolverOptions(reuse_factorization=True)             # cuDSS
krylov = fe.KrylovSolverOptions(solver="cg", tol=1e-8, atol=1e-11,    # Krylov CG
                                maxiter=2000)                         #   + fe.jacobi()
amg_cg = fe.KrylovSolverOptions(solver="cg", tol=1e-8, atol=1e-10,    # AMG-CG
                                maxiter=500)                          #   + fe.amg(...)
P = fe.amg(problem, bc, tp, fe.AMGOptions(near_nullspace="rigid_body"),
           traced_structure=ts)                                       # hierarchy, built once

solver = fe.create_solver(problem, bc, solver_options=direct, linear=True,
                          traced_params=tp, traced_structure=ts)
# Krylov CG: solver_options=krylov, preconditioner=fe.jacobi()
# AMG-CG:    solver_options=amg_cg, preconditioner=P  (and preconditioner=P per call
#            inside jit, so the hierarchy is a jit argument, not a constant)
u = solver(tp, fe.zero_like_initial_guess(problem, bc), traced_structure=ts)

eager runs this op-by-op, jit wraps it in jax.jit, and the vmap modes map it over a batch — the material for vmap·lhs (matrix A varies), or the load for vmap·rhs (A fixed, factor-once / solve-many).

The rows are embedded in the dashboard itself (docs/static/benchmarks/dashboard.html) and kept as tracked CSVs under examples/benchmark/results/. New runs append to the local examples/benchmark/bench.csv (git-ignored): produce it with examples/benchmark/bench_linear_elasticity.py or the run_guarded.sh launcher (see --help), move the rows you want to keep into results/, then run examples/benchmark/sync_docs_dashboard.py to re-embed everything.

Benchmark setup​

Physical problem​

Solvers​

Reading the comparison fairly​

Code​

Benchmark setup

Physical problem

Solvers

Reading the comparison fairly

Code