#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.'''
@@ -614,20 +615,13 @@ def _validate_operates_on_dof(self, need_evaluator):
         self._validate_no_mesh_properties()
         self._validate_not_intergrid()
 
-        for arg in self._arg_descriptors:
-            # No vector arguments are permitted
-            if arg.vector_size > 1:
-                raise ParseError(
-                    f"Kernel '{self.name}' operates on 'dof' but has a "
-                    f"vector argument '{arg.argument_type}*{arg.vector_size}'."
-                    f" This is not permitted in the LFRic API.")
-
         # Check function spaces are the same
         # list out all function spaces
         arg_fs_names = set()
         for arg in self._arg_descriptors:
             for fs_name in arg.function_spaces:
                 arg_fs_names.add(fs_name)
+
         # 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
@@ -130,32 +131,6 @@ def test_dof_kernel_invalid_arg():
             in str(excinfo.value))
 
 
-def test_dof_kernel_invalid_field_vector():
-    '''
-    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*3, gh_real, gh_write, w1), &
-                      arg_type(gh_scalar, gh_real, gh_read)  &
-        """,
-        1)
-    ast = fpapi.parse(code, ignore_comments=False)
-    name = "testkern_dofs_type"
-    with pytest.raises(ParseError) as excinfo:
-        _ = LFRicKernMetadata(ast, name=name)
-    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_upper_bounds(monkeypatch, annexed, dist_mem, tmpdir):
     '''
     Checks that the correct upper bound is generated for a dof-kernel for all
@@ -209,7 +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)")
+                "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
@@ -293,7 +269,8 @@ def test_multi_invoke_cell_dof_builtin(tmpdir, monkeypatch, annexed, dist_mem):
     # Consistent declarations
     output = (
         "      REAL(KIND=r_def), intent(in) :: scalar_arg, a\n"
-        "      TYPE(field_type), intent(in) :: f1, f2, f3, f4, m1, m2\n"
+        "      TYPE(field_type), intent(in) :: f1, f2, f3, f4, "
+        "field_vec(3), m1, m2\n"
         "      INTEGER(KIND=i_def) cell\n"
         "      INTEGER(KIND=i_def) df\n"
         "      INTEGER(KIND=i_def) loop2_start, loop2_stop\n"
@@ -302,6 +279,8 @@ def test_multi_invoke_cell_dof_builtin(tmpdir, monkeypatch, annexed, dist_mem):
         "      INTEGER(KIND=i_def) nlayers_f1\n"
         "      REAL(KIND=r_def), pointer, dimension(:) :: m2_data => null()\n"
         "      REAL(KIND=r_def), pointer, dimension(:) :: m1_data => null()\n"
+        "      REAL(KIND=r_def), pointer, dimension(:) :: field_vec_1_data =>"
+        " null(), field_vec_2_data => null(), field_vec_3_data => null()\n"
         "      REAL(KIND=r_def), pointer, dimension(:) :: f4_data => null()\n"
         "      REAL(KIND=r_def), pointer, dimension(:) :: f3_data => null()\n"
         "      REAL(KIND=r_def), pointer, dimension(:) :: f2_data => null()\n"
@@ -313,7 +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)\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
@@ -326,7 +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)\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 "
@@ -341,7 +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)\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,10 +45,11 @@ program single_invoke_dofs
   implicit none
 
   type(field_type) :: f1, f2, f3, f4
+  type(field_type) :: field_vec(3)
   real(kind=r_def) :: scalar_arg
 
-  call invoke(                            &
-       testkern_dofs_type(f1, f2, f3, f4, scalar_arg) &
+  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,12 +47,13 @@ program multikernel_invokes_cell_dof_builtin
   implicit none
 
   type(field_type) :: f1, f2, f3, f4, m1, m2
+  type(field_type) :: field_vec(3)
   real(kind=r_def) :: a, scalar_arg
 
-  call invoke(                                         &
-       testkern_dofs_type(f1, f2, f3, f4, scalar_arg), &
-       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
 
@@ -42,26 +42,33 @@ module testkern_dofs_mod
   implicit none
 
   type, extends(kernel_type) :: testkern_dofs_type
-     type(arg_type), dimension(5) :: 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)        &
+     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_field*3, gh_real, gh_read,  w1), &
+             arg_type(gh_scalar,  gh_real, gh_read)       &
            /)
+
      integer :: operates_on = DOF
    contains
      procedure, nopass :: code => testkern_dofs_code
   end type testkern_dofs_type
 
 contains
 
-  subroutine testkern_dofs_code(a, b, c, d, scalar_arg)
+  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
+    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