#2711 DoF Kernel Code Generation

doc/user_guide/dynamo0p3.rst

@@ -796,9 +796,6 @@ support for i-first kernels
 point the looping (and associated parallelisation) will be put
 back into the PSy layer.
 
-.. note:: Support for DoF kernels have not yet been implemented in PSyclone
-          (see PSyclone issue #1351 for progress).
-
 .. _lfric-user-kernel-rules:
 
 Rules for all User-Supplied Kernels that Operate on Cell-Columns
@@ -981,9 +978,6 @@ on a ``CELL_COLUMN`` without CMA Operators. Specifically:
 Rules for all User-Supplied Kernels that Operate on DoFs (DoF Kernels)
 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 
-.. note:: Support for DoF kernels have not yet been implemented in PSyclone
-          (see PSyclone issue #1351 for progress).
-
 Kernels that have ``operates_on = DOF`` and
 :ref:`LFRic Built-ins<lfric-built-ins>` overlap significantly in their
 scope, and the conventions that DoF Kernels must follow are influenced
@@ -1015,8 +1009,9 @@ The list of rules for DoF Kernels is as follows:
    to do this.) Any scalar arguments must therefore be declared in the metadata
    as `GH_READ` - see :ref:`below<lfric-kernel-valid-access>`
 
-6) Kernels must be written to operate on a single DoF, such that single DoFs
-   can be provided to the Kernel within a loop over the DoFs of a field.
+6) Kernels must be written to operate on a single DoF, such that field values
+   at the same dof location/index can be provided to the Kernel within a loop
+   over the DoFs of the function space of the field that is being updated.
 
 .. _lfric-api-kernel-metadata:
 
@@ -1889,16 +1884,14 @@ operates_on
 The fourth type of metadata provided is ``OPERATES_ON``. This
 specifies that the Kernel has been written with the assumption that it
 is supplied with the specified data for each field/operator argument.
-For user-supplied kernels this is currently only permitted to be
-``CELL_COLUMN`` or ``DOMAIN``. The possible values for ``OPERATES_ON``
-and their interpretation are summarised in the following table:
+The possible values for ``OPERATES_ON`` and their interpretation are
+summarised in the following table:
 
 ===========  =========================================================
 operates_on  Data passed for each field/operator argument
 ===========  =========================================================
 cell_column  Single column of cells
-dof          Single DoF (currently :ref:`built-ins` only, but see PSyclone
-             issue #1351)
+dof          Single DoF
 domain       All columns of cells
 ===========  =========================================================
 
@@ -2477,9 +2470,6 @@ as the second argument to the kernel (after ``nlayers``).
 Rules for DoF Kernels
 #####################
 
-.. note:: Support for DoF kernels have not yet been implemented in PSyclone
-          (see PSyclone issue #1351 for progress).
-
 The rules for kernels that have ``operates_on = DOF`` are similar to those for
 general-purpose kernels but, due to the restriction that only fields and
 scalars can be passed to them, are much fewer. The full set of rules, along

src/psyclone/domain/lfric/function_space.py

@@ -238,6 +238,13 @@ def undf_name(self):
         '''
         return "undf_" + self.mangled_name
 
+    @property
+    def bare_undf_name(self):
+        ''':returns: undf with no FunctionSpace object name.
+        :rtype: str
+        '''
+        return "undf"
+
     def get_basis_name(self, qr_var=None, on_space=None):
         '''
         Returns a name for the basis function on this FunctionSpace. If

src/psyclone/domain/lfric/kern_call_arg_list.py

@@ -48,6 +48,7 @@
 from psyclone.core import AccessType, Signature
 from psyclone.domain.lfric import ArgOrdering, LFRicConstants
 # Avoid circular import:
+from psyclone.domain.lfric.lfric_builtins import LFRicBuiltIn
 from psyclone.domain.lfric.lfric_types import LFRicTypes
 from psyclone.errors import GenerationError, InternalError
 from psyclone.psyir.nodes import ArrayReference, Reference, StructureReference
@@ -218,9 +219,9 @@ def mesh_height(self, var_accesses=None):
         '''Add mesh height (nlayers) to the argument list and if supplied
         stores this access in var_accesses.
 
-        :param var_accesses: optional VariablesAccessInfo instance to store
+        :param var_accesses: optional VariablesAccessInfo instance to store \
             the information about variable accesses.
-        :type var_accesses:
+        :type var_accesses: \
             :py:class:`psyclone.core.VariablesAccessInfo`
 
         '''
@@ -380,11 +381,23 @@ def field(self, arg, var_accesses=None):
         # Look-up the name of the variable that stores the reference to
         # the data in this field.
         sym = self._symtab.lookup_with_tag(f"{arg.name}:{suffix}")
-        # Add the field data array as being read.
-        self.append(sym.name, var_accesses, var_access_name=sym.name,
-                    mode=arg.access, metadata_posn=arg.metadata_index)
 
-        self.psyir_append(Reference(sym))
+        if self._kern.iterates_over == "dof" and not isinstance(
+          self._kern, LFRicBuiltIn):
+            # If dof kernel, add access to the field by dof ref
+            dof_sym = self._symtab.find_or_create_integer_symbol(
+                "df", tag="dof_loop_idx")
+            self.append_array_reference(sym.name, [Reference(dof_sym)],
+                                        ScalarType.Intrinsic.INTEGER,
+                                        symbol=sym)
+            # Then append our symbol
+            name = f"{sym.name}({Reference(dof_sym)})"
+            self.append(name, var_accesses, var_access_name=sym.name)
+        else:
+            # Add the field data array as being read.
+            self.append(sym.name, var_accesses, var_access_name=sym.name,
+                        mode=arg.access, metadata_posn=arg.metadata_index)
+            self.psyir_append(Reference(sym))
 
     def stencil_unknown_extent(self, arg, var_accesses=None):
         '''Add stencil information to the argument list associated with the
@@ -611,7 +624,10 @@ def fs_compulsory_field(self, function_space, var_accesses=None):
             :py:class:`psyclone.core.VariablesAccessInfo`
 
         '''
-        sym = self.append_integer_reference(function_space.undf_name)
+        if self._kern.iterates_over == "dof":
+            sym = self.append_integer_reference(function_space.bare_undf_name)
+        else:
+            sym = self.append_integer_reference(function_space.undf_name)
         self.append(sym.name, var_accesses)
 
         map_name = function_space.map_name
@@ -621,6 +637,9 @@ def fs_compulsory_field(self, function_space, var_accesses=None):
             sym = self.append_array_reference(map_name, [":", ":"],
                                               ScalarType.Intrinsic.INTEGER)
             self.append(sym.name, var_accesses, var_access_name=sym.name)
+        elif self._kern.iterates_over == "dof":
+            # Dofmaps are not required for DoF kernels
+            pass
         else:
             # Pass the dofmap for the cell column
             cell_name, cell_ref = self.cell_ref_name(var_accesses)

src/psyclone/domain/lfric/kern_stub_arg_list.py

@@ -328,6 +328,7 @@ def fs_compulsory_field(self, function_space, var_accesses=None):
 
         '''
         self.append(function_space.undf_name, var_accesses)
+
         self.append(function_space.map_name, var_accesses)
 
     def basis(self, function_space, var_accesses=None):

src/psyclone/domain/lfric/lfric_kern.py

@@ -582,8 +582,8 @@ def gen_stub(self):
         const = LFRicConstants()
         supported_operates_on = const.USER_KERNEL_ITERATION_SPACES[:]
         # TODO #925 Add support for 'domain' kernels
-        # TODO #1351 Add suport for 'dof' kernels
         supported_operates_on.remove("domain")
+        # TODO #1351 Add support for 'dof' kernels
         supported_operates_on.remove("dof")
 
         # Check operates-on (iteration space) before generating code

src/psyclone/domain/lfric/lfric_kern_call_factory.py

@@ -70,6 +70,9 @@ def create(call, parent=None):
             # We still need a loop object though as that is where the logic
             # for handling halo exchanges is currently implemented.
             loop_type = "null"
+        elif call.ktype.iterates_over == "dof":
+            # Loop over dofs within a field.
+            loop_type = "dof"
         else:
             # Loop over cells, indicated by an empty string.
             loop_type = ""

src/psyclone/domain/lfric/lfric_loop.py

@@ -225,9 +225,8 @@ def load(self, kern):
         # Loop bounds
         self.set_lower_bound("start")
         const = LFRicConstants()
-        if isinstance(kern, LFRicBuiltIn):
-            # If the kernel is a built-in/pointwise operation
-            # then this loop must be over DoFs
+        if kern.iterates_over == "dof":
+            # This loop must be over DoFs
             if Config.get().api_conf("lfric").compute_annexed_dofs \
                and Config.get().distributed_memory \
                and not kern.is_reduction:
@@ -475,14 +474,20 @@ def _upper_bound_fortran(self):
         if self._upper_bound_name in ["ndofs", "nannexed"]:
             if Config.get().distributed_memory:
                 if self._upper_bound_name == "ndofs":
-                    result = (f"{self.field.proxy_name_indexed}%"
-                              f"{self.field.ref_name()}%get_last_dof_owned()")
+                    result = (
+                        f"{self.field.proxy_name_indexed}%"
+                        f"{self.field.ref_name()}%get_last_dof_owned()")
                 else:  # nannexed
                     result = (
                         f"{self.field.proxy_name_indexed}%"
                         f"{self.field.ref_name()}%get_last_dof_annexed()")
             else:
-                result = self._kern.undf_name
+                if isinstance(self._kern, LFRicBuiltIn):
+                    result = self._kern.undf_name
+                else:
+                    # User-defined dof kernel has undf_name in a different
+                    # location
+                    result = self._field_space.bare_undf_name
             return result
         if self._upper_bound_name == "ncells":
             if Config.get().distributed_memory:

src/psyclone/dynamo0p3.py

@@ -1215,10 +1215,16 @@ def __init__(self, kern_or_invoke):
             # field proxy and undf is not required.
             if self._invoke and self._invoke.field_on_space(function_space):
                 if not (self._dofs_only and Config.get().distributed_memory):
-                    self._var_list.append(function_space.undf_name)
+                    if not self._invoke.operates_on_dofs_only:
+                        self._var_list.append(function_space.undf_name)
+                    else:
+                        self._var_list.append(function_space.bare_undf_name)
             elif self._kernel and \
                     function_space.field_on_space(self._kernel.arguments):
-                self._var_list.append(function_space.undf_name)
+                if not self._kernel.iterates_over == "dof":
+                    self._var_list.append(function_space.undf_name)
+                else:
+                    self._var_list.append(function_space.bare_undf_name)
 
     def _stub_declarations(self, parent):
         '''
@@ -1295,7 +1301,10 @@ def initialise(self, parent):
             # from the field proxy and undf is not required.
             if not (self._dofs_only and Config.get().distributed_memory):
                 if self._invoke.field_on_space(function_space):
-                    undf_name = function_space.undf_name
+                    if self._invoke.operates_on_dofs_only:
+                        undf_name = function_space.bare_undf_name
+                    else:
+                        undf_name = function_space.undf_name
                     parent.add(AssignGen(parent, lhs=undf_name,
                                          rhs=name + "%" +
                                          arg.ref_name(function_space) +