diff --git a/src/psyclone/domain/lfric/kern_call_arg_list.py b/src/psyclone/domain/lfric/kern_call_arg_list.py index 27071ad28a..d693f7e6b1 100644 --- a/src/psyclone/domain/lfric/kern_call_arg_list.py +++ b/src/psyclone/domain/lfric/kern_call_arg_list.py @@ -43,9 +43,10 @@ ''' from dataclasses import dataclass +from typing import Optional from psyclone import psyGen -from psyclone.core import AccessType, Signature +from psyclone.core import AccessType, Signature, VariablesAccessInfo from psyclone.domain.lfric.arg_ordering import ArgOrdering from psyclone.domain.lfric.lfric_constants import LFRicConstants # Avoid circular import: @@ -950,18 +951,20 @@ def ndf_positions(self): "before the ndf_positions() method") return self._ndf_positions - def cell_ref_name(self, var_accesses=None): - '''Utility routine which determines whether to return the cell value - or the colourmap lookup value. If supplied it also stores this access - in var_accesses. + def cell_ref_name( + self, var_accesses: Optional[VariablesAccessInfo] = None + ) -> tuple[str, Reference]: + ''' Utility routine which determines whether to return the cell + reference or the colourmap lookup array reference. If supplied with + a "var_accesses" it also stores the Variables Access information. - :param var_accesses: optional VariablesAccessInfo instance to store \ + :param var_accesses: optional VariablesAccessInfo instance to store the information about variable accesses. - :type var_accesses: \ - :py:class:`psyclone.core.VariablesAccessInfo` - :returns: the Fortran code needed to access the current cell index. - :rtype: Tuple[str, py:class:`psyclone.psyir.nodes.Reference`] + :returns: the reference needed to access the current cell index. + + TODO #2874: The name, argument, and first tuple component of this + and similar methods should be refactored. ''' cell_sym = self._symtab.find_or_create_integer_symbol( @@ -969,18 +972,7 @@ def cell_ref_name(self, var_accesses=None): if self._kern.is_coloured(): colour_sym = self._symtab.find_or_create_integer_symbol( "colour", tag="colours_loop_idx") - if self._kern.is_intergrid: - tag = None - else: - # If there is only one colourmap we need to specify the tag - # to make sure we get the right symbol. - tag = "cmap" - symbol = self._symtab.find_or_create( - self._kern.colourmap.name, symbol_type=DataSymbol, - datatype=ArrayType( - LFRicTypes("LFRicIntegerScalarDataType")(), - [ArrayType.Extent.DEFERRED, ArrayType.Extent.DEFERRED]), - tag=tag) + symbol = self._kern.colourmap array_ref = ArrayReference.create( symbol, [Reference(colour_sym), Reference(cell_sym)]) @@ -993,8 +985,7 @@ def cell_ref_name(self, var_accesses=None): AccessType.READ, self._kern, ["colour", "cell"]) - return (self._kern.colourmap.name + "(colour,cell)", - array_ref) + return (array_ref.debug_string(), array_ref) if var_accesses is not None: var_accesses.add_access(Signature("cell"), AccessType.READ, diff --git a/src/psyclone/domain/lfric/lfric_cell_iterators.py b/src/psyclone/domain/lfric/lfric_cell_iterators.py index b27fe9db98..fac5310e5b 100644 --- a/src/psyclone/domain/lfric/lfric_cell_iterators.py +++ b/src/psyclone/domain/lfric/lfric_cell_iterators.py @@ -110,6 +110,8 @@ def _stub_declarations(self, cursor): if self._kernel.cma_operation not in ["apply", "matrix-matrix"]: for name in self._nlayers_names: sym = self.symtab.lookup(name) + # Symbols are created thinking for the Invoke context, so make + # sure the are arguments when we are in a Stub context. sym.interface = ArgumentInterface( ArgumentInterface.Access.READ) self.symtab.append_argument(sym) @@ -121,6 +123,7 @@ def initialise(self, cursor): :param int cursor: position where to add the next initialisation statements. + :returns: Updated cursor value. :rtype: int diff --git a/src/psyclone/domain/lfric/lfric_collection.py b/src/psyclone/domain/lfric/lfric_collection.py index 71844f1d16..2cd0ba83ed 100644 --- a/src/psyclone/domain/lfric/lfric_collection.py +++ b/src/psyclone/domain/lfric/lfric_collection.py @@ -41,9 +41,12 @@ related entities within an Invoke or Kernel stub.''' import abc +from typing import List + from psyclone.domain.lfric.lfric_invoke import LFRicInvoke from psyclone.domain.lfric.lfric_kern import LFRicKern from psyclone.errors import InternalError +from psyclone.psyGen import Kern class LFRicCollection(): @@ -65,12 +68,10 @@ def __init__(self, node): # We are handling declarations/initialisations for an Invoke self._invoke = node self._kernel = None - # The list of Kernel calls we are responsible for elif isinstance(node, LFRicKern): # We are handling declarations for a Kernel stub self._invoke = None self._kernel = node - # We only have a single Kernel call in this case else: raise InternalError(f"LFRicCollection takes only an LFRicInvoke " f"or an LFRicKern but got: {type(node)}") @@ -94,10 +95,10 @@ def symtab(self): return self._kernel._stub_symbol_table @property - def _calls(self): + def kernel_calls(self) -> List[Kern]: ''' - :returns: associated kernels. - :rtype: List[:py:class:`psyclone.psyGen.kern`] + :returns: associated kernels calls. + ''' if self._invoke: return self._invoke.schedule.kernels() @@ -112,6 +113,7 @@ def declarations(self, cursor): :param int cursor: position where to add the next initialisation statements. + :returns: Updated cursor value. :rtype: int diff --git a/src/psyclone/domain/lfric/lfric_dofmaps.py b/src/psyclone/domain/lfric/lfric_dofmaps.py index a405ae7f7b..8b139c1922 100644 --- a/src/psyclone/domain/lfric/lfric_dofmaps.py +++ b/src/psyclone/domain/lfric/lfric_dofmaps.py @@ -88,7 +88,7 @@ def __init__(self, node): # "argument" and "direction" entries. self._unique_indirection_maps = OrderedDict() - for call in self._calls: + for call in self.kernel_calls: # We only need a dofmap if the kernel operates on cells # rather than dofs. if call.iterates_over != "dof": diff --git a/src/psyclone/domain/lfric/lfric_fields.py b/src/psyclone/domain/lfric/lfric_fields.py index a24c7b2d49..1d6f758b31 100644 --- a/src/psyclone/domain/lfric/lfric_fields.py +++ b/src/psyclone/domain/lfric/lfric_fields.py @@ -113,14 +113,13 @@ def _invoke_declarations(self, cursor: int) -> int: # new entry field_datatype_map[(arg.data_type, arg.module_name)] = [arg] - symtab = self._invoke.schedule.symbol_table # Add the Invoke subroutine argument declarations for the # different fields types. They are declared as intent "in" as # they contain a pointer to the data that is modified. for fld_type, fld_mod in field_datatype_map: args = field_datatype_map[(fld_type, fld_mod)] for arg in args: - arg_symbol = symtab.lookup(arg.name) + arg_symbol = self.symtab.lookup(arg.name) arg_symbol.interface.access = ArgumentInterface.Access.READ return cursor diff --git a/src/psyclone/domain/lfric/lfric_halo_depths.py b/src/psyclone/domain/lfric/lfric_halo_depths.py index 2658f2581f..173a65a5fd 100644 --- a/src/psyclone/domain/lfric/lfric_halo_depths.py +++ b/src/psyclone/domain/lfric/lfric_halo_depths.py @@ -66,7 +66,7 @@ def __init__(self, node): # No distributed memory so there are no halo regions. return depth_names = set() - for kern in self._calls: + for kern in self.kernel_calls: if not kern.halo_depth: continue if not isinstance(kern.halo_depth, Literal): diff --git a/src/psyclone/domain/lfric/lfric_invoke.py b/src/psyclone/domain/lfric/lfric_invoke.py index 40e9f6f841..d1390d52fe 100644 --- a/src/psyclone/domain/lfric/lfric_invoke.py +++ b/src/psyclone/domain/lfric/lfric_invoke.py @@ -272,17 +272,12 @@ def field_on_space(self, func_space): return None def declare(self): - ''' Declare and initialise all symbols associated this this Invoke. + ''' Declare and initialise all symbols associated this Invoke. Generates LFRic-specific invocation code (the subroutine called by the associated Invoke call in the algorithm layer). This consists of the PSy invocation subroutine and the declaration of its arguments. - :param parent: the parent node in the AST (of the code to be \ - generated) to which the node describing the PSy \ - subroutine will be added. - :type parent: :py:class:`psyclone.f2pygen.ModuleGen` - ''' cursor = 0 diff --git a/src/psyclone/domain/lfric/lfric_kern.py b/src/psyclone/domain/lfric/lfric_kern.py index a379274dfa..5ec7a97927 100644 --- a/src/psyclone/domain/lfric/lfric_kern.py +++ b/src/psyclone/domain/lfric/lfric_kern.py @@ -384,9 +384,11 @@ def _setup(self, ktype, module_name, args, parent, check=True): qr_arg = args[idx] quad_map = const.QUADRATURE_TYPE_MAP[shape] - # Use the InvokeSchedule symbol_table to create a unique symbol - # name for the whole Invoke. + # Use the InvokeSchedule or Stub symbol_table that we obtained + # earlier to create a unique symbol name if qr_arg.varname: + # If we have a name for the qr argument, we are dealing with + # an Invoke tag = "AlgArgs_" + qr_arg.text qr_sym = symtab.find_or_create( qr_arg.varname, tag=tag, symbol_type=DataSymbol, @@ -399,8 +401,6 @@ def _setup(self, ktype, module_name, args, parent, check=True): else: # If we don't have a name then we must be doing kernel-stub # generation so create a suitable name. - # TODO #719 we don't yet have a symbol table to prevent - # clashes. qr_name = "qr_"+shape.split("_")[-1] # Append the name of the qr argument to the names of the qr-related @@ -469,12 +469,9 @@ def is_intergrid(self): return self._intergrid_ref is not None @property - def colourmap(self): + def colourmap(self: DataSymbol): ''' - Getter for the name of the colourmap associated with this kernel call. - - :returns: name of the colourmap (Fortran array). - :rtype: str + :returns: the symbol representing the colourmap for this kernekl call. :raises InternalError: if this kernel is not coloured. @@ -653,12 +650,11 @@ def argument_kinds(self): return self._argument_kinds @property - def gen_stub(self): + def gen_stub(self) -> Container: ''' Create the PSyIR for a kernel stub. :returns: the kernel stub root Container. - :rtype: :py:class:`psyclone.psyir.nodes.Container` :raises GenerationError: if the supplied kernel stub does not operate on a supported subset of the domain (currently only those that @@ -707,15 +703,17 @@ def gen_stub(self): DynFunctionSpaces, DynCMAOperators, DynBoundaryConditions, DynLMAOperators, LFRicMeshProperties, DynBasisFunctions, DynReferenceElement) + cursor = 0 for entities in [LFRicCellIterators, LFRicDofmaps, DynFunctionSpaces, DynCMAOperators, LFRicScalarArgs, LFRicFields, DynLMAOperators, LFRicStencils, DynBasisFunctions, DynBoundaryConditions, DynReferenceElement, LFRicMeshProperties]: - entities(self).declarations(stub_routine) + cursor = entities(self).declarations(cursor) - # The declarations above are not in order, we need to use the - # KernStubArgList to generate a list of strings with the correct order + # TODO #2874: The declarations above are not in order, we need to use + # the KernStubArgList to generate a list of strings with the correct + # order create_arg_list = KernStubArgList(self) create_arg_list._forced_symtab = stub_routine.symbol_table create_arg_list.generate() diff --git a/src/psyclone/domain/lfric/lfric_loop.py b/src/psyclone/domain/lfric/lfric_loop.py index 8ab7b5051d..2c81372f73 100644 --- a/src/psyclone/domain/lfric/lfric_loop.py +++ b/src/psyclone/domain/lfric/lfric_loop.py @@ -103,18 +103,20 @@ def __init__(self, loop_type="", **kwargs): # Initialise loop bounds ischedule = self.ancestor(InvokeSchedule) - if ischedule: - idx = len(ischedule.loops()) - start_name = f"loop{idx}_start" - stop_name = f"loop{idx}_stop" - lbound = ischedule.symbol_table.find_or_create_integer_symbol( - start_name, tag=start_name) - ubound = ischedule.symbol_table.find_or_create_integer_symbol( - stop_name, tag=stop_name) - else: - # Only for testing, PS-layer creation always has an invoke - lbound = DataSymbol("undeclared_loop_start", datatype=INTEGER_TYPE) - ubound = DataSymbol("undeclared_loop_stop", datatype=INTEGER_TYPE) + if not ischedule: + raise InternalError( + "LFRic loops can only be inside a InvokeSchedule, a parent " + "argument is mandatory when they are created.") + # The loop bounds names are given by the number of previous LFRic loops + # already present in the Schedule. Since this are inserted in order it + # will produce sequencially ascending loop bound names. + idx = len(ischedule.loops()) + start_name = f"loop{idx}_start" + stop_name = f"loop{idx}_stop" + lbound = ischedule.symbol_table.find_or_create_integer_symbol( + start_name, tag=start_name) + ubound = ischedule.symbol_table.find_or_create_integer_symbol( + stop_name, tag=stop_name) self.addchild(Reference(lbound)) # start self.addchild(Reference(ubound)) # stop self.addchild(Literal("1", INTEGER_TYPE, parent=self)) # step diff --git a/src/psyclone/domain/lfric/lfric_scalar_args.py b/src/psyclone/domain/lfric/lfric_scalar_args.py index 80d9ab0cb9..680cc2d3a9 100644 --- a/src/psyclone/domain/lfric/lfric_scalar_args.py +++ b/src/psyclone/domain/lfric/lfric_scalar_args.py @@ -158,7 +158,7 @@ def _stub_declarations(self, cursor: int) -> int: ''' # Extract all scalar arguments - for arg in self._calls[0].arguments.args: + for arg in self.kernel_calls[0].arguments.args: if arg.is_scalar: self._scalar_args[arg.intent].append(arg) diff --git a/src/psyclone/domain/lfric/lfric_stencils.py b/src/psyclone/domain/lfric/lfric_stencils.py index cfafc5d926..a727d1e581 100644 --- a/src/psyclone/domain/lfric/lfric_stencils.py +++ b/src/psyclone/domain/lfric/lfric_stencils.py @@ -75,7 +75,7 @@ def __init__(self, node): self._unique_extent_args = [] extent_names = [] # pylint: disable=too-many-nested-blocks - for call in self._calls: + for call in self.kernel_calls: for arg in call.arguments.args: if arg.stencil: # Check for the existence of arg.extent here as in @@ -96,7 +96,7 @@ def __init__(self, node): # argument names are removed. self._unique_direction_args = [] direction_names = [] - for call in self._calls: + for call in self.kernel_calls: for idx, arg in enumerate(call.arguments.args): if arg.stencil and arg.stencil.direction_arg: if arg.stencil.direction_arg.is_literal(): @@ -115,7 +115,7 @@ def __init__(self, node): # List of stencil args with an extent variable passed in. The same # field name may occur more than once here from different kernels. self._kern_args = [] - for call in self._calls: + for call in self.kernel_calls: for arg in call.arguments.args: if arg.stencil: if not arg.stencil.extent: diff --git a/src/psyclone/dynamo0p3.py b/src/psyclone/dynamo0p3.py index 11a6b38634..f69c5ee40e 100644 --- a/src/psyclone/dynamo0p3.py +++ b/src/psyclone/dynamo0p3.py @@ -386,7 +386,7 @@ def __init__(self, node): # kernel stub. self._properties = [] - for call in self._calls: + for call in self.kernel_calls: if call.mesh: self._properties += [prop for prop in call.mesh.properties if prop not in self._properties] @@ -631,7 +631,7 @@ def initialise(self, cursor: int) -> int: # it now, rather than when this class was first constructed. need_colour_limits = False need_colour_halo_limits = False - for call in self._calls: + for call in self.kernel_calls: if call.is_coloured() and not call.is_intergrid: loop = call.parent.parent # Record whether or not this coloured loop accesses the halo. @@ -742,7 +742,7 @@ def __init__(self, node): self._properties = [] self._nfaces_h_required = False - for call in self._calls: + for call in self.kernel_calls: if call.reference_element: self._properties.extend(call.reference_element.properties) if call.mesh and call.mesh.properties: @@ -1124,7 +1124,7 @@ def __init__(self, kern_or_invoke): if self._invoke: self._function_spaces = self._invoke.unique_fss()[:] else: - self._function_spaces = self._calls[0].arguments.unique_fss + self._function_spaces = self.kernel_calls[0].arguments.unique_fss self._var_list = [] @@ -1681,7 +1681,7 @@ def __init__(self, node): # You can't index into an OrderedDict so we keep a separate ref # to the first CMA argument we find. self._first_cma_arg = None - for call in self._calls: + for call in self.kernel_calls: if call.cma_operation: # Get a list of all of the CMA arguments to this call cma_args = psyGen.args_filter( @@ -2580,7 +2580,7 @@ def __init__(self, node): # DynKernelArgument) tuples. self._eval_targets = OrderedDict() - for call in self._calls: + for call in self.kernel_calls: if isinstance(call, LFRicBuiltIn) or not call.eval_shapes: # Skip this kernel if it doesn't require basis/diff basis fns @@ -3599,7 +3599,7 @@ def __init__(self, node): # pylint: disable=import-outside-toplevel from psyclone.domain.lfric.metadata_to_arguments_rules import ( MetadataToArgumentsRules) - for call in self._calls: + for call in self.kernel_calls: if MetadataToArgumentsRules.bc_kern_regex.match(call.name): bc_fs = None for fspace in call.arguments.unique_fss: diff --git a/src/psyclone/tests/domain/lfric/lfric_loop_test.py b/src/psyclone/tests/domain/lfric/lfric_loop_test.py index 377b9476ef..2b80627484 100644 --- a/src/psyclone/tests/domain/lfric/lfric_loop_test.py +++ b/src/psyclone/tests/domain/lfric/lfric_loop_test.py @@ -47,7 +47,8 @@ from psyclone.configuration import Config from psyclone.core import AccessType from psyclone.domain.lfric import (LFRicConstants, LFRicSymbolTable, - LFRicKern, LFRicKernMetadata, LFRicLoop) + LFRicKern, LFRicKernMetadata, LFRicLoop, + LFRicInvokeSchedule) from psyclone.errors import GenerationError, InternalError from psyclone.parse.algorithm import parse from psyclone.psyGen import PSyFactory @@ -66,6 +67,28 @@ TEST_API = "lfric" +def test_constructor_loop_bound_names(): + ''' Check that the constructor creates the appropriate loop bound + references (with names with a sequencially ascending index) + ''' + with pytest.raises(InternalError) as err: + _ = LFRicLoop(loop_type="null") + assert ("LFRic loops can only be inside a InvokeSchedule, a parent " + "argument is mandatory when they are created." in str(err.value)) + + schedule = LFRicInvokeSchedule.create("test") + schedule.addchild(LFRicLoop(parent=schedule)) + schedule.addchild(LFRicLoop(parent=schedule)) + schedule.addchild(LFRicLoop(parent=schedule)) + loops = schedule.loops() + assert loops[0].start_expr.name == "loop0_start" + assert loops[1].start_expr.name == "loop1_start" + assert loops[2].start_expr.name == "loop2_start" + assert loops[0].stop_expr.name == "loop0_stop" + assert loops[1].stop_expr.name == "loop1_stop" + assert loops[2].stop_expr.name == "loop2_stop" + + def test_constructor_invalid_loop_type(monkeypatch): ''' Check that the constructor raises the expected errors when an invalid loop type is specified. @@ -95,7 +118,7 @@ def test_set_lower_bound_functions(monkeypatch): # TODO #1954: Remove the protected access using a factory monkeypatch.setattr(ScopingNode, "_symbol_table_class", LFRicSymbolTable) - schedule = Schedule() + schedule = LFRicInvokeSchedule.create("test") my_loop = LFRicLoop(parent=schedule) schedule.children = [my_loop] with pytest.raises(GenerationError) as excinfo: @@ -116,7 +139,7 @@ def test_set_upper_bound_functions(monkeypatch): # TODO #1954: Remove the protected access using a factory monkeypatch.setattr(ScopingNode, "_symbol_table_class", LFRicSymbolTable) - schedule = Schedule() + schedule = LFRicInvokeSchedule.create("test") my_loop = LFRicLoop(parent=schedule) schedule.children = [my_loop] with pytest.raises(GenerationError) as excinfo: @@ -718,7 +741,7 @@ def test_null_loop(): check in the 'load()' method behaves as expected. ''' - loop = LFRicLoop(loop_type="null") + loop = LFRicLoop(loop_type="null", parent=LFRicInvokeSchedule.create("a")) assert loop.loop_type == "null" assert loop.node_str(colour=False) == "Loop[type='null']" @@ -753,7 +776,7 @@ def test_loop_independent_iterations(monkeypatch, dist_mem): '''Tests for the independent_iterations() method.''' # A 'null' loop cannot be parallelised (because there's nothing to # parallelise). - loop = LFRicLoop(loop_type="null") + loop = LFRicLoop(loop_type="null", parent=LFRicInvokeSchedule.create("a")) assert not loop.independent_iterations() # A loop over all columns that contains a kernel that increments a field # on a continuous function space does not have independent iterations.