MG: prolong working in parallel for single refinement

Author: pbrubeck

firedrake/function.py

@@ -382,7 +382,8 @@ def vector(self):
         return vector.Vector(self)
 
     def nodal_dat(self):
-        return op3.HierarchicalArray(self.function_space().nodal_axes, data=self.dat.data_rw_with_halos)
+        return op3.HierarchicalArray(self.function_space().nodal_axes,
+                                     data=self.dat.data_rw_with_halos)
 
     @PETSc.Log.EventDecorator()
     def interpolate(

firedrake/functionspaceimpl.py

@@ -712,15 +712,13 @@ def owned_cell_node_list(self):
 
     @utils.cached_property
     def nodes(self):
-        if self.value_size > 1:
-            raise NotImplementedError("TODO")
         ax = self.block_axes
         return op3.Axis([op3.AxisComponent((ax.owned.size, ax.size),
                          "XXX", rank_equal=False)], "nodes", numbering=None, sf=ax.sf)
 
     @utils.cached_property
     def nodal_axes(self):
-        return op3.AxisTree(self.nodes)
+        return op3.AxisTree.from_iterable([self.nodes, self.value_size])
 
     @utils.cached_property
     def topological(self):

firedrake/mg/interface.py

@@ -83,17 +83,13 @@ def prolong(coarse, fine):
         for d in [coarse, coarse_coords]:
             d.dat.assemble()
 
-        #op2.par_loop(kernel, next.node_set,
-        #             next.dat(op2.WRITE),
-        #             coarse.dat(op2.READ, fine_to_coarse),
-        #             node_locations.dat(op2.READ),
-        #             coarse_coords.dat(op2.READ, fine_to_coarse_coords))
+        id_map = utils.owned_node_map(Vf)
         op3.do_loop(
-            n := Vf.nodes.index(),
+            n := Vf.nodes.owned.index(),
             kernel(
-                next.nodal_dat()[n],
+                next.nodal_dat()[id_map(n)],
                 coarse.nodal_dat()[fine_to_coarse(n)],
-                node_locations.nodal_dat()[n],
+                node_locations.nodal_dat()[id_map(n)],
                 coarse_coords.nodal_dat()[fine_to_coarse_coords(n)],
             ),
         )
@@ -153,17 +149,14 @@ def restrict(fine_dual, coarse_dual):
         for d in [coarse_coords]:
             d.dat.assemble()
         kernel = kernels.restrict_kernel(Vf, Vc)
-        #op2.par_loop(kernel, fine_dual.node_set,
-        #             next.dat(op2.INC, fine_to_coarse),
-        #             fine_dual.dat(op2.READ),
-        #             node_locations.dat(op2.READ),
-        #             coarse_coords.dat(op2.READ, fine_to_coarse_coords))
+
+        id_map = utils.owned_node_map(Vf)
         op3.do_loop(
-            n := Vf.nodes.index(),
+            n := Vf.nodes.owned.index(),
             kernel(
                 next.nodal_dat()[fine_to_coarse(n)],
-                fine_dual.nodal_dat()[n],
-                node_locations.nodal_dat()[n],
+                fine_dual.nodal_dat()[id_map(n)],
+                node_locations.nodal_dat()[id_map(n)],
                 coarse_coords.nodal_dat()[fine_to_coarse_coords(n)],
             ),
         )
@@ -235,16 +228,13 @@ def inject(fine, coarse):
             # this expanded stencil
             for d in [fine, fine_coords]:
                 d.dat.assemble()
-            #op2.par_loop(kernel, next.node_set,
-            #             next.dat(op2.INC),
-            #             node_locations.dat(op2.READ),
-            #             fine.dat(op2.READ, coarse_node_to_fine_nodes),
-            #             fine_coords.dat(op2.READ, coarse_node_to_fine_coords))
+
+            id_map = utils.owned_node_map(Vc)
             op3.do_loop(
-                n := Vc.nodes.index(),
+                n := Vc.nodes.owned.index(),
                 kernel(
-                    next.nodal_dat()[n],
-                    node_locations.nodal_dat()[n],
+                    next.nodal_dat()[id_map(n)],
+                    node_locations.nodal_dat()[id_map(n)],
                     fine.nodal_dat()[coarse_node_to_fine_nodes(n)],
                     fine_coords.nodal_dat()[coarse_node_to_fine_coords(n)],
                 ),
@@ -264,7 +254,7 @@ def inject(fine, coarse):
             #             fine_coords.dat(op2.READ, coarse_cell_to_fine_coords),
             #             coarse_coords.dat(op2.READ, coarse_coords.cell_node_map()))
             op3.do_loop(
-                c := Vc.mesh().cells.index(),
+                c := Vc.mesh().cells.owned.index(),
                 kernel(
                     next.dat[c],
                     fine.nodal_dat()[coarse_cell_to_fine_nodes(c)],

firedrake/mg/mesh.py

@@ -8,7 +8,6 @@
 from firedrake.utils import cached_property
 from firedrake.cython import mgimpl as impl
 from .utils import set_level
-from firedrake.petsc import PETSc
 
 __all__ = ("HierarchyBase", "MeshHierarchy", "ExtrudedMeshHierarchy", "NonNestedHierarchy",
            "SemiCoarsenedExtrudedHierarchy")

firedrake/mg/utils.py

@@ -18,6 +18,31 @@ def coarse_to_fine_cell_map(coarse_mesh, fine_mesh, coarse_to_fine_data):
     return op3.Map(connectivity)
 
 
+def create_node_map(iterset, toset, arity=1, values=None):
+    axes = op3.AxisTree.from_iterable([iterset, arity])
+    if values is None:
+        values = numpy.arange(iterset.size, dtype=IntType)
+    dat = op3.HierarchicalArray(axes, data=values)
+    return op3.Map({
+        freeze({"nodes": iterset.owned.component.label}): [
+            op3.TabulatedMapComponent("nodes", toset.component.label, dat)
+        ]
+    })
+
+
+def owned_node_map(V):
+    """ This should not be necessary
+    """
+    mesh = V.mesh()
+    key = entity_dofs_key(V.finat_element.entity_dofs())
+    cache = mesh._shared_data_cache["owned_node_map"]
+    try:
+        return cache[key]
+    except KeyError:
+        node_map = create_node_map(V.nodes.owned, V.nodes)
+        return cache.setdefault(key, node_map)
+
+
 def fine_node_to_coarse_node_map(Vf, Vc):
     if len(Vf) > 1:
         assert len(Vf) == len(Vc)
@@ -51,15 +76,9 @@ def fine_node_to_coarse_node_map(Vf, Vc):
 
         fine_to_coarse = hierarchy.fine_to_coarse_cells[levelf]
         fine_to_coarse_nodes = impl.fine_to_coarse_nodes(Vf, Vc, fine_to_coarse)
-
-        axes = op3.AxisTree.from_iterable([Vf.nodes, fine_to_coarse_nodes.shape[1]])
-        fine_to_coarse_node_dat = op3.HierarchicalArray(axes, data=fine_to_coarse_nodes)
-        fine_to_coarse_node_map = op3.Map({
-            freeze({"nodes": "XXX"}): [
-                op3.TabulatedMapComponent("nodes", "XXX", fine_to_coarse_node_dat)
-            ]
-        })
-        return cache.setdefault(key, fine_to_coarse_node_map)
+        return cache.setdefault(key, create_node_map(Vf.nodes, Vc.nodes,
+                                                     arity=fine_to_coarse_nodes.shape[1],
+                                                     values=fine_to_coarse_nodes))
 
 
 def coarse_node_to_fine_node_map(Vc, Vf):
@@ -95,15 +114,9 @@ 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)
-
-        axes = op3.AxisTree.from_iterable([Vc.nodes, coarse_to_fine_nodes.shape[1]])
-        coarse_to_fine_node_dat = op3.HierarchicalArray(axes, data=coarse_to_fine_nodes)
-        coarse_to_fine_node_map = op3.Map({
-            freeze({"nodes": "XXX"}): [
-                op3.TabulatedMapComponent("nodes", "XXX", coarse_to_fine_node_dat)
-            ]
-        })
-        return cache.setdefault(key, coarse_to_fine_node_map)
+        return cache.setdefault(key, create_node_map(Vc.nodes, Vf.nodes,
+                                                     arity=coarse_to_fine_nodes.shape[1],
+                                                     values=coarse_to_fine_nodes))
 
 
 def coarse_cell_to_fine_node_map(Vc, Vf):
@@ -153,8 +166,8 @@ def coarse_cell_to_fine_node_map(Vc, Vf):
         axes = op3.AxisTree.from_iterable([Vc.nodes, arity*level_ratio])
         coarse_cell_to_fine_node_dat = op3.HierarchicalArray(axes, data=coarse_to_fine_nodes)
         coarse_cell_to_fine_node_map = op3.Map({
-            freeze({Vc.mesh().topology.name: Vc.mesh().cell_label}): [
-                op3.TabulatedMapComponent("nodes", "XXX", coarse_cell_to_fine_node_dat)
+            freeze({Vc.mesh().topology.name: Vc.mesh().cells.owned.label}): [
+                op3.TabulatedMapComponent("nodes", Vf.nodes.component.label, coarse_cell_to_fine_node_dat)
             ]
         })
         return cache.setdefault(key, coarse_cell_to_fine_node_map)

tests/multigrid/test_grid_transfer.py

@@ -21,24 +21,23 @@ def space(request, cell):
         return request.param
 
 
-@pytest.fixture(params=[1, 2], scope="module")
-def refinements_per_level(request):
+@pytest.fixture(params=[(1, 1), (2, 1), (1, 2)], scope="module")
+def nref_refinements_per_level(request):
     return request.param
 
 
 @pytest.fixture(scope="module")
-def hierarchy(cell, refinements_per_level):
+def hierarchy(cell, nref_refinements_per_level):
     if cell == "interval":
         mesh = UnitIntervalMesh(3)
-        return MeshHierarchy(mesh, 2)
     elif cell in {"triangle", "triangle-nonnested", "prism"}:
         mesh = UnitSquareMesh(3, 3, quadrilateral=False)
     elif cell in {"quadrilateral", "hexahedron"}:
         mesh = UnitSquareMesh(3, 3, quadrilateral=True)
     elif cell == "tetrahedron":
         mesh = UnitCubeMesh(2, 2, 2)
 
-    nref = {2: 1, 1: 2}[refinements_per_level]
+    nref, refinements_per_level = nref_refinements_per_level
     hierarchy = MeshHierarchy(mesh, nref, refinements_per_level=refinements_per_level)
 
     if cell in {"prism", "hexahedron"}:
@@ -328,6 +327,5 @@ def test_grid_transfer_periodic(periodic_hierarchy, periodic_space):
 
 if __name__ == "__main__":
     bmesh = UnitIntervalMesh(3)
-    mh = MeshHierarchy(bmesh, 1)
-
-    run_restriction(mh, False, "CG", [1])
+    mh = MeshHierarchy(bmesh, 2, refinements_per_level=1)
+    run_prolongation(mh, False, "CG", [1])