Stop DG injection reading past a coarse cell's children

firedrake/mg/interface.py

@@ -281,7 +281,8 @@ def inject(fine, coarse):
                          coarse.dat(op2.INC, coarse.cell_node_map()),
                          fine.dat(op2.READ, compose_map(fine)),
                          fine_coords.dat(op2.READ, compose_map(fine_coords)),
-                         coarse_coords.dat(op2.READ, coarse_coords.cell_node_map()))
+                         coarse_coords.dat(op2.READ, coarse_coords.cell_node_map()),
+                         utils.coarse_cell_child_count(Vc, Vf)(op2.READ))
 
         if needs_quadrature:
             # Transfer to the actual target space

firedrake/mg/kernels.py

@@ -421,7 +421,9 @@ def dg_injection_kernel(Vf, Vc, ncell):
     from firedrake.slate.slac import compile_expression
     if complex_mode:
         raise NotImplementedError("In complex mode we are waiting for Slate")
-    macro_builder = MacroKernelBuilder(ScalarType, ncell)
+    # The kernel integrates over one micro-cell per call. The outer kernel
+    # below calls it once for each real child of the coarse cell.
+    macro_builder = MacroKernelBuilder(ScalarType, 1)
     macro_builder._domain_integral_type_map = {Vf.mesh(): "cell"}
     macro_builder._entity_ids = {Vf.mesh(): (0,)}
     f = ufl.Coefficient(Vf)
@@ -586,7 +588,7 @@ def name_multiindex(multiindex, name):
         lp.TemporaryVariable(local_tensor.name, shape=local_tensor.shape, dtype=local_tensor.dtype))
     depends_on |= {"zero"}
 
-    # 2. Fill the local tensor
+    # 2. Fill the local tensor, one micro-cell at a time
     macro_coordinates_arg = macro_builder.generate_arg_from_expression(
         macro_builder.coefficient_map[macro_builder.domain_coordinate[Vf.mesh()]])
     coarse_coordinates_arg = coarse_builder.generate_arg_from_expression(
@@ -604,9 +606,26 @@ def name_multiindex(multiindex, name):
         ScalarType, kernel_name="pyop2_kernel_evaluate", index_names=index_names)
     subkernels.append(eval_kernel)
 
+    # The macro arguments arrive holding every child slot of the coarse cell,
+    # back to back. The callee takes one slot, so each call gets the slice
+    # that starts at this child.
+    macro_args = [*macro_builder.kernel_args, macro_coordinates_arg]
+    macro_names = {arg.name for arg in macro_args}
+    entity = pym.var("entity")
+    offsets = [entity * arg.shape[0] if arg.name in macro_names else None
+               for arg in eval_args]
+
+    # Stop at this coarse cell's own children, so the kernel never reads the
+    # padded slots of its coarse_cell_to_fine_node_map row.
+    nchild_arg = lp.GlobalArg("nchild", dtype=IntType, shape=(1,))
+    domains.append(f"{{ [entity]: 0 <= entity < {ncell} }}")
     fill_insn, extra_domains = _generate_call_insn(
         "pyop2_kernel_evaluate", eval_args, iname_prefix="fill", id="fill",
-        depends_on=depends_on, within_inames_is_final=True)
+        offsets=offsets, depends_on=depends_on,
+        within_inames=frozenset({"entity"}), within_inames_is_final=True,
+        predicates=frozenset({
+            pym.primitives.Comparison(
+                entity, "<", pym.subscript(pym.var(nchild_arg.name), (0,)))}))
     instructions.append(fill_insn)
     domains.extend(extra_domains)
     depends_on |= {fill_insn.id}
@@ -616,9 +635,14 @@ def name_multiindex(multiindex, name):
     retarg = lp.GlobalArg(
         "R", dtype=ScalarType, shape=local_tensor.shape, is_output=True)
 
+    # The caller holds every child slot, so its macro arguments are ncell
+    # times as long as the ones the callee takes.
+    outer_macro_args = [
+        lp.GlobalArg(arg.name, dtype=arg.dtype, shape=(arg.shape[0] * ncell,))
+        for arg in macro_args]
     kernel_data = [
-        retarg, *macro_builder.kernel_args, macro_coordinates_arg,
-        coarse_coordinates_arg, *kernel_data]
+        retarg, *outer_macro_args, coarse_coordinates_arg, nchild_arg,
+        *kernel_data]
 
     u = TrialFunction(Vc)
     v = TestFunction(Vc)
@@ -645,7 +669,7 @@ def name_multiindex(multiindex, name):
         headers=Ainv.headers, events=Ainv.events)
 
 
-def _generate_call_insn(name, args, *, iname_prefix=None, **kwargs):
+def _generate_call_insn(name, args, *, iname_prefix=None, offsets=None, **kwargs):
     """Create an appropriate loopy call instruction from its arguments.
 
     This function is useful because :class:`loopy.CallInstruction` are a
@@ -661,6 +685,11 @@ def _generate_call_insn(name, args, *, iname_prefix=None, **kwargs):
         vector shape.
     iname_prefix : str, optional
         Prefix to the autogenerated inames, defaults to ``name``.
+    offsets : iterable of pymbolic.primitives.Expression, optional
+        One offset per argument, or `None` for no offset. The call passes the
+        slice of that argument which starts at the offset and is as long as
+        the callee expects. Use this to hand a callee one block of a caller
+        array that holds several.
     kwargs
         All other keyword arguments are passed to the
         :class:`loopy.CallInstruction` constructor.
@@ -675,12 +704,14 @@ def _generate_call_insn(name, args, *, iname_prefix=None, **kwargs):
     """
     if not iname_prefix:
         iname_prefix = name
+    if offsets is None:
+        offsets = (None,) * len(args)
 
     domains = []
     assignees = []
     parameters = []
     swept_iname_counter = 0
-    for arg in args:
+    for arg, offset in zip(args, offsets):
         try:
             shape, = arg.shape
         except ValueError:
@@ -690,8 +721,12 @@ def _generate_call_insn(name, args, *, iname_prefix=None, **kwargs):
         swept_iname_counter += 1
         domains.append(f"{{ [{swept_iname}]: 0 <= {swept_iname} < {shape} }}")
         swept_index = (pym.var(swept_iname),)
+        if offset is None:
+            outer_index = swept_index
+        else:
+            outer_index = (offset + pym.var(swept_iname),)
         param = lp.symbolic.SubArrayRef(
-            swept_index, pym.subscript(pym.var(arg.name), swept_index))
+            swept_index, pym.subscript(pym.var(arg.name), outer_index))
         parameters.append(param)
         if arg.is_output:
             assignees.append(param)

firedrake/mg/mesh.py

@@ -42,7 +42,9 @@ class HierarchyBase(object):
         List of meshes (coarse to fine).
     coarse_to_fine_cells :
         List of numpy arrays for each level pair, mapping each coarse cell
-        into fine cells it intersects.
+        into fine cells it intersects. Every row is as wide as the busiest
+        coarse cell's count, so a coarse cell with fewer fine cells has its
+        row right-padded with -1.
     fine_to_coarse_cells :
         List of numpy arrays for each level pair, mapping each fine cell into
         coarse cells it intersects.

firedrake/mg/utils.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import numpy
 from fractions import Fraction
 from pyop2 import op2
@@ -73,27 +75,20 @@ def coarse_node_to_fine_node_map(Vc, Vf):
 
         coarse_to_fine = hierarchy.coarse_to_fine_cells[levelc]
         coarse_to_fine_nodes = impl.coarse_to_fine_nodes(Vc, Vf, coarse_to_fine)
-        # Under adaptive refinement, coarse cells have varying numbers of
-        # fine descendants, so coarse_to_fine (and hence coarse_to_fine_nodes)
-        # is right-padded with -1 up to the busiest coarse cell's count.
-        # op2.Map cannot hold negative indices, and every *owned* coarse
-        # node needs at least one real candidate to inject from; but padding
-        # slots on rows that do have candidates can safely be filled with a
-        # duplicate of one of that row's real entries; the injection kernel
-        # below only ever reads (op2.READ) through this map and picks the
-        # candidate matching the coarse node's physical location, so a
-        # repeated valid entry is just redundantly (harmlessly) considered.
+        # op2.Map cannot hold the -1 that pads a short row, so fill each
+        # padded slot with a real entry from its own row. The injection
+        # kernel picks the candidate that matches the coarse node's physical
+        # location, so a repeated entry changes nothing.
         valid = coarse_to_fine_nodes >= 0
-        if not valid.all():
-            nonempty = valid.any(axis=1)
-            if not nonempty[:Vc.node_set.size].all():
-                raise RuntimeError("Adaptive coarse-to-fine map has empty node candidates")
-            replacement = numpy.zeros(coarse_to_fine_nodes.shape[0],
-                                      dtype=coarse_to_fine_nodes.dtype)
-            rows = numpy.nonzero(nonempty)[0]
-            replacement[rows] = coarse_to_fine_nodes[rows, valid[rows].argmax(axis=1)]
-            coarse_to_fine_nodes = numpy.where(valid, coarse_to_fine_nodes,
-                                               replacement[:, None])
+        nonempty = valid.any(axis=1)
+        if not nonempty[:Vc.node_set.size].all():
+            raise RuntimeError("Adaptive coarse-to-fine map has empty node candidates")
+        replacement = numpy.zeros(coarse_to_fine_nodes.shape[0],
+                                  dtype=coarse_to_fine_nodes.dtype)
+        rows = numpy.nonzero(nonempty)[0]
+        replacement[rows] = coarse_to_fine_nodes[rows, valid[rows].argmax(axis=1)]
+        coarse_to_fine_nodes = numpy.where(valid, coarse_to_fine_nodes,
+                                           replacement[:, None])
         return cache.setdefault(key, op2.Map(Vc.node_set, Vf.node_set,
                                              coarse_to_fine_nodes.shape[1],
                                              values=coarse_to_fine_nodes))
@@ -155,6 +150,66 @@ def coarse_cell_to_fine_node_map(Vc, Vf):
                                              offset=offset))
 
 
+def coarse_cell_child_count(
+    Vc: firedrake.functionspaceimpl.WithGeometry,
+    Vf: firedrake.functionspaceimpl.WithGeometry,
+) -> op2.Dat:
+    """Count the fine cells that each coarse cell was refined into.
+
+    A row of `HierarchyBase.coarse_to_fine_cells` is as wide as the busiest
+    coarse cell's count, so its width overstates how many children most cells
+    have. The DG injection kernel reads this count to stop at a coarse cell's
+    own children, and so leaves the padding alone.
+
+    Parameters
+    ----------
+    Vc : firedrake.functionspaceimpl.WithGeometry
+        The coarse function space.
+    Vf : firedrake.functionspaceimpl.WithGeometry
+        The fine function space, on the next level of the same hierarchy.
+
+    Returns
+    -------
+    pyop2.types.dat.Dat
+        One count per cell of ``Vc``'s mesh, over that mesh's cell set. Halo
+        cells are left at zero: a par_loop visits the core and owned parts
+        only, so the kernel never reads them.
+
+    """
+    mesh = Vc.mesh()
+    assert hasattr(mesh, "_shared_data_cache")
+    hierarchyf, levelf = get_level(Vf.mesh())
+    hierarchyc, levelc = get_level(Vc.mesh())
+
+    if hierarchyc != hierarchyf:
+        raise ValueError("Can't map across hierarchies")
+
+    hierarchy = hierarchyf
+    increment = Fraction(1, hierarchyf.refinements_per_level)
+    if levelc + increment != levelf:
+        raise ValueError("Can't map between level %s and level %s" % (levelc, levelf))
+
+    key = (levelc, Vc.extruded and (Vf.mesh().layers, Vc.mesh().layers))
+    cache = mesh._shared_data_cache["hierarchy_coarse_cell_child_count"]
+    try:
+        return cache[key]
+    except KeyError:
+        if Vc.extruded:
+            level_ratio = (Vf.mesh().layers - 1) // (Vc.mesh().layers - 1)
+        else:
+            level_ratio = 1
+        coarse_to_fine = hierarchy.coarse_to_fine_cells[levelc]
+        iterset = mesh.cell_set
+        counts = numpy.zeros(iterset.total_size, dtype=IntType)
+        # Each child of a coarse cell becomes level_ratio cells once extruded.
+        counts[:iterset.size] = (coarse_to_fine[:iterset.size] >= 0).sum(axis=1) * level_ratio
+        # A count belongs to a base cell, and every layer of that cell shares
+        # it. An ExtrudedSet holds no data of its own, so hang the counts off
+        # the base set that it was built on.
+        dset = op2.DataSet(iterset.parent if Vc.extruded else iterset, 1)
+        return cache.setdefault(key, op2.Dat(dset, counts, dtype=IntType))
+
+
 def physical_node_locations(V):
     element = V.ufl_element()
     if V.value_shape:

tests/firedrake/multigrid/test_adaptive_multigrid.py

@@ -34,10 +34,12 @@ def _linear_expr(mesh):
 def coarse_mesh(request):
     dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 1)}
     mesher = request.param
+    # Big enough that refining part of it leaves untouched cells behind, and
+    # that a coarse cell's child count varies widely across the mesh.
     if mesher == "firedrake-square":
-        return UnitSquareMesh(1, 1, distribution_parameters=dparams)
+        return UnitSquareMesh(4, 4, distribution_parameters=dparams)
     elif mesher == "firedrake-cube":
-        return UnitCubeMesh(1, 1, 1, distribution_parameters=dparams)
+        return UnitCubeMesh(2, 2, 2, distribution_parameters=dparams)
     elif mesher == "netgen-square":
         from netgen.occ import WorkPlane, OCCGeometry
         wp = WorkPlane()
@@ -362,6 +364,73 @@ def test_DG0(mh, operator):
         assert errornorm(stepc, u_coarse) <= 1e-12
 
 
+def _coarse_cell_integrals(mh, level, u_coarse, u_fine):
+    """Integrate a coarse and a fine function over each owned coarse cell.
+
+    Both returned arrays hold one entry per owned cell of ``mh[level]``. The
+    first is the integral of ``u_coarse`` over that cell. The second is the
+    integral of ``u_fine`` over that cell's fine children.
+    """
+    coarse_mesh = mh[level]
+    fine_mesh = mh[level + 1]
+
+    # A DG0 test function integrates over one cell per entry.
+    W_coarse = FunctionSpace(coarse_mesh, "DG", 0)
+    mass_coarse = assemble(TestFunction(W_coarse) * u_coarse * dx).dat.data_ro
+    W_fine = FunctionSpace(fine_mesh, "DG", 0)
+    mass_per_child = assemble(TestFunction(W_fine) * u_fine * dx).dat.data_ro
+
+    # Refinement acts on each rank's own plex, so the children of an owned
+    # coarse cell are owned fine cells. Summing the owned children of each
+    # owned coarse cell therefore needs no halo exchange.
+    children = mh.coarse_to_fine_cells[level][:coarse_mesh.cell_set.size]
+    valid = children >= 0
+    assert (children[valid] < fine_mesh.cell_set.size).all()
+    mass_fine = np.where(valid, mass_per_child[children], 0).sum(axis=1)
+    return mass_coarse[:coarse_mesh.cell_set.size], mass_fine
+
+
+@pytest.mark.skipcomplex
+@pytest.mark.parallel([1, 2, 4])
+@pytest.mark.parametrize("family, degree", [("DG", 0), ("DG", 1), ("DG", 2)])
+def test_dg_injection_conserves_mass(mh, family, degree):
+    """DG injection conserves mass on every coarse cell.
+
+    Injection into a DG space is a cellwise L2 projection. Every DG space
+    holds the constants. Test that projection against the constant 1, and
+    the integral of the injected function over a coarse cell must equal the
+    integral of the fine function over that cell's children.
+
+    A random fine function makes this test bite. The step function that
+    `test_DG0` injects is constant on a unit domain. Injecting a constant
+    only checks that the children's volumes add up to the coarse cell's
+    volume. It passes even when the kernel integrates over the wrong set
+    of children.
+    """
+    rg = RandomGenerator(PCG64(seed=0))
+    padded = False
+    for level in range(len(mh) - 1):
+        # A coarse cell that the refinement left alone has one child, and a
+        # refined one has several. The macro-cell map pads the short rows.
+        # Only the levels that leave some cells alone exercise that padding.
+        padded |= bool((mh.coarse_to_fine_cells[level] < 0).any())
+
+        V_coarse = FunctionSpace(mh[level], family, degree)
+        V_fine = FunctionSpace(mh[level + 1], family, degree)
+
+        u_fine = rg.uniform(V_fine)
+
+        u_coarse = Function(V_coarse)
+        inject(u_fine, u_coarse)
+
+        mass_coarse, mass_fine = _coarse_cell_integrals(mh, level, u_coarse, u_fine)
+        assert np.allclose(mass_coarse, mass_fine, rtol=1e-12, atol=1e-14)
+
+    # The padded rows are the point of this test. A hierarchy that refines
+    # every cell of every level says nothing about them.
+    assert mh[0].comm.allreduce(padded, MPI.LOR)
+
+
 @pytest.mark.parallel([1, 2, 4])
 @pytest.mark.parametrize("operator", ["prolong", "inject"])
 def test_CG1(mh, operator):