Merge pull request #2389 from stfc/928_field_bcs_kernel

(Towards #928) add support for  field bcs kernel in kernel_types.py

changelog

@@ -57,7 +57,7 @@
 	21) PR #2362 for #2243 and #2312. Add support for Fortran save statements.
 
 	22) PR #2379 for #2370. Adds a Node.get_sibling_lists() method.
-	
+
 	23) PR #2366 towards #342. Improve symbols and tighten-up ACCRoutineTrans
 	validation.
 
@@ -77,6 +77,10 @@
 	29) PR #2406 for #2404. Fixes bug in setop_random kernel (was only
 	filling the lowest vertical level).
 
+	30) PR #2389 toward #928. Adds support for field_bcs_kernel when
+	parsing kernels using kernel_interface.py. This is required when
+	validating LFRic kernel arguments.
+
 release 2.4.0 29th of September 2023
 
 	1) PR #1758 for #1741. Splits the PSyData read functionality into a

src/psyclone/domain/lfric/arg_ordering.py

@@ -47,6 +47,8 @@
 # a circular dependency.
 from psyclone.domain.lfric import LFRicConstants
 from psyclone.domain.lfric.lfric_symbol_table import LFRicSymbolTable
+from psyclone.domain.lfric.metadata_to_arguments_rules import (
+    MetadataToArgumentsRules)
 from psyclone.errors import GenerationError, InternalError
 from psyclone.psyir.nodes import ArrayReference, Reference
 from psyclone.psyir.symbols import ScalarType
@@ -472,8 +474,8 @@ def generate(self, var_accesses=None):
                     self.diff_basis(unique_fs, var_accesses=var_accesses)
             # Fix for boundary_dofs array to the boundary condition
             # kernel (enforce_bc_kernel) arguments
-            if self._kern.name.lower() == "enforce_bc_code" and \
-               unique_fs.orig_name.lower() == "any_space_1":
+            if (MetadataToArgumentsRules.bc_kern_regex.match(self._kern.name)
+                    and unique_fs.orig_name.lower() == "any_space_1"):
                 self.field_bcs_kernel(unique_fs, var_accesses=var_accesses)
 
         # Add boundary dofs array to the operator boundary condition

src/psyclone/domain/lfric/kernel_interface.py

@@ -586,20 +586,54 @@ def diff_basis(self, function_space, var_accesses=None):
                            basis_name_func, first_dim_value_func)
 
     def field_bcs_kernel(self, function_space, var_accesses=None):
-        '''Not implemented.
+        '''Create the boundary-dofs mask argument required for the
+        enforce_bc_code kernel. Adds it to the symbol table and the argument
+        list.
 
         :param function_space: the function space for this boundary condition.
         :type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
-        :param var_accesses: an unused optional argument that stores \
+        :param var_accesses: an unused optional argument that stores
             information about variable accesses.
-        :type var_accesses: :\
-            py:class:`psyclone.core.VariablesAccessInfo`
+        :type var_accesses: :py:class:`psyclone.core.VariablesAccessInfo`
 
-        :raises NotImplementedError: as this method is not implemented.
+        :raises InternalError: if the kernel does not have a single field as
+                               argument.
+        :raises InternalError: if the field argument is not on the
+                               'ANY_SPACE_1' function space.
 
         '''
-        raise NotImplementedError(
-            "TODO #928: field_bcs_kernel not implemented")
+        # Sanity check - expect the enforce_bc_code to have a single argument
+        # that is a field.
+        if (len(self._kern.arguments.args) != 1 or
+                not self._kern.arguments.args[0].is_field):
+            const = LFRicConstants()
+            raise InternalError(
+                f"Kernel '{self._kern.name}' applies boundary conditions to a "
+                f"field and therefore should have a single, field argument "
+                f"(one of {const.VALID_FIELD_NAMES}) but got "
+                f"{[arg.argument_type for arg in self._kern.arguments.args]}")
+        farg = self._kern.arguments.args[0]
+        fspace = farg.function_space
+        # Sanity check - expect the field argument to the enforce_bc_code
+        # kernel to be on the ANY_SPACE_1 space.
+        if fspace.orig_name != "any_space_1":
+            raise InternalError(
+                f"Kernel '{self._kern.name}' applies boundary conditions to a "
+                f"field but the supplied argument, '{farg.name}', is on "
+                f"'{fspace.orig_name}' rather than the expected 'ANY_SPACE_1'")
+
+        ndf_symbol = self._symtab.find_or_create_tag(
+            f"ndf_{fspace.orig_name}", fs=fspace.orig_name,
+            symbol_type=LFRicTypes("NumberOfDofsDataSymbol"),
+            interface=self._read_access)
+
+        # Add the boundary-dofs array argument.
+        sym = self._symtab.find_or_create_tag(
+            f"boundary_dofs_{farg.name}",
+            interface=self._read_access,
+            symbol_type=LFRicTypes("VerticalBoundaryDofMaskDataSymbol"),
+            dims=[Reference(ndf_symbol), 2])
+        self._arglist.append(sym)
 
     def operator_bcs_kernel(self, function_space, var_accesses=None):
         '''Not implemented.

src/psyclone/domain/lfric/lfric_types.py

@@ -31,13 +31,15 @@
 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 # POSSIBILITY OF SUCH DAMAGE.
 # -----------------------------------------------------------------------------
-# Author: J. Henrichs, Bureau of Meteorology
-# Modified: O. Brunt, Met Office
+# Authors: J. Henrichs, Bureau of Meteorology
+#          A. R. Porter, STFC Daresbury Lab
+#          O. Brunt, Met Office
 
 '''This module contains a singleton class that manages LFRic types. '''
 
 
 from collections import namedtuple
+from dataclasses import dataclass
 
 from psyclone.configuration import Config
 from psyclone.domain.lfric.lfric_constants import LFRicConstants
@@ -126,6 +128,7 @@ def init():
         LFRicTypes._create_lfric_dimension()
         LFRicTypes._create_specific_scalars()
         LFRicTypes._create_fields()
+
         # Generate LFRic vector-field-data symbols as subclasses of
         # field-data symbols
         const = LFRicConstants()
@@ -271,30 +274,30 @@ def _create_lfric_dimension():
         # The actual class:
         class LFRicDimension(Literal):
             '''An LFRic-specific scalar integer that captures a literal array
-            dimension which can either have the value 1 or 3. This is used for
-            one of the dimensions in basis and differential basis
-            functions.
+            dimension which can have a value between 1 and 3, inclusive. This
+            is used for one of the dimensions in basis and differential basis
+            functions and also for the vertical-boundary dofs mask.
 
             :param str value: the value of the scalar integer.
 
-            :raises ValueError: if the supplied value is not '1 or '3'.
+            :raises ValueError: if the supplied value is not '1', '2' or '3'.
 
             '''
             # pylint: disable=undefined-variable
             def __init__(self, value):
                 super().__init__(value,
                                  LFRicTypes("LFRicIntegerScalarDataType")())
-                if value not in ['1', '3']:
-                    raise ValueError(f"An LFRic dimension object must be '1' "
-                                     f"or '3', but found '{value}'.")
+                if value not in ['1', '2', '3']:
+                    raise ValueError(f"An LFRic dimension object must be '1', "
+                                     f"'2' or '3', but found '{value}'.")
         # --------------------------------------------------------------------
 
         # Create the required entries in the dictionary
-        LFRicTypes._name_to_class["LFRicDimension"] = LFRicDimension
-        LFRicTypes._name_to_class["LFRIC_SCALAR_DIMENSION"] = \
-            LFRicDimension("1")
-        LFRicTypes._name_to_class["LFRIC_VECTOR_DIMENSION"] = \
-            LFRicDimension("3")
+        LFRicTypes._name_to_class.update({
+            "LFRicDimension": LFRicDimension,
+            "LFRIC_SCALAR_DIMENSION": LFRicDimension("1"),
+            "LFRIC_VERTICAL_BOUNDARIES_DIMENSION": LFRicDimension("2"),
+            "LFRIC_VECTOR_DIMENSION": LFRicDimension("3")})
 
     # ------------------------------------------------------------------------
     @staticmethod
@@ -404,14 +407,24 @@ def _create_fields():
         # list because they are used to create vector field datatypes and
         # symbols.
 
-        # The Array namedtuple has 4 properties: the first determines the
-        # names of the resultant datatype and datasymbol classes, the second
-        # references the generic scalar type classes declared above, the third
-        # property is a textual description of each of the dimensions.
-        # The fourth specifies any additional class properties that should be
-        # declared in the generated datasymbol class.
-        Array = namedtuple('Array',
-                           ["name", "scalar_type", "dims", "properties"])
+        @dataclass(frozen=True)
+        class Array:
+            '''
+            Holds the properties of an LFRic array type, used when generating
+            DataSymbol and DataSymbolType classes.
+
+            :param name: the base name to use for the datatype and datasymbol.
+            :param scalar_type: the name of the LFRic scalar type that this is
+                                an array of.
+            :param dims: textual description of each of the dimensions.
+            :param properties: names of additional class properties that should
+                               be declared in the generated datasymbol class.
+            '''
+            name: str
+            scalar_type: str
+            dims: list        # list[str] notation supported in Python 3.9+
+            properties: list  # ditto
+
         field_datatypes = [
             Array("RealField", "LFRicRealScalarDataType",
                   ["number of unique dofs"], ["fs"]),
@@ -471,7 +484,9 @@ def _create_fields():
             Array("QrWeightsInFaces", "LFRicRealScalarDataType",
                   ["number of qr points in faces"], []),
             Array("QrWeightsInEdges", "LFRicRealScalarDataType",
-                  ["number of qr points in edges"], [])
+                  ["number of qr points in edges"], []),
+            Array("VerticalBoundaryDofMask", "LFRicIntegerScalarDataType",
+                  ["number of dofs", LFRicTypes("LFRicDimension")], [])
             ]
 
         for array_type in array_datatypes + field_datatypes:

src/psyclone/domain/lfric/metadata_to_arguments_rules.py

@@ -38,6 +38,7 @@
 
 '''
 from collections import OrderedDict
+import re
 
 from psyclone.domain.lfric import LFRicConstants
 from psyclone.domain.lfric.kernel import (
@@ -67,6 +68,12 @@ class MetadataToArgumentsRules():
     '''
     _metadata = None
     _info = None
+    # Regex used to identify the special 'enforce_bc_code' kernel that
+    # applies boundary conditions to a field. Allows for the renaming
+    # performed by PSyclone when performing kernel transformations.
+    # TODO #487 - this can be removed when we have metadata to specify
+    # that a kernel applies boundary conditinos.
+    bc_kern_regex = re.compile(r"enforce_bc_(\d+_)?code", flags=re.I)
 
     @classmethod
     def mapping(cls, metadata, info=None):
@@ -538,8 +545,8 @@ def _generate(cls):
 
         # The boundary condition kernel (enforce_bc_kernel) is a
         # special case.
-        if cls._metadata.procedure_name and \
-                cls._metadata.procedure_name.lower() == "enforce_bc_code":
+        if (cls._metadata.procedure_name and
+                cls.bc_kern_regex.match(cls._metadata.procedure_name)):
             cls._field_bcs_kernel()
 
         # The operator boundary condition kernel

src/psyclone/dynamo0p3.py

@@ -5098,8 +5098,11 @@ def __init__(self, node):
         # Check through all the kernel calls to see whether any of them
         # require boundary conditions. Currently this is done by recognising
         # the kernel name.
+        # pylint: disable=import-outside-toplevel
+        from psyclone.domain.lfric.metadata_to_arguments_rules import (
+            MetadataToArgumentsRules)
         for call in self._calls:
-            if call.name.lower() == "enforce_bc_code":
+            if MetadataToArgumentsRules.bc_kern_regex.match(call.name):
                 bc_fs = None
                 for fspace in call.arguments.unique_fss:
                     if fspace.orig_name == "any_space_1":

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

@@ -775,14 +775,78 @@ def test_diff_basis():
     assert diff_basis_symbol.shape[3].upper.symbol is nqpv_symbol
 
 
-@pytest.mark.xfail(reason="Issue #928: this callback is not yet implemented")
-def test_field_bcs_kernel():
+def test_field_bcs_kernel(monkeypatch):
     '''Test that the KernelInterface class field_bcs_kernel method adds the
     expected symbols to the symbol table and the _arglist list.
 
     '''
-    kernel_interface = KernelInterface(None)
+    _, invoke_info = parse(os.path.join(
+        BASE_PATH, "12.2_enforce_bc_kernel.f90"),
+                           api="dynamo0.3")
+    psy = PSyFactory("dynamo0.3",
+                     distributed_memory=False).create(invoke_info)
+    schedule = psy.invokes.invoke_list[0].schedule
+    kernel = schedule[0].loop_body[0]
+    kernel_interface = KernelInterface(kernel)
     kernel_interface.field_bcs_kernel(None)
+    fld_name = kernel.arguments.args[0].name
+    fspace = kernel.arguments.unique_fss[0]
+    fs_name = fspace.orig_name
+    # ndf declared
+    ndf_symbol = kernel_interface._symtab.lookup(f"ndf_{fs_name}")
+    assert isinstance(ndf_symbol, LFRicTypes("NumberOfDofsDataSymbol"))
+    assert isinstance(ndf_symbol.interface, ArgumentInterface)
+    assert (ndf_symbol.interface.access ==
+            kernel_interface._read_access.access)
+    # vertical-boundary dofs mask declared
+    mask_sym = kernel_interface._symtab.lookup(f"boundary_dofs_{fld_name}")
+    assert isinstance(mask_sym,
+                      LFRicTypes("VerticalBoundaryDofMaskDataSymbol"))
+    assert isinstance(mask_sym.interface, ArgumentInterface)
+    assert mask_sym.interface.access == kernel_interface._read_access.access
+    assert len(mask_sym.shape) == 2
+    assert isinstance(mask_sym.shape[0].upper, Reference)
+    assert mask_sym.shape[0].upper.symbol is ndf_symbol
+    assert isinstance(mask_sym.shape[1].upper, Literal)
+    assert mask_sym.shape[1].upper.value == "2"
+
+
+def test_field_bcs_kernel_errors(monkeypatch):
+    '''
+    Test that the field_bcs_kernel method raises the expected errors if the
+    kernel does not have exactly one argument that is itself a field on the
+    'ANY_SPACE_1' function space.
+
+    '''
+    _, invoke_info = parse(os.path.join(
+        BASE_PATH, "1_single_invoke.f90"), api="dynamo0.3")
+    psy = PSyFactory("dynamo0.3",
+                     distributed_memory=False).create(invoke_info)
+    schedule = psy.invokes.invoke_list[0].schedule
+    kernel = schedule[0].loop_body[0]
+    kernel_interface = KernelInterface(kernel)
+    with pytest.raises(InternalError) as err:
+        kernel_interface.field_bcs_kernel(None)
+    assert ("Kernel 'testkern_code' applies boundary conditions to a field "
+            "and therefore should have a single, field argument (one of "
+            "['gh_field']) but got ['gh_scalar', 'gh_field'" in str(err.value))
+    # Repeat for a kernel that *does* have an argument on ANY_SPACE_1.
+    _, invoke_info = parse(os.path.join(
+        BASE_PATH, "12.2_enforce_bc_kernel.f90"), api="dynamo0.3")
+    psy = PSyFactory("dynamo0.3",
+                     distributed_memory=False).create(invoke_info)
+    schedule = psy.invokes.invoke_list[0].schedule
+    kernel = schedule[0].loop_body[0]
+    # Monkeypatch the argument so that it appears to be on the wrong space.
+    monkeypatch.setattr(
+        kernel.arguments._args[0]._function_spaces[0],
+        "_orig_name", "ANY_SPACE_2")
+    kernel_interface = KernelInterface(kernel)
+    with pytest.raises(InternalError) as err:
+        kernel_interface.field_bcs_kernel(None)
+    assert ("Kernel 'enforce_bc_code' applies boundary conditions to a field "
+            "but the supplied argument, 'a', is on 'ANY_SPACE_2' rather than "
+            "the expected 'ANY_SPACE_1'" in str(err.value))
 
 
 @pytest.mark.xfail(reason="Issue #928: this callback is not yet implemented")

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

@@ -146,8 +146,8 @@ def test_scalar_literals():
     assert isinstance(LFRicTypes("LFRicDimension")("3"),
                       lfric_dim_class)
     with pytest.raises(ValueError) as info:
-        LFRicTypes("LFRicDimension")("2")
-    assert ("An LFRic dimension object must be '1' or '3', but found '2'."
+        LFRicTypes("LFRicDimension")("4")
+    assert ("An LFRic dimension object must be '1', '2' or '3', but found '4'."
             in str(info.value))
     # LFRIC_SCALAR_DIMENSION instance
     assert isinstance(LFRicTypes("LFRIC_SCALAR_DIMENSION"), lfric_dim_class)

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

@@ -70,6 +70,24 @@ def check_called(monkeypatch, function, method_name, metadata):
     assert method_name in str(info.value)
 
 
+def test_bc_kern_regex():
+    '''
+    Test the regular expression used to identify the boundary-condition kernel
+    and its transformed forms.
+
+    TODO #487 - this test should be removed once metadata is used to identify
+    the boundary-condition kernel.
+
+    '''
+    cls = MetadataToArgumentsRules
+    assert cls.bc_kern_regex.match("enforce_bc_code")
+    assert cls.bc_kern_regex.match("enforce_BC_code")
+    assert cls.bc_kern_regex.match("enforce_BC_1099_code")
+    assert not cls.bc_kern_regex.match("other_bc_code")
+    assert not cls.bc_kern_regex.match("enforce_bc_a1_code")
+    assert not cls.bc_kern_regex.match("enforce_bc_1a_code")
+
+
 def test_mapping(monkeypatch):
     '''Test the MetadataToArgumentRules class mapping method works as
     expected.