Make serial submesh construction linear in the number of points

firedrake/cython/dmcommon.pyx

@@ -3988,7 +3988,7 @@ def create_halo_exchange_sf(PETSc.DM dm):
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def submesh_create(PETSc.DM dm,
-                   PetscInt subdim,
+                   subdim,
                    label_name,
                    subdomain_id,
                    PetscBool ignore_label_halo,
@@ -3999,12 +3999,13 @@ def submesh_create(PETSc.DM dm,
     ----------
     dm : PETSc.DM
         DMPlex representing the mesh topology
-    subdim : int
-        Topological dimension of the submesh
-    label_name : str
-        Name of the label
-    subdomain_id : int | Sequence
-        Values in the label
+    subdim : int | None
+        Topological dimension of the submesh, or None to be inferred from other kwargs.
+        See :func:`~.mesh.Submesh`.
+    label_name : str | None
+        Name of the label, or `None` to select every cell
+    subdomain_id : int | Sequence | None
+        Values in the label, unused if ``label_name`` is `None`
     ignore_label_halo : bool
         If labeled points in the halo are ignored.
     comm : PETSc.Comm | None
@@ -4018,6 +4019,36 @@ def submesh_create(PETSc.DM dm,
         PetscInt pStart, pEnd, p, i, stratum_size = 0, label_value = 1
         const PetscInt *stratum_indices = NULL
 
+    # Parse default subdim, label_name, and subdomain_id
+    dim = dm.getDimension()
+    if subdomain_id == "on_boundary":
+        if subdim is None:
+            subdim = dim - 1
+        elif subdim != dim - 1:
+            raise ValueError('subdomain_id="on_boundary" requires subdim=dim-1')
+        if label_name is None:
+            label_name = "exterior_facets"
+        elif label_name != "exterior_facets":
+            raise ValueError('subdomain_id="on_boundary" requires label_name="exterior_facets"')
+        subdomain_id = 1
+
+    if subdim is None:
+        subdim = dim
+    if subdim not in {dim, dim - 1}:
+        raise NotImplementedError(f"Submesh construction is only implemented for codimension 0 or 1. "
+                                   "Found submesh dim ({subdim}) and parent dim ({dim})")
+    if subdomain_id is None:
+        if label_name is not None:
+            raise ValueError("subdomain_id=None requires label_name=None.")
+        # Select all entities
+        label_name = "depth"
+        subdomain_id = subdim
+    elif label_name is None:
+        if subdim == dim:
+            label_name = CELL_SETS_LABEL
+        elif subdim == dim - 1:
+            label_name = FACE_SETS_LABEL
+
     # Cast subdomain_id into an iterable
     if isinstance(subdomain_id, str) or not isinstance(subdomain_id, Sequence):
         subdomain_id = (subdomain_id,)
@@ -4056,6 +4087,8 @@ def submesh_create(PETSc.DM dm,
                                              ignoreHalo=ignore_label_halo,
                                              sanitizeSubMesh=PETSC_TRUE,
                                              comm=comm)
+    if subdm.getDimension() != subdim:
+        raise RuntimeError(f"Found subplex dim ({subdm.getDimension()}) != expected ({subdim})")
     # Destroy temp_label.
     dm.removeLabel(temp_label_name)
     subdm.removeLabel(temp_label_name)
@@ -4087,6 +4120,7 @@ def submesh_correct_entity_classes(PETSc.DM dm,
         const PetscInt *ilocal = NULL
         const PetscSFNode *iremote = NULL
         PETSc.IS subpoint_is
+        PETSc.IS all_points
         const PetscInt *subpoint_indices = NULL
         np.ndarray ownership_loss
         np.ndarray ownership_gain
@@ -4109,17 +4143,13 @@ def submesh_correct_entity_classes(PETSc.DM dm,
 
     if subdm.comm.size == 1:
         # Undistributed case: relabel every point as core
-        for subp in range(subpStart, subpEnd):
-            CHKERR(DMLabelHasPoint(lbl_core, subp, &has))
-            if has:
-                continue
-            CHKERR(DMLabelHasPoint(lbl_ghost, subp, &has))
-            if has:
-                CHKERR(DMLabelClearValue(lbl_ghost, subp, 1))
-            CHKERR(DMLabelHasPoint(lbl_owned, subp, &has))
-            if has:
-                CHKERR(DMLabelClearValue(lbl_owned, subp, 1))
-            CHKERR(DMLabelSetValue(lbl_core, subp, 1))
+        all_points = PETSc.IS().createStride(subpEnd - subpStart,
+                                             first=subpStart, step=1,
+                                             comm=PETSc.COMM_SELF)
+        CHKERR(DMLabelClearStratum(lbl_owned, 1))
+        CHKERR(DMLabelClearStratum(lbl_ghost, 1))
+        CHKERR(DMLabelSetStratumIS(lbl_core, 1, (<PETSc.IS>all_points).iset))
+        all_points.destroy()
     else:
         ownership_loss = np.zeros(pEnd - pStart, dtype=IntType)
         ownership_gain = np.zeros(pEnd - pStart, dtype=IntType)

firedrake/cython/petschdr.pxi

@@ -95,6 +95,8 @@ cdef extern from "petscdmlabel.h" nogil:
     PetscErrorCode DMLabelClearValue(DMLabel, PetscInt, PetscInt)
     PetscErrorCode DMLabelGetStratumSize(DMLabel, PetscInt, PetscInt*)
     PetscErrorCode DMLabelGetStratumIS(DMLabel, PetscInt, PETSc.PetscIS*)
+    PetscErrorCode DMLabelSetStratumIS(DMLabel, PetscInt, PETSc.PetscIS)
+    PetscErrorCode DMLabelClearStratum(DMLabel, PetscInt)
 
 cdef extern from "petscdm.h" nogil:
     PetscErrorCode DMCreateLabel(PETSc.PetscDM,char[])

firedrake/mesh.py

@@ -5003,41 +5003,12 @@ def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ig
     elif isinstance(mesh.topology, VertexOnlyMeshTopology):
         raise NotImplementedError("Can not create a submesh of a ``VertexOnlyMesh``")
 
-    if subdomain_id == "on_boundary":
-        if subdim is None:
-            subdim = mesh.topological_dimension - 1
-        elif subdim != mesh.topological_dimension - 1:
-            raise ValueError('subdomain_id="on_boundary" requires subdim=dim-1')
-        if label_name is None:
-            label_name = "exterior_facets"
-        elif label_name != "exterior_facets":
-            raise ValueError('subdomain_id="on_boundary" requires label_name="exterior_facets"')
-        subdomain_id = 1
-
-    if subdim is None:
-        subdim = mesh.topological_dimension
-    plex = mesh.topology_dm
-    dim = plex.getDimension()
-    if subdim not in {dim, dim - 1}:
-        raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})")
-    if subdomain_id is None:
-        if label_name is not None:
-            raise ValueError("subdomain_id=None requires label_name=None.")
-        # Select all entities
-        label_name = "depth"
-        subdomain_id = subdim
-    elif label_name is None:
-        if subdim == dim:
-            label_name = dmcommon.CELL_SETS_LABEL
-        elif subdim == dim - 1:
-            label_name = dmcommon.FACE_SETS_LABEL
-    subplex = dmcommon.submesh_create(plex, subdim, label_name, subdomain_id, ignore_halo, comm=comm)
+    subplex = dmcommon.submesh_create(mesh.topology_dm, subdim, label_name, subdomain_id, ignore_halo, comm=comm)
 
     comm = comm or mesh.comm
     name = name or _generate_default_submesh_name(mesh.name)
     subplex.setName(_generate_default_mesh_topology_name(name))
     if subplex.getDimension() != subdim:
-        raise RuntimeError(f"Found subplex dim ({subplex.getDimension()}) != expected ({subdim})")
     if reorder is None:
         # Ideally we should set perm_is = mesh._dm_renumbering[label_indices]
         reorder = mesh._did_reordering

tests/firedrake/submesh/test_submesh_comm.py

@@ -44,6 +44,27 @@ def test_create_submesh_comm_self(reorder, ignore_halo):
     assert np.allclose(submesh.coordinates.dat.data_ro, x.dat.data_ro)
 
 
+@pytest.mark.parallel([1, 3])
+@pytest.mark.parametrize("ignore_halo", [False, True])
+def test_submesh_comm_self_entity_classes(ignore_halo):
+    """A submesh on COMM_SELF must own every point that it holds.
+
+    The parent mesh divides its points into the pyop2 classes core, owned and
+    ghost. A submesh on COMM_SELF has no neighbour, so it has no ghost points
+    and no owned points either. Every point is core.
+    """
+    mesh = UnitSquareMesh(
+        8, 8, distribution_parameters={
+            "overlap_type": (DistributedMeshOverlapType.VERTEX, 1)})
+    submesh = Submesh(mesh, ignore_halo=ignore_halo, comm=COMM_SELF)
+
+    plex = submesh.topology_dm
+    pStart, pEnd = plex.getChart()
+    assert plex.getStratumSize("pyop2_core", 1) == pEnd - pStart
+    assert plex.getStratumSize("pyop2_owned", 1) == 0
+    assert plex.getStratumSize("pyop2_ghost", 1) == 0
+
+
 @pytest.mark.parallel([1, 3])
 @pytest.mark.parametrize("family,degree", [("DG", 0), ("CG", 1)])
 @pytest.mark.parametrize("reorder", [False, True])