diff --git a/check_simple.py b/check_simple.py new file mode 100644 index 0000000000..47493ec58f --- /dev/null +++ b/check_simple.py @@ -0,0 +1,17 @@ +from firedrake import * + + +mesh = UnitSquareMesh(2, 2) +x, _ = SpatialCoordinate(mesh) + +assert np.allclose(assemble(1 * dx(domain=mesh)), 1) +assert np.allclose(assemble(x * dx(domain=mesh)), 0.5) + +plex = mesh.topology_dm +coords_vec = plex.getCoordinatesLocal() +coords_vec *= 2 + +assert np.allclose(assemble(1 * dx(domain=mesh)), 4) +assert np.allclose(assemble(x * dx(domain=mesh)), 4) + +print("Success! Yay!") diff --git a/firedrake/cython/dmcommon.pyx b/firedrake/cython/dmcommon.pyx index b5e0777dfe..4c1bd960e4 100644 --- a/firedrake/cython/dmcommon.pyx +++ b/firedrake/cython/dmcommon.pyx @@ -3919,3 +3919,106 @@ def submesh_create_cell_closure_cell_submesh(PETSc.DM subdm, CHKERR(PetscFree(subpoint_indices_inv)) CHKERR(ISRestoreIndices(subpoint_is.iset, &subpoint_indices)) return subcell_closure + + +def make_geometry_dm(coordinates) -> PETSc.DM: + cdef: + PETSc.DM topology_dm, geometry_dm + PETSc.Section vector_section + PetscInt gdim + + topology_dm = coordinates.function_space().mesh().topology_dm + geometry_dm = topology_dm.clone() + coordinate_element = coordinates.ufl_element() + gdim = topology_dm.getCoordinateDim() + + # the existing section has the correct numbering but is scalar, expand to gdim + scalar_dm = coordinates.function_space().dm + scalar_section = scalar_dm.getLocalSection() + p_start, p_end = scalar_section.getChart() + + vector_section = PETSc.Section().create(comm=scalar_section.comm) + vector_section.setNumFields(1) + vector_section.setFieldComponents(0, gdim) + vector_section.setChart(p_start, p_end) + vector_section.setPermutation(scalar_section.getPermutation()) + + for p in range(p_start, p_end): + scalar_ndof = scalar_section.getDof(p) + CHKERR(PetscSectionSetDof(vector_section.sec, p, scalar_ndof*gdim)) + CHKERR(PetscSectionSetFieldDof(vector_section.sec, p, 0, scalar_ndof*gdim)) + vector_section.setUp() + + if coordinate_element.family() == "Lagrange": + # apparently gdim ignored in this call and set explicitly below + CHKERR(DMSetCoordinateSection(geometry_dm.dm, gdim, vector_section.sec)) + geometry_dm.setCoordinateDim(gdim) + geometry_dm.setCoordinatesLocal(coordinates.dat._vec) + else: + geometry_dm.setCellCoordinateDM(geometry_dm.getCoordinateDM().clone()) + geometry_dm.setCellCoordinateSection(gdim, vector_section) # gdim ignored + geometry_dm.setCoordinateDim(gdim) + geometry_dm.setCellCoordinatesLocal(coordinates.dat._vec) + + return geometry_dm + + +def dmplex_migrate(PETSc.DM dm, PETSc.SF sf) -> PETSc.DM: + cdef: + PETSc.DM migrated_dm + + migrated_dm = PETSc.DMPlex().create(comm=dm.comm) + CHKERR(DMPlexMigrate(dm.dm, sf.sf, migrated_dm.dm)) + return migrated_dm + + +# not used currently +def densify_sf(PETSc.DM topology_dm, PETSc.SF sparse_sf) -> PETSc.SF: + cdef: + PETSc.SF dense_sf + PetscInt dof_nroots, dof_nleaves + PetscInt *ilocal_new = NULL + PetscSFNode *iremote_new = NULL + PETSc.SF point_sf + PetscInt nroots, nleaves + const PetscInt *ilocal = NULL + const PetscSFNode *iremote = NULL + PETSc.Section local_sec + PetscInt pStart, pEnd, p, dof, off, m, n, i, j + np.ndarray local_offsets + np.ndarray remote_offsets + + pStart, pEnd = topology_dm.getChart() + nPoints = pEnd - pStart + CHKERR(PetscSFGetGraph(sparse_sf.sf, &nroots, &nleaves, &ilocal, &iremote)) + + CHKERR(PetscMalloc1(nPoints, &ilocal_new)) + CHKERR(PetscMalloc1(nPoints, &iremote_new)) + for p in range(pStart, pEnd): + ilocal_new[p] = p + iremote_new[p].rank = topology_dm.comm.rank + iremote_new[p].index = p + + for i in range(nleaves): + p = ilocal[i] if ilocal else i + iremote_new[p].rank = iremote[i].rank + iremote_new[p].index = iremote[i].index + + dense_sf = PETSc.SF().create(comm=sparse_sf.comm) + CHKERR(PetscSFSetGraph(dense_sf.sf, nPoints, nPoints, ilocal_new, PETSC_OWN_POINTER, iremote_new, PETSC_OWN_POINTER)) + return dense_sf + + +# not used currently +def dmplex_create_overlap_migration_sf(PETSc.DM topology_dm, PETSc.SF overlap_sf) -> PETSc.SF: + cdef: + PETSc.SF migration_sf + + migration_sf = PETSc.SF().create(comm=topology_dm.comm) + CHKERR(DMPlexCreateOverlapMigrationSF(topology_dm.dm, overlap_sf.sf, &migration_sf.sf)) + return migration_sf + + +# not used currently +def set_label(PETSc.DM dm, PETSc.DMLabel label) -> None: + CHKERR(DMSetLabel(dm.dm, label.dmlabel)) diff --git a/firedrake/cython/petschdr.pxi b/firedrake/cython/petschdr.pxi index 9a0bff609d..19b239b84d 100644 --- a/firedrake/cython/petschdr.pxi +++ b/firedrake/cython/petschdr.pxi @@ -56,6 +56,11 @@ cdef extern from "petscdmplex.h" nogil: int DMPlexGetSubpointMap(PETSc.PetscDM,PETSc.PetscDMLabel*) int DMPlexSetSubpointMap(PETSc.PetscDM,PETSc.PetscDMLabel) + int DMSetCoordinateSection(PETSc.PetscDM,PetscInt,PETSc.PetscSection) + int DMPlexMigrate(PETSc.PetscDM,PETSc.PetscSF,PETSc.PetscDM) + int DMPlexCreateOverlapMigrationSF(PETSc.PetscDM,PETSc.PetscSF,PETSc.PetscSF*) + + cdef extern from "petscdmlabel.h" nogil: struct _n_DMLabel ctypedef _n_DMLabel* DMLabel "DMLabel" @@ -72,6 +77,7 @@ cdef extern from "petscdmlabel.h" nogil: cdef extern from "petscdm.h" nogil: int DMCreateLabel(PETSc.PetscDM,char[]) int DMGetLabel(PETSc.PetscDM,char[],DMLabel*) + int DMSetLabel(PETSc.PetscDM,PETSc.PetscDMLabel) int DMGetPointSF(PETSc.PetscDM,PETSc.PetscSF*) int DMSetLabelValue(PETSc.PetscDM,char[],PetscInt,PetscInt) int DMGetLabelValue(PETSc.PetscDM,char[],PetscInt,PetscInt*) diff --git a/firedrake/functionspace.py b/firedrake/functionspace.py index 509f7a2863..0af27f6431 100644 --- a/firedrake/functionspace.py +++ b/firedrake/functionspace.py @@ -50,6 +50,7 @@ def make_scalar_element(mesh, family, degree, vfamily, vdegree, variant): the provided mesh, by calling :meth:`.AbstractMeshTopology.init` (or :meth:`.MeshGeometry.init`) as appropriate. """ + mesh.init() topology = mesh.topology cell = topology.ufl_cell() if isinstance(family, finat.ufl.FiniteElementBase): diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 8a0f8ec8a9..c813ca0496 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -543,17 +543,15 @@ def _from_cell_list(dim, cells, coords, comm, name=None): return plex -class AbstractMeshTopology(object, metaclass=abc.ABCMeta): +class AbstractMeshTopology(metaclass=abc.ABCMeta): """A representation of an abstract mesh topology without a concrete PETSc DM implementation""" - def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm): + def __init__(self, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm): """Initialise a mesh topology. Parameters ---------- - topology_dm : PETSc.DMPlex or PETSc.DMSwarm - `PETSc.DMPlex` or `PETSc.DMSwarm` representing the mesh topology. name : str Name of the mesh topology. reorder : bool @@ -577,9 +575,6 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, """ utils._init() - dmcommon.validate_mesh(topology_dm) - topology_dm.setFromOptions() - self.topology_dm = topology_dm r"The PETSc DM representation of the mesh topology." self.sfBC = None r"The PETSc SF that pushes the input (naive) plex to current (good) plex." @@ -591,19 +586,11 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, # Internal comm self._comm = internal_comm(self.user_comm, self) - dmcommon.label_facets(self.topology_dm) - self._distribute() self._grown_halos = False def callback(self): """Finish initialisation.""" del self._callback - if self.comm.size > 1: - self._add_overlap() - if self.sfXB is not None: - self.sfXC = sfXB.compose(self.sfBC) if self.sfBC else self.sfXB - dmcommon.label_facets(self.topology_dm) - dmcommon.complete_facet_labels(self.topology_dm) # TODO: Allow users to set distribution name if they want to save # conceptually the same mesh but with different distributions, # e.g., those generated by different partitioners. @@ -642,9 +629,6 @@ def callback(self): self._facet_ordering = dmcommon.get_facet_ordering(self.topology_dm, facet_numbering) self._callback = callback self.name = name - # Set/Generate names to be used when checkpointing. - self._distribution_name = distribution_name or _generate_default_mesh_topology_distribution_name(self.topology_dm.comm.size, self._distribution_parameters) - self._permutation_name = permutation_name or _generate_default_mesh_topology_permutation_name(reorder) # A cache of shared function space data on this mesh self._shared_data_cache = defaultdict(dict) # Cell subsets for integration over subregions @@ -657,22 +641,15 @@ def callback(self): # submesh self.submesh_parent = None + self._distribution_name = distribution_name + self._permutation_name = permutation_name + layers = None """No layers on unstructured mesh""" variable_layers = False """No variable layers on unstructured mesh""" - @abc.abstractmethod - def _distribute(self): - """Distribute the mesh toplogy.""" - pass - - @abc.abstractmethod - def _add_overlap(self): - """Add overlap.""" - pass - @abc.abstractmethod def _mark_entity_classes(self): """Mark entities with pyop2 classes.""" @@ -1112,13 +1089,11 @@ class MeshTopology(AbstractMeshTopology): """A representation of mesh topology implemented on a PETSc DMPlex.""" @PETSc.Log.EventDecorator("CreateMesh") - def __init__(self, plex, name, reorder, distribution_parameters, sfXB=None, perm_is=None, distribution_name=None, permutation_name=None, comm=COMM_WORLD): + def __init__(self, name, reorder, distribution_parameters, sfXB=None, perm_is=None, distribution_name=None, permutation_name=None, comm=COMM_WORLD): """Initialise a mesh topology. Parameters ---------- - plex : PETSc.DMPlex - `PETSc.DMPlex` representing the mesh topology. name : str Name of the mesh topology. reorder : bool @@ -1144,61 +1119,7 @@ def __init__(self, plex, name, reorder, distribution_parameters, sfXB=None, perm """ self._distribution_parameters = {} - distribute = distribution_parameters.get("partition") - if distribute is None: - distribute = True - self._distribution_parameters["partition"] = distribute - partitioner_type = distribution_parameters.get("partitioner_type") - self._set_partitioner(plex, distribute, partitioner_type) - self._distribution_parameters["partitioner_type"] = self._get_partitioner(plex).getType() - self._distribution_parameters["overlap_type"] = distribution_parameters.get("overlap_type", - (DistributedMeshOverlapType.FACET, 1)) - # Disable auto distribution and reordering before setFromOptions is called. - plex.distributeSetDefault(False) - plex.reorderSetDefault(PETSc.DMPlex.ReorderDefaultFlag.FALSE) - super().__init__(plex, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm) - - def _distribute(self): - # Distribute/redistribute the dm to all ranks - distribute = self._distribution_parameters["partition"] - if self.comm.size > 1 and distribute: - plex = self.topology_dm - # We distribute with overlap zero, in case we're going to - # refine this mesh in parallel. Later, when we actually use - # it, we grow the halo. - original_name = plex.getName() - sfBC = plex.distribute(overlap=0) - plex.setName(original_name) - self.sfBC = sfBC - # plex carries a new dm after distribute, which - # does not inherit partitioner from the old dm. - # It probably makes sense as chaco does not work - # once distributed. - - def _add_overlap(self): - overlap_type, overlap = self._distribution_parameters["overlap_type"] - if overlap < 0: - raise ValueError("Overlap depth must be >= 0") - if overlap_type == DistributedMeshOverlapType.NONE: - if overlap > 0: - raise ValueError("Can't have NONE overlap with overlap > 0") - elif overlap_type == DistributedMeshOverlapType.FACET: - dmcommon.set_adjacency_callback(self.topology_dm) - original_name = self.topology_dm.getName() - sfBC = self.topology_dm.distributeOverlap(overlap) - self.topology_dm.setName(original_name) - self.sfBC = self.sfBC.compose(sfBC) if self.sfBC else sfBC - dmcommon.clear_adjacency_callback(self.topology_dm) - self._grown_halos = True - elif overlap_type == DistributedMeshOverlapType.VERTEX: - # Default is FEM (vertex star) adjacency. - original_name = self.topology_dm.getName() - sfBC = self.topology_dm.distributeOverlap(overlap) - self.topology_dm.setName(original_name) - self.sfBC = self.sfBC.compose(sfBC) if self.sfBC else sfBC - self._grown_halos = True - else: - raise ValueError("Unknown overlap type %r" % overlap_type) + super().__init__(name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm) def _mark_entity_classes(self): dmcommon.mark_entity_classes(self.topology_dm) @@ -1404,52 +1325,6 @@ def cell_set(self): size = list(self._entity_classes[self.cell_dimension(), :]) return op2.Set(size, "Cells", comm=self._comm) - @PETSc.Log.EventDecorator() - def _set_partitioner(self, plex, distribute, partitioner_type=None): - """Set partitioner for (re)distributing underlying plex over comm. - - :arg distribute: Boolean or (sizes, points)-tuple. If (sizes, point)- - tuple is given, it is used to set shell partition. If Boolean, no-op. - :kwarg partitioner_type: Partitioner to be used: "chaco", "ptscotch", "parmetis", - "shell", or `None` (unspecified). Ignored if the distribute parameter - specifies the distribution. - """ - if plex.comm.size == 1 or distribute is False: - return - partitioner = plex.getPartitioner() - if distribute is True: - if partitioner_type: - if partitioner_type not in ["chaco", "ptscotch", "parmetis", "simple", "shell"]: - raise ValueError( - f"Unexpected partitioner_type: {partitioner_type}") - if partitioner_type in ["chaco", "ptscotch", "parmetis"] and \ - partitioner_type not in get_external_packages(): - raise ValueError( - f"Unable to use {partitioner_type} as PETSc is not " - f"installed with {partitioner_type}." - ) - if partitioner_type == "chaco" and plex.isDistributed(): - raise ValueError( - "Unable to use 'chaco' mesh partitioner, 'chaco' is a " - "serial partitioner, but the mesh is distributed." - ) - else: - partitioner_type = DEFAULT_PARTITIONER - - partitioner.setType({ - "chaco": partitioner.Type.CHACO, - "ptscotch": partitioner.Type.PTSCOTCH, - "parmetis": partitioner.Type.PARMETIS, - "shell": partitioner.Type.SHELL, - "simple": partitioner.Type.SIMPLE - }[partitioner_type]) - else: - sizes, points = distribute - partitioner.setType(partitioner.Type.SHELL) - partitioner.setShellPartition(plex.comm.size, sizes, points) - # Command line option `-petscpartitioner_type ` overrides. - # partitioner.setFromOptions() is called from inside plex.setFromOptions(). - @PETSc.Log.EventDecorator() def _get_partitioner(self, plex): """Get partitioner actually used for (re)distributing underlying plex over comm.""" @@ -1934,15 +1809,9 @@ def __init__(self, swarm, parentmesh, name, reorder, input_ordering_swarm=None, "partitioner_type": None, "overlap_type": (DistributedMeshOverlapType.NONE, 0)} self.input_ordering_swarm = input_ordering_swarm - super().__init__(swarm, name, reorder, None, perm_is, distribution_name, permutation_name, parentmesh.comm) + super().__init__(name, reorder, None, perm_is, distribution_name, permutation_name, parentmesh.comm) self._parent_mesh = parentmesh - def _distribute(self): - pass - - def _add_overlap(self): - pass - def _mark_entity_classes(self): if self.input_ordering_swarm: assert isinstance(self._parent_mesh, MeshTopology) @@ -2242,6 +2111,7 @@ def __new__(cls, element, comm): cargo = MeshGeometryCargo(uid) assert isinstance(element, finat.ufl.FiniteElementBase) ufl.Mesh.__init__(mesh, element, ufl_id=mesh.uid, cargo=cargo) + return mesh @MeshGeometryMixin._ad_annotate_init @@ -2253,9 +2123,19 @@ def __init__(self, coordinates): coordinates : CoordinatelessFunction The `CoordinatelessFunction` containing the coordinates. + Notes + ----- + At the time this constructor is called the mesh is fully distributed. + """ topology = coordinates.function_space().mesh() + if isinstance(topology, MeshTopology): + self.geometry_dm = dmcommon.make_geometry_dm(coordinates) + else: + assert isinstance(topology, VertexOnlyMeshTopology) + self.geometry_dm = topology.topology_dm + # this is codegen information so we attach it to the MeshGeometry rather than its cargo self.extruded = isinstance(topology, ExtrudedMeshTopology) self.variable_layers = self.extruded and topology.variable_layers @@ -2272,6 +2152,10 @@ def __init__(self, coordinates): self._spatial_index = None self._saved_coordinate_dat_version = coordinates.dat.dat_version + @property + def cell_closure_list(self): + return self.dm.getCoordinateDM().getAttr("cell_closure_list") + def _ufl_signature_data_(self, *args, **kwargs): return (type(self), self.extruded, self.variable_layers, super()._ufl_signature_data_(*args, **kwargs)) @@ -2284,7 +2168,7 @@ def init(self): if hasattr(self, '_callback'): self._callback(self) - def _init_topology(self, topology): + def _init_topology(self, topology, geometry_dm): """Initialise the topology. :arg topology: The :class:`.MeshTopology` object. @@ -2298,14 +2182,48 @@ def _init_topology(self, topology): import firedrake.function as function self._topology = topology + self._input_geometry_dm = geometry_dm def callback(self): """Finish initialisation.""" del self._callback - # Finish the initialisation of mesh topology + + if isinstance(self.topology, MeshTopology): + # Add overlap to the mesh + geometry_dm = self._input_geometry_dm + sfBC = None + grown_halos = False + if self.comm.size > 1: + overlap_type, overlap = self.topology._distribution_parameters["overlap_type"] + sfBC, grown_halos = _add_overlap_to_dm(geometry_dm, overlap_type, overlap) + self.topology._grown_halos = grown_halos + + if sfBC: + topology.sfBC = topology.sfBC.compose(sfBC) if topology.sfBC else sfBC + if topology.sfXB is not None: + topology.sfXC = topology.sfXB.compose(topology.sfBC) if topology.sfBC else topology.sfXB + + # Finish the initialisation of mesh topology + topology_dm = geometry_dm.clone() + # effectively 'topology_dm.setCoordinates(NULL)' + topology_dm.setCoordinates(geometry_dm.getCoordinateDM().getCoordinates()) + else: + assert isinstance(self.topology, VertexOnlyMeshTopology) + geometry_dm = self._input_geometry_dm + topology_dm = self._input_geometry_dm + + self.topology.topology_dm = topology_dm self.topology.init() + coordinates_fs = functionspace.FunctionSpace(self.topology, self.ufl_coordinate_element()) - coordinates_data = dmcommon.reordered_coords(topology.topology_dm, coordinates_fs.dm.getDefaultSection(), + + # reordered_coords only needs to reorder the coordinates + # *some of the time*. If the input is a DM that has already + # been reordered to respect Firedrake then we don't have to do + # anything. Otherwise, we have the set the correct section. + # TODO: Ask Matt if there is a nice way to do this. + + coordinates_data = dmcommon.reordered_coords(geometry_dm, coordinates_fs.dm.getDefaultSection(), (self.num_vertices(), self.geometric_dimension())) coordinates = function.CoordinatelessFunction(coordinates_fs, val=coordinates_data, @@ -2733,7 +2651,7 @@ def input_ordering(self): """ if not isinstance(self.topology, VertexOnlyMeshTopology): raise AttributeError("Input ordering is only defined for vertex-only meshes.") - _input_ordering = make_vom_from_vom_topology(self.topology.input_ordering, self.name + "_input_ordering") + _input_ordering = make_vom_from_vom_topology(self.topology.input_ordering, self.geometry_dm, self.name + "_input_ordering") if _input_ordering: _input_ordering._parent_mesh = self _input_ordering.init() @@ -2868,8 +2786,9 @@ def make_mesh_from_coordinates(coordinates, name, tolerance=0.5): return mesh -def make_mesh_from_mesh_topology(topology, name, tolerance=0.5): - """Make mesh from tpology. +# Should be make_mesh_from_geometry_dm +def make_mesh_from_mesh_topology(topology, geometry_dm, name, tolerance=0.5): + """Make mesh from topology. Parameters ---------- @@ -2886,23 +2805,36 @@ def make_mesh_from_mesh_topology(topology, name, tolerance=0.5): The mesh. """ + if topology.comm.size > 1: + assert geometry_dm.isDistributed() + # Construct coordinate element - # TODO: meshfile might indicates higher-order coordinate element - cell = topology.ufl_cell() - geometric_dim = topology.topology_dm.getCoordinateDim() - if not topology.topology_dm.getCoordinatesLocalized(): + tdim = geometry_dm.getDimension() + # Allow empty local meshes on a process + cStart, cEnd = geometry_dm.getHeightStratum(0) # cells + if cStart == cEnd: + nfacets = -1 + else: + nfacets = geometry_dm.getConeSize(cStart) + + # TODO: this needs to be updated for mixed-cell meshes. + nfacets = topology._comm.allreduce(nfacets, op=MPI.MAX) + + cell = ufl.Cell(_cells[tdim][nfacets]) + geometric_dim = geometry_dm.getCoordinateDim() + if not geometry_dm.getCoordinatesLocalized(): element = finat.ufl.VectorElement("Lagrange", cell, 1, dim=geometric_dim) else: element = finat.ufl.VectorElement("DQ" if cell in [ufl.quadrilateral, ufl.hexahedron] else "DG", cell, 1, dim=geometric_dim, variant="equispaced") # Create mesh object mesh = MeshGeometry.__new__(MeshGeometry, element, topology.comm) - mesh._init_topology(topology) + mesh._init_topology(topology, geometry_dm) mesh.name = name mesh._tolerance = tolerance return mesh -def make_vom_from_vom_topology(topology, name, tolerance=0.5): +def make_vom_from_vom_topology(topology, dm, name, tolerance=0.5): """Make `VertexOnlyMesh` from a mesh topology. Parameters @@ -2924,16 +2856,16 @@ def make_vom_from_vom_topology(topology, name, tolerance=0.5): import firedrake.functionspace as functionspace import firedrake.function as function - gdim = topology.topology_dm.getCoordinateDim() + gdim = dm.getCoordinateDim() cell = topology.ufl_cell() element = finat.ufl.VectorElement("DG", cell, 0, dim=gdim) vmesh = MeshGeometry.__new__(MeshGeometry, element, topology.comm) - vmesh._init_topology(topology) + vmesh._init_topology(topology, dm) # Save vertex reference coordinate (within reference cell) in function parent_tdim = topology._parent_mesh.ufl_cell().topological_dimension() if parent_tdim > 0: reference_coordinates_fs = functionspace.VectorFunctionSpace(topology, "DG", 0, dim=parent_tdim) - reference_coordinates_data = dmcommon.reordered_coords(topology.topology_dm, reference_coordinates_fs.dm.getDefaultSection(), + reference_coordinates_data = dmcommon.reordered_coords(dm, reference_coordinates_fs.dm.getDefaultSection(), (topology.num_vertices(), parent_tdim), reference_coord=True) reference_coordinates = function.CoordinatelessFunction(reference_coordinates_fs, @@ -3094,17 +3026,65 @@ def Mesh(meshfile, **kwargs): raise RuntimeError("Mesh file %s has unknown format '%s'." % (meshfile, ext[1:])) plex.setName(_generate_default_mesh_topology_name(name)) - # Create mesh topology - topology = MeshTopology(plex, name=plex.getName(), reorder=reorder, + + # Create mesh topology - note that this is in a poorly defined + # state until the mesh is finalised + topology = MeshTopology(name=plex.getName(), reorder=reorder, distribution_parameters=distribution_parameters, distribution_name=kwargs.get("distribution_name"), permutation_name=kwargs.get("permutation_name"), comm=user_comm) + + # Distribute the plex + # from MeshTopology.__init__ + distribute = distribution_parameters.get("partition") + if distribute is None: # NOTE: Use a sane default in .get? + distribute = True + topology._distribution_parameters["partition"] = distribute + topology._distribution_parameters["overlap_type"] = \ + distribution_parameters.get("overlap_type", (DistributedMeshOverlapType.FACET, 1)) + partitioner_type = distribution_parameters.get("partitioner_type") + _set_dmplex_partitioner(plex, distribute, partitioner_type) + topology._distribution_parameters["partitioner_type"] = plex.getPartitioner().getType() + + # Disable auto distribution and reordering before setFromOptions is called. + plex.distributeSetDefault(False) + plex.reorderSetDefault(PETSc.DMPlex.ReorderDefaultFlag.FALSE) + + # from AbstractMeshTopology.__init__ + dmcommon.validate_mesh(plex) + plex.setFromOptions() + dmcommon.label_facets(plex) + + # Set/Generate names to be used when checkpointing. + topology._distribution_name = topology._distribution_name or _generate_default_mesh_topology_distribution_name(user_comm.size, topology._distribution_parameters) + topology._permutation_name = topology._permutation_name or _generate_default_mesh_topology_permutation_name(reorder) + + if user_comm.size > 1 and distribution_parameters["partition"]: + # We distribute with overlap zero, in case we're going to + # refine this mesh in parallel. Later, when we actually use + # it, we grow the halo. + original_name = plex.getName() + sfBC = plex.distribute(overlap=0) + plex.setName(original_name) + # NOTE: plex carries a new dm after distribute, which + # does not inherit partitioner from the old dm. + # It probably makes sense as chaco does not work + # once distributed. + topology.sfBC = sfBC + + # from callback + dmcommon.label_facets(plex) + dmcommon.complete_facet_labels(plex) + if netgen and isinstance(meshfile, netgen.libngpy._meshing.Mesh): + # for compat with netgen + topology.topology_dm = plex + netgen_firedrake_mesh.createFromTopology(topology, name=name, comm=user_comm) mesh = netgen_firedrake_mesh.firedrakeMesh else: - mesh = make_mesh_from_mesh_topology(topology, name) + mesh = make_mesh_from_mesh_topology(topology, plex, name) mesh._tolerance = tolerance return mesh @@ -3396,7 +3376,8 @@ def VertexOnlyMesh(mesh, vertexcoords, reorder=None, missing_points_behaviour='e reorder=reorder, input_ordering_swarm=input_ordering_swarm, ) - vmesh_out = make_vom_from_vom_topology(topology, name, tolerance) + topology.topology_dm = swarm + vmesh_out = make_vom_from_vom_topology(topology, swarm, name, tolerance) vmesh_out._parent_mesh = mesh vmesh_out.init() return vmesh_out @@ -4528,7 +4509,7 @@ def RelabeledMesh(mesh, indicator_functions, subdomain_ids, **kwargs): if not isinstance(subid, numbers.Integral): raise TypeError(f"subdomain id must be an integer: got {subid}") name1 = kwargs.get("name", DEFAULT_MESH_NAME) - plex = tmesh.topology_dm + plex = mesh.geometry_dm # Clone plex: plex1 will share topology with plex. plex1 = plex.clone() plex1.setName(_generate_default_mesh_topology_name(name1)) @@ -4565,13 +4546,13 @@ def RelabeledMesh(mesh, indicator_functions, subdomain_ids, **kwargs): distribution_parameters_noop = {"partition": False, "overlap_type": (DistributedMeshOverlapType.NONE, 0)} reorder_noop = None - tmesh1 = MeshTopology(plex1, name=plex1.getName(), reorder=reorder_noop, + tmesh1 = MeshTopology(name=plex1.getName(), reorder=reorder_noop, distribution_parameters=distribution_parameters_noop, perm_is=tmesh._dm_renumbering, distribution_name=tmesh._distribution_name, permutation_name=tmesh._permutation_name, comm=tmesh.comm) - return make_mesh_from_mesh_topology(tmesh1, name1) + return make_mesh_from_mesh_topology(tmesh1, plex1, name1) @PETSc.Log.EventDecorator() @@ -4679,3 +4660,75 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None): submesh.submesh_parent = mesh submesh.init() return submesh + + +def _set_dmplex_partitioner(plex, distribute, partitioner_type=None): + """Set partitioner for (re)distributing underlying plex over comm. + + :arg distribute: Boolean or (sizes, points)-tuple. If (sizes, point)- + tuple is given, it is used to set shell partition. If Boolean, no-op. + :kwarg partitioner_type: Partitioner to be used: "chaco", "ptscotch", "parmetis", + "shell", or `None` (unspecified). Ignored if the distribute parameter + specifies the distribution. + """ + if plex.comm.size == 1 or distribute is False: + return + partitioner = plex.getPartitioner() + if distribute is True: + if partitioner_type: + if partitioner_type not in ["chaco", "ptscotch", "parmetis", "simple", "shell"]: + raise ValueError( + f"Unexpected partitioner_type: {partitioner_type}") + if partitioner_type in ["chaco", "ptscotch", "parmetis"] and \ + partitioner_type not in get_external_packages(): + raise ValueError( + f"Unable to use {partitioner_type} as PETSc is not " + f"installed with {partitioner_type}." + ) + if partitioner_type == "chaco" and plex.isDistributed(): + raise ValueError( + "Unable to use 'chaco' mesh partitioner, 'chaco' is a " + "serial partitioner, but the mesh is distributed." + ) + else: + partitioner_type = DEFAULT_PARTITIONER + + partitioner.setType({ + "chaco": partitioner.Type.CHACO, + "ptscotch": partitioner.Type.PTSCOTCH, + "parmetis": partitioner.Type.PARMETIS, + "shell": partitioner.Type.SHELL, + "simple": partitioner.Type.SIMPLE + }[partitioner_type]) + else: + sizes, points = distribute + partitioner.setType(partitioner.Type.SHELL) + partitioner.setShellPartition(plex.comm.size, sizes, points) + # Command line option `-petscpartitioner_type ` overrides. + # partitioner.setFromOptions() is called from inside plex.setFromOptions(). + + +def _add_overlap_to_dm(dm, overlap_type, overlap): + sfBC = None + grown_halos = False + if overlap < 0: + raise ValueError("Overlap depth must be >= 0") + if overlap_type == DistributedMeshOverlapType.NONE: + if overlap > 0: + raise ValueError("Can't have NONE overlap with overlap > 0") + elif overlap_type == DistributedMeshOverlapType.FACET: + dmcommon.set_adjacency_callback(dm) + original_name = dm.getName() + sfBC = dm.distributeOverlap(overlap) + dm.setName(original_name) + dmcommon.clear_adjacency_callback(dm) + grown_halos = True + elif overlap_type == DistributedMeshOverlapType.VERTEX: + # Default is FEM (vertex star) adjacency. + original_name = dm.getName() + sfBC = dm.distributeOverlap(overlap) + dm.setName(original_name) + grown_halos = True + else: + raise ValueError(f"Unknown overlap type {overlap_type}") + return sfBC, grown_halos diff --git a/firedrake/mg/mesh.py b/firedrake/mg/mesh.py index 42ea985666..045e7db1ff 100644 --- a/firedrake/mg/mesh.py +++ b/firedrake/mg/mesh.py @@ -131,7 +131,7 @@ def MeshHierarchy(mesh, refinement_levels, raise RuntimeError("Cannot create a NetgenHierarchy from a mesh that has not been generated by\ Netgen.") - cdm = mesh.topology_dm + cdm = mesh._input_geometry_dm cdm.setRefinementUniform(True) dms = [] if mesh.comm.size > 1 and mesh._grown_halos: @@ -180,10 +180,10 @@ def MeshHierarchy(mesh, refinement_levels, lgmaps = [] for i, m in enumerate(meshes): - no = impl.create_lgmap(m.topology_dm) + no = impl.create_lgmap(m._input_geometry_dm) m.init() - o = impl.create_lgmap(m.topology_dm) - m.topology_dm.setRefineLevel(i) + o = impl.create_lgmap(m._input_geometry_dm) + m._input_geometry_dm.setRefineLevel(i) lgmaps.append((no, o)) coarse_to_fine_cells = [] diff --git a/firedrake/utility_meshes.py b/firedrake/utility_meshes.py index 99d687449e..bd9eff7f82 100644 --- a/firedrake/utility_meshes.py +++ b/firedrake/utility_meshes.py @@ -80,7 +80,7 @@ def _postprocess_periodic_mesh(coords, comm, distribution_parameters, reorder, name, distribution_name, permutation_name): - dm = coords.function_space().mesh().topology.topology_dm + dm = coords.function_space().mesh().geometry_dm dm.removeLabel("pyop2_core") dm.removeLabel("pyop2_owned") dm.removeLabel("pyop2_ghost")