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

src/psyclone/domain/lfric/lfric_kern_metadata.py

@@ -614,20 +614,19 @@ 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)
+
+        # 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

@@ -130,32 +130,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 +183,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), scalar_arg, "
+                "field_vec_1_data, field_vec_2_data, field_vec_3_data)")
 
     assert expected in code
     # Check compilation
@@ -293,7 +268,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 +278,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 +291,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), scalar_arg, field_vec_1_data, "
+        "field_vec_2_data, field_vec_3_data)\n"
         "      END DO\n"
     )
     assert output in code
@@ -326,7 +305,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), scalar_arg, field_vec_1_data, "
+                "field_vec_2_data, field_vec_3_data)\n"
                 "      END DO\n"
                 "      !\n"
                 "      ! Set halos dirty/clean for fields modified in the "
@@ -341,7 +321,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), scalar_arg, field_vec_1_data, "
+                "field_vec_2_data, field_vec_3_data)\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

@@ -46,9 +46,10 @@ program single_invoke_dofs
 
   type(field_type) :: f1, f2, f3, f4
   real(kind=r_def) :: scalar_arg
+  type(field_type) :: field_vec(3)
 
   call invoke(                            &
-       testkern_dofs_type(f1, f2, f3, f4, scalar_arg) &
+       testkern_dofs_type(f1, f2, f3, f4, scalar_arg, field_vec) &
           )
 
 end program single_invoke_dofs

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

@@ -47,9 +47,10 @@ program multikernel_invokes_cell_dof_builtin
 
   type(field_type) :: f1, f2, f3, f4, m1, m2
   real(kind=r_def) :: a, scalar_arg
+  type(field_type) :: field_vec(3)
 
   call invoke(                                         &
-       testkern_dofs_type(f1, f2, f3, f4, scalar_arg), &
+       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)                &
           )

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

@@ -42,26 +42,30 @@ module testkern_dofs_mod
   implicit none
 
   type, extends(kernel_type) :: testkern_dofs_type
-     type(arg_type), dimension(5) :: meta_args =          &
+     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_scalar, gh_real, gh_read),       &
+             arg_type(gh_field*3, gh_real, gh_read,  &
+              ANY_SPACE_1)                                &
            /)
+
      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, scalar_arg, field_vec)
     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)
 
   end subroutine testkern_dofs_code