Make serial submesh construction linear in the number of points

firedrake/cython/dmcommon.pyx

@@ -4001,10 +4001,10 @@ def submesh_create(PETSc.DM 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
+    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 +4018,10 @@ def submesh_create(PETSc.DM dm,
         PetscInt pStart, pEnd, p, i, stratum_size = 0, label_value = 1
         const PetscInt *stratum_indices = NULL
 
+    if label_name is None:
+        # Default to all entities of the given dimension.
+        label_name = "depth"
+        subdomain_id = subdim
     # Cast subdomain_id into an iterable
     if isinstance(subdomain_id, str) or not isinstance(subdomain_id, Sequence):
         subdomain_id = (subdomain_id,)
@@ -4087,6 +4091,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 +4114,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

@@ -5023,9 +5023,6 @@ def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ig
     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

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])