# -----------------------------------------------------------------------------
# BSD 3-Clause License
#
# Copyright (c) 2017-2025, 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 and L. Turner, 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 psyclone.core import VariablesAccessMap
from psyclone.domain.lfric.arg_ordering import ArgOrdering
from psyclone.domain.lfric.lfric_constants import LFRicConstants
from psyclone.errors import InternalError
[docs]
class KernStubArgList(ArgOrdering):
# pylint: disable=too-many-public-methods, abstract-method
# TODO: #845 Check that all implicit variables have the right type.
'''Creates the argument list required to create and declare the
required arguments for a kernel subroutine. The ordering and type
of the arguments is captured by the base class.
:param kern: Kernel for which to create argument list.
:type kern: :py:class:`psyclone.domain.lfric.LFRicKern`
:raises NotImplementedError: if the kernel is inter-grid.
:raises NotImplementedError: if the kernel requires properties of the \
reference element.
'''
def __init__(self, kern):
ArgOrdering.__init__(self, kern)
[docs]
def cell_position(self, var_accesses=None):
'''Adds a cell argument to the argument list and if supplied stores
this access in var_accesses.
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
self.append("cell", var_accesses)
[docs]
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 VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
self.append("nlayers", var_accesses)
def _mesh_ncell2d(self, var_accesses=None):
'''Add the number of columns in the mesh to the argument list and if
supplied stores this access in var_accesses.
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
self.append("ncell_2d", var_accesses)
[docs]
def cma_operator(self, arg, var_accesses=None):
'''Add the CMA operator and associated scalars to the argument
list and optionally add them to the variable access
information.
:param arg: the CMA operator argument.
:type arg: :py:class:`psyclone.lfric.LFRicKernelArgument`
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
# The CMA operator itself
self.append(arg.name, var_accesses)
# Associated scalar parameters
nrow = arg.name + "_nrow"
_local_args = [nrow]
if arg.function_space_to.orig_name != \
arg.function_space_from.orig_name:
# If the to- and from-spaces are different then so are ncol and
# nrow so we pass both of them. If they are the same then we
# could pass either but choose to pass nrow and not ncol.
ncol = arg.name + "_ncol"
_local_args.append(ncol)
bandwidth = arg.name + "_bandwidth"
alpha = arg.name + "_alpha"
beta = arg.name + "_beta"
gamma_m = arg.name + "_gamma_m"
gamma_p = arg.name + "_gamma_p"
_local_args += [bandwidth, alpha, beta, gamma_m, gamma_p]
self.extend(_local_args, var_accesses)
[docs]
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
var_access object.
:param argvect: the field vector to add.
:type argvect: :py:class:`psyclone.lfric.LFRicKernelArgument`
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
# the range function below returns values from
# 1 to the vector size which is what we
# require in our Fortran code
for idx in range(1, argvect.vector_size+1):
text = (argvect.name + "_" +
argvect.function_space.mangled_name +
"_v" + str(idx))
self.append(text, var_accesses)
[docs]
def field(self, arg, var_accesses=None):
'''Add the field array associated with the argument 'arg' to the
argument list. If supplied it also stores this access in var_accesses.
:param arg: the field to be added.
:type arg: :py:class:`psyclone.lfric.LFRicKernelArgument`
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
text = arg.name + "_" + arg.function_space.mangled_name
self.append(text, var_accesses, mode=arg.access)
[docs]
def stencil_unknown_extent(self, arg, var_accesses=None):
'''Add stencil information to the argument list associated with the
argument 'arg' if the extent is unknown. If supplied it also stores
this access in var_accesses.
:param arg: the kernel argument with which the stencil is associated.
:type arg: :py:class:`psyclone.lfric.LFRicKernelArgument`
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
# Import here to avoid circular dependency
# pylint: disable=import-outside-toplevel
from psyclone.domain.lfric.lfric_stencils import LFRicStencils
name = LFRicStencils.dofmap_size_symbol(self._symtab, arg).name
self.append(name, var_accesses)
[docs]
def stencil_unknown_direction(self, arg, var_accesses=None):
'''Add stencil information to the argument list associated with the
argument 'arg' if the direction is unknown. If supplied it also stores
this access in var_accesses.
:param arg: the kernel argument with which the stencil is associated.
:type arg: :py:class:`psyclone.lfric.LFRicKernelArgument`
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
# Import here to avoid circular dependency
# pylint: disable=import-outside-toplevel
from psyclone.domain.lfric.lfric_stencils import LFRicStencils
name = LFRicStencils.direction_name(self._symtab, arg).name
self.append(name, var_accesses)
[docs]
def stencil(self, arg, var_accesses=None):
'''Add general stencil information associated with the argument 'arg'
to the argument list. If supplied it also stores this access in
var_accesses.
:param arg: the meta-data description of the kernel \
argument with which the stencil is associated.
:type arg: :py:class:`psyclone.lfric.LFRicKernelArgument`
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
# Import here to avoid circular dependency
# pylint: disable=import-outside-toplevel
from psyclone.domain.lfric.lfric_stencils import LFRicStencils
var_name = LFRicStencils.dofmap_symbol(self._symtab, arg).name
self.append(var_name, var_accesses)
[docs]
def stencil_2d_max_extent(self, arg, var_accesses=None):
'''Add the maximum branch extent for a 2D stencil associated with the
argument 'arg' to the argument list. If supplied it also stores this
in var_accesses.
:param arg: the kernel argument with which the stencil is associated.
:type arg: :py:class:`psyclone.lfric.LFRicKernelArgument`
:param var_accesses: optional `AccessSequence` \
instance to store the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.AccessSequence`
'''
# The maximum branch extent is not specified in the metadata so pass
# the value in.
# Import here to avoid circular dependency
# pylint: disable=import-outside-toplevel
from psyclone.domain.lfric.lfric_stencils import LFRicStencils
name = LFRicStencils.max_branch_length(self._symtab, arg).name
self.append(name, var_accesses)
[docs]
def stencil_2d_unknown_extent(self, arg, var_accesses=None):
'''Add 2D stencil information to the argument list associated with the
argument 'arg' if the extent is unknown. If supplied it also stores
this access in var_accesses.
:param arg: the kernel argument with which the stencil is associated.
:type arg: :py:class:`psyclone.lfric.LFRicKernelArgument`
:param var_accesses: optional `VariablesAccessMap` instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
# Import here to avoid circular dependency
# pylint: disable=import-outside-toplevel
from psyclone.domain.lfric.lfric_stencils import LFRicStencils
name = LFRicStencils.dofmap_size_symbol(self._symtab, arg).name
self.append(name, var_accesses)
[docs]
def stencil_2d(self, arg, var_accesses=None):
'''Add general 2D stencil information associated with the argument
'arg' to the argument list. If supplied it also stores this access in
var_accesses.
:param arg: the meta-data description of the kernel \
argument with which the stencil is associated.
:type arg: :py:class:`psyclone.lfric.LFRicKernelArgument`
:param var_accesses: optional `VariablesAccessMap` instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
# The stencil_2D differs from the stencil in that the direction
# of the branch is baked into the stencil_dofmap array.
# The array dimensions are thus (dof_in_cell, cell_in_branch,
# branch_in_stencil) where the branch_in_stencil is always ordered
# West, South, East, North which is standard in LFRic. This allows
# for knowledge of what direction a stencil cell is in relation
# to the centre even when the stencil is truncated at boundaries.
# Import here to avoid circular dependency
# pylint: disable=import-outside-toplevel
from psyclone.domain.lfric.lfric_stencils import LFRicStencils
var_name = LFRicStencils.dofmap_symbol(self._symtab, arg).name
self.append(var_name, var_accesses)
[docs]
def operator(self, arg, var_accesses=None):
'''Add the operator arguments to the argument list. If supplied it
also stores this access in var_accesses.
:param arg: the meta-data description of the operator.
:type arg: :py:class:`psyclone.lfric.LFRicKernelArgument`
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
size = arg.name + "_ncell_3d"
self.append(size, var_accesses)
self.append(arg.name, var_accesses)
[docs]
def fs_compulsory_field(self, function_space, var_accesses=None):
''' Provide compulsory arguments if there is a field on this
function space. If supplied it also stores this access in
var_accesses.
:param function_space: the function space for which the compulsory \
arguments are added.
:type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
self.append(function_space.undf_name, var_accesses)
self.append(function_space.map_name, var_accesses)
[docs]
def basis(self, function_space, var_accesses=None):
'''Add basis function information for this function space to the
argument list and optionally to the variable access information.
There can be more than one if this is an evaluator and/or multiple
'gh_shape's have been requested in the kernel metadata. If supplied
it also stores these accesses in var_accesses.
:param function_space: the function space for which to provide \
the basis functions
:type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
:raises InternalError: if the evaluator shape is not recognised.
'''
const = LFRicConstants()
for shape in self._kern.eval_shapes:
if shape in const.VALID_QUADRATURE_SHAPES:
# A kernel stub won't have a name for the corresponding
# quadrature argument so we create one by appending the last
# part of the shape name to "qr_".
basis_name = function_space.get_basis_name(
qr_var="qr_"+shape.split("_")[-1])
self.append(basis_name, var_accesses)
elif shape in const.VALID_EVALUATOR_SHAPES:
# Need a basis array for each target space upon which the basis
# functions have been evaluated. _kern.eval_targets is a dict
# where the values are 2-tuples of (FunctionSpace, argument).
for _, target in self._kern.eval_targets.items():
basis_name = \
function_space.get_basis_name(on_space=target[0])
self.append(basis_name, var_accesses)
else:
raise InternalError(
f"Unrecognised evaluator shape ('{shape}'). Expected one "
f"of: {const.VALID_EVALUATOR_SHAPES}")
[docs]
def diff_basis(self, function_space, var_accesses=None):
'''Add differential basis information for the function space to the
argument list. If supplied it also stores this access in
var_accesses.
:param function_space: the function space for which the differential \
basis functions are required.
:type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
:raises InternalError: if the evaluator shape is not recognised.
'''
const = LFRicConstants()
for shape in self._kern.eval_shapes:
if shape in const.VALID_QUADRATURE_SHAPES:
# We need differential basis functions for quadrature. A
# kernel stub won't have a name for the corresponding
# quadrature argument so we create one by appending the
# last part of the shape name to "qr_".
diff_basis_name = function_space.get_diff_basis_name(
qr_var="qr_"+shape.split("_")[-1])
self.append(diff_basis_name, var_accesses)
elif shape in const.VALID_EVALUATOR_SHAPES:
# We need differential basis functions for an evaluator,
# potentially for multiple target spaces. _kern.eval_targets is
# a dict where the values are 2-tuples of
# (FunctionSpace, argument).
for _, target in self._kern.eval_targets.items():
diff_basis_name = function_space.get_diff_basis_name(
on_space=target[0])
self.append(diff_basis_name, var_accesses)
else:
raise InternalError(f"Unrecognised evaluator shape "
f"('{shape}'). Expected one of: "
f"{const.VALID_EVALUATOR_SHAPES}")
[docs]
def field_bcs_kernel(self, function_space, var_accesses=None):
'''Implement the boundary_dofs array fix for a field. If supplied it
also stores this access in var_accesses.
:param function_space: the function space for which boundary dofs \
are required.
:type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
arg = self._kern.arguments.get_arg_on_space(function_space)
name = "boundary_dofs_"+arg.name
self.append(name, var_accesses)
[docs]
def operator_bcs_kernel(self, function_space, var_accesses=None):
'''Supply necessary additional arguments for the kernel that
applies boundary conditions to a LMA operator. If supplied it
also stores this access in var_accesses.
:param function_space: the 'to' function space of the operator.
:type function_space: :py:class:`psyclone.lfric.FunctionSpace`
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
self.field_bcs_kernel(function_space, var_accesses=var_accesses)
[docs]
def mesh_properties(self, var_accesses=None):
'''Provide the kernel arguments required for the mesh properties
specified in the kernel metadata. If supplied it also stores this
access in var_accesses.
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
if self._kern.mesh.properties:
# Avoid circular import
# pylint: disable=import-outside-toplevel
from psyclone.lfric import LFRicMeshProperties
self.extend(LFRicMeshProperties(self._kern).
kern_args(stub=True, var_accesses=var_accesses))
[docs]
def quad_rule(self, var_accesses=None):
'''Add quadrature-related information to the kernel argument list.
Adds the necessary arguments to the argument list, and optionally
adds variable access information to the var_accesses object.
:param var_accesses: optional VariablesAccessMap instance to store \
the information about variable accesses.
:type var_accesses: \
:py:class:`psyclone.core.VariablesAccessMap`
'''
for rule in self._kern.qr_rules.values():
self.extend(rule.kernel_args, var_accesses)
[docs]
def indirection_dofmap(
self, function_space, operator=None,
var_accesses: Optional[VariablesAccessMap] = None
):
'''Add indirection dofmap required when applying a CMA operator. If
supplied it also stores this access in var_accesses.
:param function_space: the function space for which the indirect
dofmap is required.
:type function_space: :py:class:`psyclone.domain.lfric.FunctionSpace`
:param operator: the CMA operator.
:type operator: :py:class:`psyclone.lfric.LFRicKernelArgument`
:param var_accesses: optional VariablesAccessMap instance to store
the information about variable accesses.
:raises InternalError: if no kernel argument is supplied.
:raises InternalError: if the supplied kernel argument is not a
CMA operator.
'''
if not operator:
raise InternalError("No CMA operator supplied.")
if operator.argument_type != "gh_columnwise_operator":
raise InternalError(
f"A CMA operator (gh_columnwise_operator) must "
f"be supplied but got '{operator.argument_type}'.")
super().indirection_dofmap(function_space, operator, var_accesses)
# ============================================================================
# For automatic documentation creation:
__all__ = ["KernStubArgList"]
