Merge pull request #376 from ecmwf-ifs/naml-inline-extract-reorg

Transformations: Re-organise `inline` and `extract` sub-packages

loki/transformations/extract/__init__.py

@@ -0,0 +1,81 @@
+# (C) Copyright 2018- ECMWF.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+"""
+Transformations sub-package that provides various forms of
+source-code extraction into standalone :any:`Subroutine` objects.
+
+The various extractions mechanisms are provided as standalone utility
+methods, or via the :any:`ExtractTransformation` class for for batch
+processing.
+
+These utilities represent the conceptual inverse operation to
+"inlining", as done by the :any:`InlineTransformation`.
+"""
+
+from loki.transformations.extract.internal import * # noqa
+from loki.transformations.extract.outline import * # noqa
+
+from loki.batch import Transformation
+
+
+__all__ = ['ExtractTransformation']
+
+
+class ExtractTransformation(Transformation):
+    """
+    :any:`Transformation` class to apply several types of source
+    extraction when batch-processing large source trees via the
+    :any:`Scheduler`.
+
+    Parameters
+    ----------
+    extract_internals : bool
+        Extract internal procedure (see :any:`extract_internal_procedures`);
+        default: False.
+    outline_regions : bool
+        Outline pragma-annotated code regions to :any:`Subroutine` objects.
+        (see :any:`outline_pragma_regions`); default: True.
+    """
+    def __init__(self, extract_internals=False, outline_regions=True):
+        self.extract_internals = extract_internals
+        self.outline_regions = outline_regions
+
+    def transform_module(self, module, **kwargs):
+        """
+        Extract internals procedures and marked subroutines and add
+        them to the given :any:`Module`.
+        """
+
+        # Extract internal (contained) procedures into standalone ones
+        if self.extract_internals:
+            for routine in module.subroutines:
+                new_routines = extract_internal_procedures(routine)
+                module.contains.append(new_routines)
+
+        # Extract pragma-marked code regions into standalone subroutines
+        if self.outline_regions:
+            for routine in module.subroutines:
+                new_routines = outline_pragma_regions(routine)
+                module.contains.append(new_routines)
+
+    def transform_file(self, sourcefile, **kwargs):
+        """
+        Extract internals procedures and marked subroutines and add
+        them to the given :any:`Sourcefile`.
+        """
+
+        # Extract internal (contained) procedures into standalone ones
+        if self.extract_internals:
+            for routine in sourcefile.subroutines:
+                new_routines = extract_internal_procedures(routine)
+                sourcefile.ir.append(new_routines)
+
+        # Extract pragma-marked code regions into standalone subroutines
+        if self.outline_regions:
+            for routine in sourcefile.subroutines:
+                new_routines = outline_pragma_regions(routine)
+                sourcefile.ir.append(new_routines)

loki/transformations/extract.py → loki/transformations/extract/internal.py

@@ -14,21 +14,24 @@
 from loki.types import DerivedType
 
 
-__all__ = ['extract_contained_procedures', 'extract_contained_procedure']
+__all__ = [
+    'extract_contained_procedures', 'extract_contained_procedure',
+    'extract_internal_procedures', 'extract_internal_procedure'
+]
 
 
-def extract_contained_procedures(procedure):
+def extract_internal_procedures(procedure):
     """
     This transform creates "standalone" :any:`Subroutine`s
-    from the contained procedures (subroutines or functions) of ``procedure``.
+    from the internal procedures (subroutines or functions) of ``procedure``.
 
-    A list of :any:`Subroutine`s corresponding to each contained subroutine of
+    A list of :any:`Subroutine`s corresponding to each internal subroutine of
     ``procedure`` is returned and ``procedure`` itself is
     modified (see below).
     This function does the following transforms:
-    1. all global bindings from the point of view of the contained procedures(s) are introduced
-    as imports or dummy arguments to the modified contained procedures(s) to make them standalone.
-    2. all calls or invocations of the contained procedures in parent are modified accordingly.
+    1. all global bindings from the point of view of the internal procedures(s) are introduced
+    as imports or dummy arguments to the modified internal procedures(s) to make them standalone.
+    2. all calls or invocations of the internal procedures in parent are modified accordingly.
     3. All procedures are removed from the CONTAINS block of ``procedure``.
 
     As a basic example of this transformation, the Fortran subroutine:
@@ -79,32 +82,33 @@ def extract_contained_procedures(procedure):
     """
     new_procedures = []
     for r in procedure.subroutines:
-        new_procedures += [extract_contained_procedure(procedure, r.name)]
+        new_procedures += [extract_internal_procedure(procedure, r.name)]
 
     # Remove all subroutines (or functions) from the CONTAINS section.
     newbody = tuple(r for r in procedure.contains.body if not isinstance(r, Subroutine))
     procedure.contains = procedure.contains.clone(body=newbody)
     return new_procedures
 
-def extract_contained_procedure(procedure, name):
+
+def extract_internal_procedure(procedure, name):
     """
-    Extract a single contained procedure with name ``name`` from the parent procedure ``procedure``.
+    Extract a single internal procedure with name ``name`` from the parent procedure ``procedure``.
 
     This function does the following transforms:
-    1. all global bindings from the point of view of the contained procedure are introduced
-    as imports or dummy arguments to the modified contained procedure returned from this function.
-    2. all calls or invocations of the contained procedure in the parent are modified accordingly.
+    1. all global bindings from the point of view of the internal procedure are introduced
+    as imports or dummy arguments to the modified internal procedure returned from this function.
+    2. all calls or invocations of the internal procedure in the parent are modified accordingly.
 
-    See also the "driver" function ``extract_contained_procedures``, which applies this function to each
-    contained procedure of a parent procedure and additionally empties the CONTAINS section of subroutines.
+    See also the "driver" function ``extract_internal_procedures``, which applies this function to each
+    internal procedure of a parent procedure and additionally empties the CONTAINS section of subroutines.
     """
     inner = procedure.subroutine_map[name] # Fetch the subprocedure to extract (or crash with 'KeyError').
 
     # Check if there are variables that don't have a scope. This means that they are not defined anywhere
     # and execution cannot continue.
     undefined = tuple(v for v in FindVariables().visit(inner.body) if not v.scope)
     if undefined:
-        msg = f"The following variables appearing in the contained procedure '{inner.name}' are undefined "
+        msg = f"The following variables appearing in the internal procedure '{inner.name}' are undefined "
         msg += f"in both '{inner.name}' and the parent procedure '{procedure.name}': "
         for u in undefined:
             msg += f"{u.name}, "
@@ -152,7 +156,7 @@ def extract_contained_procedure(procedure, name):
     # Produce kinds appearing in `vars_to_resolve` or in `inner.spec` that need to be resolved
     # from imports of `procedure`.
     kind_imports_to_add = tuple(v.type.kind for v in vars_to_resolve + inner_spec_vars \
-        if v.type.kind and v.type.kind.scope is procedure)
+        if v.type.kind and hasattr(v.type.kind, 'scope') and v.type.kind.scope is procedure)
 
     # Produce all imports to add.
     # Here the imports are also tidied to only import what is strictly necessary, and with single
@@ -204,3 +208,8 @@ def extract_contained_procedure(procedure, name):
         procedure.body = Transformer(call_map).visit(procedure.body)
 
     return inner
+
+
+# Aliases to the original names
+extract_contained_procedures = extract_internal_procedures
+extract_contained_procedure = extract_internal_procedure

loki/transformations/extract/outline.py

@@ -0,0 +1,136 @@
+# (C) Copyright 2018- ECMWF.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+from loki.analyse import dataflow_analysis_attached
+from loki.expression import Variable
+from loki.ir import (
+    CallStatement, Import, PragmaRegion, Section, FindNodes,
+    FindVariables, MaskedTransformer, Transformer, is_loki_pragma,
+    get_pragma_parameters, pragma_regions_attached
+)
+from loki.logging import info
+from loki.subroutine import Subroutine
+from loki.tools import as_tuple, CaseInsensitiveDict
+
+
+__all__ = ['outline_pragma_regions']
+
+
+def outline_pragma_regions(routine):
+    """
+    Convert regions annotated with ``!$loki outline`` pragmas to subroutine calls.
+
+    The pragma syntax for regions to convert to subroutines is
+    ``!$loki outline [name(...)] [in(...)] [out(...)] [inout(...)]``
+    and ``!$loki end outline``.
+
+    A new subroutine is created with the provided name (or an auto-generated default name
+    derived from the current subroutine name) and the content of the pragma region as body.
+
+    Variables provided with the ``in``, ``out`` and ``inout`` options are used as
+    arguments in the routine with the corresponding intent, all other variables used in this
+    region are assumed to be local variables.
+
+    The pragma region in the original routine is replaced by a call to the new subroutine.
+
+    Parameters
+    ----------
+    routine : :any:`Subroutine`
+        The routine from which to extract marked pragma regions.
+
+    Returns
+    -------
+    list of :any:`Subroutine`
+        the list of newly created subroutines.
+
+    """
+    counter = 0
+    routines, starts, stops = [], [], []
+    imports = {var for imprt in FindNodes(Import).visit(routine.spec) for var in imprt.symbols}
+    mask_map = {}
+    with pragma_regions_attached(routine):
+        with dataflow_analysis_attached(routine):
+            for region in FindNodes(PragmaRegion).visit(routine.body):
+                if not is_loki_pragma(region.pragma, starts_with='outline'):
+                    continue
+
+                # Name the external routine
+                parameters = get_pragma_parameters(region.pragma, starts_with='outline')
+                name = parameters.get('name', f'{routine.name}_outlined_{counter}')
+                counter += 1
+
+                # Create the external subroutine containing the routine's imports and the region's body
+                spec = Section(body=Transformer().visit(FindNodes(Import).visit(routine.spec)))
+                body = Section(body=Transformer().visit(region.body))
+                region_routine = Subroutine(name, spec=spec, body=body)
+
+                # Use dataflow analysis to find in, out and inout variables to that region
+                # (ignoring any symbols that are external imports)
+                region_in_args = region.uses_symbols - region.defines_symbols - imports
+                region_inout_args = region.uses_symbols & region.defines_symbols - imports
+                region_out_args = region.defines_symbols - region.uses_symbols - imports
+
+                # Remove any parameters from in args
+                region_in_args = {arg for arg in region_in_args if not arg.type.parameter}
+
+                # Extract arguments given in pragma annotations
+                region_var_map = CaseInsensitiveDict(
+                    (v.name, v.clone(dimensions=None))
+                    for v in FindVariables().visit(region.body) if v.clone(dimensions=None) not in imports
+                )
+                pragma_in_args = {region_var_map[v.lower()] for v in parameters.get('in', '').split(',') if v}
+                pragma_inout_args = {region_var_map[v.lower()] for v in parameters.get('inout', '').split(',') if v}
+                pragma_out_args = {region_var_map[v.lower()] for v in parameters.get('out', '').split(',') if v}
+
+                # Override arguments according to pragma annotations
+                region_in_args = (region_in_args - (pragma_inout_args | pragma_out_args)) | pragma_in_args
+                region_inout_args = (region_inout_args - (pragma_in_args | pragma_out_args)) | pragma_inout_args
+                region_out_args = (region_out_args - (pragma_in_args | pragma_inout_args)) | pragma_out_args
+
+                # Now fix the order
+                region_inout_args = as_tuple(region_inout_args)
+                region_in_args = as_tuple(region_in_args)
+                region_out_args = as_tuple(region_out_args)
+
+                # Set the list of variables used in region routine (to create declarations)
+                # and put all in the new scope
+                region_routine_variables = {v.clone(dimensions=v.type.shape or None)
+                                            for v in FindVariables().visit(region_routine.body)
+                                            if v.name in region_var_map}
+                region_routine.variables = as_tuple(region_routine_variables)
+                region_routine.rescope_symbols()
+
+                # Build the call signature
+                region_routine_var_map = region_routine.variable_map
+                region_routine_arguments = []
+                for intent, args in zip(('in', 'inout', 'out'), (region_in_args, region_inout_args, region_out_args)):
+                    for arg in args:
+                        local_var = region_routine_var_map[arg.name]
+                        local_var = local_var.clone(type=local_var.type.clone(intent=intent))
+                        region_routine_var_map[arg.name] = local_var
+                        region_routine_arguments += [local_var]
+
+                # We need to update the list of variables again to avoid duplicate declarations
+                region_routine.variables = as_tuple(region_routine_var_map.values())
+                region_routine.arguments = as_tuple(region_routine_arguments)
+
+                # insert into list of new routines
+                routines.append(region_routine)
+
+                # Register start and end nodes in transformer mask for original routine
+                starts += [region.pragma_post]
+                stops += [region.pragma]
+
+                # Replace end pragma by call in original routine
+                call_arguments = region_in_args + region_inout_args + region_out_args
+                call = CallStatement(name=Variable(name=name), arguments=call_arguments)
+                mask_map[region.pragma_post] = call
+
+    routine.body = MaskedTransformer(active=True, start=starts, stop=stops, mapper=mask_map).visit(routine.body)
+    info('%s: converted %d region(s) to calls', routine.name, counter)
+
+    return routines