Merge pull request #2678 from stfc/2636_nlevels_per_kernel

(Closes #2636) alter LFRic PSy-layer to lookup nlevels for each kernel

changelog

@@ -225,6 +225,9 @@
 
 	77) PR #2708 for #2663. Add LFRic OpenACC in the integration tests.
 
+	78) PR #2678 for #2636. Alter LFRic PSy-layer to lookup nlevels for
+	each kernel
+
 release 2.5.0 14th of February 2024
 
 	1) PR #2199 for #2189. Fix bugs with missing maps in enter data

doc/user_guide/dynamo0p3.rst

@@ -1934,7 +1934,9 @@ conventions, are:
 1) If an LMA operator is passed then include the ``cells`` argument.
    ``cells`` is an ``integer`` of kind ``i_def`` and has intent ``in``.
 2) Include ``nlayers``, the number of layers in a column. ``nlayers``
-   is an ``integer`` of kind ``i_def`` and has intent ``in``.
+   is an ``integer`` of kind ``i_def`` and has intent ``in``. PSyclone
+   will obtain the value of ``nlayers`` to use for a particular kernel
+   from the first field (in the argument list) that is written to.
 3) For each scalar/field/vector_field/operator in the order specified by
    the meta_args metadata:
 

src/psyclone/domain/lfric/__init__.py

@@ -51,6 +51,7 @@
 from psyclone.domain.lfric.kern_call_invoke_arg_list import \
     KernCallInvokeArgList
 from psyclone.domain.lfric.kernel_interface import KernelInterface
+from psyclone.domain.lfric.lfric_cell_iterators import LFRicCellIterators
 from psyclone.domain.lfric.lfric_extract_driver_creator import \
     LFRicExtractDriverCreator
 from psyclone.domain.lfric.lfric_symbol_table import LFRicSymbolTable
@@ -84,6 +85,7 @@
     'KernelInterface',
     'KernStubArgList',
     'LFRicArgDescriptor',
+    'LFRicCellIterators',
     'LFRicCollection',
     'LFRicConstants',
     'LFRicDofmaps',

src/psyclone/domain/lfric/kern_call_arg_list.py

@@ -218,15 +218,16 @@ def mesh_height(self, var_accesses=None):
         '''Add mesh height (nlayers) to the argument list and if supplied
         stores this access in var_accesses.
 
-        :param var_accesses: optional VariablesAccessInfo instance to store \
+        :param var_accesses: optional VariablesAccessInfo instance to store
             the information about variable accesses.
-        :type var_accesses: \
+        :type var_accesses:
             :py:class:`psyclone.core.VariablesAccessInfo`
 
         '''
         if self._kern.iterates_over not in ["cell_column", "domain"]:
             return
-        nlayers_symbol = self.append_integer_reference("nlayers")
+        name = f"nlayers_{self._kern.arguments.iteration_space_arg().name}"
+        nlayers_symbol = self.append_integer_reference(name, tag=name)
         self.append(nlayers_symbol.name, var_accesses)
         self._nlayers_positions.append(self.num_args)
 

src/psyclone/domain/lfric/kern_stub_arg_list.py

@@ -39,7 +39,9 @@
 for a kernel subroutine.
 '''
 
-from psyclone.domain.lfric import ArgOrdering, LFRicConstants, LFRicSymbolTable
+from psyclone.domain.lfric.arg_ordering import ArgOrdering
+from psyclone.domain.lfric.lfric_constants import LFRicConstants
+from psyclone.domain.lfric.lfric_symbol_table import LFRicSymbolTable
 from psyclone.errors import InternalError
 
 

src/psyclone/domain/lfric/lfric_cell_iterators.py

@@ -0,0 +1,157 @@
+# -----------------------------------------------------------------------------
+# BSD 3-Clause License
+#
+# Copyright (c) 2017-2024, Science and Technology Facilities Council.
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+# -----------------------------------------------------------------------------
+# Authors R. W. Ford, A. R. Porter and S. Siso, STFC Daresbury Lab
+# 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
+
+''' This module implements the LFRicCellIterators collection which handles
+    the requirements of kernels that operator on cells.'''
+
+from psyclone.configuration import Config
+from psyclone.domain.lfric.lfric_collection import LFRicCollection
+from psyclone.domain.lfric.lfric_kern import LFRicKern
+from psyclone.domain.lfric.lfric_types import LFRicTypes
+from psyclone.errors import GenerationError
+from psyclone.f2pygen import AssignGen, CommentGen, DeclGen
+
+
+class LFRicCellIterators(LFRicCollection):
+    '''
+    Handles all entities required by kernels that operate on cell-columns.
+
+    :param kern_or_invoke: the Kernel or Invoke for which to manage cell
+                           iterators.
+    :type kern_or_invoke: :py:class:`psyclone.domain.lfric.LFRicKern` |
+                          :py:class:`psyclone.dynamo0p3.LFRicInvoke`
+
+    :raises GenerationError: if an Invoke has no field or operator arguments.
+
+    '''
+    def __init__(self, kern_or_invoke):
+        super().__init__(kern_or_invoke)
+
+        # Dictionary to hold the names of the various nlayers variables and
+        # (for invokes) the kernel argument to which each corresponds.
+        self._nlayers_names = {}
+
+        if not self._invoke:
+            # We are dealing with a single Kernel so there is only one
+            # 'nlayers' variable and we don't need to store the associated
+            # argument.
+            self._nlayers_names[self._symbol_table.find_or_create_tag(
+                "nlayers",
+                symbol_type=LFRicTypes("MeshHeightDataSymbol")).name] = None
+            # We're not generating a PSy layer so we're done here.
+            return
+
+        # Each kernel that operates on either the domain or cell-columns needs
+        # an 'nlayers' obtained from the first written field/operator argument.
+        for kern in self._invoke.schedule.walk(LFRicKern):
+            if kern.iterates_over in ["cell_column", "domain"]:
+                arg = kern.arguments.iteration_space_arg()
+                self._nlayers_names[self._symbol_table.find_or_create_tag(
+                    f"nlayers_{arg.name}",
+                    symbol_type=LFRicTypes("MeshHeightDataSymbol")).name] = arg
+
+        first_var = None
+        for var in self._invoke.psy_unique_vars:
+            if not var.is_scalar:
+                first_var = var
+                break
+        if not first_var:
+            raise GenerationError(
+                "Cannot create an Invoke with no field/operator arguments.")
+        self._first_var = first_var
+
+    def _invoke_declarations(self, parent):
+        '''
+        Declare entities required for iterating over cells in the Invoke.
+
+        :param parent: the f2pygen node representing the PSy-layer routine.
+        :type parent: :py:class:`psyclone.f2pygen.SubroutineGen`
+
+        '''
+        api_config = Config.get().api_conf("lfric")
+
+        # Declare the number of layers in the mesh for each kernel that
+        # operates on cell-columns or the domain.
+        name_list = list(self._nlayers_names.keys())
+        if name_list:
+            name_list.sort()  # Purely for test reproducibility.
+            parent.add(DeclGen(parent, datatype="integer",
+                               kind=api_config.default_kind["integer"],
+                               entity_decls=name_list))
+
+    def _stub_declarations(self, parent):
+        '''
+        Declare entities required for a kernel stub that operates on
+        cell-columns.
+
+        :param parent: the f2pygen node representing the Kernel stub.
+        :type parent: :py:class:`psyclone.f2pygen.SubroutineGen`
+
+        '''
+        api_config = Config.get().api_conf("lfric")
+
+        if self._kernel.cma_operation not in ["apply", "matrix-matrix"]:
+            parent.add(DeclGen(parent, datatype="integer",
+                               kind=api_config.default_kind["integer"],
+                               intent="in",
+                               entity_decls=list(self._nlayers_names.keys())))
+
+    def initialise(self, parent):
+        '''
+        Look-up the number of vertical layers in the mesh for each user-
+        supplied kernel that operates on cell columns.
+
+        :param parent: the f2pygen node representing the PSy-layer routine.
+        :type parent: :py:class:`psyclone.f2pygen.SubroutineGen`
+
+        '''
+        if not self._nlayers_names or not self._invoke:
+            return
+
+        parent.add(CommentGen(parent, ""))
+        parent.add(CommentGen(parent, " Initialise number of layers"))
+        parent.add(CommentGen(parent, ""))
+        # Sort for test reproducibility
+        sorted_names = list(self._nlayers_names.keys())
+        sorted_names.sort()
+        for name in sorted_names:
+            var = self._nlayers_names[name]
+            parent.add(AssignGen(
+                parent, lhs=name,
+                rhs=(f"{var.proxy_name_indexed}%{var.ref_name()}%"
+                     f"get_nlayers()")))

src/psyclone/domain/lfric/lfric_collection.py

@@ -42,8 +42,9 @@
 
 # Imports
 import abc
-from psyclone.domain.lfric import (LFRicSymbolTable, LFRicInvoke,
-                                   LFRicKern)
+from psyclone.domain.lfric.lfric_invoke import LFRicInvoke
+from psyclone.domain.lfric.lfric_kern import LFRicKern
+from psyclone.domain.lfric.lfric_symbol_table import LFRicSymbolTable
 from psyclone.errors import InternalError
 
 

src/psyclone/domain/lfric/lfric_invoke.py

@@ -97,11 +97,10 @@ def __init__(self, alg_invocation, idx, invokes):
                                         DynLMAOperators, DynReferenceElement,
                                         DynCMAOperators, DynBasisFunctions,
                                         DynMeshes, DynBoundaryConditions,
-                                        DynProxies, DynCellIterators,
-                                        LFRicMeshProperties)
-        from psyclone.domain.lfric import (LFRicLoopBounds, LFRicRunTimeChecks,
-                                           LFRicScalarArgs, LFRicFields,
-                                           LFRicDofmaps, LFRicStencils)
+                                        DynProxies, LFRicMeshProperties)
+        from psyclone.domain.lfric import (
+            LFRicCellIterators, LFRicLoopBounds, LFRicRunTimeChecks,
+            LFRicScalarArgs, LFRicFields, LFRicDofmaps, LFRicStencils)
 
         self.scalar_args = LFRicScalarArgs(self)
 
@@ -144,7 +143,7 @@ def __init__(self, alg_invocation, idx, invokes):
         self.run_time_checks = LFRicRunTimeChecks(self)
 
         # Information required by kernels that operate on cell-columns
-        self.cell_iterators = DynCellIterators(self)
+        self.cell_iterators = LFRicCellIterators(self)
 
         # Information on the required properties of the reference element
         self.reference_element_properties = DynReferenceElement(self)

src/psyclone/domain/lfric/lfric_kern.py

@@ -44,8 +44,10 @@
 
 from psyclone.configuration import Config
 from psyclone.core import AccessType
-from psyclone.domain.lfric import (KernCallArgList, KernStubArgList,
-                                   KernelInterface, LFRicConstants)
+from psyclone.domain.lfric.kern_call_arg_list import KernCallArgList
+from psyclone.domain.lfric.lfric_constants import LFRicConstants
+from psyclone.domain.lfric.kern_stub_arg_list import KernStubArgList
+from psyclone.domain.lfric.kernel_interface import KernelInterface
 from psyclone.errors import GenerationError, InternalError, FieldNotFoundError
 from psyclone.f2pygen import ModuleGen, SubroutineGen, UseGen
 from psyclone.parse.algorithm import Arg, KernelCall
@@ -601,13 +603,14 @@ def gen_stub(self):
         # Add all the declarations
         # Import here to avoid circular dependency
         # pylint: disable=import-outside-toplevel
-        from psyclone.domain.lfric import (LFRicScalarArgs, LFRicFields,
-                                           LFRicDofmaps, LFRicStencils)
-        from psyclone.dynamo0p3 import (DynCellIterators, DynFunctionSpaces,
+        from psyclone.domain.lfric import (
+            LFRicCellIterators, LFRicScalarArgs, LFRicFields,
+            LFRicDofmaps, LFRicStencils)
+        from psyclone.dynamo0p3 import (DynFunctionSpaces,
                                         DynCMAOperators, DynBoundaryConditions,
                                         DynLMAOperators, LFRicMeshProperties,
                                         DynBasisFunctions, DynReferenceElement)
-        for entities in [DynCellIterators, LFRicDofmaps, DynFunctionSpaces,
+        for entities in [LFRicCellIterators, LFRicDofmaps, DynFunctionSpaces,
                          DynCMAOperators, LFRicScalarArgs, LFRicFields,
                          DynLMAOperators, LFRicStencils, DynBasisFunctions,
                          DynBoundaryConditions, DynReferenceElement,