Solvers in pure C++/CUDA: relaxation, CG, and full multigrid (FMG)

[balbasty]

Umbrella for implementing the linear-system solvers (relaxation, conjugate gradient, and full multigrid) for spatially-regularised fields in pure C++/CUDA, launching CPU/CUDA kernels directly from the algorithm — rather than orchestrating them from Python.

Why now (the nitorch → jitfields → fastfields arc)

• nitorch implemented the full multigrid + relaxation solvers in pure C++/CUDA (nitorch/_C/src/: fmg.cpp/fmg_common.cpp, multires.cpp/multires_common.cpp, relax.cpp/relax_grid.cpp + *_common.cpp, precond.cpp/precond_grid.cpp, resize.cpp/resize_common.cpp; exposed relax, relax_grid, precond, precond_grid, resize, prolongation, restriction, fmg).
• jitfields couldn't do this: cupy-JIT only lets you write kernels, not host code, so the algorithm had to move to Python, built from the jitfields ingredients — see nitorch/nitorch-solvers (solve_field_fmg / solve_grid_fmg, CG + Jacobi preconditioning or checkerboard relaxation, in field and grid=flow flavours).
• fastfields launches its own kernels (host launchers in cpu-impl / cuda-impl), so we can move the algorithm back into pure C++/CUDA and drive it from the hub, matching nitorch's design.

Ingredients already present in the stack

• Relaxation (Gauss-Seidel): flow_relax landed; field relaxation exists in fastfields-cpu-impl/reg_field.h (relax_*) and just needs wiring (tracked in Field regulariser: full jitfields parity (kernel / precond / forward / relax / accumulate) #33).
• Prolong / restrict: resample / restriction (spline prolongation + its adjoint).
• Diagonal / preconditioner: flow_diag/flow_precond, field_diag/field_precond.
• Regulariser matvec / forward: *_matvec, *_forward, plus the compact-symmetric sym_matvec/sym_solve for the data-term Hessian.

So the missing piece is the solver driver itself (V-cycle / FMG, CG/PCG), which can now be genuine C++/CUDA.

Proposed breakdown (field + flow/"grid" flavours)

• prolong/restrict for solvers — thin C++ helpers over resample/restriction with the grid bookkeeping (make_grid, per-level shapes) nitorch's _FMG does in Python
• per-level operator apply + relaxation — (M + R) x and a Gauss-Seidel/Jacobi sweep at each level (reuse *_forward / *_relax)
• conjugate gradient (CG/PCG) — Jacobi-preconditioned CG solver in C++ (the jitfields default inner solver)
• V-cycle + FMG driver — the multigrid recursion in C++, mirroring fmg_common.cpp / multires_common.cpp
• expose through the stack — solve_field / solve_flow (relax / cg / fmg) in cpu-lib → cuda-lib → lib → bind-py → numpy/torch/cupy → fastfields.any
• CPU test gate — validate each solver reduces the residual (M+R)x − g to tolerance vs a dense reference (mirrors the existing relax residual tests)

References

• nitorch C++: nitorch/_C/src/{fmg,multires,relax,relax_grid,precond,precond_grid,resize}*.cpp
• nitorch-solvers (Python, the jitfields-era port): https://github.com/nitorch/nitorch-solvers
• Related in-flight parity: flow regulariser Flow regulariser: full jitfields parity (+ Lamé self-adjointness fix) #25, field regulariser Field regulariser: full jitfields parity (kernel / precond / forward / relax / accumulate) #33 (relaxation wiring feeds directly into this).

Depends on the field_relax wiring from #33; otherwise the ingredients are in place.

Workstream: claude-jitfields-to-fastfields

[balbasty]

Triage for the 0.0.dev0 release: confirmed non-blocking, deferred as post-release scope.

This is a net-new feature umbrella (CG/PCG, V-cycle, FMG solver drivers in pure C++/CUDA) building on ingredients that already exist and ship (field_relax/flow_relax, resample/restriction, field_diag/flow_diag, sym_matvec/sym_solve). Nothing here is a defect in what's currently exposed through the stack — it's asking for capability the project doesn't claim to have yet. The checklist itself is still all-unchecked and explicitly depends on #33's field_relax wiring, which is a separate in-flight item. No part of a correct 0.0.dev0 wheel build/publish depends on this landing first. Leaving open for post-release prioritization.

---

Generated by Claude Code

[balbasty]

Done: conjugate gradient (CG/PCG), field flavour

Third checkbox on the list — Jacobi-preconditioned CG in C++ — is implemented for the field family across three PRs. Leaving this issue open: it is one piece of the umbrella, not the whole thing.

• fastfields-cpu-impl#71 — solve_field::dot / axpby_, the two BLAS-1 element loops CG needs
• fastfields-cpu-lib#86 — ff::cpu::field_cg, the driver + the test gate
• fastfields-lib#75 — ff::field_cg, device dispatch

What it does

field_cg solves (H + L) x = g — the same system field_relax attacks with Gauss-Seidel sweeps — at a single level, no multigrid machinery involved. As the issue anticipated, the ingredients were all already in the stack, so the new code is genuinely just the solver driver plus the two vector primitives that were missing:

CG needs	already existed
apply (H + L) x	field_forward (= sym_matvec + field_addmatvec_)
Jacobi preconditioner	field_diag + posdef::sym_solve (i.e. field_precond's body)
dot / axpby	new (solve_field.h in cpu-impl)

One deliberate deviation from a naive field_forward + field_precond composition: the regulariser diagonal depends only on the penalties, not on the iterate, so it is hoisted out of the loop and computed once. A CG iteration costs one field_forward + one sym_solve.

Validation

A manufactured solution rather than a residual check: build g = (H + L) x_true, then require CG to recover x_true. The right-hand side is built from sym_matvec + field_matvec separately, not through field_forward, so it does not share the solver's own operator composition. A solver that quietly dropped the Hessian or the regulariser would still shrink its own residual but would not land on x_true.

Representative numbers (double, target 1e-10 relative residual):

case	iters	rel. residual	rel. error
2D 8×8 C=2 membrane (128 unknowns)	18	5.2e-11	4.0e-11
2D 8×8 C=2 bending	46	4.6e-11	2.8e-11
3D 8×8×8 C=2 membrane	20	9.5e-11	1.2e-10
2D 8×8 C=1 stiff bending (h=0.05, bend=10)	74	7.9e-12	4.0e-12

452 checks across absolute/membrane/bending, Zero/DCT2/DST2/DFT bounds, 1D/2D/3D, C=1..3, float and double, stiff systems, warm starts and channel-major strides. The full cpu-lib gate is green with every existing suite at unchanged counts, under both clang++ and g++.

CG has no coloring/parity constraints the way relaxation does — it is not a colored sweep — so the folding-incompatible boundary conditions are unremarkable here; DFT converges in the same ballpark as DCT2.

Still open on this issue

• flow / "grid" CG — structurally the same driver over flow_forward / flow_diag; the natural next slice
• prolong/restrict helpers for solvers and per-level operator apply
• V-cycle + FMG driver — the actual multigrid recursion, by far the largest remaining piece
• CUDA — field_cg is CPU-only. The blocker is specific: the dot products need a device-side reduction the cuda backend does not expose yet. CUDA tensors currently get an explicit "not implemented on CUDA yet" rather than silently falling through to the host path
• Python wiring — bind-py → numpy/torch/cupy → fastfields.any. Deliberately not touched, matching this issue's "pure C++/CUDA ... rather than orchestrating them from Python" framing

Relaxation (field_relax / flow_relax) was already wired before this work, so the first two solver flavours (relax, CG) now exist in C++ for the field family; FMG is the remaining algorithmic piece.

---

Generated by Claude Code