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

@@ -419,7 +419,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)
@@ -569,7 +571,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(
@@ -587,9 +589,28 @@ 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]
+
+    # A coarse cell that adaptive refinement left alone has fewer children
+    # than the busiest cell of the level, and coarse_cell_to_fine_node_map
+    # pads its row out to that width. Stop at this cell's own children, so
+    # the padding is never read.
+    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}
@@ -599,9 +620,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)
@@ -628,7 +654,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
@@ -644,6 +670,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.
@@ -658,12 +689,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:
@@ -673,8 +706,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/utils.py

@@ -9,6 +9,32 @@
 
 
 def fine_node_to_coarse_node_map(Vf, Vc):
+    """Map each fine node to the nodes of its parent coarse cell.
+
+    A fine cell has exactly one parent. This holds for uniform refinement
+    and for adaptive refinement. ``hierarchy.fine_to_coarse_cells`` is
+    therefore one column wide, and it holds no padding. This map needs no
+    special handling for adaptive refinement.
+
+    `coarse_node_to_fine_node_map` runs in the other direction. There the
+    number of children varies, and padding does appear.
+
+    `prolong` and `restrict` both read through this map.
+
+    Parameters
+    ----------
+    Vf : firedrake.functionspaceimpl.WithGeometry
+        The fine function space.
+    Vc : firedrake.functionspaceimpl.WithGeometry
+        The coarse function space, on the previous level of the same
+        hierarchy.
+
+    Returns
+    -------
+    pyop2.types.map.Map
+        A map from the nodes of ``Vf`` to the nodes of ``Vc``.
+
+    """
     if len(Vf) > 1:
         assert len(Vf) == len(Vc)
         return op2.MixedMap(map(fine_node_to_coarse_node_map, Vf, Vc))
@@ -44,6 +70,40 @@ def fine_node_to_coarse_node_map(Vf, Vc):
 
 
 def coarse_node_to_fine_node_map(Vc, Vf):
+    """Map each coarse node to the fine nodes it could come from.
+
+    A coarse node gets one candidate per fine cell that descends from it.
+    Uniform refinement gives every coarse cell the same number of children.
+    Every row is then full. Adaptive refinement gives them different
+    numbers. ``hierarchy.coarse_to_fine_cells`` then right-pads its short
+    rows with -1, out to the busiest coarse cell's count.
+
+    This map fills that padding, because `op2.Map` cannot hold a negative
+    index. `fine_node_to_coarse_node_map` runs in the other direction. One
+    parent per fine cell makes padding impossible there.
+
+    Injection into a space with a pointwise dual basis reads through this
+    map. Injection into a DG space uses `coarse_cell_to_fine_node_map`
+    instead. That map handles its padding a different way.
+
+    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.map.Map
+        A map from the nodes of ``Vc`` to the nodes of ``Vf``.
+
+    Raises
+    ------
+    RuntimeError
+        If an owned coarse node has no fine node to inject from.
+
+    """
     if len(Vf) > 1:
         assert len(Vf) == len(Vc)
         return op2.MixedMap(map(coarse_node_to_fine_node_map, Vf, Vc))
@@ -73,33 +133,66 @@ 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.
+        # Adaptive refinement gives coarse cells different numbers of fine
+        # descendants. Each row of coarse_to_fine_nodes is therefore padded
+        # with -1, out to the busiest coarse cell's count, and op2.Map cannot
+        # hold a negative index. Fill each padded slot with a duplicate of a
+        # real entry from its own row. The injection kernel only reads
+        # through this map, and picks the candidate that matches the coarse
+        # node's physical location, so a repeated entry changes nothing.
+        #
+        # Every rank runs this fill. The partition decides which rows hold
+        # padding, so the fill must not depend on a rank-local test.
+        # The check covers the owned rows alone. coarse_to_fine_cells spans
+        # the owned coarse cells, so most halo rows hold no candidate at all.
+        # Those rows keep the zero filler, and injection visits owned nodes
+        # only, so nothing reads them.
         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))
 
 
 def coarse_cell_to_fine_node_map(Vc, Vf):
+    """Map each coarse cell to the fine nodes of all its children.
+
+    Each row holds one block of nodes per child cell. The blocks are stored
+    back to back. Uniform refinement gives every coarse cell the same number
+    of children, so every row is full. Adaptive refinement gives them
+    different numbers. ``hierarchy.coarse_to_fine_cells`` then right-pads
+    its short rows with -1, and that padding reaches this map.
+
+    The padding stays as it is. The DG injection kernel reads
+    `coarse_cell_child_count`. It stops at a coarse cell's own children, and
+    so never reads a padded block.
+
+    `coarse_node_to_fine_node_map` fills its padding instead. The kernel
+    that reads that map visits every candidate.
+
+    The DG branch of `inject` reads through this map.
+
+    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.map.Map
+        A map from the cells of ``Vc``'s mesh to the nodes of ``Vf``.
+
+    """
     if len(Vf) > 1:
         assert len(Vf) == len(Vc)
         return op2.MixedMap(coarse_cell_to_fine_node_map(f, c) for f, c in zip(Vf, Vc))
@@ -155,6 +248,65 @@ def coarse_cell_to_fine_node_map(Vc, Vf):
                                              offset=offset))
 
 
+def coarse_cell_child_count(Vc, Vf):
+    """Count the fine cells that each coarse cell was refined into.
+
+    Uniform refinement gives every coarse cell the same number of children.
+    Every count then equals the width of a `coarse_cell_to_fine_node_map`
+    row. Adaptive refinement leaves some coarse cells alone, and splits
+    others. A coarse cell then has from one child up to the busiest cell's
+    count. The map pads its short rows out to that busiest count.
+
+    The DG injection kernel reads this count. It stops at a coarse cell's
+    own children, and so leaves that 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.
+
+    """
+    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: