#2739 Update PSyKAl transformation signature to accept PSYIR node (an…

…d add FileContainer.invokes for backwards compatibility)

examples/gocean/eg6/backends_transform.py

@@ -63,7 +63,7 @@ def trans(psy):
     fvisitor = FortranWriter()
     print("")
     print("FortranWriter code:")
-    print(fvisitor(psy.container))
+    print(fvisitor(schedule.root))
 
     # This PSyclone call should terminate gracefully here
     sys.exit(0)

examples/lfric/eg2/loop_fuse_trans.py

@@ -56,8 +56,6 @@ def trans(psy):
     :rtype: :py:class:`psyclone.dynamo0p3.DynamoPSy`
 
     '''
-    print(psy.gen)
-
     print(psy.invokes.names)
 
     schedule = psy.invokes.get('invoke_0').schedule

examples/lfric/eg3/colouring_and_omp.py

@@ -36,45 +36,39 @@
 # Modified by J. Henrichs, Bureau of Meteorology
 
 
-''' File containing a PSyclone transformation script for the Dynamo0p3
+''' File containing a PSyclone transformation script for the LFRic
 API to apply colouring and OpenMP generically. This can be applied via
 the -s option in the "psyclone" script. '''
-from __future__ import print_function
 from psyclone.transformations import Dynamo0p3ColourTrans, \
     DynamoOMPParallelLoopTrans
-from psyclone.psyir.nodes import Loop
+from psyclone.psyir.nodes import Loop, Routine
 from psyclone.domain.lfric import LFRicConstants
 
 
-def trans(psy):
+def trans(psyir):
     ''' PSyclone transformation script for the dynamo0p3 api to apply
     colouring and OpenMP generically.'''
     ctrans = Dynamo0p3ColourTrans()
     otrans = DynamoOMPParallelLoopTrans()
     const = LFRicConstants()
 
     # Loop over all of the Invokes in the PSy object
-    for invoke in psy.invokes.invoke_list:
+    for subroutine in psyir.walk(Routine):
 
-        print("Transforming invoke '"+invoke.name+"'...")
-        schedule = invoke.schedule
+        print(f"Transforming invoke '{subroutine.name}' ...")
 
         # Colour all of the loops over cells unless they are on
         # discontinuous spaces
-        for child in schedule.children:
+        for child in subroutine.children:
             if isinstance(child, Loop) \
                and child.field_space.orig_name \
                not in const.VALID_DISCONTINUOUS_NAMES \
                and child.iteration_space == "cell_column":
                 ctrans.apply(child)
         # Then apply OpenMP to each of the colour loops
-        for child in schedule.children:
+        for child in subroutine.children:
             if isinstance(child, Loop):
                 if child.loop_type == "colours":
                     otrans.apply(child.loop_body[0])
                 else:
                     otrans.apply(child)
-
-        print(schedule.view())
-
-    return psy

examples/lfric/eg4/backends_transform.py

@@ -57,9 +57,6 @@ def trans(psy):
         print("DSL level view:")
         print(schedule.view())
 
-    print("f2pygen code:")
-    print(str(psy.gen))
-
     # TODO #1010: This script should terminate here until LFRic declares
     # all its symbols to the symbol table.
     sys.exit(0)

src/psyclone/domain/gocean/transformations/gocean_opencl_trans.py

@@ -49,7 +49,7 @@
 from psyclone.psyir.frontend.fortran import FortranReader
 from psyclone.psyir.nodes import Routine, Call, Reference, Literal, \
     Assignment, IfBlock, ArrayReference, Schedule, BinaryOperation, \
-    StructureReference, FileContainer, CodeBlock, IntrinsicCall
+    StructureReference, FileContainer, CodeBlock, IntrinsicCall, Container
 from psyclone.psyir.symbols import (
     ArrayType, DataSymbol, RoutineSymbol, ContainerSymbol,
     UnsupportedFortranType, ArgumentInterface, ImportInterface,
@@ -264,15 +264,16 @@ def apply(self, node, options=None):
 
         # Insert, if they don't already exist, the necessary OpenCL helper
         # subroutines in the root Container.
-        psy_init = self._insert_opencl_init_routine(node.root)
-        init_grid = self._insert_initialise_grid_buffers(node.root)
-        write_grid_buf = self._insert_write_grid_buffers(node.root)
-        self._insert_ocl_read_from_device_function(node.root)
-        self._insert_ocl_write_to_device_function(node.root)
-        init_buf = self._insert_ocl_initialise_buffer(node.root)
+        module = node.ancestor(Container)
+        psy_init = self._insert_opencl_init_routine(module)
+        init_grid = self._insert_initialise_grid_buffers(module)
+        write_grid_buf = self._insert_write_grid_buffers(module)
+        self._insert_ocl_read_from_device_function(module)
+        self._insert_ocl_write_to_device_function(module)
+        init_buf = self._insert_ocl_initialise_buffer(module)
 
         for kern in node.coded_kernels():
-            self._insert_ocl_arg_setter_routine(node.root, kern)
+            self._insert_ocl_arg_setter_routine(module, kern)
 
         # Insert fortcl, clfortran and c_iso_binding import statement
         fortcl = ContainerSymbol("fortcl")

src/psyclone/domain/lfric/kern_call_acc_arg_list.py

@@ -45,6 +45,7 @@
 from psyclone import psyGen
 from psyclone.domain.lfric import KernCallArgList
 from psyclone.errors import InternalError
+from psyclone.psyir.nodes import Container
 
 
 class KernCallAccArgList(KernCallArgList):
@@ -110,8 +111,8 @@ def stencil(self, arg, var_accesses=None):
         # 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._kern.root.symbol_table,
-                                               arg).name
+        var_name = LFRicStencils.dofmap_symbol(
+            self._kern.ancestor(Container).symbol_table, arg).name
         self.append(var_name, var_accesses)
 
     def stencil_2d(self, arg, var_accesses=None):
@@ -148,8 +149,8 @@ def stencil_unknown_extent(self, arg, var_accesses=None):
         # 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._kern.root.symbol_table,
-                                                arg).name
+        name = LFRicStencils.dofmap_size_symbol(
+                    self._kern.ancestor(Container).symbol_table, arg).name
         self.append(name, var_accesses)
 
     def stencil_2d_unknown_extent(self, arg, var_accesses=None):

src/psyclone/dynamo0p3.py

@@ -70,7 +70,8 @@
 from psyclone.psyir.frontend.fortran import FortranReader
 from psyclone.psyir.nodes import (
     Reference, ACCEnterDataDirective, ScopingNode, ArrayOfStructuresReference,
-    StructureReference, Literal, IfBlock, Call, BinaryOperation, IntrinsicCall)
+    StructureReference, Literal, IfBlock, Call, BinaryOperation, IntrinsicCall,
+    Container)
 from psyclone.psyir.symbols import (INTEGER_TYPE, DataSymbol, ScalarType,
                                     UnresolvedType, DataTypeSymbol,
                                     ContainerSymbol, ImportInterface,
@@ -6037,12 +6038,13 @@ def infer_datatype(self, proxy=False):
         :raises NotImplementedError: if an unsupported argument type is found.
 
         '''
-        # We want to put any Container symbols in the outermost scope so find
-        # the corresponding symbol table.
-        symbol_table = self._call.scope.symbol_table
-        root_table = symbol_table
-        while root_table.parent_symbol_table():
-            root_table = root_table.parent_symbol_table()
+        scope = self._call.ancestor(Container)
+        if scope is None:
+            # Prefer the module scope, but some tests that are disconnected can
+            # use the current scope
+            scope = self._call.scope
+
+        symtab = scope.symbol_table
 
         def _find_or_create_type(mod_name, type_name):
             '''
@@ -6058,11 +6060,11 @@ def _find_or_create_type(mod_name, type_name):
             :rtype: :py:class:`psyclone.psyir.symbols.DataTypeSymbol`
 
             '''
-            return root_table.find_or_create(
+            return symtab.find_or_create(
                     type_name,
                     symbol_type=DataTypeSymbol,
                     datatype=UnresolvedType(),
-                    interface=ImportInterface(root_table.find_or_create(
+                    interface=ImportInterface(symtab.find_or_create(
                         mod_name,
                         symbol_type=ContainerSymbol)
                         ))
@@ -6081,22 +6083,22 @@ def _find_or_create_type(mod_name, type_name):
 
             kind_name = self.precision
             try:
-                kind_symbol = symbol_table.lookup(kind_name)
+                kind_symbol = symtab.lookup(kind_name)
             except KeyError:
                 mod_map = LFRicConstants().UTILITIES_MOD_MAP
                 const_mod = mod_map["constants"]["module"]
                 try:
-                    constants_container = symbol_table.lookup(const_mod)
+                    constants_container = symtab.lookup(const_mod)
                 except KeyError:
                     # TODO Once #696 is done, we should *always* have a
                     # symbol for this container at this point so should
                     # raise an exception if we haven't.
                     constants_container = LFRicTypes(const_mod)
-                    root_table.add(constants_container)
+                    symtab.add(constants_container)
                 kind_symbol = DataSymbol(
                     kind_name, INTEGER_TYPE,
                     interface=ImportInterface(constants_container))
-                root_table.add(kind_symbol)
+                symtab.add(kind_symbol)
             return ScalarType(prim_type, kind_symbol)
 
         if self.is_field or self.is_operator:

src/psyclone/generator.py

@@ -159,12 +159,9 @@ def generate(filename, api="", kernel_paths=None, script_name=None,
     # pylint: disable=too-many-arguments, too-many-statements
     # pylint: disable=too-many-branches, too-many-locals
     '''Takes a PSyclone algorithm specification as input and outputs the
-    associated generated algorithm and psy codes suitable for
+    associated generated algorithm and psy-layer codes suitable for
     compiling with the specified kernel(s) and support
-    infrastructure. Uses the :func:`parse.algorithm.parse` function to
-    parse the algorithm specification, the :class:`psyGen.PSy` class
-    to generate the PSy code and the :class:`alg_gen.Alg` class to
-    generate the modified algorithm code.
+    infrastructure.
 
     :param str filename: the file containing the algorithm specification.
     :param str api: the name of the API to use. Defaults to empty string.
@@ -252,7 +249,7 @@ def generate(filename, api="", kernel_paths=None, script_name=None,
         if script_name is not None:
             # Apply provided recipe to PSyIR
             recipe, _ = load_script(script_name)
-            recipe(psy)
+            recipe(psy.container.root)
         alg_gen = None
 
     elif api in GOCEAN_API_NAMES or (api in LFRIC_API_NAMES and LFRIC_TESTING):
@@ -380,7 +377,7 @@ def generate(filename, api="", kernel_paths=None, script_name=None,
         if script_name is not None:
             # Call the optimisation script for psy-layer optimisations
             recipe, _ = load_script(script_name)
-            recipe(psy)
+            recipe(psy.container.root)
 
     # TODO issue #1618 remove Alg class and tests from PSyclone
     if api in LFRIC_API_NAMES and not LFRIC_TESTING:

src/psyclone/gocean1p0.py

@@ -188,12 +188,12 @@ def gen_code(self, parent):
 
             # Lower the GOcean PSyIR to language level so it can be visited
             # by the backends
-            invoke.schedule.root.lower_to_language_level()
+            invoke.schedule.parent.lower_to_language_level()
             # Then insert it into a f2pygen AST as a PSyIRGen node.
             # Note that other routines besides the Invoke could have been
             # inserted during the lowering (e.g. module-inlined kernels),
             # so have to iterate over all current children of root.
-            for child in invoke.schedule.root.children:
+            for child in invoke.schedule.parent.children:
                 parent.add(PSyIRGen(parent, child))
 
 
@@ -1498,19 +1498,26 @@ def infer_datatype(self):
                                property is not 'go_r_scalar' or 'go_i_scalar'.
 
         '''
+        scope = self._call.ancestor(Container)
+        if scope is None:
+            # Prefer the module scope, but some tests that are disconnected can
+            # use the current scope
+            scope = self._call.scope
+
+        symtab = scope.symbol_table
         # All GOcean fields are r2d_field
         if self.argument_type == "field":
             # r2d_field can have UnresolvedType and UnresolvedInterface because
             # it is an unnamed import from a module.
-            type_symbol = self._call.root.symbol_table.find_or_create_tag(
+            type_symbol = symtab.find_or_create_tag(
                 "r2d_field", symbol_type=DataTypeSymbol,
                 datatype=UnresolvedType(), interface=UnresolvedInterface())
             return type_symbol
 
         # Gocean scalars can be REAL or INTEGER
         if self.argument_type == "scalar":
             if self.space.lower() == "go_r_scalar":
-                go_wp = self._call.root.symbol_table.find_or_create_tag(
+                go_wp = symtab.find_or_create_tag(
                     "go_wp", symbol_type=DataSymbol, datatype=UnresolvedType(),
                     interface=UnresolvedInterface())
                 return ScalarType(ScalarType.Intrinsic.REAL, go_wp)
@@ -2150,8 +2157,18 @@ def _read_from_device_routine(self):
 
         :returns: the symbol representing the read_from_device routine.
         :rtype: :py:class:`psyclone.psyir.symbols.symbol`
+
+        :raises GenerationError: if this class is lowered from a location where
+            a Routine can not be inserted.
         '''
-        symtab = self.root.symbol_table
+        # Insert the routine as a child of the ancestor Container
+        if not self.ancestor(Container):
+            raise GenerationError(
+                f"The GOACCEnterDataDirective can only be generated/lowered "
+                f"inside a Container in order to insert a sibling "
+                f"subroutine, but '{self}' is not inside a Container.")
+
+        symtab = self.ancestor(Container).symbol_table
         try:
             return symtab.lookup_with_tag("openacc_read_func")
         except KeyError:
@@ -2180,12 +2197,6 @@ def _read_from_device_routine(self):
         # Add an ACCUpdateDirective inside the subroutine
         subroutine.addchild(ACCUpdateDirective([Signature("to")], "host",
                                                if_present=False))
-        # Insert the routine as a child of the ancestor Container
-        if not self.ancestor(Container):
-            raise GenerationError(
-                f"The GOACCEnterDataDirective can only be generated/lowered "
-                f"inside a Container in order to insert a sibling "
-                f"subroutine, but '{self}' is not inside a Container.")
 
         self.ancestor(Container).symbol_table.add(subroutine_symbol,
                                                   tag="openacc_read_func")

src/psyclone/psyGen.py

@@ -54,7 +54,7 @@
 from psyclone.psyir.backend.fortran import FortranWriter
 from psyclone.psyir.nodes import (ArrayReference, Call, Container, Literal,
                                   Loop, Node, OMPDoDirective, Reference,
-                                  Routine, Schedule, Statement)
+                                  Routine, Schedule, Statement, FileContainer)
 from psyclone.psyir.symbols import (ArgumentInterface, ArrayType,
                                     ContainerSymbol, DataSymbol,
                                     UnresolvedType,
@@ -230,9 +230,11 @@ class PSy():
     def __init__(self, invoke_info):
         self._name = invoke_info.name
         self._invokes = None
-        # create an empty PSy layer container
+        # Create an empty PSy layer PSyIR file with a Container (module) inside
         # TODO 1010: Alternatively the PSy object could be a Container itself
-        self._container = Container(self.name)
+        module = Container(self.name)
+        FileContainer(self.name, children=[module])
+        self._container = module
 
     @property
     def container(self):

src/psyclone/psyir/nodes/file_container.py

@@ -38,6 +38,7 @@
 ''' This module contains the FileContainer node implementation.'''
 
 from psyclone.psyir.nodes.container import Container
+from psyclone.errors import GenerationError
 
 
 class FileContainer(Container):
@@ -55,6 +56,40 @@ class FileContainer(Container):
     def __str__(self):
         return f"FileContainer[name='{self.name}']\n"
 
+    @property
+    def invokes(self):
+        ''' Return the Invokes object associated to this FileContainer.
+        This is for compatibility with old psyclone transformation scripts.
+        Before the entry point was PSy, and the script had to find the list
+        of InvokeSchedules, now the entry point is the root FileContainer:
+
+        before: PSy -> Invokes -> Invoke -> InvokeSchedule
+        now:                FileContainer --^
+
+        This method creates a shortcut:
+        PSy -> Invokes -> Invoke -> InvokeSchedule
+               ^--- FileContainer --^
+
+        So that previous:
+            def trans(psy):
+                psy.invokes.get_invoke('name').schedule
+
+        Still work as expected. However, it still exposes the PSy hierachy to
+        users scripts, so this will be deprecated after a grace period.
+
+        :return: the associated Invokes object.
+        :rtype: :py:class:`psyclone.psyGen.Invokes`
+        :raises GenerationError: if no InvokeSchedule was found.
+        '''
+        # pylint: disable=import-outside-toplevel
+        from psyclone.psyGen import InvokeSchedule
+        invokes = self.walk(InvokeSchedule, stop_type=InvokeSchedule)
+        if not invokes:
+            raise GenerationError(
+                f"No InvokeSchedule found in {self.name}, does it come from"
+                f" a PSyKAl file?")
+        return invokes[0].invoke.invokes
+
 
 # For AutoAPI documentation generation
 __all__ = ['FileContainer']