GPU support for Firedrake

.github/workflows/build.yml

@@ -83,6 +83,8 @@ jobs:
             --install defcon \
             --install gadopt \
             --install asQ \
+            --package-branch tsfc gpu --package-branch pyop2 gpu \
+
             || (cat firedrake-install.log && /bin/false)
       - name: Install test dependencies
         run: |

firedrake/__init__.py

@@ -118,6 +118,7 @@
 from firedrake.ensemble import *
 from firedrake.randomfunctiongen import *
 from firedrake.external_operators import *
+from firedrake.offload import *
 from firedrake.progress_bar import ProgressBar  # noqa: F401
 
 from firedrake.logging import *

firedrake/assemble.py

@@ -1109,7 +1109,7 @@ def allocate(self):
         # Getting the comm attribute of a form isn't straightforward
         # form.ufl_domains()[0]._comm seems the most robust method
         # revisit in a refactor
-        return op2.Global(
+        return op2.compute_backend.Global(
             1,
             [0.0],
             dtype=utils.ScalarType,
@@ -1450,7 +1450,7 @@ def _apply_bc(self, tensor, bc):
     def _apply_bcs_mat_real_block(op2tensor, i, j, component, node_set):
         dat = op2tensor[i, j].handle.getPythonContext().dat
         if component is not None:
-            dat = op2.DatView(dat, component)
+            dat = op2.compute_backend.DatView(dat, component)
         dat.zero(subset=node_set)
 
     def _check_tensor(self, tensor):
@@ -1535,7 +1535,9 @@ def _global_kernel_cache_key(form, local_knl, subdomain_id, all_integer_subdomai
     return ((sig, subdomain_id)
             + tuple(subdomain_key)
             + tuplify(all_integer_subdomain_ids)
-            + cachetools.keys.hashkey(local_knl, **kwargs))
+            + cachetools.keys.hashkey(local_knl, **kwargs)
+            + op2.compute_backend.cache_key
+            )
 
 
 @cachetools.cached(cache={}, key=_global_kernel_cache_key)
@@ -1595,14 +1597,14 @@ def build(self):
         extruded_periodic = self._mesh.extruded_periodic
         constant_layers = extruded and not self._mesh.variable_layers
 
-        return op2.GlobalKernel(self._kinfo.kernel,
-                                kernel_args,
-                                iteration_region=iteration_region,
-                                pass_layer_arg=self._kinfo.pass_layer_arg,
-                                extruded=extruded,
-                                extruded_periodic=extruded_periodic,
-                                constant_layers=constant_layers,
-                                subset=self._needs_subset)
+        return op2.compute_backend.GlobalKernel(self._kinfo.kernel,
+                                                kernel_args,
+                                                iteration_region=iteration_region,
+                                                pass_layer_arg=self._kinfo.pass_layer_arg,
+                                                extruded=extruded,
+                                                extruded_periodic=extruded_periodic,
+                                                constant_layers=constant_layers,
+                                                subset=self._needs_subset)
 
     @property
     def _integral_type(self):
@@ -1831,7 +1833,7 @@ def build(self, tensor):
             unroll=self.needs_unrolling()
         )
         try:
-            return op2.Parloop(_global_knl, self._iterset, parloop_args)
+            return op2.compute_backend.Parloop(_global_knl, self._iterset, parloop_args)
         except MapValueError:
             raise RuntimeError("Integral measure does not match measure of all "
                                "coefficients/arguments")
@@ -2060,7 +2062,7 @@ def _as_parloop_arg_cell_facet(_, self):
 
 @_as_parloop_arg.register(LayerCountKernelArg)
 def _as_parloop_arg_layer_count(_, self):
-    glob = op2.Global(
+    glob = op2.compute_backend.Global(
         (1,),
         self._iterset.layers-2,
         dtype=numpy.int32,

firedrake/bcs.py

@@ -170,7 +170,7 @@ def node_set(self):
         '''The subset corresponding to the nodes at which this
         boundary condition applies.'''
 
-        return op2.Subset(self._function_space.node_set, self.nodes)
+        return op2.op2.compute_backend.Subset(self._function_space.node_set, self.nodes)
 
     @PETSc.Log.EventDecorator()
     def zero(self, r):

firedrake/constant.py

@@ -75,7 +75,7 @@ def __new__(cls, value, domain=None, name=None, count=None):
                 "create a Function in the Real space.", FutureWarning
             )
 
-            dat, rank, shape = _create_dat(op2.Global, value, domain._comm)
+            dat, rank, shape = _create_dat(op2.compute_backend.Global, value, domain._comm)
 
             if not isinstance(domain, ufl.AbstractDomain):
                 cell = ufl.as_cell(domain)
@@ -102,7 +102,7 @@ def __init__(self, value, domain=None, name=None, count=None):
         # Init also called in mesh constructor, but constant can be built without mesh
         utils._init()
 
-        self.dat, rank, self._ufl_shape = _create_dat(op2.Constant, value, None)
+        self.dat, rank, self._ufl_shape = _create_dat(op2.compute_backend.Constant, value, None)
 
         super().__init__()
         Counted.__init__(self, count, Counted)

firedrake/extrusion_utils.py

@@ -65,7 +65,7 @@ def make_extruded_coords(extruded_topology, base_coords, ext_coords,
         layer_height = numpy.cumsum(numpy.concatenate(([0], layer_height)))
 
     layer_heights = layer_height.size
-    layer_height = op2.Global(layer_heights, layer_height, dtype=RealType, comm=extruded_topology._comm)
+    layer_height = op2.compute_backend.Global(layer_heights, layer_height, dtype=RealType, comm=extruded_topology._comm)
 
     if kernel is not None:
         op2.ParLoop(kernel,

firedrake/functionspacedata.py

@@ -104,7 +104,7 @@ def get_node_set(mesh, key):
     global_numbering, constrained_size = get_global_numbering(mesh, key)
     node_classes = mesh.node_classes(nodes_per_entity, real_tensorproduct=real_tensorproduct)
     halo = halo_mod.Halo(mesh.topology_dm, global_numbering, comm=mesh.comm)
-    node_set = op2.Set(node_classes, halo=halo, comm=mesh.comm, constrained_size=constrained_size)
+    node_set = op2.compute_backend.Set(node_classes, halo=halo, comm=mesh.comm, constrained_size=constrained_size)
     extruded = mesh.cell_set._extruded
 
     assert global_numbering.getStorageSize() == node_set.total_size
@@ -509,12 +509,12 @@ def get_map(self, V, entity_set, map_arity, name, offset, offset_quotient):
         entity_node_list = self.entity_node_lists[entity_set]
         val = self.map_cache[entity_set]
         if val is None:
-            val = op2.Map(entity_set, self.node_set,
-                          map_arity,
-                          entity_node_list,
-                          ("%s_"+name) % (V.name),
-                          offset=offset,
-                          offset_quotient=offset_quotient)
+            val = op2.compute_backend.Map(entity_set, self.node_set,
+                                          map_arity,
+                                          entity_node_list,
+                                          ("%s_"+name) % (V.name),
+                                          offset=offset,
+                                          offset_quotient=offset_quotient)
 
             self.map_cache[entity_set] = val
         return val

firedrake/functionspaceimpl.py

@@ -695,7 +695,7 @@ def dim(self):
     def make_dat(self, val=None, valuetype=None, name=None):
         r"""Return a newly allocated :class:`pyop2.types.dat.Dat` defined on the
         :attr:`dof_dset` of this :class:`.Function`."""
-        return op2.Dat(self.dof_dset, val, valuetype, name)
+        return op2.compute_backend.Dat(self.dof_dset, val, valuetype, name)
 
     def entity_node_map(self, source_mesh, source_integral_type, source_subdomain_id, source_all_integer_subdomain_ids):
         r"""Return entity node map rebased on ``source_mesh``.
@@ -1073,7 +1073,7 @@ def node_set(self):
         :class:`FunctionSpace`\s this :class:`MixedFunctionSpace` is
         composed of one or (for VectorFunctionSpaces) more degrees of freedom
         are stored at each node."""
-        return op2.MixedSet(s.node_set for s in self._spaces)
+        return op2.compute_backend.MixedSet(s.node_set for s in self._spaces)
 
     @utils.cached_property
     def dof_dset(self):
@@ -1082,7 +1082,7 @@ def dof_dset(self):
         :attr:`FunctionSpace.dof_dset`\s of the underlying
         :class:`FunctionSpace`\s of which this :class:`MixedFunctionSpace` is
         composed."""
-        return op2.MixedDataSet(s.dof_dset for s in self._spaces)
+        return op2.compute_backend.MixedDataSet(s.dof_dset for s in self._spaces)
 
     def entity_node_map(self, source_mesh, source_integral_type, source_subdomain_id, source_all_integer_subdomain_ids):
         r"""Return entity node map rebased on ``source_mesh``.
@@ -1114,17 +1114,17 @@ def cell_node_map(self):
         :attr:`FunctionSpace.cell_node_map`\s of the underlying
         :class:`FunctionSpace`\s of which this :class:`MixedFunctionSpace` is
         composed."""
-        return op2.MixedMap(s.cell_node_map() for s in self._spaces)
+        return op2.compute_backend.MixedMap(s.cell_node_map() for s in self._spaces)
 
     def interior_facet_node_map(self):
         r"""Return the :class:`pyop2.types.map.MixedMap` from interior facets to
         function space nodes."""
-        return op2.MixedMap(s.interior_facet_node_map() for s in self)
+        return op2.compute_backend.MixedMap(s.interior_facet_node_map() for s in self)
 
     def exterior_facet_node_map(self):
         r"""Return the :class:`pyop2.types.map.Map` from exterior facets to
         function space nodes."""
-        return op2.MixedMap(s.exterior_facet_node_map() for s in self)
+        return op2.compute_backend.MixedMap(s.exterior_facet_node_map() for s in self)
 
     def local_to_global_map(self, bcs):
         r"""Return a map from process local dof numbering to global dof numbering.
@@ -1139,8 +1139,8 @@ def make_dat(self, val=None, valuetype=None, name=None):
             assert len(val) == len(self)
         else:
             val = [None for _ in self]
-        return op2.MixedDat(s.make_dat(v, valuetype, "%s[cmpt-%d]" % (name, i))
-                            for i, (s, v) in enumerate(zip(self._spaces, val)))
+        return op2.compute_backend.MixedDat(s.make_dat(v, valuetype, "%s[cmpt-%d]" % (name, i))
+                                            for i, (s, v) in enumerate(zip(self._spaces, val)))
 
     @utils.cached_property
     def dm(self):
@@ -1339,12 +1339,12 @@ def set_shared_data(self):
         pass
 
     def make_dof_dset(self):
-        return op2.GlobalDataSet(self.make_dat())
+        return op2.compute_backend.GlobalDataSet(self.make_dat())
 
     def make_dat(self, val=None, valuetype=None, name=None):
         r"""Return a newly allocated :class:`pyop2.types.glob.Global` representing the
         data for a :class:`.Function` on this space."""
-        return op2.Global(self.value_size, val, valuetype, name, self._comm)
+        return op2.compute_backend.Global(self.value_size, val, valuetype, name, self._comm)
 
     def entity_node_map(self, source_mesh, source_integral_type, source_subdomain_id, source_all_integer_subdomain_ids):
         return None

firedrake/interpolation.py

@@ -1112,7 +1112,7 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
 
     if tensor in set((c.dat for c in coefficients)):
         output = tensor
-        tensor = op2.Dat(tensor.dataset)
+        tensor = op2.compute_backend.Dat(tensor.dataset)
         if access is not op2.WRITE:
             copyin = (partial(output.copy, tensor), )
         else:

firedrake/matrix.py

@@ -116,7 +116,7 @@ def __init__(self, a, bcs, mat_type, *args, **kwargs):
         # sets self._a, self._bcs, and self._mat_type
         MatrixBase.__init__(self, a, bcs, mat_type)
         options_prefix = kwargs.pop("options_prefix")
-        self.M = op2.Mat(*args, mat_type=mat_type, **kwargs)
+        self.M = op2.compute_backend.Mat(*args, mat_type=mat_type, **kwargs)
         self.petscmat = self.M.handle
         self.petscmat.setOptionsPrefix(options_prefix)
         self.mat_type = mat_type

firedrake/mesh.py

@@ -189,15 +189,15 @@ def __init__(self, mesh, facets, classes, kind, facet_cell, local_facet_number,
         self.facet_cell = facet_cell
 
         if isinstance(self.set, op2.ExtrudedSet):
-            dset = op2.DataSet(self.set.parent, self._rank)
+            dset = op2.compute_backend.DataSet(self.set.parent, self._rank)
         else:
-            dset = op2.DataSet(self.set, self._rank)
+            dset = op2.compute_backend.DataSet(self.set, self._rank)
 
         # Dat indicating which local facet of each adjacent cell corresponds
         # to the current facet.
-        self.local_facet_dat = op2.Dat(dset, local_facet_number, np.uintc,
-                                       "%s_%s_local_facet_number" %
-                                       (self.mesh.name, self.kind))
+        self.local_facet_dat = op2.compute_backend.Dat(dset, local_facet_number, np.uintc,
+                                                       "%s_%s_local_facet_number" %
+                                                       (self.mesh.name, self.kind))
 
         self.unique_markers = [] if unique_markers is None else unique_markers
         self._subsets = {}
@@ -209,17 +209,17 @@ def set(self):
             label = "%s_facets" % self.kind
             layers = self.mesh.entity_layers(1, label)
             base = getattr(self.mesh._base_mesh, label).set
-            return op2.ExtrudedSet(base, layers=layers)
-        return op2.Set(size, "%sFacets" % self.kind.capitalize()[:3],
-                       comm=self.mesh.comm)
+            return op2.compute_backend.ExtrudedSet(base, layers=layers)
+        return op2.compute_backend.Set(size, "%sFacets" % self.kind.capitalize()[:3],
+                                       comm=self.mesh.comm)
 
     @utils.cached_property
     def _null_subset(self):
         '''Empty subset for the case in which there are no facets with
         a given marker value. This is required because not all
         markers need be represented on all processors.'''
 
-        return op2.Subset(self.set, [])
+        return op2.compute_backend.Subset(self.set, [])
 
     @PETSc.Log.EventDecorator()
     def measure_set(self, integral_type, subdomain_id,
@@ -260,7 +260,7 @@ def measure_set(self, integral_type, subdomain_id,
             except KeyError:
                 unmarked_points = self._collect_unmarked_points(all_integer_subdomain_ids)
                 _, indices, _ = np.intersect1d(self.facets, unmarked_points, return_indices=True)
-                return self._subsets.setdefault(key, op2.Subset(self.set, indices))
+                return self._subsets.setdefault(key, op2.compute_backend.Subset(self.set, indices))
         else:
             return self.subset(subdomain_id)
 
@@ -295,7 +295,7 @@ def subset(self, markers):
                         marked_points_list.append(self.mesh.topology_dm.getStratumIS(dmcommon.FACE_SETS_LABEL, i).indices)
             if marked_points_list:
                 _, indices, _ = np.intersect1d(self.facets, np.concatenate(marked_points_list), return_indices=True)
-                return self._subsets.setdefault(markers, op2.Subset(self.set, indices))
+                return self._subsets.setdefault(markers, op2.compute_backend.Subset(self.set, indices))
             else:
                 return self._subsets.setdefault(markers, self._null_subset)
 
@@ -314,8 +314,8 @@ def _collect_unmarked_points(self, markers):
     @utils.cached_property
     def facet_cell_map(self):
         """Map from facets to cells."""
-        return op2.Map(self.set, self.mesh.cell_set, self._rank, self.facet_cell,
-                       "facet_to_cell_map")
+        return op2.compute_backend.Map(self.set, self.mesh.cell_set, self._rank, self.facet_cell,
+                                       "facet_to_cell_map")
 
 
 @PETSc.Log.EventDecorator()
@@ -889,7 +889,7 @@ def cell_subset(self, subdomain_id, all_integer_subdomain_ids=None):
                 indices = dmcommon.get_cell_markers(self.topology_dm,
                                                     self._cell_numbering,
                                                     subdomain_id)
-            return self._subsets.setdefault(key, op2.Subset(self.cell_set, indices))
+            return self._subsets.setdefault(key, op2.compute_backend.Subset(self.cell_set, indices))
 
     @PETSc.Log.EventDecorator()
     def measure_set(self, integral_type, subdomain_id,
@@ -1329,9 +1329,9 @@ def cell_to_facets(self):
                                                    self._cell_numbering,
                                                    self.cell_closure)
         if isinstance(self.cell_set, op2.ExtrudedSet):
-            dataset = op2.DataSet(self.cell_set.parent, dim=cell_facets.shape[1:])
+            dataset = op2.compute_backend.DataSet(self.cell_set.parent, dim=cell_facets.shape[1:])
         else:
-            dataset = op2.DataSet(self.cell_set, dim=cell_facets.shape[1:])
+            dataset = op2.compute_backend.DataSet(self.cell_set, dim=cell_facets.shape[1:])
         return op2.Dat(dataset, cell_facets, dtype=cell_facets.dtype,
                        name="cell-to-local-facet-dat")
 
@@ -1362,7 +1362,7 @@ def num_entities(self, d):
     @utils.cached_property
     def cell_set(self):
         size = list(self._entity_classes[self.cell_dimension(), :])
-        return op2.Set(size, "Cells", comm=self._comm)
+        return op2.compute_backend.Set(size, "Cells", comm=self._comm)
 
     @PETSc.Log.EventDecorator()
     def _set_partitioner(self, plex, distribute, partitioner_type=None):
@@ -1691,7 +1691,7 @@ def __init__(self, mesh, layers, periodic=False, name=None):
             """
         else:
             self.variable_layers = False
-        self.cell_set = op2.ExtrudedSet(mesh.cell_set, layers=layers, extruded_periodic=periodic)
+        self.cell_set = op2.computed_backend.ExtrudedSet(mesh.cell_set, layers=layers, extruded_periodic=periodic)
         # submesh
         self.submesh_parent = None
 
@@ -2016,7 +2016,7 @@ def num_entities(self, d):
     @utils.cached_property  # TODO: Recalculate if mesh moves
     def cell_set(self):
         size = list(self._entity_classes[self.cell_dimension(), :])
-        return op2.Set(size, "Cells", comm=self.comm)
+        return op2.compute_backend.Set(size, "Cells", comm=self.comm)
 
     @utils.cached_property  # TODO: Recalculate if mesh moves
     def cell_parent_cell_list(self):
@@ -4477,7 +4477,7 @@ def SubDomainData(geometric_expr):
 
     # Create cell subset
     indices, = np.nonzero(f.dat.data_ro_with_halos > 0.5)
-    return op2.Subset(m.cell_set, indices)
+    return op2.compute_backend.Subset(m.cell_set, indices)
 
 
 def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None):