#2711 Loops now use undf_name if dofkern

src/psyclone/domain/lfric/lfric_kern.py

@@ -110,6 +110,8 @@ def __init__(self):
         self._reference_element = None
         # The mesh properties required by this kernel
         self._mesh_properties = None
+        # An identifier for a user-defined dof kernel
+        self._is_dofkern = None
         # Initialise kinds (precisions) of all kernel arguments (start
         # with 'real' and 'integer' kinds)
         api_config = Config.get().api_conf("lfric")
@@ -289,6 +291,9 @@ def _setup(self, ktype, module_name, args, parent, check=True):
             intergrid = DynInterGrid(fine_args[0], coarse_args[0])
             self._intergrid_ref = intergrid
 
+        # Set the identifier for a dof kernels
+        self._is_dofkern = ktype.is_user_dofkern
+
         const = LFRicConstants()
         # Check that all specified evaluator shapes are recognised
         invalid_shapes = set(self._eval_shapes) \
@@ -394,6 +399,14 @@ def is_intergrid(self):
         '''
         return self._intergrid_ref is not None
 
+    @property
+    def is_dofkern(self):
+        '''
+        :return: True if it is a user-defined dof kernel, False otherwise
+        :rtype: bool
+        '''
+        return self._is_dofkern
+
     @property
     def colourmap(self):
         '''

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_kern_metadata.py

@@ -95,6 +95,10 @@ def __init__(self, ast, name=None):
         # set to True if all the checks in '_validate_inter_grid()' pass.
         self._is_intergrid = False
 
+        # Whether or not this is a user-defined kernel operating on dofs. This
+        # is set to True if the `_validate_operates_on_dof` method succeeds.
+        self._is_user_dofkern = False
+
         # Parse the 'arg_type' metadata
         self._arg_descriptors = []
         for idx, arg_type in enumerate(self._inits):
@@ -636,6 +640,8 @@ def _validate_operates_on_dof(self, need_evaluator):
                 f"fields on different function spaces: {sorted(arg_fs_names)}."
                 f" This is not permitted in the LFRic API.")
 
+        self._is_user_dofkern = True
+
     def _validate_only_scalars_and_fields(self):
         '''
         Checks metadata arguments are either fields or scalars by comparing
@@ -770,6 +776,16 @@ def is_intergrid(self):
         '''
         return self._is_intergrid
 
+    @property
+    def is_user_dofkern(self):
+        '''
+        Returns whether or not this is a user-defined kernel operating on dofs.
+
+        :return: True if kernel is a user-defined dof kernel, False otherwise
+        :rtype: bool
+        '''
+        return self._is_user_dofkern
+
 
 # ---------- Documentation utils -------------------------------------------- #
 # The list of module members that we wish AutoAPI to generate

src/psyclone/domain/lfric/lfric_loop.py

@@ -234,6 +234,10 @@ def load(self, kern):
                 self.set_upper_bound("nannexed")
             else:
                 self.set_upper_bound("ndofs")
+        elif kern.is_dofkern:
+            # If the kernel is user-defined and operats on dofs, always
+            # set upper bound to 'ndofs'
+            self.set_upper_bound("ndofs")
         else:
             if Config.get().distributed_memory:
                 if self._field.is_operator:
@@ -473,16 +477,21 @@ def _upper_bound_fortran(self):
             sym = sym_tab.find_or_create_tag(root_name)
             return f"{sym.name}(colour, {depth})"
         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()")
-                else:  # nannexed
-                    result = (
-                        f"{self.field.proxy_name_indexed}%"
-                        f"{self.field.ref_name()}%get_last_dof_annexed()")
+            if isinstance(self._kern, LFRicBuiltIn):
+                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()")
+                    else:  # nannexed
+                        result = (
+                            f"{self.field.proxy_name_indexed}%"
+                            f"{self.field.ref_name()}%get_last_dof_annexed()")
+                else:
+                    result = self._kern.undf_name
             else:
-                result = self._kern.undf_name
+                # User-defined dof kernel has undf_name in a different location
+                # so use that here instead of calling `get_undf` again
+                result = self._field_space.undf_name
             return result
         if self._upper_bound_name == "ncells":
             if Config.get().distributed_memory: