#2711 dof access added

src/psyclone/domain/lfric/kern_call_arg_list.py

@@ -379,11 +379,20 @@ 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.is_dofkern:
+            # 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")
+            dof_sym = Reference(dof_sym)
+            self.append_array_reference(sym.name, [dof_sym],
+                                        ScalarType.Intrinsic.INTEGER,
+                                        symbol=sym)
+        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
@@ -620,6 +629,10 @@ 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.is_dofkern:
+            # Dofmaps are not compatible with user-defined DoF kernels, so
+            # just pass without creating a dofmap.
+            pass
         else:
             # Pass the dofmap for the cell column
             cell_name, cell_ref = self.cell_ref_name(var_accesses)

src/psyclone/domain/lfric/lfric_loop.py

@@ -225,7 +225,7 @@ def load(self, kern):
         # Loop bounds
         self.set_lower_bound("start")
         const = LFRicConstants()
-        if isinstance(kern, LFRicBuiltIn):
+        if isinstance(kern, LFRicBuiltIn) or kern.is_dofkern:
             # If the kernel is a built-in/pointwise operation
             # then this loop must be over DoFs
             if Config.get().api_conf("lfric").compute_annexed_dofs \
@@ -234,10 +234,6 @@ 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:
@@ -477,22 +473,22 @@ 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 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
+            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:
-                # 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
+                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.undf_name
             return result
         if self._upper_bound_name == "ncells":
             if Config.get().distributed_memory:

src/psyclone/tests/domain/lfric/dofkern_test.py

@@ -37,13 +37,15 @@
 This module tests the metadata validation in LFRicKernMetadata of
 user-supplied kernels operating on degrees of freedom (dofs)
 '''
-
+import os
 import pytest
 
 from fparser import api as fpapi
 
 from psyclone.configuration import Config
-from psyclone.domain.lfric import LFRicKernMetadata
+from psyclone.domain.lfric import LFRicKernMetadata, LFRicKern, LFRicLoop
+from psyclone.parse.algorithm import parse
+from psyclone.psyGen import PSyFactory
 from psyclone.parse.utils import ParseError
 
 
@@ -54,6 +56,9 @@ def setup():
     yield
     Config._instance = None
 
+BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(
+                os.path.abspath(__file__)))), "test_files", "dynamo0p3")
+TEST_API = "lfric"
 
 CODE = '''
         module testkern_dofs_mod
@@ -142,3 +147,37 @@ def test_dof_kernel_invalid_field_vector():
     assert ("Kernel 'testkern_dofs_type' operates on 'dof' but has a vector "
             "argument 'gh_field*3'. This is not permitted in the LFRic API."
             in str(excinfo.value))
+
+
+def test_dof_kernel():
+    ''' Check that we raise an exception if we encounter metadata
+    for a dof kernel with a field vector.
+
+    '''
+    # Substitute field for field vector
+    code = CODE.replace(
+        """
+                    (/ arg_type(gh_field, gh_real, gh_write, w1),  &
+                        arg_type(gh_field, gh_real, gh_read, w2)   &
+        """,
+        """
+                    (/ arg_type(gh_field, gh_real, gh_write, w1),  &
+                        arg_type(gh_field, gh_real, gh_read, w1)   &
+        """,
+        1)
+    ast = fpapi.parse(code, ignore_comments=False)
+    name = "testkern_dofs_type"
+    md = LFRicKernMetadata(ast, name=name)
+    kern = LFRicKern()
+    kern.load_meta(ktype=md)
+
+
+def test_gen():
+    _, invoke_info = parse(os.path.join(BASE_PATH,
+                                        "1.14_single_invoke_dofs.f90"),
+                           api=TEST_API)
+    print(invoke_info)
+    psy = PSyFactory(TEST_API, distributed_memory=True).create(invoke_info)
+    sched = psy.invokes.invoke_list[0].schedule
+    print(sched.view())
+    print(psy.gen)