#2892 Allow user-defined DoF kernels to act on field vectors

doc/user_guide/dynamo0p3.rst

@@ -1002,12 +1002,14 @@ Kernels should follow those for General-Purpose Kernels.
 
 The list of rules for DoF Kernels is as follows:
 
-1) A DoF Kernel must have at least one argument that is a field.
-   This rule reflects that a Kernel operates on some subset of the
-   whole domain and is therefore designed to be called from within
-   a loop that iterates over those subsets of the domain. Only fields
-   (as opposed to e.g. field vectors or operators) are accepted for DoF
-   Kernels because only they have a single value at each DoF.
+1) A DoF Kernel must have at least one argument that is a field. This rule
+   reflects that a Kernel operates on some subset of the whole domain
+   and is therefore designed to be called from within a loop that iterates
+   over those subsets of the domain. Fields (as opposed to e.g. operators)
+   are accepted for DoF Kernels because only they have a single value at
+   each DoF. Field vectors can be represented in a DoF kernel as a
+   collection of field arguments, each one corresponding to an index in the
+   field vector.
 
 2) All Kernel arguments must be either fields or scalars (`real-` and/or
    `integer`-valued). DoF Kernels cannot accept operators.
@@ -2533,6 +2535,12 @@ with PSyclone's naming conventions, are:
          passed in separately. Again, the intent is determined from the
          metadata (see :ref:`meta_args <lfric-api-meta-args>`).
 
+   3) For each field vector in the order specified by the meta_args metadata,
+      there needs to be an equivalent number of arguments in the kernel as
+      the dimension of the field vector. The dimension is specified in the
+      metadata. The arguments must be ordered following the indexing of the
+      field vector.
+
 .. _lfric-kernel-arg-intents:
 
 Argument Intents
@@ -2663,11 +2671,11 @@ metadata for the Built-ins that update a ``real``-valued field,
   type, public, extends(kernel_type) :: aX_plus_bY
      private
      type(arg_type) :: meta_args(5) = (/                              &
-          arg_type(GH_FIELD,  GH_REAL  GH_WRITE, ANY_SPACE_1),        &
+          arg_type(GH_FIELD,  GH_REAL, GH_WRITE, ANY_SPACE_1),        &
           arg_type(GH_SCALAR, GH_REAL, GH_READ              ),        &
           arg_type(GH_FIELD,  GH_REAL, GH_READ,  ANY_SPACE_1),        &
           arg_type(GH_SCALAR, GH_REAL, GH_READ              ),        &
-          arg_type(GH_FIELD,  GH_REAL  GH_READ,  ANY_SPACE_1)         &
+          arg_type(GH_FIELD,  GH_REAL, GH_READ,  ANY_SPACE_1)         &
           /)
      integer :: operates_on = DOF
    contains

src/psyclone/domain/lfric/kern_call_arg_list.py

@@ -32,7 +32,7 @@
 # POSSIBILITY OF SUCH DAMAGE.
 # -----------------------------------------------------------------------------
 # Authors R. W. Ford, A. R. Porter and S. Siso, STFC Daresbury Lab
-# Modified I. Kavcic, A. Coughtrie and L. Turner, Met Office
+# Modified I. Kavcic, A. Coughtrie, L. Turner, and A. Pirrie, Met Office
 # Modified J. Henrichs, Bureau of Meteorology
 
 '''This module implements a class that manages the argument for a kernel
@@ -358,14 +358,29 @@ def field_vector(self, argvect, var_accesses=None):
         '''
         suffix = LFRicConstants().ARG_TYPE_SUFFIX_MAPPING[
             argvect.argument_type]
+
         # The range function below returns values from
         # 1 to the vector size which is what we
         # require in our Fortran code
         for idx in range(1, argvect.vector_size + 1):
+            # Look-up the name of the variable that stores the reference to
+            # the data in this field.
             cmpt_sym = self._symtab.lookup_with_tag(
                 f"{argvect.name}_{idx}:{suffix}")
-            self.psyir_append(Reference(cmpt_sym))
-            text = cmpt_sym.name
+            if self._kern.iterates_over == "dof":
+                # 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(cmpt_sym.name,
+                                            [Reference(dof_sym)],
+                                            ScalarType.Intrinsic.INTEGER,
+                                            symbol=cmpt_sym)
+                # Append the dof symbol
+                text = f"{cmpt_sym.name}({dof_sym.name})"
+            else:
+                self.psyir_append(Reference(cmpt_sym))
+                text = cmpt_sym.name
             self.append(text, metadata_posn=argvect.metadata_index)
 
         if var_accesses is not None:

src/psyclone/domain/lfric/lfric_kern_metadata.py

@@ -35,6 +35,7 @@
 # Modified I. Kavcic, A. Coughtrie, L. Turner and O. Brunt, Met Office
 # Modified J. Henrichs, Bureau of Meteorology
 # Modified A. B. G. Chalk and N. Nobre, STFC Daresbury Lab
+# Modified A. Pirrie, Met Office
 
 ''' This module implements the PSyclone LFRic API by capturing the Kernel
 subroutine code and metadata describing the subroutine for the LFRic API.'''
@@ -621,12 +622,6 @@ def _validate_operates_on_dof(self, need_evaluator):
             for fs_name in arg.function_spaces:
                 arg_fs_names.add(fs_name)
 
-        # Previously the field vector test fails before this
-        # In those cases stop the test before the function space check
-        for arg in self._arg_descriptors:
-            if arg.vector_size > 1:
-                return
-
         # dof kernels should only have one function space so a set of fs
         # names should be of length 1
         if len(arg_fs_names) > 1:

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

@@ -32,6 +32,7 @@
 # POSSIBILITY OF SUCH DAMAGE.
 # -----------------------------------------------------------------------------
 # Author O. Brunt, Met Office
+# Modified A. Pirrie, Met Office
 
 '''
 This module tests metadata validation and code generation of
@@ -183,8 +184,8 @@ def test_indexed_field_args(tmpdir):
     code = str(psy.gen)
 
     expected = ("CALL testkern_dofs_code(f1_data(df), f2_data(df), "
-                "f3_data(df), f4_data(df), scalar_arg, "
-                "field_vec_1_data, field_vec_2_data, field_vec_3_data)")
+                "f3_data(df), f4_data(df), field_vec_1_data(df), "
+                "field_vec_2_data(df), field_vec_3_data(df), scalar_arg)")
 
     assert expected in code
     # Check compilation
@@ -291,8 +292,8 @@ def test_multi_invoke_cell_dof_builtin(tmpdir, monkeypatch, annexed, dist_mem):
     output = (
         "      DO df = loop0_start, loop0_stop, 1\n"
         "        CALL testkern_dofs_code(f1_data(df), f2_data(df), "
-        "f3_data(df), f4_data(df), scalar_arg, field_vec_1_data, "
-        "field_vec_2_data, field_vec_3_data)\n"
+        "f3_data(df), f4_data(df), field_vec_1_data(df), "
+        "field_vec_2_data(df), field_vec_3_data(df), scalar_arg)\n"
         "      END DO\n"
     )
     assert output in code
@@ -305,8 +306,8 @@ def test_multi_invoke_cell_dof_builtin(tmpdir, monkeypatch, annexed, dist_mem):
             output = (
                 "      DO df = loop0_start, loop0_stop, 1\n"
                 "        CALL testkern_dofs_code(f1_data(df), f2_data(df), "
-                "f3_data(df), f4_data(df), scalar_arg, field_vec_1_data, "
-                "field_vec_2_data, field_vec_3_data)\n"
+                "f3_data(df), f4_data(df), field_vec_1_data(df), "
+                "field_vec_2_data(df), field_vec_3_data(df), scalar_arg)\n"
                 "      END DO\n"
                 "      !\n"
                 "      ! Set halos dirty/clean for fields modified in the "
@@ -321,8 +322,8 @@ def test_multi_invoke_cell_dof_builtin(tmpdir, monkeypatch, annexed, dist_mem):
             output = (
                 "      DO df = loop0_start, loop0_stop, 1\n"
                 "        CALL testkern_dofs_code(f1_data(df), f2_data(df), "
-                "f3_data(df), f4_data(df), scalar_arg, field_vec_1_data, "
-                "field_vec_2_data, field_vec_3_data)\n"
+                "f3_data(df), f4_data(df), field_vec_1_data(df), "
+                "field_vec_2_data(df), field_vec_3_data(df), scalar_arg)\n"
                 "      END DO\n"
                 "      !\n"
                 "      ! Set halos dirty/clean for fields modified in the "

src/psyclone/tests/test_files/dynamo0p3/1.14_single_invoke_dofs.f90

@@ -32,7 +32,7 @@
 ! POSSIBILITY OF SUCH DAMAGE.
 ! -----------------------------------------------------------------------------
 ! Author I. Kavcic, Met Office
-! Modified O. Brunt, Met Office
+! Modified O. Brunt, and A. Pirrie, Met Office
 
 program single_invoke_dofs
 
@@ -45,11 +45,11 @@ program single_invoke_dofs
   implicit none
 
   type(field_type) :: f1, f2, f3, f4
-  real(kind=r_def) :: scalar_arg
   type(field_type) :: field_vec(3)
+  real(kind=r_def) :: scalar_arg
 
-  call invoke(                            &
-       testkern_dofs_type(f1, f2, f3, f4, scalar_arg, field_vec) &
+  call invoke(                                                   &
+       testkern_dofs_type(f1, f2, f3, f4, field_vec, scalar_arg) &
           )
 
 end program single_invoke_dofs

src/psyclone/tests/test_files/dynamo0p3/4.17_multikernel_invokes_cell_dof_builtin.f90

@@ -32,6 +32,7 @@
 ! POSSIBILITY OF SUCH DAMAGE.
 ! -----------------------------------------------------------------------------
 ! Author O. Brunt, Met Office
+! Modified A. Pirrie, Met Office
 
 program multikernel_invokes_cell_dof_builtin
 
@@ -46,13 +47,13 @@ program multikernel_invokes_cell_dof_builtin
   implicit none
 
   type(field_type) :: f1, f2, f3, f4, m1, m2
-  real(kind=r_def) :: a, scalar_arg
   type(field_type) :: field_vec(3)
+  real(kind=r_def) :: a, scalar_arg
 
-  call invoke(                                         &
-       testkern_dofs_type(f1, f2, f3, f4, scalar_arg, field_vec), &
-       testkern_type(a, f1, f2, m1, m2),               &
-       inc_aX_plus_Y(0.5_r_def, f1, f2)                &
+  call invoke(                                                    &
+       testkern_dofs_type(f1, f2, f3, f4, field_vec, scalar_arg), &
+       testkern_type(a, f1, f2, m1, m2),                          &
+       inc_aX_plus_Y(0.5_r_def, f1, f2)                           &
           )
 
 end program multikernel_invokes_cell_dof_builtin

src/psyclone/tests/test_files/dynamo0p3/testkern_dofs_mod.f90

@@ -30,7 +30,7 @@
 ! OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ! -----------------------------------------------------------------------------
 ! Author R. W. Ford, STFC Daresbury Lab
-! Modified I. Kavcic, O. Brunt, Met Office
+! Modified I. Kavcic, O. Brunt, and A. Pirrie, Met Office
 
 module testkern_dofs_mod
 
@@ -43,13 +43,12 @@ module testkern_dofs_mod
 
   type, extends(kernel_type) :: testkern_dofs_type
      type(arg_type), dimension(6) :: meta_args =          &
-          (/ arg_type(gh_field, gh_real, gh_write, w1),   &
-             arg_type(gh_field, gh_real, gh_read,  w1),   &
-             arg_type(gh_field, gh_real, gh_read,  w1),   &
-             arg_type(gh_field, gh_real, gh_read,  w1),   &
-             arg_type(gh_scalar, gh_real, gh_read),       &
-             arg_type(gh_field*3, gh_real, gh_read,  &
-              ANY_SPACE_1)                                &
+          (/ arg_type(gh_field,   gh_real, gh_write, w1), &
+             arg_type(gh_field,   gh_real, gh_read,  w1), &
+             arg_type(gh_field,   gh_real, gh_read,  w1), &
+             arg_type(gh_field,   gh_real, gh_read,  w1), &
+             arg_type(gh_field*3, gh_real, gh_read,  w1), &
+             arg_type(gh_scalar,  gh_real, gh_read)       &
            /)
 
      integer :: operates_on = DOF
@@ -59,13 +58,17 @@ module testkern_dofs_mod
 
 contains
 
-  subroutine testkern_dofs_code(a, b, c, d, scalar_arg, field_vec)
+  subroutine testkern_dofs_code(a, b, c, d,  &
+                                field_vec_1, &
+                                field_vec_2, &
+                                field_vec_3, &
+                                scalar_arg)
     implicit none
 
-    real(kind=r_def),    intent(inout) :: a
-    real(kind=r_def),    intent(in)    :: b, c, d
-    real(kind=r_def),    intent(in)    :: scalar_arg
-    type(field_type),    intent(in)    :: field_vec(3)
+    real(kind=r_def), intent(inout) :: a
+    real(kind=r_def), intent(in)    :: b, c, d
+    real(kind=r_def), intent(in)    :: field_vec_1, field_vec_2, field_vec_3
+    real(kind=r_def), intent(in)    :: scalar_arg
 
   end subroutine testkern_dofs_code