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

.github/workflows/lfric_test.yml

@@ -87,7 +87,7 @@ jobs:
         export LFRIC_DIR=${HOME}/LFRic/lfric_apps
         cd ${LFRIC_DIR}
         # Clean previous version and compile again
-        rm -rf appplications/gungho_model/working
+        rm -rf applications/gungho_model/working
         ./build/local_build.py -a gungho_model
         # Run
         cd applications/gungho_model/example
@@ -131,7 +131,7 @@ jobs:
         mkdir -p applications/gungho_model/optimisation/psyclone-test
         cp ${PSYCLONE_LFRIC_DIR}/everything_everywhere_all_at_once.py ${OPT_DIR}/global.py
         # Clean previous version and compile again
-        rm -rf appplications/gungho_model/working
+        rm -rf applications/gungho_model/working
         ./build/local_build.py -a gungho_model -p psyclone-test
         # Run
         cd applications/gungho_model/example

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/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/lfric_invoke.py

@@ -97,7 +97,7 @@ def __init__(self, alg_invocation, idx, invokes):
                                         DynLMAOperators, DynReferenceElement,
                                         DynCMAOperators, DynBasisFunctions,
                                         DynMeshes, DynBoundaryConditions,
-                                        DynProxies, DynCellIterators,
+                                        DynProxies, LFRicCellIterators,
                                         LFRicMeshProperties)
         from psyclone.domain.lfric import (LFRicLoopBounds, LFRicRunTimeChecks,
                                            LFRicScalarArgs, LFRicFields,
@@ -144,7 +144,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

@@ -601,11 +601,11 @@ def gen_stub(self):
         # pylint: disable=import-outside-toplevel
         from psyclone.domain.lfric import (LFRicScalarArgs, LFRicFields,
                                            LFRicDofmaps, LFRicStencils)
-        from psyclone.dynamo0p3 import (DynCellIterators, DynFunctionSpaces,
+        from psyclone.dynamo0p3 import (LFRicCellIterators, 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,

src/psyclone/dynamo0p3.py

@@ -36,8 +36,8 @@
 # Modified J. Henrichs, Bureau of Meteorology
 # Modified A. B. G. Chalk and N. Nobre, STFC Daresbury Lab
 
-''' This module implements the PSyclone Dynamo 0.3 API by 1)
-    specialising the required base classes in parser.py (KernelType) and
+''' This module implements the PSyclone LFRic API by 1) specialising the
+    required base classes in parser.py (KernelType) and
     adding a new class (DynFuncDescriptor03) to capture function descriptor
     metadata and 2) specialising the required base classes in psyGen.py
     (PSy, Invokes, Invoke, InvokeSchedule, Loop, Kern, Inf, Arguments and
@@ -1620,29 +1620,44 @@ def initialise(self, parent):
                         f"handled in DynProxies.initialise()")
 
 
-class DynCellIterators(LFRicCollection):
+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 \
+    :param kern_or_invoke: the Kernel or Invoke for which to manage cell
                            iterators.
-    :type kern_or_invoke: :py:class:`psyclone.domain.lfric.LFRicKern` or \
+    :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.
+    :raises GenerationError: if an Invoke has no field or operator arguments.
 
     '''
     def __init__(self, kern_or_invoke):
         super().__init__(kern_or_invoke)
 
-        self._nlayers_name = self._symbol_table.find_or_create_tag(
-            "nlayers", symbol_type=LFRicTypes("MeshHeightDataSymbol")).name
-
-        # Store a reference to the first field/operator object that
-        # we can use to look-up nlayers in the PSy layer.
         if not self._invoke:
+            # We are dealing with a single Kernel so there is only one
+            # 'nlayers' variable.
+            self._nlayers_names = {
+                self._symbol_table.find_or_create_tag(
+                    "nlayers",
+                    symbol_type=LFRicTypes("MeshHeightDataSymbol")).name}
             # We're not generating a PSy layer so we're done here.
             return
+
+        # Each kernel needs an 'nlayers' obtained from the first written
+        # field/operator argument.
+        self._nlayers_names = {}
+        for kern in self._invoke.schedule.walk(LFRicKern):
+            if kern.iterates_over not in ["cell_column", "domain"]:
+                # Only user-supplied kernels that operate on either the domain
+                # or cell-columns require nlayers.
+                continue
+            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:
@@ -1663,12 +1678,14 @@ def _invoke_declarations(self, parent):
         '''
         api_config = Config.get().api_conf("lfric")
 
-        # We only need the number of layers in the mesh if we are calling
-        # one or more kernels that operate on cell-columns.
-        if not self._dofs_only:
+        # 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=[self._nlayers_name]))
+                               entity_decls=name_list))
 
     def _stub_declarations(self, parent):
         '''
@@ -1684,24 +1701,33 @@ def _stub_declarations(self, parent):
         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=[self._nlayers_name]))
+                               intent="in",
+                               entity_decls=list(self._nlayers_names)))
 
     def initialise(self, parent):
         '''
-        Look-up the number of vertical layers in the mesh in the PSy layer.
+        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._dofs_only:
-            parent.add(CommentGen(parent, ""))
-            parent.add(CommentGen(parent, " Initialise number of layers"))
-            parent.add(CommentGen(parent, ""))
+        if not self._nlayers_names:
+            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=self._nlayers_name,
-                rhs=self._first_var.proxy_name_indexed + "%" +
-                self._first_var.ref_name() + "%get_nlayers()"))
+                parent, lhs=name,
+                rhs=(f"{var.proxy_name_indexed}%{var.ref_name()}%"
+                     f"get_nlayers()")))
 
 
 class DynLMAOperators(LFRicCollection):
@@ -6216,7 +6242,7 @@ def data_on_device(self, _):
     'DynamoPSy',
     'DynFunctionSpaces',
     'DynProxies',
-    'DynCellIterators',
+    'LFRicCellIterators',
     'DynLMAOperators',
     'DynCMAOperators',
     'DynMeshes',

src/psyclone/tests/core/variables_access_info_test.py

@@ -452,7 +452,8 @@ def test_variables_access_info_domain_loop():
     vai = VariablesAccessInfo(invoke.schedule)
     assert str(vai) == ("a: READ, b: READ, f1_data: READWRITE, f2_data: "
                         "READWRITE, map_w3: READ, ncell_2d_no_halos: READ, "
-                        "ndf_w3: READ, nlayers: READ, undf_w3: READ")
+                        "ndf_w3: READ, nlayers_f1: READ, nlayers_f2: READ, "
+                        "undf_w3: READ")
 
 
 # -----------------------------------------------------------------------------
@@ -472,7 +473,7 @@ def test_lfric_access_info():
             "diff_basis_w2_qr: READ, diff_basis_w3_qr: READ, f1_data: "
             "READ+WRITE, f2_data: READ, loop0_start: READ, loop0_stop: READ, "
             "m1_data: READ, m2_data: READ, map_w1: READ, map_w2: READ, map_w3:"
-            " READ, ndf_w1: READ, ndf_w2: READ, ndf_w3: READ, nlayers: READ, "
-            "np_xy_qr: READ, np_z_qr: READ, undf_w1: READ, undf_w2: READ, "
-            "undf_w3: READ, weights_xy_qr: READ, weights_z_qr: READ"
+            " READ, ndf_w1: READ, ndf_w2: READ, ndf_w3: READ, nlayers_f1: "
+            "READ, np_xy_qr: READ, np_z_qr: READ, undf_w1: READ, undf_w2: "
+            "READ, undf_w3: READ, weights_xy_qr: READ, weights_z_qr: READ"
             == str(vai))

src/psyclone/tests/dependency_test.py

@@ -294,7 +294,7 @@ def test_lfric():
         "loop0_start: READ, loop0_stop: READ, m1_data: READ, "
         "m2_data: READ, map_w1: READ, map_w2: READ, "
         "map_w3: READ, ndf_w1: READ, ndf_w2: READ, ndf_w3: READ, "
-        "nlayers: READ, undf_w1: READ, undf_w2: READ, undf_w3: READ")
+        "nlayers_f1: READ, undf_w1: READ, undf_w2: READ, undf_w3: READ")
 
 
 def test_lfric_kern_cma_args():

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

@@ -255,7 +255,7 @@ def test_arg_ordering_generate_domain_kernel(dist_mem, fortran_writer):
     assert not create_arg_list._psyir_arglist
     create_arg_list.generate()
     assert create_arg_list._arglist == [
-        'nlayers', 'ncell_2d_no_halos', 'b', 'f1_data', 'ndf_w3',
+        'nlayers_f1', 'ncell_2d_no_halos', 'b', 'f1_data', 'ndf_w3',
         'undf_w3', 'map_w3']
 
     check_psyir_results(create_arg_list, fortran_writer)
@@ -280,7 +280,7 @@ def test_arg_ordering_generate_cma_kernel(dist_mem, fortran_writer):
     assert not create_arg_list._arglist
     create_arg_list.generate()
     assert create_arg_list._arglist == [
-        'cell', 'nlayers', 'ncell_2d', 'lma_op1_proxy%ncell_3d',
+        'cell', 'nlayers_cma_op1', 'ncell_2d', 'lma_op1_proxy%ncell_3d',
         'lma_op1_local_stencil', 'cma_op1_cma_matrix', 'cma_op1_nrow',
         'cma_op1_ncol', 'cma_op1_bandwidth', 'cma_op1_alpha', 'cma_op1_beta',
         'cma_op1_gamma_m', 'cma_op1_gamma_p', 'ndf_adspc1_lma_op1',

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

@@ -110,7 +110,7 @@ def test_cellmap_intergrid(dist_mem, fortran_writer):
     assert Signature("cell_map_field2") in vai
 
     assert create_arg_list._arglist == [
-        'nlayers', 'cell_map_field2(:,:,cell)', 'ncpc_field1_field2_x',
+        'nlayers_field2', 'cell_map_field2(:,:,cell)', 'ncpc_field1_field2_x',
         'ncpc_field1_field2_y', 'ncell_field1', 'field1_data',
         'field2_data', 'ndf_w1', 'undf_w1', 'map_w1', 'undf_w2',
         'map_w2(:,cell)']
@@ -134,7 +134,7 @@ def test_kerncallarglist_face_xyoz(dist_mem, fortran_writer):
 
     create_arg_list.generate()
     assert create_arg_list._arglist == [
-        'nlayers', 'f1_data', 'f2_1_data', 'f2_2_data',
+        'nlayers_f1', 'f1_data', 'f2_1_data', 'f2_2_data',
         'f2_3_data', 'f3_data', 'istp', 'ndf_w2', 'undf_w2',
         'map_w2(:,cell)', 'basis_w2_qr_xyoz', 'basis_w2_qr_face', 'ndf_wchi',
         'undf_wchi', 'map_wchi(:,cell)', 'diff_basis_wchi_qr_xyoz',
@@ -174,7 +174,7 @@ def test_kerncallarglist_face_edge(dist_mem, fortran_writer):
 
     create_arg_list.generate()
     assert create_arg_list._arglist == [
-        'nlayers', 'f1_data', 'f2_data', 'm1_data',
+        'nlayers_f1', 'f1_data', 'f2_data', 'm1_data',
         'm2_data', 'ndf_w1', 'undf_w1', 'map_w1(:,cell)',
         'basis_w1_qr_face', 'basis_w1_qr_edge', 'ndf_w2', 'undf_w2',
         'map_w2(:,cell)', 'diff_basis_w2_qr_face', 'diff_basis_w2_qr_edge',
@@ -208,7 +208,7 @@ def test_kerncallarglist_colouring(dist_mem, fortran_writer):
     create_arg_list = KernCallArgList(schedule.kernels()[0])
     create_arg_list.generate()
     assert create_arg_list._arglist == [
-        'nlayers', 'rdt', 'h_data', 'f_data', 'c_data',
+        'nlayers_h', 'rdt', 'h_data', 'f_data', 'c_data',
         'd_data', 'ndf_w1', 'undf_w1', 'map_w1(:,cmap(colour,cell))',
         'ndf_w2', 'undf_w2', 'map_w2(:,cmap(colour,cell))', 'ndf_w3',
         'undf_w3', 'map_w3(:,cmap(colour,cell))']
@@ -233,15 +233,15 @@ def test_kerncallarglist_mesh_properties(fortran_writer):
     assert str(var_info) == ("a: READ, adjacent_face: READ, cell: READ, "
                              "cmap: READ, colour: READ, f1_data: READ+WRITE, "
                              "map_w1: READ, ndf_w1: READ, nfaces_re_h: "
-                             "READ, nlayers: READ, undf_w1: READ")
+                             "READ, nlayers_f1: READ, undf_w1: READ")
     # Tests that multiple reads are reported as expected:
     assert str(var_info[Signature("cell")]) == "cell:READ(0),READ(0)"
     assert str(var_info[Signature("colour")]) == "colour:READ(0),READ(0)"
     assert str(var_info[Signature("cmap")]) == "cmap:READ(0),READ(0)"
     assert str(var_info[Signature("adjacent_face")]) == "adjacent_face:READ(0)"
 
     assert create_arg_list._arglist == [
-        'nlayers', 'a', 'f1_data', 'ndf_w1', 'undf_w1',
+        'nlayers_f1', 'a', 'f1_data', 'ndf_w1', 'undf_w1',
         'map_w1(:,cmap(colour,cell))', 'nfaces_re_h',
         'adjacent_face(:,cmap(colour,cell))']
 
@@ -262,7 +262,7 @@ def test_kerncallarglist_evaluator(fortran_writer):
     create_arg_list = KernCallArgList(schedule.kernels()[0])
     create_arg_list.generate()
     assert create_arg_list._arglist == [
-        'nlayers', 'f0_data', 'cmap_data', 'ndf_w0', 'undf_w0',
+        'nlayers_f0', 'f0_data', 'cmap_data', 'ndf_w0', 'undf_w0',
         'map_w0(:,cmap_1(colour,cell))', 'basis_w0_on_w0', 'ndf_w1', 'undf_w1',
         'map_w1(:,cmap_1(colour,cell))', 'diff_basis_w1_on_w0']
 
@@ -284,7 +284,7 @@ def test_kerncallarglist_stencil(fortran_writer):
     create_arg_list.generate()
 
     assert create_arg_list._arglist == [
-        'nlayers', 'f1_data', 'f2_data',
+        'nlayers_f1', 'f1_data', 'f2_data',
         'f2_stencil_size(cmap(colour,cell))',
         'f2_stencil_dofmap(:,:,cmap(colour,cell))', 'f3_data',
         'f3_stencil_size(cmap(colour,cell))', 'f3_direction',
@@ -316,7 +316,7 @@ def test_kerncallarglist_cross2d_stencil(fortran_writer):
     create_arg_list.generate()
 
     assert create_arg_list._arglist == [
-        'nlayers', 'f1_data', 'f2_data',
+        'nlayers_f1', 'f1_data', 'f2_data',
         'f2_stencil_size(:,cell)', 'f2_max_branch_length',
         'f2_stencil_dofmap(:,:,:,cell)', 'f3_data', 'f4_data',
         'ndf_w1', 'undf_w1', 'map_w1(:,cell)', 'ndf_w2', 'undf_w2',
@@ -336,7 +336,7 @@ def test_kerncallarglist_bcs(fortran_writer, monkeypatch):
     create_arg_list = KernCallArgList(schedule.kernels()[0])
     create_arg_list.generate()
     assert create_arg_list._arglist == [
-        'nlayers', 'a_data', 'ndf_aspc1_a', 'undf_aspc1_a',
+        'nlayers_a', 'a_data', 'ndf_aspc1_a', 'undf_aspc1_a',
         'map_aspc1_a(:,cell)', 'boundary_dofs_a']
 
     check_psyir_results(create_arg_list, fortran_writer)
@@ -370,7 +370,7 @@ def test_kerncallarglist_bcs_operator(fortran_writer):
     access_info = VariablesAccessInfo()
     create_arg_list.generate(access_info)
     assert create_arg_list._arglist == [
-        'cell', 'nlayers', 'op_a_proxy%ncell_3d', 'op_a_local_stencil',
+        'cell', 'nlayers_op_a', 'op_a_proxy%ncell_3d', 'op_a_local_stencil',
         'ndf_aspc1_op_a', 'ndf_aspc2_op_a', 'boundary_dofs_op_a']
 
     check_psyir_results(create_arg_list, fortran_writer)
@@ -485,7 +485,7 @@ def test_kerncallarglist_scalar_literal(fortran_writer):
     assert "1.0" not in str(vai)
 
     assert create_arg_list._arglist == [
-        'nlayers', 'f1_data', 'f2_data', 'm1_data',
+        'nlayers_f1', 'f1_data', 'f2_data', 'm1_data',
         '1.0_r_def', 'm2_data', '2_i_def', 'ndf_w1', 'undf_w1',
         'map_w1(:,cell)', 'basis_w1_qr', 'ndf_w2', 'undf_w2', 'map_w2(:,cell)',
         'diff_basis_w2_qr', 'ndf_w3', 'undf_w3', 'map_w3(:,cell)',
@@ -611,12 +611,12 @@ def test_ref_element_handling(fortran_writer):
     create_arg_list.generate(vai)
 
     assert (create_arg_list._arglist == [
-        'nlayers', 'f1_data', 'ndf_w1', 'undf_w1', 'map_w1(:,cell)',
+        'nlayers_f1', 'f1_data', 'ndf_w1', 'undf_w1', 'map_w1(:,cell)',
         'nfaces_re_h', 'nfaces_re_v', 'normals_to_horiz_faces',
         'normals_to_vert_faces'])
 
     assert ("cell: READ, f1_data: READ+WRITE, map_w1: READ, ndf_w1: READ, "
-            "nfaces_re_h: READ, nfaces_re_v: READ, nlayers: READ, "
+            "nfaces_re_h: READ, nfaces_re_v: READ, nlayers_f1: READ, "
             "normals_to_horiz_faces: READ, normals_to_vert_faces: READ, "
             "undf_w1: READ" == str(vai))