Merge pull request #154 from ecmwf-ifs/naml-transformation-manifest

Static Transformation properties (manifest)

loki/bulk/scheduler.py

@@ -614,7 +614,7 @@ def item_successors(self, item):
                 successors += [self.item_map[child.name]] + self.item_successors(child)
         return successors
 
-    def process(self, transformation, reverse=False, item_filter=SubroutineItem, use_file_graph=False):
+    def process(self, transformation):
         """
         Process all :attr:`items` in the scheduler's graph
 
@@ -638,16 +638,16 @@ def process(self, transformation, reverse=False, item_filter=SubroutineItem, use
         log = f'[Loki::Scheduler] Applied transformation <{trafo_name}>' + ' in {:.2f}s'
         with Timer(logger=info, text=log):
 
-            if use_file_graph:
-                graph = self.file_graph
-            else:
-                graph = self.item_graph
+            # Extract the graph iteration properties from the transformation
+            graph = self.file_graph if transformation.traverse_file_graph else self.item_graph
+            item_filter = as_tuple(transformation.item_filter)
 
+            # Construct the actual graph to traverse
             traversal = nx.topological_sort(graph)
-            if reverse:
+            if transformation.reverse_traversal:
                 traversal = reversed(list(traversal))
 
-            if use_file_graph:
+            if transformation.traverse_file_graph:
                 for node in traversal:
                     items = graph.nodes[node]['items']
 
@@ -656,6 +656,7 @@ def process(self, transformation, reverse=False, item_filter=SubroutineItem, use
                         if not items:
                             continue
 
+                    _item = items[0]
                     transformation.apply(items[0].source, items=items)
             else:
                 for item in traversal:

loki/lint/linter.py

@@ -248,6 +248,9 @@ class LinterTransformation(Transformation):
 
     _key = 'LinterTransformation'
 
+    # This transformation is applied over the file graph
+    traverse_file_graph = True
+
     def __init__(self, linter, key=None, **kwargs):
         self.linter = linter
         self.counter = 0
@@ -272,7 +275,7 @@ def lint_files_scheduler(linter, basedir, config):
     """
     scheduler = Scheduler(paths=[basedir], config=SchedulerConfig.from_dict(config))
     transformation = LinterTransformation(linter=linter)
-    scheduler.process(transformation=transformation, use_file_graph=True)
+    scheduler.process(transformation=transformation)
     return transformation.counter
 
 

loki/transform/build_system_transform.py

@@ -13,6 +13,7 @@
 
 from loki.logging import info
 from loki.transform.transformation import Transformation
+from loki.bulk.item import SubroutineItem, GlobalVarImportItem
 
 __all__ = ['CMakePlanner', 'FileWriteTransformation']
 
@@ -129,12 +130,33 @@ class FileWriteTransformation(Transformation):
         omitted, it will preserve the original file type.
     cuf : bool, optional
         Use CUF (CUDA Fortran) backend instead of Fortran backend.
+    include_module_var_imports : bool, optional
+        Flag to force the :any:`Scheduler` traversal graph to recognise
+        module variable imports and write the modified module files.
     """
-    def __init__(self, builddir=None, mode='loki', suffix=None, cuf=False):
+
+    # This transformation is applied over the file graph
+    traverse_file_graph = True
+
+    def __init__(
+            self, builddir=None, mode='loki', suffix=None, cuf=False,
+            include_module_var_imports=False
+    ):
         self.builddir = Path(builddir)
         self.mode = mode
         self.suffix = suffix
         self.cuf = cuf
+        self.include_module_var_imports = include_module_var_imports
+
+    @property
+    def item_filter(self):
+        """
+        Override ``item_filter`` to configure whether module variable
+        imports are honoured in the :any:`Scheduler` traversal.
+        """
+        if self.include_module_var_imports:
+            return (SubroutineItem, GlobalVarImportItem)
+        return SubroutineItem
 
     def transform_file(self, sourcefile, **kwargs):
         item = kwargs.get('item', None)

loki/transform/dependency_transform.py

@@ -62,6 +62,15 @@ class DependencyTransformation(Transformation):
         in the ``ignore``. Default is ``True``.
     """
 
+    # This transformation is applied over the file graph
+    traverse_file_graph = True
+
+    # This transformation recurses from the Sourcefile down
+    recurse_to_modules = True
+    recurse_to_procedures = True
+    recurse_to_internal_procedures = False
+
+
     def __init__(self, suffix, mode='module', module_suffix=None, include_path=None,
                  replace_ignore_items=True):
         self.suffix = suffix
@@ -151,18 +160,6 @@ def transform_file(self, sourcefile, **kwargs):
         if role == 'kernel' and self.mode == 'module':
             self.module_wrap(sourcefile, **kwargs)
 
-        for module in sourcefile.modules:
-            # Recursion into contained modules using the sourcefile's "role"
-            self.transform_module(module, role=role, targets=targets, **kwargs)
-
-        if items:
-            # Recursion into all subroutine items in the current file
-            for item in items:
-                self.transform_subroutine(item.routine, item=item, role=item.role, targets=item.targets, **kwargs)
-        else:
-            for routine in sourcefile.all_subroutines:
-                self.transform_subroutine(routine, role=role, targets=targets, **kwargs)
-
     def rename_calls(self, routine, **kwargs):
         """
         Update calls to actively transformed subroutines.

loki/transform/transform_hoist_variables.py

@@ -46,7 +46,7 @@
 .. code-block:: python
 
     # Transformation: Analysis
-    scheduler.process(transformation=HoistTemporaryArraysAnalysis(), reverse=True)
+    scheduler.process(transformation=HoistTemporaryArraysAnalysis())
     # Transformation: Synthesis
     scheduler.process(transformation=HoistVariablesTransformation())
 
@@ -73,10 +73,10 @@
     .. code-block:: python
 
         key = "UniqueKey"
-        scheduler.process(transformation=HoistTemporaryArraysAnalysis(dim_vars=('b',), key=key), reverse=True)
+        scheduler.process(transformation=HoistTemporaryArraysAnalysis(dim_vars=('b',), key=key))
         scheduler.process(transformation=HoistTemporaryArraysTransformation(key=key))
         key = "AnotherUniqueKey"
-        scheduler.process(transformation=HoistTemporaryArraysAnalysis(dim_vars=('a',), key=key), reverse=True)
+        scheduler.process(transformation=HoistTemporaryArraysAnalysis(dim_vars=('a',), key=key))
         scheduler.process(transformation=HoistTemporaryArraysTransformationAllocatable(key=key))
 """
 from loki.expression import FindVariables, SubstituteExpressions
@@ -100,9 +100,6 @@ class HoistVariablesAnalysis(Transformation):
     Create a derived class and override :func:`find_variables<HoistVariablesAnalysis.find_variables>`
     to define which variables to be hoisted.
 
-    .. note::
-        To be applied **reversed**, in order to recursively find all variables to be hoisted.
-
     Parameters
     ----------
     key : str
@@ -112,6 +109,9 @@ class HoistVariablesAnalysis(Transformation):
 
     _key = 'HoistVariablesTransformation'
 
+    # Apply in reverse order to recursively find all variables to be hoisted.
+    reverse_traversal = True
+
     def __init__(self, key=None):
         if key is not None:
             self._key = key
@@ -325,6 +325,9 @@ class HoistTemporaryArraysAnalysis(HoistVariablesAnalysis):
         for the array dimensions.
     """
 
+    # Apply in reverse order to recursively find all variables to be hoisted.
+    reverse_traversal = True
+
     def __init__(self, key=None, dim_vars=None, **kwargs):
         super().__init__(key=key, **kwargs)
         self.dim_vars = dim_vars

loki/transform/transformation.py

@@ -11,6 +11,7 @@
 from loki.module import Module
 from loki.sourcefile import Sourcefile
 from loki.subroutine import Subroutine
+from loki.bulk.item import SubroutineItem
 
 
 __all__ = ['Transformation']
@@ -35,8 +36,53 @@ class Transformation:
 
     Note that in :any:`Sourcefile` objects, all :any:`Module` members will be
     traversed before standalone :any:`Subroutine` objects.
+
+    Classes inheriting from :any:`Transformation` may configure the
+    invocation and behaviour during batch processing via a predefined
+    set of class attributes. These flags determine the underlying
+    graph traversal when processing complex call trees and determine
+    how the transformations are invoked for a given type of scheduler
+    :any:`Item`.
+
+    Attributes
+    ----------
+    reverse_traversal : bool
+        Forces scheduler traversal in reverse order from the leaf
+        nodes upwards (default: ``False``).
+    traverse_file_graph : bool
+         Apply :any:`Transformation` to the :any:`Sourcefile` object
+         corresponding to the :any:`Item` being processed, instead of
+         the program unit in question (default: ``False``).
+    item_filter : bool
+        Filter by graph node types to prune the graph and change connectivity.
+        By default, only calls to :any:`Subroutine` items are used to construct
+        the graph.
+    recurse_to_modules : bool
+        Apply transformation to all :any:`Module` objects when processing
+        a :any:`Sourcefile` (default ``False``)
+    recurse_to_procedures : bool
+        Apply transformation to all :any:`Subroutine` objects when processing
+        :any:`Sourcefile` or :any:``Module`` objects (default ``False``)
+    recurse_to_internal_procedures : bool
+        Apply transformation to all internal :any:`Subroutine` objects
+        when processing :any:`Subroutine` objects (default ``False``)
     """
 
+    # Forces scheduler traversal in reverse order from the leaf nodes upwards
+    reverse_traversal = False
+
+    # Traverse a graph of Sourcefile options corresponding to scheduler items
+    traverse_file_graph = False
+
+    # Filter certain graph nodes to prune the graph and change connectivity
+    item_filter = SubroutineItem  # This can also be a tuple of types
+
+    # Recursion behaviour when invoking transformations via ``trafo.apply()``
+    recurse_to_modules = False  # Recurse from Sourcefile to Module
+    recurse_to_procedures = False  # Recurse from Sourcefile/Module to subroutines and functions
+    recurse_to_internal_procedures = False  # Recurse to subroutines in ``contains`` clause
+
+
     def transform_subroutine(self, routine, **kwargs):
         """
         Defines the transformation to apply to :any:`Subroutine` items.
@@ -120,6 +166,13 @@ def apply_file(self, sourcefile, **kwargs):
 
         This calls :meth:`transform_file`.
 
+        If the :attr:`recurse_to_modules` class property is set, it
+        will also invoke :meth:`apply` on all :any:`Module` objects in
+        this :any:`Sourcefile`. Likewise, if
+        :attr:`recurse_to_procedures` is set, it will invoke
+        :meth:`apply` on all free :any:`Subroutine` objects in this
+        :any:`Sourcefile`.
+
         Parameters
         ----------
         sourcefile : :any:`Sourcefile`
@@ -133,15 +186,55 @@ def apply_file(self, sourcefile, **kwargs):
         if sourcefile._incomplete:
             raise RuntimeError('Transformation.apply_file requires Sourcefile to be complete')
 
+        item = kwargs.pop('item', None)
+        items = kwargs.pop('items', None)
+        role = kwargs.pop('role', None)
+        targets = kwargs.pop('targets', None)
+
+        if items:
+            # TODO: This special logic is required for the
+            # DependencyTransformation to capture certain corner
+            # cases. Once the module wrapping is split into its
+            # own transformation, we can probably simplify this.
+
+            # We consider the sourcefile to be a "kernel" file if all items are kernels
+            role = 'kernel' if all(item.role == 'kernel' for item in items) else 'driver'
+
+            if targets is None:
+                # We collect the targets for file/module-level imports from all items
+                targets = [target for item in items for target in item.targets]
+
         # Apply file-level transformations
-        self.transform_file(sourcefile, **kwargs)
+        self.transform_file(sourcefile, item=item, role=role, targets=targets, items=items, **kwargs)
+
+        # Recurse to modules, if configured
+        if self.recurse_to_modules:
+            for module in sourcefile.modules:
+                self.transform_module(module, item=item, role=role, targets=targets, items=items, **kwargs)
+
+        # Recurse into procedures, if configured
+        if self.recurse_to_procedures:
+            if items:
+                # Recursion into all subroutine items in the current file
+                for item in items:
+                    self.transform_subroutine(
+                        item.routine, item=item, role=item.role, targets=item.targets, **kwargs
+                    )
+            else:
+                for routine in sourcefile.all_subroutines:
+                    self.transform_subroutine(routine, item=item, role=role, targets=targets, **kwargs)
 
     def apply_subroutine(self, subroutine, **kwargs):
         """
         Apply transformation to a given :any:`Subroutine` object and its members.
 
         This calls :meth:`transform_subroutine`.
 
+        If the :attr:`recurse_to_member_procedures` class property is
+        set, it will also invoke :meth:`apply` on all
+        :any:`Subroutine` objects in the ``contains`` clause of this
+        :any:`Subroutine`.
+
         Parameters
         ----------
         subroutine : :any:`Subroutine`
@@ -158,12 +251,21 @@ def apply_subroutine(self, subroutine, **kwargs):
         # Apply the actual transformation for subroutines
         self.transform_subroutine(subroutine, **kwargs)
 
+        # Recurse to internal procedures
+        if self.recurse_to_internal_procedures:
+            for routine in subroutine.subroutines:
+                self.apply_subroutine(routine, **kwargs)
+
     def apply_module(self, module, **kwargs):
         """
         Apply transformation to a given :any:`Module` object and its members.
 
         This calls :meth:`transform_module`.
 
+        If the :attr:`recurse_to_procedures` class property is set,
+        it will also invoke :meth:`apply` on all :any:`Subroutine`
+        objects in the ``contains`` clause of this :any:`Module`.
+
         Parameters
         ----------
         module : :any:`Module`
@@ -180,6 +282,11 @@ def apply_module(self, module, **kwargs):
         # Apply the actual transformation for modules
         self.transform_module(module, **kwargs)
 
+        # Recurse to procedures contained in this module
+        if self.recurse_to_procedures:
+            for routine in module.subroutines:
+                self.apply_subroutine(routine, **kwargs)
+
     def post_apply(self, source, rescope_symbols=False):
         """
         Dispatch method for actions to be carried out after applying a transformation