#1312 ScalarArray code generation

.github/workflows/lfric_test.yml

@@ -42,7 +42,7 @@ on:
   push
 
 env:
-  LFRIC_APPS_HASH: 0b2283814583258e940f54c9beca826cdf8b72d7
+  LFRIC_APPS_HASH: bc2249045eeac3c332123b0031930681f2cc2445
   PYTHON_VERSION: 3.14
   GNU_TOOLCHAIN: gnu14_openmpi
   NVHPC_VERSION: 25.1

changelog

@@ -1,3 +1,6 @@
+   31) PR #3101 for #1312. Implements support for ScalarArray arguments to
+   LFRic kernels.
+
    30) PR #3280 towards #2971. Move OMPParallelTrans to psyir.transformations. 
 
    29) PR #3275 for #3216. Allows the specification of private variables

external/lfric_infrastructure/src/kernel_metadata/argument_mod.f90

@@ -60,6 +60,7 @@ module argument_mod
   integer, public, parameter :: GH_FIELD               = 507
   integer, public, parameter :: GH_OPERATOR            = 735
   integer, public, parameter :: GH_COLUMNWISE_OPERATOR = 841
+  integer, public, parameter :: GH_SCALAR_ARRAY        = 973
   !> @}
 
   !> @defgroup data_type Enumeration of argument data type property descriptors.

src/psyclone/domain/lfric/__init__.py

@@ -32,7 +32,7 @@
 # POSSIBILITY OF SUCH DAMAGE.
 # -----------------------------------------------------------------------------
 # Author J. Henrichs, Bureau of Meteorology
-# Modified: I. Kavcic, L. Turner and O. Brunt, Met Office
+# Modified: I. Kavcic, L. Turner, O. Brunt and A. Pirrie, Met Office
 #           R. W. Ford and A. R. Porter, STFC Daresbury Lab
 
 '''Module for the LFRic domain.
@@ -71,6 +71,7 @@
 from psyclone.domain.lfric.lfric_run_time_checks import LFRicRunTimeChecks
 from psyclone.domain.lfric.lfric_invokes import LFRicInvokes
 from psyclone.domain.lfric.lfric_scalar_args import LFRicScalarArgs
+from psyclone.domain.lfric.lfric_scalar_array_args import LFRicScalarArrayArgs
 from psyclone.domain.lfric.lfric_loop_bounds import LFRicLoopBounds
 from psyclone.domain.lfric.lfric_kern_metadata import LFRicKernMetadata
 from psyclone.domain.lfric.lfric_psy import LFRicPSy
@@ -105,5 +106,6 @@
     'LFRicPSy',
     'LFRicRunTimeChecks',
     'LFRicScalarArgs',
+    'LFRicScalarArrayArgs',
     'LFRicStencils',
     'LFRicSymbolTable']

src/psyclone/domain/lfric/arg_ordering.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 the base class for managing arguments to
@@ -459,7 +459,7 @@ def generate(self, var_accesses=None):
                 self.operator(arg, var_accesses=var_accesses)
             elif arg.argument_type == "gh_columnwise_operator":
                 self.cma_operator(arg, var_accesses=var_accesses)
-            elif arg.is_scalar:
+            elif arg.is_scalar or arg.is_scalar_array:
                 self.scalar(arg, var_accesses=var_accesses)
             else:
                 raise GenerationError(
@@ -770,11 +770,11 @@ def scalar(self, scalar_arg, var_accesses=None):
 
         '''
         const = LFRicConstants()
-        if not scalar_arg.is_scalar:
+        if not (scalar_arg.is_scalar or scalar_arg.is_scalar_array):
             raise InternalError(
                 f"Expected argument type to be one of "
-                f"{const.VALID_SCALAR_NAMES} but got "
-                f"'{scalar_arg.argument_type}'")
+                f"{const.VALID_SCALAR_NAMES + const.VALID_ARRAY_NAMES}"
+                f" but got '{scalar_arg.argument_type}'")
 
         if scalar_arg.is_literal:
             # If we have a literal, do not add it to the variable access
@@ -785,7 +785,8 @@ def scalar(self, scalar_arg, var_accesses=None):
                 var_accesses.add_access(Signature(scalar_arg.precision),
                                         AccessType.CONSTANT, self._kern)
         else:
-            self.append(scalar_arg.name, var_accesses, mode=scalar_arg.access,
+            self.append(scalar_arg.name, var_accesses,
+                        mode=scalar_arg.access,
                         metadata_posn=scalar_arg.metadata_index)
 
     def fs_common(self, function_space, var_accesses=None):

src/psyclone/domain/lfric/kern_call_acc_arg_list.py

@@ -261,7 +261,8 @@ def fs_intergrid(self, function_space, var_accesses=None):
     def scalar(self, scalar_arg, var_accesses=None):
         '''
         Override the default implementation as there's no need to specify
-        scalars for an OpenACC data region.
+        scalar values for an OpenACC data region. But there is a need for
+        ScalarArrays.
 
         :param scalar_arg: the kernel argument.
         :type scalar_arg: :py:class:`psyclone.lfric.LFRicKernelArgument`
@@ -270,7 +271,14 @@ def scalar(self, scalar_arg, var_accesses=None):
         :type var_accesses: Optional[
             :py:class:`psyclone.core.VariablesAccessMap`]
 
+        :raises NotImplementedError: OpenACC data regions are not currently
+            supported for arrays of scalars.
+
         '''
+        if scalar_arg.is_scalar_array:
+            raise NotImplementedError(
+                "OpenACC data regions are not currently supported for arrays"
+                " of scalars.")
 
 
 # ============================================================================

src/psyclone/domain/lfric/kern_call_arg_list.py

@@ -252,6 +252,12 @@ def scalar(self, scalar_arg,
         if scalar_arg.is_literal:
             self.psyir_append(scalar_arg.psyir_expression())
         else:
+            if scalar_arg.is_scalar_array:
+                # If it's a ScalarArray we need to add the dimensions
+                # array to the call
+                dims_sym = self._symtab.lookup_with_tag(
+                    "dims_" + scalar_arg.name)
+                self.psyir_append(Reference(dims_sym))
             sym = self._symtab.lookup(scalar_arg.name)
             self.psyir_append(Reference(sym))
 

src/psyclone/domain/lfric/kern_stub_arg_list.py

@@ -32,16 +32,18 @@
 # 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 creates the argument list
 for a kernel subroutine.
 '''
 
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 from psyclone.core import VariablesAccessMap
+if TYPE_CHECKING:
+    from psyclone.lfric import LFRicKernelArgument
 from psyclone.domain.lfric.arg_ordering import ArgOrdering
 from psyclone.domain.lfric.lfric_constants import LFRicConstants
 from psyclone.errors import InternalError
@@ -133,6 +135,34 @@ def cma_operator(self, arg, var_accesses=None):
         _local_args += [bandwidth, alpha, beta, gamma_m, gamma_p]
         self.extend(_local_args, var_accesses)
 
+    def scalar(self,
+               scalar_arg: 'LFRicKernelArgument',
+               var_accesses: Optional[VariablesAccessMap] = None):
+        '''Add the name associated with the scalar argument to the argument
+        list and optionally add this scalar to the variable access
+        information.
+
+        :param scalar_arg: the kernel argument.
+        :param var_accesses: optional VariablesAccessMap instance that
+            stores information about variable accesses.
+
+        :raises InternalError: if the argument is not a recognised scalar type.
+
+        '''
+        const = LFRicConstants()
+        if not (scalar_arg.is_scalar or scalar_arg.is_scalar_array):
+            raise InternalError(
+                f"Expected argument type to be one of "
+                f"{const.VALID_SCALAR_NAMES + const.VALID_ARRAY_NAMES}"
+                f" but got '{scalar_arg.argument_type}'")
+
+        if scalar_arg.is_scalar:
+            self.append(scalar_arg.name, var_accesses)
+        else:
+            # ScalarArray
+            self.append("dims_" + scalar_arg.name, var_accesses)
+            self.append(scalar_arg.name, var_accesses)
+
     def field_vector(self, argvect, var_accesses=None):
         '''Add the field vector associated with the argument 'argvect' to the
         argument list. If supplied it also stores these accesses to the

src/psyclone/domain/lfric/lfric_invoke.py

@@ -32,7 +32,8 @@
 # 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 I. Kavcic, A. Coughtrie, L. Turner, O. Brunt
+# and A. Pirrie, Met Office
 # Modified J. Henrichs, Bureau of Meteorology
 # Modified A. B. G. Chalk and N. Nobre, STFC Daresbury Lab
 
@@ -93,10 +94,13 @@ def __init__(self, alg_invocation, idx, invokes):
         from psyclone.domain.lfric import (
             LFRicCellIterators, LFRicHaloDepths, LFRicLoopBounds,
             LFRicRunTimeChecks,
-            LFRicScalarArgs, LFRicFields, LFRicDofmaps, LFRicStencils)
+            LFRicScalarArgs, LFRicScalarArrayArgs, LFRicFields, LFRicDofmaps,
+            LFRicStencils)
 
         self.scalar_args = LFRicScalarArgs(self)
 
+        self.scalar_array_args = LFRicScalarArrayArgs(self)
+
         # Initialise our Invoke stencil information
         self.stencil = LFRicStencils(self)
 
@@ -279,7 +283,8 @@ def setup_psy_layer_symbols(self):
 
         '''
         # Declare all quantities required by this PSy routine (Invoke)
-        for entities in [self.scalar_args, self.fields, self.lma_ops,
+        for entities in [self.scalar_args, self. scalar_array_args,
+                         self.fields, self.lma_ops,
                          self.stencil, self.meshes,
                          self.function_spaces, self.dofmaps, self.cma_ops,
                          self.boundary_conditions, self.evaluators,

src/psyclone/domain/lfric/lfric_kern.py

@@ -32,7 +32,8 @@
 # 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 I. Kavcic, A. Coughtrie, L. Turner, O. Brunt
+# and A. Pirrie, Met Office
 # Modified J. Henrichs, Bureau of Meteorology
 # Modified A. B. G. Chalk and N. Nobre, STFC Daresbury Lab
 
@@ -196,7 +197,8 @@ def load_meta(self, ktype):
             elif descriptor.argument_type.lower() == "gh_field":
                 pre = "field_"
             elif (descriptor.argument_type.lower() in
-                  const.VALID_SCALAR_NAMES):
+                  (const.VALID_SCALAR_NAMES +
+                   const.VALID_ARRAY_NAMES)):
                 if descriptor.data_type.lower() == "gh_real":
                     pre = "rscalar_"
                 elif descriptor.data_type.lower() == "gh_integer":
@@ -208,6 +210,9 @@ def load_meta(self, ktype):
                         f"Expected one of {const.VALID_SCALAR_DATA_TYPES} "
                         f"data types for a scalar argument but found "
                         f"'{descriptor.data_type}'.")
+                if (descriptor.argument_type.lower() in
+                        const.VALID_ARRAY_NAMES):
+                    pre += "array_"
             else:
                 raise GenerationError(
                     f"LFRicKern.load_meta() expected one of "
@@ -773,15 +778,16 @@ def gen_stub(self) -> Container:
         # Import here to avoid circular dependency
         # pylint: disable=import-outside-toplevel
         from psyclone.domain.lfric import (
-            LFRicCellIterators, LFRicScalarArgs, LFRicFields,
-            LFRicDofmaps, LFRicStencils)
+            LFRicCellIterators, LFRicScalarArgs, LFRicScalarArrayArgs,
+            LFRicFields, LFRicDofmaps, LFRicStencils)
         from psyclone.lfric import (
             LFRicFunctionSpaces, LFRicCMAOperators, LFRicBoundaryConditions,
             LFRicLMAOperators, LFRicMeshProperties, LFRicBasisFunctions,
             LFRicReferenceElement)
         for entities in [LFRicCellIterators, LFRicDofmaps, LFRicFunctionSpaces,
-                         LFRicCMAOperators, LFRicScalarArgs, LFRicFields,
-                         LFRicLMAOperators, LFRicStencils, LFRicBasisFunctions,
+                         LFRicCMAOperators, LFRicScalarArgs,
+                         LFRicScalarArrayArgs, LFRicFields, LFRicLMAOperators,
+                         LFRicStencils, LFRicBasisFunctions,
                          LFRicBoundaryConditions, LFRicReferenceElement,
                          LFRicMeshProperties]:
             entities(self).stub_declarations()

src/psyclone/domain/lfric/lfric_loop.py

@@ -697,7 +697,7 @@ def _halo_read_access(self, arg):
 
         '''
         const = LFRicConstants()
-        if arg.is_scalar or arg.is_operator:
+        if arg.is_scalar or arg.is_operator or arg.is_scalar_array:
             # Scalars and operators do not have halos
             return False
         if arg.is_field: