diff --git a/firedrake/adapt.py b/firedrake/adapt.py index f76c8622bd..96cd6f8861 100644 --- a/firedrake/adapt.py +++ b/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 diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 8b0cc4504a..533036669b 100644 --- a/firedrake/mesh.py +++ b/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): diff --git a/firedrake/mg/adaptive_hierarchy.py b/firedrake/mg/adaptive_hierarchy.py index 6d4b282c4b..2b4c7b1094 100644 --- a/firedrake/mg/adaptive_hierarchy.py +++ b/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] diff --git a/firedrake/mg/mesh.py b/firedrake/mg/mesh.py index 2c3f87b927..8fbafe5095 100644 --- a/firedrake/mg/mesh.py +++ b/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,7 +184,8 @@ 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, @@ -185,19 +193,39 @@ def MeshHierarchy(mesh, refinement_levels, 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) diff --git a/firedrake/redist.py b/firedrake/redist.py new file mode 100644 index 0000000000..68108471eb --- /dev/null +++ b/firedrake/redist.py @@ -0,0 +1,112 @@ +import firedrake.cython.dmcommon as dmcommon +from firedrake.cython.dmcommon import DistributedMeshOverlapType +from firedrake.halo import _get_mtype as get_mpi_type +from firedrake.mesh import MeshTopology +from pyop2.mpi import MPI + + +class RedistributedMeshTransfer: + """Transfer data between a parent-owned mesh and its redistributed mesh.""" + + def __init__(self, orig, redist, point_sf=None): + self.orig = orig + self.redist = redist + self.point_sf = point_sf if point_sf is not None else redist.sfBC + + def _section_sf(self, root, leaf): + point_sf = self.point_sf + root_section = root.function_space().dm.getDefaultSection() + leaf_section = leaf.function_space().dm.getDefaultSection() + remote_offsets, _ = point_sf.distributeSection(root_section, leaf_section) + return point_sf.createSectionSF(root_section, remote_offsets, leaf_section) + + def orig2redist(self, source, target): + section_sf = self._section_sf(source, target) + dtype, _ = get_mpi_type(source.dat) + section_sf.bcastBegin(dtype, + source.dat.data_ro_with_halos, + target.dat.data_wo_with_halos, + MPI.REPLACE) + section_sf.bcastEnd(dtype, + source.dat.data_ro_with_halos, + target.dat.data_wo_with_halos, + MPI.REPLACE) + + def redist2orig(self, source, target): + section_sf = self._section_sf(target, source) + dtype, _ = get_mpi_type(source.dat) + section_sf.reduceBegin(dtype, + source.dat.data_ro_with_halos, + target.dat.data_wo_with_halos, + MPI.REPLACE) + section_sf.reduceEnd(dtype, + source.dat.data_ro_with_halos, + target.dat.data_wo_with_halos, + MPI.REPLACE) + + +def distribute_overlap(dm, parameters): + overlap_type, overlap = parameters.get( + "overlap_type", (DistributedMeshOverlapType.FACET, 1) + ) + if overlap_type == DistributedMeshOverlapType.NONE: + if overlap > 0: + raise ValueError("Cannot have NONE overlap with overlap > 0") + return None + elif overlap_type in [DistributedMeshOverlapType.FACET, + DistributedMeshOverlapType.RIDGE]: + dmcommon.set_adjacency_callback(dm, overlap_type) + sf = dm.distributeOverlap(overlap) + dmcommon.clear_adjacency_callback(dm) + return sf + elif overlap_type == DistributedMeshOverlapType.VERTEX: + return dm.distributeOverlap(overlap) + else: + raise ValueError("Unknown overlap type %r" % (overlap_type,)) + + +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 + + +def redistribute_dm(dm, parameters, grow_overlap=True): + """Distribute ``dm`` and optionally add the requested overlap. + + Returns the SF from the partitioning step and the composed SF from + partitioning plus overlap growth. If ``grow_overlap`` is false, the + unoverlapped redistributed DM is left in place and the partition SF is + returned for both entries. + """ + dm.removeLabel("pyop2_core") + dm.removeLabel("pyop2_owned") + dm.removeLabel("pyop2_ghost") + + distribute = parameters.get("partition", True) + partitioner_type = parameters.get("partitioner_type") + MeshTopology._set_partitioner(dm, distribute, partitioner_type) + point_sf_orig = dm.distribute(overlap=0) + if not grow_overlap: + return point_sf_orig, point_sf_orig + overlap_sf = distribute_overlap(dm, parameters) + if overlap_sf is None: + point_sf = point_sf_orig + elif point_sf_orig is None: + point_sf = overlap_sf + else: + point_sf = point_sf_orig.compose(overlap_sf) + return point_sf_orig, point_sf + + +def dm_has_empty_rank(dm): + cstart, cend = dm.getHeightStratum(0) + return dm.comm.tompi4py().allreduce(cstart == cend, op=MPI.LOR) diff --git a/tests/firedrake/multigrid/test_adaptive_multigrid.py b/tests/firedrake/multigrid/test_adaptive_multigrid.py index 4285d7e1d4..32797b9980 100644 --- a/tests/firedrake/multigrid/test_adaptive_multigrid.py +++ b/tests/firedrake/multigrid/test_adaptive_multigrid.py @@ -26,6 +26,11 @@ def test_adapt_basic(): assert np.allclose(assemble(1*dx(mesh)), assemble(1*dx(base))) +def _adaptive_map_mesh(mesh): + redist = getattr(mesh, "redist", None) + return redist.orig if redist is not None else mesh + + def _linear_expr(mesh): """A linear expression in the mesh's spatial coordinates, generalizing ``x + 2*y`` to any dimension (``x + 2*y + 3*z`` in 3D, etc.).""" @@ -85,7 +90,7 @@ def test_refine_marked_elements_populates_cell_maps(coarse_mesh): fine_to_coarse = amh.fine_to_coarse_cells[1] assert coarse_to_fine.shape[0] == mesh.cell_set.size - assert fine_to_coarse.shape == (refined_mesh.cell_set.size, 1) + assert fine_to_coarse.shape == (_adaptive_map_mesh(refined_mesh).cell_set.size, 1) assert (fine_to_coarse >= -1).all() assert (fine_to_coarse >= 0).any() assert (coarse_to_fine >= 0).any() @@ -153,7 +158,7 @@ def _assert_adapt_after_uniform_refinement(mesh): fine_to_coarse = amh.fine_to_coarse_cells[1] assert coarse_to_fine.shape[0] == mesh.cell_set.size - assert fine_to_coarse.shape == (refined_mesh.cell_set.size, 1) + assert fine_to_coarse.shape == (_adaptive_map_mesh(refined_mesh).cell_set.size, 1) assert (fine_to_coarse >= 0).any() assert (coarse_to_fine >= 0).any() diff --git a/tests/firedrake/multigrid/test_redist_mesh.py b/tests/firedrake/multigrid/test_redist_mesh.py new file mode 100644 index 0000000000..274b6fdf78 --- /dev/null +++ b/tests/firedrake/multigrid/test_redist_mesh.py @@ -0,0 +1,116 @@ +import numpy as np +import pytest + +from firedrake import * + + +@pytest.mark.parallel(2) +def test_redistributed_hierarchy(): + m = UnitIntervalMesh(1) + mh = MeshHierarchy(m, 1) + + assert mh[1].cell_set.size > 0 + + +@pytest.mark.parallel(4) +def test_uniform_hierarchy_no_empty_ranks(): + dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 1)} + base = UnitSquareMesh(1, 1, distribution_parameters=dparams) + mh = MeshHierarchy(base, 2) + + for l, m in enumerate(mh[1:]): + assert m.cell_set.size > 0 + for k, v in dparams.items(): + assert m._distribution_parameters.get(k, None) == v + + Vc = FunctionSpace(mh[l], "CG", 1) + Vf = FunctionSpace(mh[l+1], "CG", 1) + + xc, yc = SpatialCoordinate(mh[l]) + xf, yf = SpatialCoordinate(mh[l+1]) + coarse_expr = xc + 2*yc + fine_expr = xf + 2*yf + + # test prolong CG1 + coarse = Function(Vc).interpolate(coarse_expr) + fine = Function(Vf) + prolong(coarse, fine) + assert errornorm(fine_expr, fine) < 1e-12 + + # test restrict CG1 + one_coarse = Function(Vc).assign(1) + one_fine = Function(Vf) + prolong(one_coarse, one_fine) + + fine_dual = assemble(conj(TestFunction(Vf))*dx) + coarse_dual = Cofunction(Vc.dual()) + restrict(fine_dual, coarse_dual) + assert np.allclose( + assemble(action(coarse_dual, one_coarse)), + assemble(action(fine_dual, one_fine)), + rtol=1e-12, + atol=1e-12, + ) + + # test inject CG1 + coarse_injected = Function(Vc) + inject(fine, coarse_injected) + assert errornorm(coarse_expr, coarse_injected) < 1e-12 + + # test inject DG0 + Qc = FunctionSpace(mh[l], "DG", 0) + Qf = FunctionSpace(mh[l+1], "DG", 0) + fine_expr = conditional(xf > 1, 1, 0) + coarse_expr = conditional(xc > 1, 1, 0) + fine = Function(Qf).interpolate(fine_expr) + coarse_injected = Function(Qc) + inject(fine, coarse_injected) + assert np.allclose( + assemble(coarse_expr * dx), + assemble(coarse_injected * dx), + rtol=1e-12, + atol=1e-12, + ) + assert l == 0 or errornorm(coarse_expr, coarse_injected) < 1e-12 + + +@pytest.mark.parallel(2) +def test_adaptive_hierarchy_redistributes_unbalanced_unitsquare(): + dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 1)} + mesh = UnitSquareMesh(4, 4, distribution_parameters=dparams) + amh = AdaptiveMeshHierarchy(mesh) + + M = FunctionSpace(mesh, "DG", 0) + markers = Function(M) + if mesh.comm.rank == 0: + markers.dat.data_wo[:] = 1 + + refined_mesh = mesh.refine_marked_elements(markers, balancing=0) + assert getattr(refined_mesh, "redist", None) is not None + amh.add_mesh(refined_mesh) + + V_coarse = FunctionSpace(mesh, "CG", 1) + V_fine = FunctionSpace(refined_mesh, "CG", 1) + xc, yc = SpatialCoordinate(mesh) + xf, yf = SpatialCoordinate(refined_mesh) + expr_coarse = xc + 2 * yc + expr_fine = xf + 2 * yf + + u_coarse = Function(V_coarse).interpolate(expr_coarse) + u_fine = Function(V_fine) + prolong(u_coarse, u_fine) + assert errornorm(expr_fine, u_fine) <= 1e-12 + + r_fine = assemble(conj(TestFunction(V_fine)) * dx) + r_coarse = Cofunction(V_coarse.dual()) + restrict(r_fine, r_coarse) + assert np.allclose( + assemble(action(r_coarse, u_coarse)), + assemble(action(r_fine, u_fine)), + rtol=1e-12, + atol=1e-12, + ) + + u_coarse_injected = Function(V_coarse) + inject(u_fine, u_coarse_injected) + assert errornorm(expr_coarse, u_coarse_injected) <= 1e-12