Utility transformation for creating standalone subroutines from contained subroutines

AUTHORS.md

@@ -1,6 +1,7 @@
 # Authors and Contributors
 
 - R. Heilemann Myhre (Met Norway)
+- S. Karppinen (FMI)
 - P. Kiepas (École polytechnique/IPSL)
 - M. Lange (ECMWF)
 - J. Legaux (CERFACS)

loki/transform/__init__.py

@@ -19,4 +19,5 @@
 from loki.transform.build_system_transform import * # noqa
 from loki.transform.transform_hoist_variables import * # noqa
 from loki.transform.transform_parametrise import * # noqa
+from loki.transform.transform_extract_contained_procedures import * # noqa
 from loki.transform.transform_sequence_association import * # noqa

loki/transform/transform_extract_contained_procedures.py

@@ -0,0 +1,198 @@
+# (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.subroutine import Subroutine
+from loki.expression import (
+    FindVariables, FindInlineCalls, SubstituteExpressions,
+    DeferredTypeSymbol, Array
+)
+from loki.ir import (
+    CallStatement, DerivedType
+)
+from loki.visitors import (
+    Transformer, FindNodes,
+)
+__all__ = ['extract_contained_procedures', 'extract_contained_procedure']
+
+def extract_contained_procedures(procedure):
+    """
+    This transform creates "standalone" :any:`Subroutine`s
+    from the contained procedures (subroutines or functions) of ``procedure``.
+
+    A list of :any:`Subroutine`s corresponding to each contained 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.
+    3. All procedures are removed from the CONTAINS block of ``procedure``.
+
+    As a basic example of this transformation, the Fortran subroutine:
+    .. code-block::
+        subroutine outer()
+            integer :: y
+            integer :: o
+            o = 0
+            y = 1
+            call inner(o)
+            contains
+            subroutine inner(o)
+               integer, intent(inout) :: o
+               integer :: x
+               x = 4
+               o = x + y ! Note, 'y' is "global" here!
+            end subroutine inner
+        end subroutine outer
+    is modified to:
+    .. code-block::
+        subroutine outer()
+            integer :: y
+            integer :: o
+            o = 0
+            y = 1
+            call inner(o, y) ! 'y' now passed as argument.
+            contains
+        end subroutine outer
+    and the (modified) child:
+    .. code-block::
+        subroutine inner(o, y)
+               integer, intent(inout) :: o
+               integer, intent(inout) :: y
+               integer :: x
+               x = 4
+               o = x + y ! Note, 'y' is no longer "global"
+        end subroutine inner
+    is returned.
+    """
+    new_procedures = []
+    for r in procedure.subroutines:
+        new_procedures += [extract_contained_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):
+    """
+    Extract a single contained 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.
+
+    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.
+    """
+    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"in both '{inner.name}' and the parent procedure '{procedure.name}': "
+        for u in undefined:
+            msg += f"{u.name}, "
+        raise RuntimeError(msg)
+
+    ## PRODUCING VARIABLES TO INTRODUCE AS DUMMY ARGUMENTS TO `inner`.
+    # Produce a list of variables defined in the scope of `procedure` that need to be resolved in `inner`'s scope
+    # by introducing them as dummy arguments.
+    # The second line drops any derived type fields, don't want them, since want to resolve the derived type itself.
+    vars_to_resolve = [v for v in FindVariables().visit(inner.body) if v.scope is procedure]
+    vars_to_resolve = [v for v in vars_to_resolve if not v.parent]
+
+    # Save any `DeferredTypeSymbol`s for later, they are in fact defined through imports in `procedure`,
+    # and therefore not to be added as arguments to `inner`. (the next step removes them from `vars_to_resolve`)
+    var_imports_to_add = tuple(v for v in vars_to_resolve if isinstance(v, DeferredTypeSymbol))
+
+    # Lookup the definition of the variables in `vars_to_resolve` from the scope of `procedure`.
+    # This provides maximal information on them.
+    vars_to_resolve = [proc_var for v in vars_to_resolve if \
+        (proc_var := procedure.variable_map.get(v.name))]
+
+    # For each array in `vars_to_resolve`, append any non-literal shape variables to `vars_to_resolve`,
+    # if not already there.
+    arr_shapes = []
+    for var in vars_to_resolve:
+        if isinstance(var, Array):
+            # Dropping variables with parents here to handle the case that the array dimension(s)
+            # are defined through the field of a derived type.
+            arr_shapes += list(v for v in FindVariables().visit(var.shape) if not v.parent)
+    for v in arr_shapes:
+        if v.name not in vars_to_resolve:
+            vars_to_resolve.append(v)
+    vars_to_resolve = tuple(vars_to_resolve)
+
+    ## PRODUCING IMPORTS TO INTRODUCE TO `inner`.
+    # Get all variables from `inner.spec`. Need to search them for resolving kinds and derived types for
+    # variables that do not need resolution.
+    inner_spec_vars = tuple(FindVariables().visit(inner.spec))
+
+    # Produce derived types appearing in `vars_to_resolve` or in `inner.spec` that need to be resolved
+    # from imports of `procedure`.
+    dtype_imports_to_add = tuple(v.type.dtype for v in vars_to_resolve + inner_spec_vars \
+        if isinstance(v.type.dtype, DerivedType))
+
+    # 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)
+
+    # Produce all imports to add.
+    # Here the imports are also tidied to only import what is strictly necessary, and with single
+    # USE statements for each module.
+    imports_to_add = []
+    to_lookup_from_imports = dtype_imports_to_add + kind_imports_to_add + var_imports_to_add
+    for val in to_lookup_from_imports:
+        imp = procedure.import_map[val.name]
+        matching_import = tuple(i for i in imports_to_add if i.module == imp.module)
+        if matching_import:
+            # Have already encountered module name, modify existing.
+            matching_import = matching_import[0]
+            imports_to_add.remove(matching_import)
+            newimport = matching_import.clone(symbols=tuple(set(matching_import.symbols + imp.symbols)))
+        else:
+            # Have not encountered the module yet, add new one.
+            newsyms = tuple(s for s in imp.symbols if s.name == val.name)
+            newimport = imp.clone(symbols=newsyms)
+        imports_to_add.append(newimport)
+
+    ## MAKING THE CHANGES TO `inner`
+    # Change `inner` to take `vars_to_resolve` as dummy arguments and add all necessary imports.
+    # Here also rescoping all variables to the scope of `inner` and specifying intent as "inout",
+    # if not set in `procedure` scope.
+    # Note: After these lines, `inner` should be self-contained or there is a bug.
+    inner.arguments += tuple(
+        v.clone(type=v.type.clone(intent=v.type.intent or 'inout'), scope=inner)
+        for v in vars_to_resolve
+    )
+    inner.spec.prepend(imports_to_add)
+
+    ## TRANSFORMING CALLS TO `inner` in `procedure`.
+    # The resolved variables are all added as keyword arguments to each call.
+    # (to avoid further modification of the call if it already happens to contain kwargs).
+    # Here any dimensions in the variables are dropped, since they should not appear in the call.
+    # Note that functions need different visitors and mappers than subroutines.
+    call_map = {}
+    if inner.is_function:
+        for call in FindInlineCalls().visit(procedure.body):
+            if call.routine is inner:
+                newkwargs = tuple((v.name, v.clone(dimensions=None, scope=procedure)) for v in vars_to_resolve)
+                call_map[call] = call.clone(kw_parameters=call.kwarguments + newkwargs)
+        procedure.body = SubstituteExpressions(call_map).visit(procedure.body)
+    else:
+        for call in FindNodes(CallStatement).visit(procedure.body):
+            if call.routine is inner:
+                newkwargs = tuple((v.name, v.clone(dimensions=None, scope=procedure)) for v in vars_to_resolve)
+                call_map[call] = call.clone(kwarguments=tuple(call.kwarguments) + newkwargs)
+        procedure.body = Transformer(call_map).visit(procedure.body)
+
+    return inner