Mesh redistribution support (split from #5215)

firedrake/adapt.py

@@ -10,6 +10,8 @@
 from firedrake.functionspace import FunctionSpace
 from firedrake.mesh import Mesh, DISTRIBUTION_PARAMETERS_NOOP
 from firedrake.netgen import _recurve_netgen_mesh
+from firedrake.redist import (RedistributedMeshTransfer, dm_has_empty_rank,
+                              redistribute_dm)
 
 
 DM_ADAPT_REFINE = 1
@@ -81,7 +83,53 @@ def _copy_adaptive_refinement_metadata(source_mesh, target_mesh):
         target_mesh.netgen_flags = source_mesh.netgen_flags
 
 
-def refine_marked_elements(mesh, cell_marker):
+def _needs_adaptive_redistribution(mesh):
+    """Return whether ``mesh`` has an empty rank."""
+    return dm_has_empty_rank(mesh.topology_dm)
+
+
+def _redistribute_adaptive_refined_mesh(coarse_mesh, transfer_mesh,
+                                        redistribute=True):
+    """Redistribute an adaptively refined mesh if the coarse mesh has empty ranks."""
+    _copy_adaptive_refinement_metadata(coarse_mesh, transfer_mesh)
+
+    needs_redist = (redistribute and coarse_mesh.comm.size > 1
+                    and _needs_adaptive_redistribution(transfer_mesh))
+    if not needs_redist:
+        return transfer_mesh
+
+    redist_parameters = dict(coarse_mesh._distribution_parameters)
+    redist_parameters["partition"] = True
+    redist_dm = transfer_mesh.topology_dm.clone()
+    _, point_sf = redistribute_dm(redist_dm, redist_parameters)
+
+    redist_topology_mesh = Mesh(
+        redist_dm,
+        dim=transfer_mesh.geometric_dimension,
+        reorder=False,
+        distribution_parameters=DISTRIBUTION_PARAMETERS_NOOP,
+        comm=transfer_mesh.comm,
+        tolerance=transfer_mesh.tolerance,
+    )
+    _copy_adaptive_refinement_metadata(transfer_mesh, redist_topology_mesh)
+
+    redist_transfer = RedistributedMeshTransfer(
+        transfer_mesh, redist_topology_mesh, point_sf
+    )
+    Vredist = transfer_mesh.coordinates.function_space().reconstruct(
+        mesh=redist_topology_mesh
+    )
+    redist_coordinates = Function(Vredist)
+    redist_transfer.orig2redist(transfer_mesh.coordinates, redist_coordinates)
+    redist_mesh = Mesh(redist_coordinates, name=transfer_mesh.name)
+    _copy_adaptive_refinement_metadata(redist_topology_mesh, redist_mesh)
+    redist_mesh.redist = RedistributedMeshTransfer(
+        transfer_mesh, redist_mesh, point_sf
+    )
+    return redist_mesh
+
+
+def refine_marked_elements(mesh, cell_marker, redistribute=True):
     """Adaptively refine a mesh using a DG0 marking function.
 
     Positive integer marker values request repeated refinement of the
@@ -95,6 +143,9 @@ def refine_marked_elements(mesh, cell_marker):
     cell_marker
         A DG0 `~firedrake.function.Function` on ``mesh``: cells with a
         positive value ``n`` are refined ``n`` times.
+    redistribute
+        If ``True``, redistribute the refined mesh when the coarse mesh
+        has empty ranks.
 
     Returns
     -------
@@ -148,5 +199,11 @@ def refine_marked_elements(mesh, cell_marker):
             final_mesh = _recurve_netgen_mesh(mesh, final_mesh, order)
 
     final_mesh._adaptive_cell_maps = (coarse_to_fine_total, fine_to_coarse_total)
-    _copy_adaptive_refinement_metadata(mesh, final_mesh)
+    final_mesh = _redistribute_adaptive_refined_mesh(
+        mesh, final_mesh, redistribute=redistribute
+    )
+    final_mesh._adaptive_cell_maps = (coarse_to_fine_total, fine_to_coarse_total)
+    redist = getattr(final_mesh, "redist", None)
+    if redist is not None:
+        redist.orig._adaptive_cell_maps = (coarse_to_fine_total, fine_to_coarse_total)
     return final_mesh

firedrake/mesh.py

@@ -2906,15 +2906,17 @@ def unique(self):
         return self
 
     @PETSc.Log.EventDecorator()
-    def refine_marked_elements(self, mark):
+    def refine_marked_elements(self, mark, redistribute=True):
         """Adaptively refine a mesh using a DG0 marking function.
 
         :arg mark: the marking function, a Firedrake DG0 function on
             this mesh; cells with a positive value are refined.
+        :arg redistribute: if ``True``, redistribute the refined mesh
+            when the coarse mesh has empty ranks.
 
         """
         from firedrake.adapt import refine_marked_elements
-        return refine_marked_elements(self, mark)
+        return refine_marked_elements(self, mark, redistribute)
 
     @PETSc.Log.EventDecorator()
     def curve_field(self, order, permutation_tol=1e-8, cg_field=None):

firedrake/mg/adaptive_hierarchy.py

@@ -20,15 +20,20 @@ class AdaptiveMeshHierarchy(HierarchyBase):
         The coarsest mesh in the hierarchy.
     nested: bool
         A flag to indicate whether the meshes are nested.
+    redistribute: bool
+        If ``True``, redistribute the adaptively-refined mesh
+        when the coarse cell has empty ranks.
 
     """
-    def __init__(self, base_mesh: MeshGeometry, nested: bool = True):
+    def __init__(self, base_mesh: MeshGeometry, nested: bool = True,
+                 redistribute: bool = True):
         self.meshes = []
         self._meshes = []
         self.coarse_to_fine_cells = {}
         self.fine_to_coarse_cells = {Fraction(0, 1): None}
         self.refinements_per_level = 1
         self.nested = nested
+        self.redistribute = redistribute
         self._shared_data_cache = defaultdict(dict)
         self.add_mesh(base_mesh)
 
@@ -51,8 +56,11 @@ def add_mesh(self, mesh: MeshGeometry,
         """
         level = len(self.meshes)
         if level > 0 and (coarse_to_fine_cells is None or fine_to_coarse_cells is None):
+            # Adaptive maps live on the parent-owned transfer mesh when redistributed.
+            redist = getattr(mesh, "redist", None)
+            map_mesh = redist.orig if redist is not None else mesh
             coarse_to_fine_cells, fine_to_coarse_cells = getattr(
-                mesh, "_adaptive_cell_maps", (None, None)
+                map_mesh, "_adaptive_cell_maps", (None, None)
             )
 
         self._meshes.append(mesh)
@@ -104,6 +112,7 @@ def adapt(self, eta: Function | Cofunction, theta: float):
         markers = Function(M)
         markers.dat.data_wo[should_refine] = 1
 
-        refined_mesh = mesh.refine_marked_elements(markers)
+        refined_mesh = mesh.refine_marked_elements(markers,
+                                                   redistribute=self.redistribute)
         self.add_mesh(refined_mesh)
         return self.meshes[-1]

firedrake/mg/mesh.py

@@ -9,29 +9,21 @@
 import firedrake
 from functools import cached_property
 
+from firedrake.mesh import DISTRIBUTION_PARAMETERS_NOOP
+from firedrake.redist import (
+    RedistributedMeshTransfer,
+    dm_has_empty_rank,
+    make_unoverlapped_dm,
+    redistribute_dm,
+)
 from firedrake import utils
 from firedrake.cython import mgimpl as impl
-import firedrake.cython.dmcommon as dmcommon
 from .utils import set_level, set_dm_refine_level
 
 __all__ = ("HierarchyBase", "MeshHierarchy", "ExtrudedMeshHierarchy", "NonNestedHierarchy",
            "SemiCoarsenedExtrudedHierarchy", "SubmeshHierarchy")
 
 
-def make_unoverlapped_dm(dm):
-    """Effectively invert dm.distributeOverlap().
-
-    The resulting plex has the identical data structure as the one before
-    distributeOverlap(). This is algorithmically guaranteed.
-    """
-    tdim = dm.getDimension()
-    dm = dmcommon.submesh_create(dm, tdim, "depth", tdim, True)
-    dm.removeLabel("pyop2_core")
-    dm.removeLabel("pyop2_owned")
-    dm.removeLabel("pyop2_ghost")
-    return dm
-
-
 class HierarchyBase(object):
     """Create an encapsulation of an hierarchy of meshes.
 
@@ -93,7 +85,8 @@ def MeshHierarchy(mesh, refinement_levels,
                   netgen_flags=False,
                   reorder=None,
                   distribution_parameters=None, callbacks=None,
-                  mesh_builder=firedrake.Mesh):
+                  mesh_builder=firedrake.Mesh,
+                  redistribute=True):
     """Build a hierarchy of meshes by uniformly refining a coarse mesh.
 
     Parameters
@@ -113,6 +106,11 @@ def MeshHierarchy(mesh, refinement_levels,
         for details.  If ``None``, use the same distribution
         parameters as were used to distribute the coarse mesh,
         otherwise, these options override the default.
+    redistribute : bool
+        If ``True``, redistribute refined meshes when this is needed to
+        avoid empty ranks.  Transfer operators use an internal
+        parent-owned mesh before moving data to or from the redistributed
+        mesh.
     reorder : bool
         optional flag indicating whether to reorder the
         refined meshes.
@@ -143,12 +141,13 @@ def MeshHierarchy(mesh, refinement_levels,
     else:
         before = after = lambda dm, i: None
 
-    # Refine an unoverlapped plex at each level. Keeping every dm here
-    # unoverlapped means overlap only ever needs to be added once, by
-    # mesh_builder below.
+    # Refine an unoverlapped plex at each level, redistributing (partition
+    # only, no overlap) whenever refinement would otherwise leave empty
+    # ranks.  Keeping every dm here unoverlapped means overlap only ever
+    # needs to be added once, by mesh_builder below.
     cdm = make_unoverlapped_dm(mesh.topology_dm)
     cdm.setRefinementUniform(True)
-    dms = [cdm]
+    dm_entries = [(cdm, cdm, None)]
     for i in range(refinement_levels*refinements_per_level):
         if i % refinements_per_level == 0:
             before(cdm, i)
@@ -165,10 +164,18 @@ def MeshHierarchy(mesh, refinement_levels,
             scale = mesh._radius / np.linalg.norm(coords, axis=1).reshape(-1, 1)
             coords *= scale
 
-        dms.append(rdm)
+        rdm_orig = rdm
+        point_sf = None
+        needs_redist = (redistribute and mesh.comm.size > 1
+                        and dm_has_empty_rank(rdm))
+        if needs_redist:
+            rdm = rdm.clone()
+            point_sf, _ = redistribute_dm(rdm, {"partition": True}, grow_overlap=False)
+        dm_entries.append((rdm, rdm_orig, point_sf))
         cdm = rdm
 
-    # Build a mesh for each level, adding overlap here.
+    # Build a mesh for each level, adding overlap (and, where needed, the
+    # transfer between the parent-owned and redistributed meshes) here.
     parameters = {}
     if distribution_parameters is not None:
         parameters.update(distribution_parameters)
@@ -177,27 +184,48 @@ def MeshHierarchy(mesh, refinement_levels,
     parameters["partition"] = False
 
     meshes = [mesh]
-    for rdm in dms[1:]:
+    mesh_origs = [mesh]
+    for rdm, rdm_orig, point_sf in dm_entries[1:]:
         fmesh = mesh_builder(
             rdm,
             dim=mesh.geometric_dimension,
             distribution_parameters=parameters,
             reorder=reorder,
             comm=mesh.comm,
         )
+        if point_sf is not None:
+            fmesh_orig = mesh_builder(
+                rdm_orig,
+                dim=mesh.geometric_dimension,
+                distribution_parameters=DISTRIBUTION_PARAMETERS_NOOP,
+                reorder=reorder,
+                comm=mesh.comm,
+            )
+            overlap_sf = fmesh.topology.sfBC
+            if overlap_sf is not None:
+                point_sf = point_sf.compose(overlap_sf)
+            fmesh.redist = RedistributedMeshTransfer(fmesh_orig, fmesh, point_sf)
+        else:
+            fmesh_orig = fmesh
         fmesh._distribution_parameters = parameters
         meshes.append(fmesh)
+        mesh_origs.append(fmesh_orig)
 
     # Build local-to-global maps and coarse/fine cell maps between
     # consecutive levels.
     lgmaps = [
         (impl.create_lgmap(dm), impl.create_lgmap(m.topology_dm))
-        for dm, m in zip(dms, meshes)
+        for (dm, _, _), m in zip(dm_entries, meshes)
+    ]
+    lgmap_origs = [
+        lgmap if point_sf is None else
+        (impl.create_lgmap(rdm_orig), impl.create_lgmap(mesh_orig.topology_dm))
+        for (lgmap, (_, rdm_orig, point_sf), mesh_orig) in zip(lgmaps, dm_entries, mesh_origs)
     ]
     coarse_to_fine_cells = []
     fine_to_coarse_cells = [None]
-    for (coarse, fine), (clgmaps, flgmaps) in zip(zip(meshes[:-1], meshes[1:]),
-                                                  zip(lgmaps[:-1], lgmaps[1:])):
+    for (coarse, fine), (clgmaps, flgmaps) in zip(zip(meshes[:-1], mesh_origs[1:]),
+                                                  zip(lgmaps[:-1], lgmap_origs[1:])):
         c2f, f2c = impl.coarse_to_fine_cells(coarse, fine, clgmaps, flgmaps)
         coarse_to_fine_cells.append(c2f)
         fine_to_coarse_cells.append(f2c)