Merge pull request #130 from ecmwf-ifs/naml-member-inlining

Member inlining utility transform

AUTHORS.md

@@ -2,6 +2,7 @@
 
 - R. Heilemann Myhre (Met Norway)
 - M. Lange (ECMWF)
+- J. Legaux (CERFACS)
 - O. Marsden (ECMWF)
 - A. Nawab (ECMWF)
 - B. Reuter (ECMWF)

cmake/loki_transform.cmake

@@ -182,7 +182,7 @@ endmacro()
 #       [OMNI_INCLUDE <omni-inc1> [<omni-inc2> ...]]
 #       [XMOD <xmod-dir1> [<xmod-dir2> ...]]
 #       [REMOVE_OPENMP] [DATA_OFFLOAD] [GLOBAL_VAR_OFFLOAD]
-#       [TRIM_VECTOR_SECTIONS] [REMOVE_DERIVED_ARGS]
+#       [TRIM_VECTOR_SECTIONS] [REMOVE_DERIVED_ARGS] [INLINE_MEMBERS]
 #   )
 #
 # Call ``loki-transform.py convert ...`` with the provided arguments.
@@ -199,7 +199,10 @@ endmacro()
 
 function( loki_transform_convert )
 
-    set( options CPP DATA_OFFLOAD REMOVE_OPENMP ASSUME_DEVICEPTR GLOBAL_VAR_OFFLOAD TRIM_VECTOR_SECTIONS REMOVE_DERIVED_ARGS )
+    set(
+        options CPP DATA_OFFLOAD REMOVE_OPENMP ASSUME_DEVICEPTR GLOBAL_VAR_OFFLOAD
+	TRIM_VECTOR_SECTIONS REMOVE_DERIVED_ARGS INLINE_MEMBERS
+    )
     set( oneValueArgs MODE DIRECTIVE FRONTEND CONFIG PATH OUTPATH )
     set( multiValueArgs OUTPUT DEPENDS INCLUDES INCLUDE HEADERS HEADER DEFINITIONS DEFINE OMNI_INCLUDE XMOD )
 
@@ -252,6 +255,10 @@ function( loki_transform_convert )
         list( APPEND _ARGS --remove-derived-args )
     endif()
 
+    if( ${_PAR_INLINE_MEMBERS} )
+        list( APPEND _ARGS --inline-members )
+    endif()
+
     _loki_transform_env_setup()
 
     add_custom_command(
@@ -588,6 +595,7 @@ endfunction()
 #       [DIRECTIVE <directive>]
 #       [CPP]
 #       [FRONTEND <frontend>]
+#       [INLINE_MEMBERS]
 #       [BUILDDIR <build-path>]
 #       [SOURCES <source1> [<source2> ...]]
 #       [HEADERS <header1> [<header2> ...]]
@@ -607,7 +615,7 @@ endfunction()
 
 function( loki_transform_command )
 
-    set( options CPP )
+    set( options CPP INLINE_MEMBERS )
     set( oneValueArgs COMMAND MODE DIRECTIVE FRONTEND CONFIG BUILDDIR )
     set( multiValueArgs OUTPUT DEPENDS SOURCES HEADERS )
 
@@ -731,7 +739,7 @@ endfunction()
 
 function( loki_transform_target )
 
-    set( options NO_PLAN_SOURCEDIR COPY_UNMODIFIED CPP CPP_PLAN )
+    set( options NO_PLAN_SOURCEDIR COPY_UNMODIFIED CPP CPP_PLAN INLINE_MEMBERS )
     set( single_value_args TARGET COMMAND MODE DIRECTIVE FRONTEND CONFIG PLAN )
     set( multi_value_args SOURCES HEADERS )
 
@@ -794,6 +802,10 @@ function( loki_transform_target )
             list( APPEND _TRANSFORM_OPTIONS CPP )
         endif()
 
+        if( _PAR_INLINE_MEMBERS )
+            list( APPEND _TRANSFORM_OPTIONS INLINE_MEMBERS )
+        endif()
+
         loki_transform_command(
             COMMAND   ${_PAR_COMMAND}
             OUTPUT    ${LOKI_SOURCES_TO_APPEND}

loki/expression/symbols.py

@@ -1069,7 +1069,7 @@ class LogicLiteral(StrCompareMixin, _Literal):
     """
 
     def __init__(self, value, **kwargs):
-        self.value = value.lower() in ('true', '.true.')
+        self.value = str(value).lower() in ('true', '.true.')
         super().__init__(**kwargs)
 
     init_arg_names = ('value', )

loki/ir.py

@@ -921,6 +921,14 @@ def arg_iter(self):
         kwargs = ((r_args[kw], arg) for kw, arg in as_tuple(self.kwarguments))
         return chain(args, kwargs)
 
+    @property
+    def arg_map(self):
+        """
+        A full map of all qualified argument matches from arguments
+        and keyword arguments.
+        """
+        return dict(self.arg_iter())
+
 
 @dataclass_strict(frozen=True)
 class _AllocationBase():

loki/transform/transform_inline.py

@@ -14,13 +14,18 @@
     FindVariables, FindInlineCalls, FindLiterals,
     SubstituteExpressions, LokiIdentityMapper
 )
-from loki.ir import Import, Comment, Assignment
+from loki.ir import Import, Comment, Assignment, VariableDeclaration, CallStatement
 from loki.expression import symbols as sym
 from loki.types import BasicType
 from loki.visitors import Transformer, FindNodes
+from loki.tools import as_tuple
+from loki.logging import warning, error
 
 
-__all__ = ['inline_constant_parameters', 'inline_elemental_functions']
+__all__ = [
+    'inline_constant_parameters', 'inline_elemental_functions',
+    'inline_member_procedures'
+]
 
 
 class InlineSubstitutionMapper(LokiIdentityMapper):
@@ -183,3 +188,132 @@ def inline_elemental_functions(routine):
         if all(hasattr(s, 'type') and s.type.dtype in removed_functions for s in im.symbols):
             import_map[im] = None
     routine.spec = Transformer(import_map).visit(routine.spec)
+
+
+def inline_member_routine(routine, member):
+    """
+    Inline an individual member :any:`Subroutine` at source level.
+
+    This will replace all :any:`Call` objects to the specified
+    subroutine with an adjusted equivalent of the member routines'
+    body. For this, argument matching, including partial dimension
+    matching for array references is performed, and all
+    member-specific declarations are hoisted to the containing
+    :any:`Subroutine`.
+
+    Parameters
+    ----------
+    routine : :any:`Subroutine`
+        The subroutine in which to inline all calls to the member routine
+    member : :any:`Subroutine`
+        The contained member subroutine to be inlined in the parent
+    """
+    # pylint: disable=import-outside-toplevel,cyclic-import
+    from loki.transform import recursive_expression_map_update
+
+    def _map_unbound_dims(var, val):
+        """
+        Maps all unbound dimension ranges in the passed array value
+        ``val`` with the indices from the local variable ``var``. It
+        returns the re-mapped symbol.
+
+        For example, mapping the passed array ``m(:,j)`` to the local
+        expression ``a(i)`` yields ``m(i,j)``.
+        """
+        new_dimensions = list(val.dimensions)
+
+        indices = [index for index, dim in enumerate(val.dimensions) if isinstance(dim, sym.Range)]
+
+        for index, dim in enumerate(var.dimensions):
+            new_dimensions[indices[index]] = dim
+
+        return val.clone(dimensions=tuple(new_dimensions))
+
+    # Prevent shadowing of member variables by renaming them a priori
+    parent_variables = routine.variable_map
+    duplicate_locals = tuple(
+        v for v in member.variables
+        if v.name in parent_variables and v.name.lower() not in member._dummies
+    )
+    shadow_mapper = SubstituteExpressions(
+        {v: v.clone(name=f'{member.name}_{v.name}') for v in duplicate_locals}
+    )
+    member.spec = shadow_mapper.visit(member.spec)
+    member.body = shadow_mapper.visit(member.body)
+
+    # Get local variable declarations and hoist them
+    decls = FindNodes(VariableDeclaration).visit(member.spec)
+    decls = tuple(d for d in decls if all(s.name.lower() not in routine._dummies for s in d.symbols))
+    decls = tuple(d for d in decls if all(s not in routine.variables for s in d.symbols))
+    routine.spec.append(decls)
+
+    call_map = {}
+    for call in FindNodes(CallStatement).visit(routine.body):
+        if call.routine == member:
+            argmap = {}
+            member_vars = FindVariables().visit(member.body)
+
+            # Match dimension indexes between the argument and the given value
+            # for all occurences of the argument in the body
+            for arg, val in call.arg_map.items():
+                if isinstance(arg, sym.Array):
+                    # Resolve implicit dimension ranges of the passed value,
+                    # eg. when passing a two-dimensional array `a` as `call(arg=a)`
+                    # Check if val is a DeferredTypeSymbol, as it does not have a `dimensions` attribute
+                    if not isinstance(val, sym.DeferredTypeSymbol) and val.dimensions:
+                        qualified_value = val
+                    else:
+                        qualified_value = val.clone(
+                            dimensions=tuple(sym.Range((None, None)) for _ in arg.shape)
+                        )
+
+                    # If sequence association (scalar-to-array argument passing) is used,
+                    # we cannot determine the right re-mapped iteration space, so we bail here!
+                    if not any(isinstance(d, sym.Range) for d in qualified_value.dimensions):
+                        error(
+                            '[Loki::TransformInline] Cannot find free dimension resolving '
+                            f' array argument for value "{qualified_value}"'
+                        )
+                        raise RuntimeError('[Loki::TransformInline] Unable to resolve member subroutine call')
+                    arg_vars = tuple(v for v in member_vars if v.name == arg.name)
+                    argmap.update((v, _map_unbound_dims(v, qualified_value)) for v in arg_vars)
+                else:
+                    argmap[arg] = val
+
+            # Recursive update of the map in case of nested variables to map
+            argmap = recursive_expression_map_update(argmap, max_iterations=10)
+
+            # Substitute argument calls into a copy of the body
+            member_body = SubstituteExpressions(argmap).visit(member.body.body)
+
+            # Inline substituted body within a pair of marker comments
+            comment = Comment(f'! [Loki] inlined member subroutine: {member.name}')
+            c_line = Comment('! =========================================')
+            call_map[call] = (comment, c_line) + as_tuple(member_body) + (c_line, )
+
+    # Replace calls to member with the member's body
+    routine.body = Transformer(call_map).visit(routine.body)
+    # Can't use transformer to replace subroutine, so strip it manually
+    contains_body = tuple(n for n in routine.contains.body if not n == member)
+    routine.contains._update(body=contains_body)
+
+
+def inline_member_procedures(routine):
+    """
+    Inline all member subroutines contained in an individual :any:`Subroutine`.
+
+    Please note that member functions are not yet supported!
+
+    Parameters
+    ----------
+    routine : :any:`Subroutine`
+        The subroutine in which to inline all member routines
+    """
+
+    # Run through all members and invoke individual inlining transforms
+    for member in routine.members:
+        if member.is_function:
+            # TODO: Implement for functions!!!
+            warning('[Loki::inline] Inlining member functions is not yet supported, only subroutines!')
+        else:
+            inline_member_routine(routine, member)

scripts/loki_transform.py

@@ -102,10 +102,12 @@ def cli(debug):
               help="Generate offload instructions for global vars imported via 'USE' statements.")
 @click.option('--remove-derived-args/--no-remove-derived-args', default=False,
               help="Remove derived-type arguments and replace with canonical arguments")
+@click.option('--inline-members/--no-inline-members', default=False,
+              help='Inline member functions for SCC-class transformations.')
 def convert(
         mode, config, build, source, header, cpp, directive, include, define, omni_include, xmod,
         data_offload, remove_openmp, assume_deviceptr, frontend, trim_vector_sections,
-        global_var_offload, remove_derived_args
+        global_var_offload, remove_derived_args, inline_members
 ):
     """
     Batch-processing mode for Fortran-to-Fortran transformations that
@@ -190,7 +192,9 @@ def convert(
         horizontal = scheduler.config.dimensions['horizontal']
         vertical = scheduler.config.dimensions['vertical']
         block_dim = scheduler.config.dimensions['block_dim']
-        transformation = (SCCBaseTransformation(horizontal=horizontal, directive=directive),)
+        transformation = (SCCBaseTransformation(
+            horizontal=horizontal, directive=directive, inline_members=inline_members
+        ),)
         transformation += (SCCDevectorTransformation(horizontal=horizontal, trim_vector_sections=trim_vector_sections),)
         transformation += (SCCDemoteTransformation(horizontal=horizontal),)
         if not 'hoist' in mode:

tests/test_expression.py

@@ -988,6 +988,9 @@ def test_string_compare():
     assert symbols.Literal('u') == 'u'
     assert symbols.Literal('u') != 'U'
     assert symbols.Literal('u') != u  # The `Variable(name='u', ...) from above
+    assert symbols.Literal('.TrUe.') == 'true'
+    # Specific test for constructor checks
+    assert symbols.LogicLiteral(value=True) == 'true'
 
 
 @pytest.mark.skipif(not HAVE_FP, reason='Fparser not available')