diff --git a/docs/developer/design/fmg-checkpoint-hierarchy.md b/docs/developer/design/fmg-checkpoint-hierarchy.md new file mode 100644 index 000000000..67bed838d --- /dev/null +++ b/docs/developer/design/fmg-checkpoint-hierarchy.md @@ -0,0 +1,117 @@ +--- +title: "Persisting the FMG mesh hierarchy across checkpoints" +--- + +# Persisting the geometric-multigrid hierarchy across checkpoints + +## Problem + +A mesh built with `refinement=N` carries a geometric refinement hierarchy in +`mesh.dm_hierarchy`, which the Stokes/scalar/vector solvers use for geometric +**Full Multigrid (FMG)** — the anisotropy-robust preconditioner of choice on +adapted meshes. But the checkpoint load path (`_from_plexh5`) reconstructs only +the single saved DMPlex, so a **reloaded mesh has `dm_hierarchy = [dm]`** (one +level) and FMG silently falls back to GAMG after a restart. This note records the +design that restores the hierarchy on reload, and the experiments that shaped it. + +## Why store the coarse levels (not reconstruct, not refine-from-label) + +Several approaches were prototyped and rejected: + +- **Reconstruct the coarse mesh from a per-node "level" label.** The coarse + *topology* is recoverable bit-exact from labelled nodes (the all-midpoint + "central" fine cells map 1:1 to coarse cells). **But** a coarse DMPlex built + from scratch (`createFromCellList`) does not reproduce the original's internal + cone-orientation / edge ordering, and PETSc's nested multigrid interpolator + needs the **canonical `refine()` numbering**. Splicing a reconstructed coarse + in throws `PETSc has generated inconsistent data` (err 77). +- **Store only the coarsest level + `refine()` back up on reload.** Works for the + hierarchy, but then the rebuilt fine has to be reconciled (numbering + field + data) with the saved deformed mesh — fragile. + +The winning insight: a **loaded** coarse DM preserves the canonical numbering +(`topologyLoad` is faithful), so it can be `setCoarseDM`-linked directly under the +working fine and FMG just works — **no refine, no node-moving, no reconstruction**. +This is exactly the live `clone_dm_hierarchy` pattern with *load* swapped for +*clone*. Validated: `refine(stored L0) == saved fine` bit-exact in 2D and 3D, and a +reloaded hierarchy drives `pc_type=mg` to convergence. + +## On-disk format: a single coarsest sidecar + +PETSc's `HDF5_PETSC` `DMView` writes to fixed top-level groups (`/topology`, +`/geometry`, `/labels`) — it is **not namespaced by DM name**. Writing a second +DMPlex into the same file (PETSc viewer append, *or* an h5py-injected subgroup +that the PETSc reader then ignores) corrupts the file (a reload BUS-errors). So +the hierarchy is stored in **one extra single-DM file** beside the main +checkpoint, holding only the **coarsest** level: + +``` +mymesh.h5 # the working/fine mesh (unchanged, fully compatible) +mymesh.hierarchy.L0.h5 # coarsest level only +``` + +The intermediate coarse levels are not stored — on reload they are rebuilt by +`refine()`-ing the coarsest `N-1` times (they come back canonically numbered, +which is all the co-located nested interpolation needs). The main file's +`metadata` group gains `hierarchy_coarse_levels = N-1` (the refinement depth). +Old checkpoints (attribute absent) and plain meshes (no hierarchy) write no +sidecar and reload exactly as before. + +## Reload and the link-free working `dm` + +On reload the coarse levels are loaded and spliced: +`dm_hierarchy = [L0, …, L_{N-2}, fine]`, linked with `setCoarseDM`. One subtlety: +the mesh's **working `self.dm` must be a link-free clone** of the finest level +(mirroring the `refinement` construction branch). If `self.dm` itself carries a +coarse-DM link, `mesh.update_lvec()`'s `createFieldDecomposition` recurses into +the 0-field coarse levels and fails (`requested fields 1 > DM fields 0`). The +linked hierarchy lives in `dm_hierarchy`; the solver clones it +(`clone_dm_hierarchy`) for its own multigrid setup. + +## Parallel: co-location via the Simple partitioner + +Works in serial **and** parallel through the same reload path. The hazard in +parallel is that the coarse sidecars and the fine reload on **independent +partitions**; linking incompatibly-partitioned levels sends the interpolator into +a cross-rank point-location spin (observed before the fix: np=2, rank 0 at 99% CPU +indefinitely, rank 1 idle). + +The fix needs no custom partition math. The fine carries the **canonical +refinement numbering** — coarse cell `c`'s children are fine cells +`c·numSubcells + r`, laid out contiguously right after `c`. So if the fine *and* +every coarse level are distributed with PETSc's **Simple** partitioner (equal +contiguous splits of `[0, Ncells)`), the fine split at `k·Nf/p` lines up with the +coarse split at `k·Nc/p` (since `Nf = numSubcells·Nc`): **each rank's coarse cells +and their fine children land on the same rank.** The multigrid interpolation is +then rank-local — no cross-partition communication, no hang — and the levels are a +genuine per-rank refinement, so the exact **nested** interpolator applies (the fine +levels are flagged via `DMPlexSetRegularRefinement`). + +Trade-off: hierarchy meshes reload with a Simple (contiguous) partition rather than +the default graph partition. For refinement meshes the canonical ordering is +reasonably coherent, and field reload is coordinate-matched (partition-agnostic), +so correctness is unaffected; partition-quality tuning can come later. Plain +(non-hierarchy) meshes are untouched — they keep the default partitioner. + +## Implementation + +All in `src/underworld3/discretisation/discretisation_mesh.py`: + +- `_hierarchy_sidecar_name()` — sidecar path convention. +- `Mesh.write()` — writes `metadata/hierarchy_coarse_levels` and one sidecar + holding the coarsest level (collective). +- `Mesh.__init__` `.h5` branch — loads the coarsest sidecar and rebuilds the + intermediate coarse levels by `refine()` (serial and parallel), stashing the + list on `self._sidecar_coarse_levels`. +- `Mesh.__init__` hierarchy section — distributes fine + coarse with the Simple + partitioner (co-location), splices them under the working dm, flags the fine + levels as regular refinements, re-establishes the link-free clone. +- `petsc_dm_{set,get}_regular_refinement` in `cython/petsc_discretisation.pyx` — + wraps `DMPlexSetRegularRefinement` (not exposed by petsc4py) so reloaded levels + take the exact nested interpolation path. + +No new user-facing surface: the same `Mesh(file)` reload transparently restores the +hierarchy when the checkpoint has one, and behaves exactly as before when it does +not. + +Tests: `tests/test_0004_checkpoint_fmg_hierarchy.py`. diff --git a/src/underworld3/cython/petsc_discretisation.pyx b/src/underworld3/cython/petsc_discretisation.pyx index 655537d65..6f48402de 100644 --- a/src/underworld3/cython/petsc_discretisation.pyx +++ b/src/underworld3/cython/petsc_discretisation.pyx @@ -163,6 +163,24 @@ def petsc_dm_filter_by_label(incoming_dm, label_name, label_value): return subdm +def petsc_dm_set_regular_refinement(dm, regular=True): + """Flag a DMPlex as having been produced by uniform (regular) refinement + of its coarse DM. + + This is the flag PETSc's geometric multigrid checks + (``DMCreateInterpolation_Plex``) to take the exact *nested* interpolation + path (``DMPlexComputeInterpolatorNested``, which maps coarse cell ``c`` to + fine children ``c*numSubcells + r``) instead of the point-location + *general* path. Only valid when the dm's point numbering follows the + canonical ``refine()`` convention and its coarse DM is set + (``setCoarseDM``). Not exposed by petsc4py, hence this wrapper. + """ + cdef DM c_dm = dm + cdef PetscBool flag = PETSC_TRUE if regular else PETSC_FALSE + CHKERRQ( DMPlexSetRegularRefinement(c_dm.dm, flag) ) + return + + # This is not cython, does it need to be here or in discretisation.py ? diff --git a/src/underworld3/cython/petsc_extras.pxi b/src/underworld3/cython/petsc_extras.pxi index 1f3f41202..23785c99c 100644 --- a/src/underworld3/cython/petsc_extras.pxi +++ b/src/underworld3/cython/petsc_extras.pxi @@ -55,6 +55,7 @@ cdef extern from "petsc.h" nogil: # PetscErrorCode DMPlexSetSNESLocalFEM( PetscDM, void *, void *, void *) # PetscErrorCode DMPlexSetSNESLocalFEM( PetscDM, PetscBool, void *) PetscErrorCode DMPlexComputeGeometryFVM( PetscDM dm, PetscVec *cellgeom, PetscVec *facegeom) + PetscErrorCode DMPlexSetRegularRefinement( PetscDM dm, PetscBool regular) PetscErrorCode MatInterpolate(PetscMat A, PetscVec x, PetscVec y) PetscErrorCode DMSetLocalSection(PetscDM, PetscSection) diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index 1436aa406..242be99e4 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -166,6 +166,20 @@ def _from_plexh5( return sf0, h5plex +def _hierarchy_sidecar_name(mesh_filename, level): + """Filename for a coarse hierarchy sidecar of a mesh checkpoint. + + The geometric-multigrid (FMG) hierarchy is persisted as a single extra + single-DM HDF5 file holding the **coarsest** level (``level=0``) beside the + main mesh checkpoint — PETSc's ``HDF5_PETSC`` format does not support several + DMPlex objects in one file. The intermediate coarse levels are rebuilt by + refinement on reload, so only ``mymesh.hierarchy.L0.h5`` is written. The + ``level`` argument is kept for forward-compatibility. + """ + base, ext = os.path.splitext(mesh_filename) + return f"{base}.hierarchy.L{level}{ext}" + + class Mesh(Stateful, uw_object): r""" Unstructured mesh with PETSc DMPlex backend. @@ -408,6 +422,35 @@ def __init__( f.close() + # Restore the geometric-multigrid (FMG) coarse hierarchy from + # sidecar files, if present. Works in serial and parallel: the + # coarse levels and the fine are co-located on reload by the + # Simple partitioner (see the hierarchy-construction branch), so + # the multigrid interpolation is rank-local and does not hang. + self._sidecar_coarse_levels = None + try: + with h5py.File(plex_or_meshfile, "r") as fh: + n_coarse = int( + fh["metadata"].attrs.get("hierarchy_coarse_levels", 0) + ) + except (KeyError, OSError): + n_coarse = 0 + if n_coarse > 0: + sidecar = _hierarchy_sidecar_name(plex_or_meshfile, 0) + if os.path.isfile(sidecar): + # Load the coarsest level and rebuild the intermediate + # coarse levels by refining it (n_coarse - 1 times). They + # come back canonically numbered — all the co-located + # nested interpolation needs — and the hierarchy branch + # then Simple-distributes each level. (refine() propagates + # the boundary labels from the coarsest.) + coarsest = _from_plexh5(sidecar, PETSc.COMM_WORLD) + levels = [coarsest] + for _ in range(n_coarse - 1): + levels[-1].setRefinementUniform() + levels.append(levels[-1].refine()) + self._sidecar_coarse_levels = levels + # Do not call setFromOptions() here. DMPlexTopologyLoad() # returns the topology SF needed to reload checkpoint fields. # setFromOptions() can repartition/reorder the DM before UW @@ -529,7 +572,55 @@ class replacement_boundaries(Enum): uw.mpi.barrier() - if not refinement is None and refinement > 0: + if getattr(self, "_sidecar_coarse_levels", None): + + # Reloaded mesh with a persisted FMG hierarchy: splice the loaded + # coarse levels under the working (fine) dm and link them so PETSc's + # geometric multigrid sees a refinement hierarchy. The loaded coarse + # DMs keep the canonical refine() numbering the nested interpolator + # needs (a *reconstructed* coarse would not — see the + # checkpoint-hierarchy design note). + # + # Co-location (serial and parallel): distribute the fine AND every + # coarse level with the Simple (contiguous-by-canonical-index) + # partitioner. Because the fine carries the canonical refinement + # numbering (coarse cell c -> fine cells c*numSubcells+r, contiguous), + # equal contiguous splits put each rank's coarse cells and their fine + # children on the same rank. The multigrid interpolation is then + # rank-local — this is what avoids the cross-partition hang that a + # default (graph) partitioner produces on independently-loaded levels. + from underworld3.cython.petsc_discretisation import ( + petsc_dm_set_regular_refinement, + ) + + for _hdm in [self.dm] + list(self._sidecar_coarse_levels): + if not _hdm.isDistributed(): + _hdm.getPartitioner().setType(PETSc.Partitioner.Type.SIMPLE) + if not self.dm.isDistributed(): + self.sf1 = self.dm.distribute() + for _cdm in self._sidecar_coarse_levels: + if not _cdm.isDistributed(): + _cdm.distribute() + + self.dm_hierarchy = list(self._sidecar_coarse_levels) + [self.dm] + for i in range(len(self.dm_hierarchy) - 1): + self.dm_hierarchy[i + 1].setCoarseDM(self.dm_hierarchy[i]) + # The fine level IS a uniform refinement of the coarse (canonical + # numbering, co-located) — flag it so PETSc builds the exact + # nested interpolator rather than falling back to point location. + petsc_dm_set_regular_refinement(self.dm_hierarchy[i + 1], True) + + self.dm_h = self.dm_hierarchy[-1] + self.dm_h.setName("uw_hierarchical_dm") + + # Working dm is a link-free clone of the finest level (mirrors the + # refinement branch). It must NOT carry a coarse-DM link or + # mesh.update_lvec()'s createFieldDecomposition recurses into the + # 0-field coarse levels and fails. + self.dm = self.dm_h.clone() + self._sidecar_coarse_levels = None + + elif not refinement is None and refinement > 0: self.dm.setRefinementUniform() @@ -3347,8 +3438,34 @@ def write(self, filename: str, index: Optional[int] = None): } g.attrs["coordinate_units"] = json.dumps(coord_units_dict) + # Number of coarse multigrid levels in the hierarchy (= number + # of refinements from the stored coarsest level up to the fine + # mesh). Used on reload to rebuild the intermediate levels. + g.attrs["hierarchy_coarse_levels"] = len(self.dm_hierarchy) - 1 + f.close() + # Persist the geometric-multigrid (FMG) hierarchy as a SINGLE sidecar + # holding the coarsest level only. On reload the intermediate coarse + # levels are rebuilt by refining it (they come back canonically numbered, + # which is all the co-located nested interpolation needs). Without this + # file a reloaded mesh has a single level and falls back to GAMG. One + # single-DM HDF5 file (PETSc's HDF5_PETSC format holds one DMPlex per + # file). Collective write. See _hierarchy_sidecar_name and the .h5 reload. + if len(self.dm_hierarchy) > 1: + coarse_dm = self.dm_hierarchy[0] + sidecar = _hierarchy_sidecar_name(filename, 0) + cviewer = PETSc.ViewerHDF5().create(sidecar, "w", comm=PETSc.COMM_WORLD) + cviewer.pushFormat(PETSc.Viewer.Format.HDF5_PETSC) + saved_name = coarse_dm.getName() + coarse_dm.setName("uw_mesh") # _from_plexh5 loads the DM named "uw_mesh" + try: + cviewer(coarse_dm) + finally: + coarse_dm.setName(saved_name) + cviewer.popFormat() + cviewer.destroy() + def vtk(self, filename: str): """ Save mesh to the specified file diff --git a/tests/parallel/mpi_runner.sh b/tests/parallel/mpi_runner.sh index 1c74260f6..c87eaaed1 100755 --- a/tests/parallel/mpi_runner.sh +++ b/tests/parallel/mpi_runner.sh @@ -38,3 +38,8 @@ echo "ptest 0010 snapshot on-disk -np 3 (uneven)" mpirun -np 3 $PYTHON ./ptest_0010_snapshot_disk.py echo "ptest 0010 snapshot on-disk -np 4" mpirun -np 4 $PYTHON ./ptest_0010_snapshot_disk.py + +echo "ptest 0004 checkpoint FMG hierarchy -np 2" +mpirun -np 2 $PYTHON ./ptest_0004_checkpoint_fmg_hierarchy.py +echo "ptest 0004 checkpoint FMG hierarchy -np 3 (uneven partition)" +mpirun -np 3 $PYTHON ./ptest_0004_checkpoint_fmg_hierarchy.py diff --git a/tests/parallel/ptest_0004_checkpoint_fmg_hierarchy.py b/tests/parallel/ptest_0004_checkpoint_fmg_hierarchy.py new file mode 100644 index 000000000..4374539b0 --- /dev/null +++ b/tests/parallel/ptest_0004_checkpoint_fmg_hierarchy.py @@ -0,0 +1,77 @@ +"""Parallel (MPI) test: FMG mesh hierarchy survives a checkpoint round-trip. + +A mesh built with ``refinement`` carries a geometric-multigrid hierarchy. On +reload from a checkpoint, ``Mesh(file)`` must transparently restore that +hierarchy in parallel too — without the cross-partition interpolation hang that +a naive (graph-partitioned) reload of independently-loaded levels produces. + +The fix co-locates the levels by distributing the fine and every coarse level +with PETSc's Simple partitioner: because the fine carries canonical refinement +numbering, equal contiguous splits put each rank's coarse cells and their fine +children on the same rank, so the multigrid interpolation is rank-local. + +Run: + + cd tests/parallel + mpirun -np 2 python ./ptest_0004_checkpoint_fmg_hierarchy.py + +Asserts (checked on rank 0): + 1. The reloaded mesh has its hierarchy restored (levels > 1). + 2. Geometric FMG converges on the reloaded mesh — proving the levels are + co-located and linked (a hang would block here, caught by the CI timeout). +""" + +import os +import glob + +import underworld3 as uw +from petsc4py import PETSc + +rank = uw.mpi.rank +size = uw.mpi.size + +fn = "/tmp/_ptest_0004_fmg_ckpt.h5" +if rank == 0: + for p in glob.glob("/tmp/_ptest_0004_fmg_ckpt*"): + os.remove(p) +uw.mpi.barrier() + +# Build a refinement mesh (3-level hierarchy) and checkpoint it. +mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), + cellSize=0.2, refinement=2, qdegree=2, +) +mesh.write(fn) +uw.mpi.barrier() + +# Reload in parallel — same call as a non-hierarchical mesh. +mesh2 = uw.discretisation.Mesh(fn) +assert len(mesh2.dm_hierarchy) == len(mesh.dm_hierarchy) == 3, ( + f"hierarchy not restored: got {len(mesh2.dm_hierarchy)} levels" +) + +# Geometric FMG must converge on the reloaded mesh (no cross-partition hang). +poisson = uw.systems.Poisson(mesh2) +poisson.constitutive_model = uw.constitutive_models.DiffusionModel +poisson.constitutive_model.Parameters.diffusivity = 1 +poisson.f = 0.0 +poisson.add_dirichlet_bc(0.0, "Bottom") +poisson.add_dirichlet_bc(1.0, "Top") +for k, v in { + "pc_type": "mg", "pc_mg_type": "full", "pc_mg_galerkin": "both", + "mg_levels_ksp_type": "chebyshev", "mg_levels_pc_type": "sor", + "mg_coarse_pc_type": "redundant", "mg_coarse_redundant_pc_type": "lu", +}.items(): + poisson.petsc_options[k] = v +poisson.solve() + +assert poisson.petsc_options.getString("pc_type") == "mg" +assert poisson.snes.getConvergedReason() > 0 + +if rank == 0: + print( + f"ptest_0004 OK (np={size}): hierarchy restored to " + f"{len(mesh2.dm_hierarchy)} levels, FMG converged in " + f"{poisson.snes.getKSP().getIterationNumber()} iters", + flush=True, + ) diff --git a/tests/test_0004_checkpoint_fmg_hierarchy.py b/tests/test_0004_checkpoint_fmg_hierarchy.py new file mode 100644 index 000000000..66b234885 --- /dev/null +++ b/tests/test_0004_checkpoint_fmg_hierarchy.py @@ -0,0 +1,73 @@ +import os +import pytest +import underworld3 as uw + +# Persisting + restoring the geometric-multigrid (FMG) hierarchy across a +# mesh checkpoint round-trip (serial). See +# docs/developer/design/fmg-checkpoint-hierarchy.md +pytestmark = [pytest.mark.level_2, pytest.mark.tier_a] + + +def _refined_box(cellSize=0.3, refinement=1): + return uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), + cellSize=cellSize, refinement=refinement, qdegree=2, + ) + + +def test_refinement_mesh_writes_single_coarsest_sidecar(tmp_path): + m = _refined_box(refinement=2) + assert len(m.dm_hierarchy) == 3 # coarse, mid, fine + fn = str(tmp_path / "mesh.h5") + m.write(fn) + # Only the coarsest level is stored; intermediate levels are rebuilt by + # refinement on reload. + assert os.path.isfile(str(tmp_path / "mesh.hierarchy.L0.h5")) + assert not os.path.isfile(str(tmp_path / "mesh.hierarchy.L1.h5")) + + +def test_reload_restores_hierarchy(tmp_path): + m = _refined_box(refinement=2) + fn = str(tmp_path / "mesh.h5") + m.write(fn) + m2 = uw.discretisation.Mesh(fn) + assert len(m2.dm_hierarchy) == len(m.dm_hierarchy) == 3 + + +def test_no_hierarchy_writes_no_sidecar(tmp_path): + # A plain mesh (no refinement) must not write sidecars and must reload + # exactly as before — regression guard for existing checkpoints. + m = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=0.25, qdegree=2) + assert len(m.dm_hierarchy) == 1 + fn = str(tmp_path / "plain.h5") + m.write(fn) + assert not os.path.isfile(str(tmp_path / "plain.hierarchy.L0.h5")) + m2 = uw.discretisation.Mesh(fn) + assert len(m2.dm_hierarchy) == 1 + + +def test_reloaded_hierarchy_drives_geometric_mg(tmp_path): + # The restored hierarchy must actually work as a geometric-multigrid + # preconditioner on the reloaded mesh. + m = _refined_box(refinement=1) + fn = str(tmp_path / "mesh.h5") + m.write(fn) + m2 = uw.discretisation.Mesh(fn) + assert len(m2.dm_hierarchy) == 2 + + poisson = uw.systems.Poisson(m2) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = 1 + poisson.f = 0.0 + poisson.add_dirichlet_bc(0.0, "Bottom") + poisson.add_dirichlet_bc(1.0, "Top") + for k, v in { + "pc_type": "mg", "pc_mg_type": "full", "pc_mg_galerkin": "both", + "mg_levels_ksp_type": "chebyshev", "mg_levels_pc_type": "sor", + "mg_coarse_pc_type": "lu", + }.items(): + poisson.petsc_options[k] = v + poisson.solve() + assert poisson.petsc_options.getString("pc_type") == "mg" + assert poisson.snes.getConvergedReason() > 0