diff --git a/firedrake/mg/interface.py b/firedrake/mg/interface.py index 4c544535cf..186e0acddd 100644 --- a/firedrake/mg/interface.py +++ b/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 diff --git a/firedrake/mg/kernels.py b/firedrake/mg/kernels.py index 70021eec11..0d5778af99 100644 --- a/firedrake/mg/kernels.py +++ b/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) diff --git a/firedrake/mg/utils.py b/firedrake/mg/utils.py index ffc3c1e8ad..87abd4b80d 100644 --- a/firedrake/mg/utils.py +++ b/firedrake/mg/utils.py @@ -73,27 +73,23 @@ 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. 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 +151,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: diff --git a/tests/firedrake/multigrid/test_adaptive_multigrid.py b/tests/firedrake/multigrid/test_adaptive_multigrid.py index 557b18889f..1ff0459afd 100644 --- a/tests/firedrake/multigrid/test_adaptive_multigrid.py +++ b/tests/firedrake/multigrid/test_adaptive_multigrid.py @@ -33,10 +33,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() @@ -356,6 +358,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):