MG: generate pyop3 kernels

Author: pbrubeck

firedrake/cofunction.py

@@ -349,14 +349,17 @@ def vector(self):
         :class:`Cofunction`"""
         return vector.Vector(self)
 
+    def nodal_dat(self):
+        return op3.HierarchicalArray(self.function_space().nodal_axes, data=self.dat.data_rw)
+
     @property
     def node_set(self):
         r"""A :class:`pyop2.types.set.Set` containing the nodes of this
         :class:`Cofunction`. One or (for rank-1 and 2
         :class:`.FunctionSpace`\s) more degrees of freedom are stored
         at each node.
         """
-        return self.function_space().node_set
+        return self.function_space().nodes
 
     def ufl_id(self):
         return self.uid

firedrake/cython/mgimpl.pyx

@@ -85,7 +85,7 @@ def coarse_to_fine_nodes(Vc, Vf, np.ndarray[PetscInt, ndim=2, mode="c"] coarse_t
     ndof = fine_per_cell * fine_cell_per_coarse_cell
     if extruded:
         ndof *= ratio
-    coarse_to_fine_map = np.full((coarse_cells,
+    coarse_to_fine_map = np.full((Vc.node_count,
                                   ndof),
                                  -1,
                                  dtype=IntType)
@@ -142,7 +142,7 @@ def fine_to_coarse_nodes(Vf, Vc, np.ndarray[PetscInt, ndim=2, mode="c"] fine_to_
     coarse_per_fine = fine_to_coarse_cells.shape[1]
     coarse_per_cell = coarse_map.shape[1]
     fine_per_cell = fine_map.shape[1]
-    fine_to_coarse_map = np.full((fine_cells,
+    fine_to_coarse_map = np.full((Vf.node_count,
                                   coarse_per_fine*coarse_per_cell),
                                  -1,
                                  dtype=IntType)

firedrake/function.py

@@ -381,6 +381,9 @@ def vector(self):
         r"""Return a :class:`.Vector` wrapping the data in this :class:`Function`"""
         return vector.Vector(self)
 
+    def nodal_dat(self):
+        return op3.HierarchicalArray(self.function_space().nodal_axes, data=self.dat.data_rw)
+
     @PETSc.Log.EventDecorator()
     def interpolate(
         self,

firedrake/functionspaceimpl.py

@@ -697,17 +697,28 @@ def _cdim(self):
 
     @utils.cached_property
     def cell_node_list(self):
+        from firedrake.parloops import pack_pyop3_tensor
         r"""A numpy array mapping mesh cells to function space nodes (includes halo)."""
         cells = self.mesh().cells
-        ncells = cells.size
-        packed_axes = self.block_axes[self.mesh().closure(cells.index())]
-
-        return packed_axes.tabulated_offsets.buffer.data.reshape((ncells, -1))
+        # Pass self.sub(0) to get nodes from the scalar version of this function space
+        packed_axes = pack_pyop3_tensor(self.block_axes, self.sub(0), cells.index(), "cell")
+        return packed_axes.tabulated_offsets.buffer.data.reshape((cells.size, -1))
 
     @utils.cached_property
     def owned_cell_node_list(self):
         r"""A numpy array mapping owned mesh cells to function space nodes."""
-        return self.cell_node_list[:self.mesh().owned_cells.size]
+        cells = self.mesh().cells
+        return self.cell_node_list[:cells.owned.size]
+
+    @utils.cached_property
+    def nodes(self):
+        if self.value_size > 1:
+            raise NotImplementedError("TODO")
+        return op3.Axis({"XXX": self.node_count}, "nodes", numbering=None)
+
+    @utils.cached_property
+    def nodal_axes(self):
+        return op3.AxisTree(self.nodes)
 
     @utils.cached_property
     def topological(self):

firedrake/mg/interface.py

@@ -1,4 +1,4 @@
-from pyop2 import op2
+import pyop3 as op3
 
 import firedrake
 from firedrake import ufl_expr
@@ -81,13 +81,21 @@ def prolong(coarse, fine):
         # Have to do this, because the node set core size is not right for
         # this expanded stencil
         for d in [coarse, coarse_coords]:
-            d.dat.global_to_local_begin(op2.READ)
-            d.dat.global_to_local_end(op2.READ)
-        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))
+            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))
+        op3.do_loop(
+            n := Vf.nodes.index(),
+            kernel(
+                next.nodal_dat()[n],
+                coarse.nodal_dat()[fine_to_coarse(n)],
+                node_locations.nodal_dat()[n],
+                coarse_coords.nodal_dat()[fine_to_coarse_coords(n)],
+            ),
+        )
         coarse = next
         Vc = Vf
     return fine
@@ -142,14 +150,22 @@ def restrict(fine_dual, coarse_dual):
         # Have to do this, because the node set core size is not right for
         # this expanded stencil
         for d in [coarse_coords]:
-            d.dat.global_to_local_begin(op2.READ)
-            d.dat.global_to_local_end(op2.READ)
+            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))
+        #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))
+        op3.do_loop(
+            n := Vf.nodes.index(),
+            kernel(
+                next.nodal_dat()[n],
+                fine_dual.nodal_dat()[fine_to_coarse(n)],
+                node_locations.nodal_dat()[n],
+                coarse_coords.nodal_dat()[fine_to_coarse_coords(n)],
+            ),
+        )
         fine_dual = next
         Vf = Vc
     return coarse_dual
@@ -217,13 +233,21 @@ def inject(fine, coarse):
             # Have to do this, because the node set core size is not right for
             # this expanded stencil
             for d in [fine, fine_coords]:
-                d.dat.global_to_local_begin(op2.READ)
-                d.dat.global_to_local_end(op2.READ)
-            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))
+                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))
+            op3.do_loop(
+                n := Vc.nodes.index(),
+                kernel(
+                    next.nodal_dat()[n],
+                    node_locations.nodal_dat()[n],
+                    fine.nodal_dat()[coarse_node_to_fine_nodes(n)],
+                    fine_coords.nodal_dat()[coarse_node_to_fine_coords(n)],
+                ),
+            )
         else:
             coarse_coords = Vc.mesh().coordinates
             fine_coords = Vf.mesh().coordinates
@@ -232,13 +256,22 @@ def inject(fine, coarse):
             # Have to do this, because the node set core size is not right for
             # this expanded stencil
             for d in [fine, fine_coords]:
-                d.dat.global_to_local_begin(op2.READ)
-                d.dat.global_to_local_end(op2.READ)
-            op2.par_loop(kernel, Vc.mesh().cell_set,
-                         next.dat(op2.INC, next.cell_node_map()),
-                         fine.dat(op2.READ, coarse_cell_to_fine_nodes),
-                         fine_coords.dat(op2.READ, coarse_cell_to_fine_coords),
-                         coarse_coords.dat(op2.READ, coarse_coords.cell_node_map()))
+                d.dat.assemble()
+            #op2.par_loop(kernel, Vc.mesh().cell_set,
+            #             next.dat(op2.INC, next.cell_node_map()),
+            #             fine.dat(op2.READ, coarse_cell_to_fine_nodes),
+            #             fine_coords.dat(op2.READ, coarse_cell_to_fine_coords),
+            #             coarse_coords.dat(op2.READ, coarse_coords.cell_node_map()))
+            op3.do_loop(
+                n := Vc.nodes.index(),
+                kernel(
+                    next.nodal_dat()[n],
+                    fine.nodal_dat()[coarse_node_to_fine_nodes(n)],
+                    fine_coords.nodal_dat()[coarse_node_to_fine_coords(n)],
+                    coarse_coords.nodal_dat()[coarse_coords(n)],
+                ),
+            )
+
         fine = next
         Vf = Vc
     return coarse