Skip to content

Latest commit

 

History

History
55 lines (48 loc) · 7.13 KB

File metadata and controls

55 lines (48 loc) · 7.13 KB

MatrixFreeOperators.jl

Lazy, composable, matrix-free operator algebra for PDEs on structured grids — device-agnostic (CPU/GPU) and differentiable (forward and reverse, including w.r.t. operator parameters).

DESIGN.md is the authoritative design document — read it before any structural change. It holds the locked decisions, the full rationale behind the invariants below, and the open/deferred choices (staggered grids, distributed backends, AMR, multigrid).

Key files

  • foundation: src/Grids.jl, src/boundaries.jl, src/Fields.jl
  • operator leaves and combinators: src/operators/
  • solver boundary (prepare, PreparedOperator, flat mul!): src/linalg.jl
  • nonlinear → linear Jacobian: src/operators/linearize.jl
  • AD / MDLA / Reactant integrations: ext/

Invariants — violating these produces silently wrong results

  • Traits default to the weaker claim. islinear / isconstant / isselfadjoint / isdiagonal / shares_exchange default to false; leaves opt in, combinators propagate explicitly. A forgotten declaration must degrade to an error, never a wrong answer.
  • Adjoints are declared, never assumed. Every linear leaf declares its adjoint including boundary contributions — BCs break self-adjointness even for the Laplacian. Check with the dot-product identity ⟨Lx,y⟩ = ⟨x,Lᵀy⟩.
  • Linear/affine split. apply! / apply_bc! enforce homogeneous BCs only, so islinear(L) ⇒ L(0) = 0. Inhomogeneous boundary data goes out separately through boundary_rhs and is folded into the solve RHS.
  • Interior-only flat vectors. Flat (Krylov) vectors span interior DOFs only. Ghost cells are scratch filled by halo_update! / apply_bc! and are never solver unknowns.
  • Element type carries tensor rank. A vector field is a Field over Array{SVector{N,T}} — operators have no rank parameter. Only the rank-changers Gradient / Divergence touch components.
  • Nonlinear operators never masquerade as linear maps. They support apply! and AD but not adjoint / prepare; linearize(F, u₀) yields the Jacobian operator that feeds Krylov — a finite-difference JVP by default, an exact forward-mode AD JVP (plus a real reverse-mode transpose) under linearize(F, u₀, EnzymeJVP()).
  • AD rules route through declared adjoints, never around them. A custom rule exists only where an exact transpose is already written and tested; it substitutes that transpose for the tape, so it must be an exact substitution, not an approximation. Rules must never fire on a path carrying operator parameters — they would silently zero coefficient-field gradients.

Two structural rules follow from these: leaf bodies stay array-level (broadcast/slicing) so they are device-agnostic and AD-friendly with no per-backend code — KernelAbstractions @kernel is the per-operator escape hatch, not the default. And halo_update! is a deliberate no-op seam in v1, so distributed/AMR work later changes only the grid and that function, never operator code.

Gotchas

  • The core depends only on Adapt, KernelAbstractions, LinearAlgebra, and StaticArrays. AD, MDLA, SciML, and Reactant integrations belong in ext/ — don't add them to [deps]. DifferentiationInterface is the documented frontend and belongs only in test/, docs/, and examples/; rules can't be routed through it.
  • Enzyme custom rules have two non-obvious constraints, both discovered by hitting them. A rule argument may not be a type mixing GC pointers with inline floats — on Julia 1.12 that is a hard CallingConventionMismatchError (EnzymeAD/Enzyme.jl#2707), and Field/BlockField/BlockForest all qualify via their embedded grid; that is why the seams are _exchange_storage!/_bc_storage! over raw storage plus a BlockLayout tag. And a rule body must not allocate — a Dict lookup and a closure in an augmented-primal body segfaulted on Linux x86_64 while passing on macOS/aarch64.
  • Enzyme behaves differently across platforms. Issue #26's EnzymeNoTypeError never reproduced on macOS/aarch64 on any Julia or Enzyme version tried; the segfault above only ever appeared on Linux x86_64. A local green run is not evidence about CI — push and read the matrix.
  • A partition_grid slab keeps the global extent; only local_range says which part it owns, and cell_center evaluates at the global cell index. Don't "fix" a slab's extent to describe its own span — that reintroduces an ulp of coordinate drift and makes a slab-assembled RHS depend on the partition count.
  • The distributability guards run once on the global operator tree, before _slab_op localizes field coefficients per slab. Re-checking a localized tree rejects it: a localized ScalingOp reports the slab grid from operator_grid while its Laplacian sibling still reports the global one.
  • A stencil that reads two arrays must not unroll dimensions with ntuple(Val(N)) do d. With one array (Laplacian) the closure inlines fine; add a coefficient array and it stops inlining, the broadcast loses vectorization, and one 256² sweep goes from 33 µs to 260 µs with identical numerics. Unroll by recursion over Val(D) instead — see _diff_axes in src/operators/diffusion.jl. Same class of trap as a Core.Box: silent, correct, and only visible in a benchmark.
  • prepare is stateful and single-threaded: call it once per concurrent solve, not once globally.
  • GPU parity tests are skipped unless MFO_TEST_GPU=true and CUDA.jl is available (test/device.jl gates test/device_gpu.jl) — a green suite does not mean GPU paths ran.
  • Running one test file needs the test env, which devs the package at path="..". Pass --check-bounds=yes — Pkg.test and therefore CI always do, and without it a local run is not comparable:
    julia --project=test --check-bounds=yes -e '
    using MatrixFreeOperators, Test, LinearAlgebra, Random, StaticArrays
    import Adapt, Enzyme, KernelAbstractions, Krylov, Mooncake
    include("test/test_utils.jl"); include("test/laplacian.jl")'
    
  • --check-bounds=yes changes what allocates, so every allocation assertion must be measured under it. Bounds checking blocks the SROA that elides a Ref or a view escaping a broadcast body: the same code measures 0 B without the flag and tens to thousands of bytes with it. This is not a platform difference — it reproduces on macOS/aarch64. A zero-allocation claim from a run without the flag is worth nothing, and a CI-only allocation failure is this before it is anything exotic.
  • quarto render docs runs the full test suite — docs/pages/coverage.qmd calls generate_coverage(...; run_test=true). Rendering docs is not cheap.
  • The docs sidebar (docs/_quarto.yml) has a fixed shape: pages/api.qmd lives alone in its own part: "API", between part: "Docs" and part: "Resources". docs/index.qmd must open with ## Overview, then ## Quickstart.

Testing conventions

Every operator gets four checks: action vs. analytic solution, the adjoint identity, composition laws, and AD gradients (field and parameters) against finite differences. test/test_utils.jl provides fd_gradient and materialize (densify a prepared operator on small grids to check structure exactly).