From df2da3218e0b3e47428d5444dadb653d95cb03dd Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Wed, 11 Sep 2024 12:36:22 +0000 Subject: [PATCH 01/12] Transformations: Separate `transformations.inline` into sub-package --- loki/transformations/inline.py | 958 ---------------------- loki/transformations/inline/__init__.py | 130 +++ loki/transformations/inline/constants.py | 88 ++ loki/transformations/inline/functions.py | 342 ++++++++ loki/transformations/inline/mapper.py | 84 ++ loki/transformations/inline/procedures.py | 395 +++++++++ 6 files changed, 1039 insertions(+), 958 deletions(-) delete mode 100644 loki/transformations/inline.py create mode 100644 loki/transformations/inline/__init__.py create mode 100644 loki/transformations/inline/constants.py create mode 100644 loki/transformations/inline/functions.py create mode 100644 loki/transformations/inline/mapper.py create mode 100644 loki/transformations/inline/procedures.py diff --git a/loki/transformations/inline.py b/loki/transformations/inline.py deleted file mode 100644 index 14827633e..000000000 --- a/loki/transformations/inline.py +++ /dev/null @@ -1,958 +0,0 @@ -# (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. - -""" -Collection of utility routines to perform code-level force-inlining. - - -""" -from collections import defaultdict, ChainMap - -from loki.batch import Transformation -from loki.ir import ( - Import, Comment, Assignment, VariableDeclaration, CallStatement, - Transformer, FindNodes, pragmas_attached, is_loki_pragma, Interface, - StatementFunction, FindVariables, FindInlineCalls, FindLiterals, - SubstituteExpressions, ExpressionFinder, Pragma -) -from loki.expression import ( - symbols as sym, LokiIdentityMapper, ExpressionRetriever -) -from loki.types import BasicType -from loki.tools import as_tuple, CaseInsensitiveDict -from loki.logging import error -from loki.subroutine import Subroutine - -from loki.transformations.remove_code import do_remove_dead_code -from loki.transformations.sanitise import transform_sequence_association_append_map -from loki.transformations.utilities import ( - single_variable_declaration, recursive_expression_map_update -) - - -__all__ = [ - 'inline_constant_parameters', 'inline_elemental_functions', - 'inline_internal_procedures', 'inline_member_procedures', - 'inline_marked_subroutines', 'InlineTransformation', - 'inline_statement_functions', 'inline_functions' -] - - -class InlineTransformation(Transformation): - """ - :any:`Transformation` class to apply several types of source inlining - when batch-processing large source trees via the :any:`Scheduler`. - - Parameters - ---------- - inline_constants : bool - Replace instances of variables with known constant values by - :any:`Literal` (see :any:`inline_constant_parameters`); default: False. - inline_elementals : bool - Replaces :any:`InlineCall` expression to elemental functions - with the called function's body (see :any:`inline_elemental_functions`); - default: True. - inline_stmt_funcs: bool - Replaces :any:`InlineCall` expression to statement functions - with the corresponding rhs of the statement function if - the statement function declaration is available; default: False. - inline_internals : bool - Inline internal procedure (see :any:`inline_internal_procedures`); - default: False. - inline_marked : bool - Inline :any:`Subroutine` objects marked by pragma annotations - (see :any:`inline_marked_subroutines`); default: True. - remove_dead_code : bool - Perform dead code elimination, where unreachable branches are - trimmed from the code (see :any:`dead_code_elimination`); default: True - allowed_aliases : tuple or list of str or :any:`Expression`, optional - List of variables that will not be renamed in the parent scope during - internal and pragma-driven inlining. - adjust_imports : bool - Adjust imports by removing the symbol of the inlined routine or adding - imports needed by the imported routine (optional, default: True) - external_only : bool, optional - Do not replace variables declared in the local scope when - inlining constants (default: True) - resolve_sequence_association: bool - Resolve sequence association for routines that contain calls to inline (default: False) - """ - - # Ensure correct recursive inlining by traversing from the leaves - reverse_traversal = True - - # This transformation will potentially change the edges in the callgraph - creates_items = False - - def __init__( - self, inline_constants=False, inline_elementals=True, - inline_stmt_funcs=False, inline_internals=False, - inline_marked=True, remove_dead_code=True, - allowed_aliases=None, adjust_imports=True, - external_only=True, resolve_sequence_association=False - ): - self.inline_constants = inline_constants - self.inline_elementals = inline_elementals - self.inline_stmt_funcs = inline_stmt_funcs - self.inline_internals = inline_internals - self.inline_marked = inline_marked - self.remove_dead_code = remove_dead_code - self.allowed_aliases = allowed_aliases - self.adjust_imports = adjust_imports - self.external_only = external_only - self.resolve_sequence_association = resolve_sequence_association - if self.inline_marked: - self.creates_items = True - - def transform_subroutine(self, routine, **kwargs): - - # Resolve sequence association in calls that are about to be inlined. - # This step runs only if all of the following hold: - # 1) it is requested by the user - # 2) inlining of "internals" or "marked" routines is activated - # 3) there is an "internal" or "marked" procedure to inline. - if self.resolve_sequence_association: - resolve_sequence_association_for_inlined_calls( - routine, self.inline_internals, self.inline_marked - ) - - # Replace constant parameter variables with explicit values - if self.inline_constants: - inline_constant_parameters(routine, external_only=self.external_only) - - # Inline elemental functions - if self.inline_elementals: - inline_elemental_functions(routine) - - # Inline Statement Functions - if self.inline_stmt_funcs: - inline_statement_functions(routine) - - # Inline internal (contained) procedures - if self.inline_internals: - inline_internal_procedures(routine, allowed_aliases=self.allowed_aliases) - - # Inline explicitly pragma-marked subroutines - if self.inline_marked: - inline_marked_subroutines( - routine, allowed_aliases=self.allowed_aliases, - adjust_imports=self.adjust_imports - ) - - # After inlining, attempt to trim unreachable code paths - if self.remove_dead_code: - do_remove_dead_code(routine) - - -class InlineSubstitutionMapper(LokiIdentityMapper): - """ - An expression mapper that defines symbolic substitution for inlining. - """ - - def map_algebraic_leaf(self, expr, *args, **kwargs): - raise NotImplementedError - - def map_scalar(self, expr, *args, **kwargs): - parent = self.rec(expr.parent, *args, **kwargs) if expr.parent is not None else None - - scope = kwargs.get('scope') or expr.scope - # We're re-scoping an imported symbol - if expr.scope != scope: - return expr.clone(scope=scope, type=expr.type.clone(), parent=parent) - return expr.clone(parent=parent) - - map_deferred_type_symbol = map_scalar - - def map_array(self, expr, *args, **kwargs): - if expr.dimensions: - dimensions = self.rec(expr.dimensions, *args, **kwargs) - else: - dimensions = None - parent = self.rec(expr.parent, *args, **kwargs) if expr.parent is not None else None - - scope = kwargs.get('scope') or expr.scope - # We're re-scoping an imported symbol - if expr.scope != scope: - return expr.clone(scope=scope, type=expr.type.clone(), parent=parent, dimensions=dimensions) - return expr.clone(parent=parent, dimensions=dimensions) - - def map_procedure_symbol(self, expr, *args, **kwargs): - parent = self.rec(expr.parent, *args, **kwargs) if expr.parent is not None else None - - scope = kwargs.get('scope') or expr.scope - # We're re-scoping an imported symbol - if expr.scope != scope: - return expr.clone(scope=scope, type=expr.type.clone(), parent=parent) - return expr.clone(parent=parent) - - def map_inline_call(self, expr, *args, **kwargs): - if expr.procedure_type is None or expr.procedure_type is BasicType.DEFERRED: - # Unkonw inline call, potentially an intrinsic - # We still need to recurse and ensure re-scoping - return super().map_inline_call(expr, *args, **kwargs) - - # if it is an inline call to a Statement Function - if isinstance(expr.routine, StatementFunction): - function = expr.routine - # Substitute all arguments through the elemental body - arg_map = dict(expr.arg_iter()) - fbody = SubstituteExpressions(arg_map).visit(function.rhs) - return fbody - - function = expr.procedure_type.procedure - v_result = [v for v in function.variables if v == function.name][0] - - # Substitute all arguments through the elemental body - arg_map = dict(expr.arg_iter()) - fbody = SubstituteExpressions(arg_map).visit(function.body) - - # Extract the RHS of the final result variable assignment - stmts = [s for s in FindNodes(Assignment).visit(fbody) if s.lhs == v_result] - assert len(stmts) == 1 - rhs = self.rec(stmts[0].rhs, *args, **kwargs) - return rhs - -def resolve_sequence_association_for_inlined_calls(routine, inline_internals, inline_marked): - """ - Resolve sequence association in calls to all member procedures (if ``inline_internals = True``) - or in calls to procedures that have been marked with an inline pragma (if ``inline_marked = True``). - If both ``inline_internals`` and ``inline_marked`` are ``False``, no processing is done. - """ - call_map = {} - with pragmas_attached(routine, node_type=CallStatement): - for call in FindNodes(CallStatement).visit(routine.body): - condition = ( - (inline_marked and is_loki_pragma(call.pragma, starts_with='inline')) or - (inline_internals and call.routine in routine.routines) - ) - if condition: - if call.routine == BasicType.DEFERRED: - # NOTE: Throwing error here instead of continuing, because the user has explicitly - # asked sequence assoc to happen with inlining, so source for routine should be - # found in calls to be inlined. - raise ValueError( - f"Cannot resolve sequence association for call to ``{call.name}`` " + - f"to be inlined in routine ``{routine.name}``, because " + - f"the ``CallStatement`` referring to ``{call.name}`` does not contain " + - "the source code of the procedure. " + - "If running in batch processing mode, please recheck Scheduler configuration." - ) - transform_sequence_association_append_map(call_map, call) - if call_map: - routine.body = Transformer(call_map).visit(routine.body) - -def inline_constant_parameters(routine, external_only=True): - """ - Replace instances of variables with known constant values by `Literals`. - - Notes - ----- - The ``.type.initial`` property is used to derive the replacement - value,a which means for symbols imported from external modules, - the parent :any:`Module` needs to be supplied in the - ``definitions`` to the constructor when creating the - :any:`Subroutine`. - - Variables that are replaced are also removed from their - corresponding import statements, with empty import statements - being removed alltogether. - - Parameters - ---------- - routine : :any:`Subroutine` - Procedure in which to inline/resolve constant parameters. - external_only : bool, optional - Do not replace variables declared in the local scope (default: True) - """ - # Find all variable instances in spec and body - variables = FindVariables().visit(routine.ir) - - # Filter out variables declared locally - if external_only: - variables = [v for v in variables if v not in routine.variables] - - def is_inline_parameter(v): - return hasattr(v, 'type') and v.type.parameter and v.type.initial - - # Create mapping for variables and imports - vmap = {v: v.type.initial for v in variables if is_inline_parameter(v)} - - # Replace kind parameters in variable types - for variable in routine.variables: - if is_inline_parameter(variable.type.kind): - routine.symbol_attrs[variable.name] = variable.type.clone(kind=variable.type.kind.type.initial) - if variable.type.initial is not None: - # Substitute kind specifier in literals in initializers (I know...) - init_map = {literal.kind: literal.kind.type.initial - for literal in FindLiterals().visit(variable.type.initial) - if is_inline_parameter(literal.kind)} - if init_map: - initial = SubstituteExpressions(init_map).visit(variable.type.initial) - routine.symbol_attrs[variable.name] = variable.type.clone(initial=initial) - - # Update imports - imprtmap = {} - substituted_names = {v.name.lower() for v in vmap} - for imprt in FindNodes(Import).visit(routine.spec): - if imprt.symbols: - symbols = tuple(s for s in imprt.symbols if s.name.lower() not in substituted_names) - if not symbols: - imprtmap[imprt] = Comment(f'! Loki: parameters from {imprt.module} inlined') - elif len(symbols) < len(imprt.symbols): - imprtmap[imprt] = imprt.clone(symbols=symbols) - - # Flush mappings through spec and body - routine.spec = Transformer(imprtmap).visit(routine.spec) - routine.spec = SubstituteExpressions(vmap).visit(routine.spec) - routine.body = SubstituteExpressions(vmap).visit(routine.body) - - # Clean up declarations that are about to become defunct - decl_map = { - decl: None for decl in routine.declarations - if all(isinstance(s, sym.IntLiteral) for s in decl.symbols) - } - routine.spec = Transformer(decl_map).visit(routine.spec) - - -def inline_elemental_functions(routine): - """ - Replaces `InlineCall` expression to elemental functions with the - called functions body. - - Parameters - ---------- - routine : :any:`Subroutine` - Procedure in which to inline functions. - """ - inline_functions(routine, inline_elementals_only=True) - - -def inline_functions(routine, inline_elementals_only=False, functions=None): - """ - Replaces `InlineCall` expression to functions with the - called functions body. Nested calls are handled/inlined through - an iterative approach calling :any:`_inline_functions`. - - Parameters - ---------- - routine : :any:`Subroutine` - Procedure in which to inline functions. - inline_elementals_only : bool, optional - Inline elemental routines/functions only (default: False). - functions : tuple, optional - Inline only functions that are provided here - (default: None, thus inline all functions). - """ - potentially_functions_to_be_inlined = True - while potentially_functions_to_be_inlined: - potentially_functions_to_be_inlined = _inline_functions(routine, inline_elementals_only=inline_elementals_only, - functions=functions) - -def _inline_functions(routine, inline_elementals_only=False, functions=None): - """ - Replaces `InlineCall` expression to functions with the - called functions body, but doesn't include nested calls! - - Parameters - ---------- - routine : :any:`Subroutine` - Procedure in which to inline functions. - inline_elementals_only : bool, optional - Inline elemental routines/functions only (default: False). - functions : tuple, optional - Inline only functions that are provided here - (default: None, thus inline all functions). - - Returns - ------- - bool - Whether inline calls are (potentially) left to be - inlined in the next call to this function. - """ - - class ExpressionRetrieverSkipInlineCallParameters(ExpressionRetriever): - """ - Expression retriever skipping parameters of inline calls. - """ - # pylint: disable=abstract-method - - def __init__(self, query, recurse_query=None, inline_elementals_only=False, - functions=None, **kwargs): - self.inline_elementals_only = inline_elementals_only - self.functions = as_tuple(functions) - super().__init__(query, recurse_query, **kwargs) - - def map_inline_call(self, expr, *args, **kwargs): - if not self.visit(expr, *args, **kwargs): - return - self.rec(expr.function, *args, **kwargs) - # SKIP parameters/args/kwargs on purpose - # under certain circumstances - if expr.procedure_type is BasicType.DEFERRED or\ - (self.inline_elementals_only and\ - not(expr.procedure_type.is_function and expr.procedure_type.is_elemental)) or\ - (self.functions and expr.routine not in self.functions): - for child in expr.parameters: - self.rec(child, *args, **kwargs) - for child in list(expr.kw_parameters.values()): - self.rec(child, *args, **kwargs) - - self.post_visit(expr, *args, **kwargs) - - class FindInlineCallsSkipInlineCallParameters(ExpressionFinder): - """ - Find inline calls but skip/ignore parameters of inline calls. - """ - retriever = ExpressionRetrieverSkipInlineCallParameters(lambda e: isinstance(e, sym.InlineCall)) - - # functions are provided, however functions is empty, thus early exit - if functions is not None and not functions: - return False - functions = as_tuple(functions) - - # Keep track of removed symbols - removed_functions = set() - - # Find and filter inline calls and corresponding nodes - function_calls = {} - # Find inline calls but skip/ignore inline calls being parameters of other inline calls - # to ensure correct ordering of inlining. Those skipped/ignored inline calls will be handled - # in the next call to this function. - retriever = ExpressionRetrieverSkipInlineCallParameters(lambda e: isinstance(e, sym.InlineCall), - inline_elementals_only=inline_elementals_only, functions=functions) - # override retriever ... - FindInlineCallsSkipInlineCallParameters.retriever = retriever - for node, calls in FindInlineCallsSkipInlineCallParameters(with_ir_node=True).visit(routine.body): - for call in calls: - if call.procedure_type is BasicType.DEFERRED or isinstance(call.routine, StatementFunction): - continue - if inline_elementals_only: - if not (call.procedure_type.is_function and call.procedure_type.is_elemental): - continue - if functions: - if call.routine not in functions: - continue - function_calls.setdefault(str(call.name).lower(),[]).append((call, node)) - - if not function_calls: - return False - - # inline functions - node_prepend_map = {} - call_map = {} - for calls_nodes in function_calls.values(): - calls, nodes = list(zip(*calls_nodes)) - for call in calls: - removed_functions.add(call.procedure_type) - # collect nodes to be appendes as well as expression replacement for inline call - inline_node_map, inline_call_map = inline_function_calls(routine, as_tuple(calls), - calls[0].routine, as_tuple(nodes)) - for node, nodes_to_prepend in inline_node_map.items(): - node_prepend_map.setdefault(node, []).extend(list(nodes_to_prepend)) - call_map.update(inline_call_map) - - # collect nodes to be prepended for each node that contains (at least one) inline call to a function - node_map = {} - for node, prepend_nodes in node_prepend_map.items(): - node_map[node] = as_tuple(prepend_nodes) + (SubstituteExpressions(call_map[node]).visit(node),) - # inline via prepending the relevant functions - routine.body = Transformer(node_map).visit(routine.body) - # We need this to ensure that symbols, as well as nested scopes - # are correctly attached to each other (eg. nested associates). - routine.rescope_symbols() - - # Remove all module imports that have become obsolete now - import_map = {} - for im in FindNodes(Import).visit(routine.spec): - if im.symbols and all(s.type.dtype in removed_functions for s in im.symbols): - import_map[im] = None - routine.spec = Transformer(import_map).visit(routine.spec) - return True - -def inline_statement_functions(routine): - """ - Replaces :any:`InlineCall` expression to statement functions with the - called statement functions rhs. - """ - # Keep track of removed symbols - removed_functions = set() - - stmt_func_decls = FindNodes(StatementFunction).visit(routine.spec) - exprmap = {} - for call in FindInlineCalls().visit(routine.body): - proc_type = call.procedure_type - if proc_type is BasicType.DEFERRED: - continue - if proc_type.is_function and isinstance(call.routine, StatementFunction): - exprmap[call] = InlineSubstitutionMapper()(call, scope=routine) - removed_functions.add(call.routine) - # Apply the map to itself to handle nested statement function calls - exprmap = recursive_expression_map_update(exprmap, max_iterations=10, mapper_cls=InlineSubstitutionMapper) - # Apply expression-level substitution to routine - routine.body = SubstituteExpressions(exprmap).visit(routine.body) - - # remove statement function declarations as well as statement function argument(s) declarations - vars_to_remove = {stmt_func.variable.name.lower() for stmt_func in stmt_func_decls} - vars_to_remove |= {arg.name.lower() for stmt_func in stmt_func_decls for arg in stmt_func.arguments} - spec_map = {stmt_func: None for stmt_func in stmt_func_decls} - for decl in routine.declarations: - if any(var in vars_to_remove for var in decl.symbols): - symbols = tuple(var for var in decl.symbols if var not in vars_to_remove) - if symbols: - decl._update(symbols=symbols) - else: - spec_map[decl] = None - routine.spec = Transformer(spec_map).visit(routine.spec) - -def map_call_to_procedure_body(call, caller, callee=None): - """ - Resolve arguments of a call and map to the called procedure body. - - Parameters - ---------- - call : :any:`CallStatment` or :any:`InlineCall` - Call object that defines the argument mapping - caller : :any:`Subroutine` - Procedure (scope) into which the callee's body gets mapped - callee : :any:`Subroutine`, optional - Procedure (scope) called. Provide if it differs from - call.routine. - """ - - 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)) - - # Get callee from the procedure type - callee = callee or call.routine - if callee is BasicType.DEFERRED: - error( - '[Loki::TransformInline] Need procedure definition to resolve ' - f'call to {call.name} from {caller}' - ) - raise RuntimeError('Procedure definition not found! ') - - argmap = {} - callee_vars = FindVariables().visit(callee.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( - f'[Loki::TransformInline] Cannot resolve procedure call to {call.name}' - ) - arg_vars = tuple(v for v in callee_vars if v.name == arg.name) - argmap.update((v, _map_unbound_dims(v, qualified_value)) for v in arg_vars) - else: - argmap[arg] = val - - # Deal with PRESENT check for optional arguments - present_checks = tuple( - check for check in FindInlineCalls().visit(callee.body) if check.function == 'PRESENT' - ) - present_map = { - check: sym.Literal('.true.') if check.arguments[0] in [arg.name for arg in call.arg_map] - else sym.Literal('.false.') - for check in present_checks - } - argmap.update(present_map) - - # 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 - callee_body = SubstituteExpressions(argmap, rebuild_scopes=True).visit( - callee.body.body, scope=caller - ) - - # Remove 'loki routine' pragmas - callee_body = Transformer( - {pragma: None for pragma in FindNodes(Pragma).visit(callee_body) - if is_loki_pragma(pragma, starts_with='routine')} - ).visit(callee_body) - - # Inline substituted body within a pair of marker comments - comment = Comment(f'! [Loki] inlined child subroutine: {callee.name}') - c_line = Comment('! =========================================') - return (comment, c_line) + as_tuple(callee_body) + (c_line, ) - - -def inline_subroutine_calls(routine, calls, callee, allowed_aliases=None): - """ - Inline a set of call to an individual :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 - calls : tuple or list of :any:`CallStatement` - callee : :any:`Subroutine` - The called target subroutine to be inlined in the parent - allowed_aliases : tuple or list of str or :any:`Expression`, optional - List of variables that will not be renamed in the parent scope, even - if they alias with a local declaration. - """ - allowed_aliases = as_tuple(allowed_aliases) - - # Ensure we process sets of calls to the same callee - assert all(call.routine == callee for call in calls) - assert isinstance(callee, Subroutine) - - # Prevent shadowing of callee's variables by renaming them a priori - parent_variables = routine.variable_map - duplicates = tuple( - v for v in callee.variables - if v.name in parent_variables and v.name.lower() not in callee._dummies - ) - # Filter out allowed aliases to prevent suffixing - duplicates = tuple(v for v in duplicates if v.symbol not in allowed_aliases) - shadow_mapper = SubstituteExpressions( - {v: v.clone(name=f'{callee.name}_{v.name}') for v in duplicates} - ) - callee.spec = shadow_mapper.visit(callee.spec) - - var_map = {} - duplicate_names = {dl.name.lower() for dl in duplicates} - for v in FindVariables(unique=False).visit(callee.body): - if v.name.lower() in duplicate_names: - var_map[v] = v.clone(name=f'{callee.name}_{v.name}') - var_map = recursive_expression_map_update(var_map) - callee.body = SubstituteExpressions(var_map).visit(callee.body) - - # Separate allowed aliases from other variables to ensure clean hoisting - if allowed_aliases: - single_variable_declaration(callee, variables=allowed_aliases) - - # Get local variable declarations and hoist them - decls = FindNodes(VariableDeclaration).visit(callee.spec) - decls = tuple(d for d in decls if all(s.name.lower() not in callee._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)) - # Rescope the declaration symbols - decls = tuple(d.clone(symbols=tuple(s.clone(scope=routine) for s in d.symbols)) for d in decls) - - # Find and apply symbol remappings for array size expressions - symbol_map = dict(ChainMap(*[call.arg_map for call in calls])) - decls = SubstituteExpressions(symbol_map).visit(decls) - - routine.spec.append(decls) - - # Resolve the call by mapping arguments into the called procedure's body - call_map = { - call: map_call_to_procedure_body(call, caller=routine) for call in calls - } - - # Replace calls to child procedure with the child's body - routine.body = Transformer(call_map).visit(routine.body) - - # We need this to ensure that symbols, as well as nested scopes - # are correctly attached to each other (eg. nested associates). - routine.rescope_symbols() - -def inline_function_calls(routine, calls, callee, nodes, allowed_aliases=None): - """ - Inline a set of call to an individual :any:`Subroutine` being functions - at source level. - - This will replace all :any:`InlineCall` 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 - calls : tuple or list of :any:`InlineCall` - Set of calls (to the same callee) to be inlined. - callee : :any:`Subroutine` - The called target function to be inlined in the parent - nodes : :any:`Node` - The corresponding nodes the functions are called from. - allowed_aliases : tuple or list of str or :any:`Expression`, optional - List of variables that will not be renamed in the parent scope, even - if they alias with a local declaration. - """ - - def rename_result_name(routine, rename): - callee = routine.clone() - var_map = {} - callee_result_var = callee.variable_map[callee.result_name.lower()] - new_callee_result_var = callee_result_var.clone(name=rename) - var_map[callee_result_var] = new_callee_result_var - callee_vars = [var for var in FindVariables().visit(callee.body) - if var.name.lower() == callee_result_var.name.lower()] - var_map.update({var: var.clone(name=rename) for var in callee_vars}) - var_map = recursive_expression_map_update(var_map) - callee.body = SubstituteExpressions(var_map).visit(callee.body) - return callee, new_callee_result_var - - allowed_aliases = as_tuple(allowed_aliases) - - # Ensure we process sets of calls to the same callee - assert all(call.routine == callee for call in calls) - assert isinstance(callee, Subroutine) - - # Prevent shadowing of callee's variables by renaming them a priori - parent_variables = routine.variable_map - duplicates = tuple( - v for v in callee.variables - if v.name.lower() != callee.result_name.lower() - and v.name in parent_variables and v.name.lower() not in callee._dummies - ) - # Filter out allowed aliases to prevent suffixing - duplicates = tuple(v for v in duplicates if v.symbol not in allowed_aliases) - shadow_mapper = SubstituteExpressions( - {v: v.clone(name=f'{callee.name}_{v.name}') for v in duplicates} - ) - callee.spec = shadow_mapper.visit(callee.spec) - - var_map = {} - duplicate_names = {dl.name.lower() for dl in duplicates} - for v in FindVariables(unique=False).visit(callee.body): - if v.name.lower() in duplicate_names: - var_map[v] = v.clone(name=f'{callee.name}_{v.name}') - - var_map = recursive_expression_map_update(var_map) - callee.body = SubstituteExpressions(var_map).visit(callee.body) - - # Separate allowed aliases from other variables to ensure clean hoisting - if allowed_aliases: - single_variable_declaration(callee, variables=allowed_aliases) - - single_variable_declaration(callee, variables=callee.result_name) - # Get local variable declarations and hoist them - decls = FindNodes(VariableDeclaration).visit(callee.spec) - decls = tuple(d for d in decls if all(s.name.lower() != callee.result_name.lower() for s in d.symbols)) - decls = tuple(d for d in decls if all(s.name.lower() not in callee._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)) - # Rescope the declaration symbols - decls = tuple(d.clone(symbols=tuple(s.clone(scope=routine) for s in d.symbols)) for d in decls) - - # Find and apply symbol remappings for array size expressions - symbol_map = dict(ChainMap(*[call.arg_map for call in calls])) - decls = SubstituteExpressions(symbol_map).visit(decls) - routine.spec.append(decls) - - # Handle result/return var/value - new_symbols = set() - result_var_map = {} - adapted_calls = [] - rename_result_var = not len(nodes) == len(set(nodes)) - for i_call, call in enumerate(calls): - callee_result_var = callee.variable_map[callee.result_name.lower()] - prefix = '' - new_callee_result_var_name = f'{prefix}result_{callee.result_name.lower()}_{i_call}'\ - if rename_result_var else f'{prefix}result_{callee.result_name.lower()}' - new_callee, new_symbol = rename_result_name(callee, new_callee_result_var_name) - adapted_calls.append(new_callee) - new_symbols.add(new_symbol) - if isinstance(callee_result_var, sym.Array): - result_var_map[(nodes[i_call], call)] = callee_result_var.clone(name=new_callee_result_var_name, - dimensions=None) - else: - result_var_map[(nodes[i_call], call)] = callee_result_var.clone(name=new_callee_result_var_name) - new_symbols = SubstituteExpressions(symbol_map).visit(as_tuple(new_symbols), recurse_to_declaration_attributes=True) - routine.variables += as_tuple([symbol.clone(scope=routine) for symbol in new_symbols]) - - # create node map to map nodes to be prepended (representing the functions) for each node - node_map = {} - call_map = {} - for i_call, call in enumerate(calls): - node_map.setdefault(nodes[i_call], []).extend( - list(map_call_to_procedure_body(call, caller=routine, callee=adapted_calls[i_call])) - ) - call_map.setdefault(nodes[i_call], {}).update({call: result_var_map[(nodes[i_call], call)]}) - return node_map, call_map - - -def inline_internal_procedures(routine, allowed_aliases=None): - """ - Inline internal subroutines contained in an individual :any:`Subroutine`. - - Please note that internal functions are not yet supported! - - Parameters - ---------- - routine : :any:`Subroutine` - The subroutine in which to inline all member routines - allowed_aliases : tuple or list of str or :any:`Expression`, optional - List of variables that will not be renamed in the parent scope, even - if they alias with a local declaration. - """ - - # Run through all members and invoke individual inlining transforms - for child in routine.members: - if child.is_function: - inline_functions(routine, functions=(child,)) - else: - calls = tuple( - call for call in FindNodes(CallStatement).visit(routine.body) - if call.routine == child - ) - inline_subroutine_calls(routine, calls, child, allowed_aliases=allowed_aliases) - - # Can't use transformer to replace subroutine/function, so strip it manually - contains_body = tuple(n for n in routine.contains.body if not n == child) - routine.contains._update(body=contains_body) - - -inline_member_procedures = inline_internal_procedures - - -def inline_marked_subroutines(routine, allowed_aliases=None, adjust_imports=True): - """ - Inline :any:`Subroutine` objects guided by pragma annotations. - - When encountering :any:`CallStatement` objects that are marked with a - ``!$loki inline`` pragma, this utility will attempt to replace the call - with the body of the called procedure and remap all passed arguments - into the calling procedures scope. - - Please note that this utility requires :any:`CallStatement` objects - to be "enriched" with external type information. - - Parameters - ---------- - routine : :any:`Subroutine` - The subroutine in which to look for pragma-marked procedures to inline - allowed_aliases : tuple or list of str or :any:`Expression`, optional - List of variables that will not be renamed in the parent scope, even - if they alias with a local declaration. - adjust_imports : bool - Adjust imports by removing the symbol of the inlined routine or adding - imports needed by the imported routine (optional, default: True) - """ - - with pragmas_attached(routine, node_type=CallStatement): - - # Group the marked calls by callee routine - call_sets = defaultdict(list) - no_call_sets = defaultdict(list) - for call in FindNodes(CallStatement).visit(routine.body): - if call.routine == BasicType.DEFERRED: - continue - - if is_loki_pragma(call.pragma, starts_with='inline'): - call_sets[call.routine].append(call) - else: - no_call_sets[call.routine].append(call) - - # Trigger per-call inlining on collected sets - for callee, calls in call_sets.items(): - if callee: # Skip the unattached calls (collected under None) - inline_subroutine_calls( - routine, calls, callee, allowed_aliases=allowed_aliases - ) - - # Remove imported symbols that have become obsolete - if adjust_imports: - callees = tuple(callee.procedure_symbol for callee in call_sets.keys()) - not_inlined = tuple(callee.procedure_symbol for callee in no_call_sets.keys()) - - import_map = {} - for impt in FindNodes(Import).visit(routine.spec): - # Remove interface header imports - if any(f'{c.name.lower()}.intfb.h' == impt.module for c in callees): - import_map[impt] = None - - if any(s.name in callees for s in impt.symbols): - new_symbols = tuple( - s for s in impt.symbols if s.name not in callees or s.name in not_inlined - ) - # Remove import if no further symbols used, otherwise clone with new symbols - import_map[impt] = impt.clone(symbols=new_symbols) if new_symbols else None - - # Remove explicit interfaces of inlined routines - for intf in routine.interfaces: - if not intf.spec: - _body = tuple( - s.type.dtype.procedure for s in intf.symbols - if s.name not in callees or s.name in not_inlined - ) - if _body: - import_map[intf] = intf.clone(body=_body) - else: - import_map[intf] = None - - # Now move any callee imports we might need over to the caller - new_imports = set() - imported_module_map = CaseInsensitiveDict((im.module, im) for im in routine.imports) - for callee in call_sets.keys(): - for impt in callee.imports: - - # Add any callee module we do not yet know - if impt.module not in imported_module_map: - new_imports.add(impt) - - # If we're importing the same module, check for missing symbols - if m := imported_module_map.get(impt.module): - _m = import_map.get(m, m) - if not all(s in _m.symbols for s in impt.symbols): - new_symbols = tuple(s.rescope(routine) for s in impt.symbols) - import_map[m] = m.clone(symbols=tuple(set(_m.symbols + new_symbols))) - - # Finally, apply the import remapping - routine.spec = Transformer(import_map).visit(routine.spec) - - # Add missing explicit interfaces from inlined subroutines - new_intfs = [] - intf_symbols = routine.interface_symbols - for callee in call_sets.keys(): - for intf in callee.interfaces: - for s in intf.symbols: - if not s in intf_symbols: - new_intfs += [s.type.dtype.procedure,] - - if new_intfs: - routine.spec.append(Interface(body=as_tuple(new_intfs))) - - # Add Fortran imports to the top, and C-style interface headers at the bottom - c_imports = tuple(im for im in new_imports if im.c_import) - f_imports = tuple(im for im in new_imports if not im.c_import) - routine.spec.prepend(f_imports) - routine.spec.append(c_imports) diff --git a/loki/transformations/inline/__init__.py b/loki/transformations/inline/__init__.py new file mode 100644 index 000000000..e1702970d --- /dev/null +++ b/loki/transformations/inline/__init__.py @@ -0,0 +1,130 @@ +# (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-level code inlining. + +The various inline mechanisms are provided as standalone utility methods, +or via the :any:`InlineTransformation` class for for batch processing. +""" + +from loki.transformations.inline.constants import * # noqa +from loki.transformations.inline.functions import * # noqa +from loki.transformations.inline.mapper import * # noqa +from loki.transformations.inline.procedures import * # noqa + +from loki.batch import Transformation +from loki.transformations.remove_code import do_remove_dead_code + + +__all__ = ['InlineTransformation'] + + +class InlineTransformation(Transformation): + """ + :any:`Transformation` class to apply several types of source inlining + when batch-processing large source trees via the :any:`Scheduler`. + + Parameters + ---------- + inline_constants : bool + Replace instances of variables with known constant values by + :any:`Literal` (see :any:`inline_constant_parameters`); default: False. + inline_elementals : bool + Replaces :any:`InlineCall` expression to elemental functions + with the called function's body (see :any:`inline_elemental_functions`); + default: True. + inline_stmt_funcs: bool + Replaces :any:`InlineCall` expression to statement functions + with the corresponding rhs of the statement function if + the statement function declaration is available; default: False. + inline_internals : bool + Inline internal procedure (see :any:`inline_internal_procedures`); + default: False. + inline_marked : bool + Inline :any:`Subroutine` objects marked by pragma annotations + (see :any:`inline_marked_subroutines`); default: True. + remove_dead_code : bool + Perform dead code elimination, where unreachable branches are + trimmed from the code (see :any:`dead_code_elimination`); default: True + allowed_aliases : tuple or list of str or :any:`Expression`, optional + List of variables that will not be renamed in the parent scope during + internal and pragma-driven inlining. + adjust_imports : bool + Adjust imports by removing the symbol of the inlined routine or adding + imports needed by the imported routine (optional, default: True) + external_only : bool, optional + Do not replace variables declared in the local scope when + inlining constants (default: True) + resolve_sequence_association: bool + Resolve sequence association for routines that contain calls to inline (default: False) + """ + + # Ensure correct recursive inlining by traversing from the leaves + reverse_traversal = True + + # This transformation will potentially change the edges in the callgraph + creates_items = False + + def __init__( + self, inline_constants=False, inline_elementals=True, + inline_stmt_funcs=False, inline_internals=False, + inline_marked=True, remove_dead_code=True, + allowed_aliases=None, adjust_imports=True, + external_only=True, resolve_sequence_association=False + ): + self.inline_constants = inline_constants + self.inline_elementals = inline_elementals + self.inline_stmt_funcs = inline_stmt_funcs + self.inline_internals = inline_internals + self.inline_marked = inline_marked + self.remove_dead_code = remove_dead_code + self.allowed_aliases = allowed_aliases + self.adjust_imports = adjust_imports + self.external_only = external_only + self.resolve_sequence_association = resolve_sequence_association + if self.inline_marked: + self.creates_items = True + + def transform_subroutine(self, routine, **kwargs): + + # Resolve sequence association in calls that are about to be inlined. + # This step runs only if all of the following hold: + # 1) it is requested by the user + # 2) inlining of "internals" or "marked" routines is activated + # 3) there is an "internal" or "marked" procedure to inline. + if self.resolve_sequence_association: + resolve_sequence_association_for_inlined_calls( + routine, self.inline_internals, self.inline_marked + ) + + # Replace constant parameter variables with explicit values + if self.inline_constants: + inline_constant_parameters(routine, external_only=self.external_only) + + # Inline elemental functions + if self.inline_elementals: + inline_elemental_functions(routine) + + # Inline Statement Functions + if self.inline_stmt_funcs: + inline_statement_functions(routine) + + # Inline internal (contained) procedures + if self.inline_internals: + inline_internal_procedures(routine, allowed_aliases=self.allowed_aliases) + + # Inline explicitly pragma-marked subroutines + if self.inline_marked: + inline_marked_subroutines( + routine, allowed_aliases=self.allowed_aliases, + adjust_imports=self.adjust_imports + ) + + # After inlining, attempt to trim unreachable code paths + if self.remove_dead_code: + do_remove_dead_code(routine) diff --git a/loki/transformations/inline/constants.py b/loki/transformations/inline/constants.py new file mode 100644 index 000000000..db7decc79 --- /dev/null +++ b/loki/transformations/inline/constants.py @@ -0,0 +1,88 @@ +# (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.ir import ( + Import, Comment, Transformer, FindNodes, FindVariables, + FindLiterals, SubstituteExpressions +) +from loki.expression import symbols as sym + + +__all__ = ['inline_constant_parameters'] + + +def inline_constant_parameters(routine, external_only=True): + """ + Replace instances of variables with known constant values by `Literals`. + + Notes + ----- + The ``.type.initial`` property is used to derive the replacement + value,a which means for symbols imported from external modules, + the parent :any:`Module` needs to be supplied in the + ``definitions`` to the constructor when creating the + :any:`Subroutine`. + + Variables that are replaced are also removed from their + corresponding import statements, with empty import statements + being removed alltogether. + + Parameters + ---------- + routine : :any:`Subroutine` + Procedure in which to inline/resolve constant parameters. + external_only : bool, optional + Do not replace variables declared in the local scope (default: True) + """ + # Find all variable instances in spec and body + variables = FindVariables().visit(routine.ir) + + # Filter out variables declared locally + if external_only: + variables = [v for v in variables if v not in routine.variables] + + def is_inline_parameter(v): + return hasattr(v, 'type') and v.type.parameter and v.type.initial + + # Create mapping for variables and imports + vmap = {v: v.type.initial for v in variables if is_inline_parameter(v)} + + # Replace kind parameters in variable types + for variable in routine.variables: + if is_inline_parameter(variable.type.kind): + routine.symbol_attrs[variable.name] = variable.type.clone(kind=variable.type.kind.type.initial) + if variable.type.initial is not None: + # Substitute kind specifier in literals in initializers (I know...) + init_map = {literal.kind: literal.kind.type.initial + for literal in FindLiterals().visit(variable.type.initial) + if is_inline_parameter(literal.kind)} + if init_map: + initial = SubstituteExpressions(init_map).visit(variable.type.initial) + routine.symbol_attrs[variable.name] = variable.type.clone(initial=initial) + + # Update imports + imprtmap = {} + substituted_names = {v.name.lower() for v in vmap} + for imprt in FindNodes(Import).visit(routine.spec): + if imprt.symbols: + symbols = tuple(s for s in imprt.symbols if s.name.lower() not in substituted_names) + if not symbols: + imprtmap[imprt] = Comment(f'! Loki: parameters from {imprt.module} inlined') + elif len(symbols) < len(imprt.symbols): + imprtmap[imprt] = imprt.clone(symbols=symbols) + + # Flush mappings through spec and body + routine.spec = Transformer(imprtmap).visit(routine.spec) + routine.spec = SubstituteExpressions(vmap).visit(routine.spec) + routine.body = SubstituteExpressions(vmap).visit(routine.body) + + # Clean up declarations that are about to become defunct + decl_map = { + decl: None for decl in routine.declarations + if all(isinstance(s, sym.IntLiteral) for s in decl.symbols) + } + routine.spec = Transformer(decl_map).visit(routine.spec) diff --git a/loki/transformations/inline/functions.py b/loki/transformations/inline/functions.py new file mode 100644 index 000000000..ae1f8476d --- /dev/null +++ b/loki/transformations/inline/functions.py @@ -0,0 +1,342 @@ +# (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 collections import ChainMap + +from loki.expression import symbols as sym, ExpressionRetriever +from loki.ir import ( + Transformer, FindNodes, FindVariables, Import, StatementFunction, + FindInlineCalls, ExpressionFinder, SubstituteExpressions, + VariableDeclaration +) +from loki.subroutine import Subroutine +from loki.types import BasicType +from loki.tools import as_tuple + +from loki.transformations.inline.mapper import InlineSubstitutionMapper +from loki.transformations.inline.procedures import map_call_to_procedure_body +from loki.transformations.utilities import ( + single_variable_declaration, recursive_expression_map_update +) + + +__all__ = [ + 'inline_elemental_functions', 'inline_functions', + 'inline_statement_functions', 'inline_function_calls' +] + + +def inline_elemental_functions(routine): + """ + Replaces `InlineCall` expression to elemental functions with the + called functions body. + + Parameters + ---------- + routine : :any:`Subroutine` + Procedure in which to inline functions. + """ + inline_functions(routine, inline_elementals_only=True) + + +def inline_functions(routine, inline_elementals_only=False, functions=None): + """ + Replaces `InlineCall` expression to functions with the + called functions body. Nested calls are handled/inlined through + an iterative approach calling :any:`_inline_functions`. + + Parameters + ---------- + routine : :any:`Subroutine` + Procedure in which to inline functions. + inline_elementals_only : bool, optional + Inline elemental routines/functions only (default: False). + functions : tuple, optional + Inline only functions that are provided here + (default: None, thus inline all functions). + """ + potentially_functions_to_be_inlined = True + while potentially_functions_to_be_inlined: + potentially_functions_to_be_inlined = _inline_functions( + routine, inline_elementals_only=inline_elementals_only, functions=functions + ) + +def _inline_functions(routine, inline_elementals_only=False, functions=None): + """ + Replaces `InlineCall` expression to functions with the + called functions body, but doesn't include nested calls! + + Parameters + ---------- + routine : :any:`Subroutine` + Procedure in which to inline functions. + inline_elementals_only : bool, optional + Inline elemental routines/functions only (default: False). + functions : tuple, optional + Inline only functions that are provided here + (default: None, thus inline all functions). + + Returns + ------- + bool + Whether inline calls are (potentially) left to be + inlined in the next call to this function. + """ + + class ExpressionRetrieverSkipInlineCallParameters(ExpressionRetriever): + """ + Expression retriever skipping parameters of inline calls. + """ + # pylint: disable=abstract-method + + def __init__(self, query, recurse_query=None, inline_elementals_only=False, + functions=None, **kwargs): + self.inline_elementals_only = inline_elementals_only + self.functions = as_tuple(functions) + super().__init__(query, recurse_query, **kwargs) + + def map_inline_call(self, expr, *args, **kwargs): + if not self.visit(expr, *args, **kwargs): + return + self.rec(expr.function, *args, **kwargs) + # SKIP parameters/args/kwargs on purpose + # under certain circumstances + if expr.procedure_type is BasicType.DEFERRED or\ + (self.inline_elementals_only and\ + not(expr.procedure_type.is_function and expr.procedure_type.is_elemental)) or\ + (self.functions and expr.routine not in self.functions): + for child in expr.parameters: + self.rec(child, *args, **kwargs) + for child in list(expr.kw_parameters.values()): + self.rec(child, *args, **kwargs) + + self.post_visit(expr, *args, **kwargs) + + class FindInlineCallsSkipInlineCallParameters(ExpressionFinder): + """ + Find inline calls but skip/ignore parameters of inline calls. + """ + retriever = ExpressionRetrieverSkipInlineCallParameters(lambda e: isinstance(e, sym.InlineCall)) + + # functions are provided, however functions is empty, thus early exit + if functions is not None and not functions: + return False + functions = as_tuple(functions) + + # Keep track of removed symbols + removed_functions = set() + + # Find and filter inline calls and corresponding nodes + function_calls = {} + # Find inline calls but skip/ignore inline calls being parameters of other inline calls + # to ensure correct ordering of inlining. Those skipped/ignored inline calls will be handled + # in the next call to this function. + retriever = ExpressionRetrieverSkipInlineCallParameters(lambda e: isinstance(e, sym.InlineCall), + inline_elementals_only=inline_elementals_only, functions=functions) + # override retriever ... + FindInlineCallsSkipInlineCallParameters.retriever = retriever + for node, calls in FindInlineCallsSkipInlineCallParameters(with_ir_node=True).visit(routine.body): + for call in calls: + if call.procedure_type is BasicType.DEFERRED or isinstance(call.routine, StatementFunction): + continue + if inline_elementals_only: + if not (call.procedure_type.is_function and call.procedure_type.is_elemental): + continue + if functions: + if call.routine not in functions: + continue + function_calls.setdefault(str(call.name).lower(),[]).append((call, node)) + + if not function_calls: + return False + + # inline functions + node_prepend_map = {} + call_map = {} + for calls_nodes in function_calls.values(): + calls, nodes = list(zip(*calls_nodes)) + for call in calls: + removed_functions.add(call.procedure_type) + # collect nodes to be appendes as well as expression replacement for inline call + inline_node_map, inline_call_map = inline_function_calls(routine, as_tuple(calls), + calls[0].routine, as_tuple(nodes)) + for node, nodes_to_prepend in inline_node_map.items(): + node_prepend_map.setdefault(node, []).extend(list(nodes_to_prepend)) + call_map.update(inline_call_map) + + # collect nodes to be prepended for each node that contains (at least one) inline call to a function + node_map = {} + for node, prepend_nodes in node_prepend_map.items(): + node_map[node] = as_tuple(prepend_nodes) + (SubstituteExpressions(call_map[node]).visit(node),) + # inline via prepending the relevant functions + routine.body = Transformer(node_map).visit(routine.body) + # We need this to ensure that symbols, as well as nested scopes + # are correctly attached to each other (eg. nested associates). + routine.rescope_symbols() + + # Remove all module imports that have become obsolete now + import_map = {} + for im in FindNodes(Import).visit(routine.spec): + if im.symbols and all(s.type.dtype in removed_functions for s in im.symbols): + import_map[im] = None + routine.spec = Transformer(import_map).visit(routine.spec) + return True + + +def inline_statement_functions(routine): + """ + Replaces :any:`InlineCall` expression to statement functions with the + called statement functions rhs. + """ + # Keep track of removed symbols + removed_functions = set() + + stmt_func_decls = FindNodes(StatementFunction).visit(routine.spec) + exprmap = {} + for call in FindInlineCalls().visit(routine.body): + proc_type = call.procedure_type + if proc_type is BasicType.DEFERRED: + continue + if proc_type.is_function and isinstance(call.routine, StatementFunction): + exprmap[call] = InlineSubstitutionMapper()(call, scope=routine) + removed_functions.add(call.routine) + # Apply the map to itself to handle nested statement function calls + exprmap = recursive_expression_map_update(exprmap, max_iterations=10, mapper_cls=InlineSubstitutionMapper) + # Apply expression-level substitution to routine + routine.body = SubstituteExpressions(exprmap).visit(routine.body) + + # remove statement function declarations as well as statement function argument(s) declarations + vars_to_remove = {stmt_func.variable.name.lower() for stmt_func in stmt_func_decls} + vars_to_remove |= {arg.name.lower() for stmt_func in stmt_func_decls for arg in stmt_func.arguments} + spec_map = {stmt_func: None for stmt_func in stmt_func_decls} + for decl in routine.declarations: + if any(var in vars_to_remove for var in decl.symbols): + symbols = tuple(var for var in decl.symbols if var not in vars_to_remove) + if symbols: + decl._update(symbols=symbols) + else: + spec_map[decl] = None + routine.spec = Transformer(spec_map).visit(routine.spec) + + +def inline_function_calls(routine, calls, callee, nodes, allowed_aliases=None): + """ + Inline a set of call to an individual :any:`Subroutine` being functions + at source level. + + This will replace all :any:`InlineCall` 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 + calls : tuple or list of :any:`InlineCall` + Set of calls (to the same callee) to be inlined. + callee : :any:`Subroutine` + The called target function to be inlined in the parent + nodes : :any:`Node` + The corresponding nodes the functions are called from. + allowed_aliases : tuple or list of str or :any:`Expression`, optional + List of variables that will not be renamed in the parent scope, even + if they alias with a local declaration. + """ + + def rename_result_name(routine, rename): + callee = routine.clone() + var_map = {} + callee_result_var = callee.variable_map[callee.result_name.lower()] + new_callee_result_var = callee_result_var.clone(name=rename) + var_map[callee_result_var] = new_callee_result_var + callee_vars = [var for var in FindVariables().visit(callee.body) + if var.name.lower() == callee_result_var.name.lower()] + var_map.update({var: var.clone(name=rename) for var in callee_vars}) + var_map = recursive_expression_map_update(var_map) + callee.body = SubstituteExpressions(var_map).visit(callee.body) + return callee, new_callee_result_var + + allowed_aliases = as_tuple(allowed_aliases) + + # Ensure we process sets of calls to the same callee + assert all(call.routine == callee for call in calls) + assert isinstance(callee, Subroutine) + + # Prevent shadowing of callee's variables by renaming them a priori + parent_variables = routine.variable_map + duplicates = tuple( + v for v in callee.variables + if v.name.lower() != callee.result_name.lower() + and v.name in parent_variables and v.name.lower() not in callee._dummies + ) + # Filter out allowed aliases to prevent suffixing + duplicates = tuple(v for v in duplicates if v.symbol not in allowed_aliases) + shadow_mapper = SubstituteExpressions( + {v: v.clone(name=f'{callee.name}_{v.name}') for v in duplicates} + ) + callee.spec = shadow_mapper.visit(callee.spec) + + var_map = {} + duplicate_names = {dl.name.lower() for dl in duplicates} + for v in FindVariables(unique=False).visit(callee.body): + if v.name.lower() in duplicate_names: + var_map[v] = v.clone(name=f'{callee.name}_{v.name}') + + var_map = recursive_expression_map_update(var_map) + callee.body = SubstituteExpressions(var_map).visit(callee.body) + + # Separate allowed aliases from other variables to ensure clean hoisting + if allowed_aliases: + single_variable_declaration(callee, variables=allowed_aliases) + + single_variable_declaration(callee, variables=callee.result_name) + # Get local variable declarations and hoist them + decls = FindNodes(VariableDeclaration).visit(callee.spec) + decls = tuple(d for d in decls if all(s.name.lower() != callee.result_name.lower() for s in d.symbols)) + decls = tuple(d for d in decls if all(s.name.lower() not in callee._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)) + # Rescope the declaration symbols + decls = tuple(d.clone(symbols=tuple(s.clone(scope=routine) for s in d.symbols)) for d in decls) + + # Find and apply symbol remappings for array size expressions + symbol_map = dict(ChainMap(*[call.arg_map for call in calls])) + decls = SubstituteExpressions(symbol_map).visit(decls) + routine.spec.append(decls) + + # Handle result/return var/value + new_symbols = set() + result_var_map = {} + adapted_calls = [] + rename_result_var = not len(nodes) == len(set(nodes)) + for i_call, call in enumerate(calls): + callee_result_var = callee.variable_map[callee.result_name.lower()] + prefix = '' + new_callee_result_var_name = f'{prefix}result_{callee.result_name.lower()}_{i_call}'\ + if rename_result_var else f'{prefix}result_{callee.result_name.lower()}' + new_callee, new_symbol = rename_result_name(callee, new_callee_result_var_name) + adapted_calls.append(new_callee) + new_symbols.add(new_symbol) + if isinstance(callee_result_var, sym.Array): + result_var_map[(nodes[i_call], call)] = callee_result_var.clone(name=new_callee_result_var_name, + dimensions=None) + else: + result_var_map[(nodes[i_call], call)] = callee_result_var.clone(name=new_callee_result_var_name) + new_symbols = SubstituteExpressions(symbol_map).visit(as_tuple(new_symbols), recurse_to_declaration_attributes=True) + routine.variables += as_tuple([symbol.clone(scope=routine) for symbol in new_symbols]) + + # create node map to map nodes to be prepended (representing the functions) for each node + node_map = {} + call_map = {} + for i_call, call in enumerate(calls): + node_map.setdefault(nodes[i_call], []).extend( + list(map_call_to_procedure_body(call, caller=routine, callee=adapted_calls[i_call])) + ) + call_map.setdefault(nodes[i_call], {}).update({call: result_var_map[(nodes[i_call], call)]}) + return node_map, call_map diff --git a/loki/transformations/inline/mapper.py b/loki/transformations/inline/mapper.py new file mode 100644 index 000000000..80439fc6c --- /dev/null +++ b/loki/transformations/inline/mapper.py @@ -0,0 +1,84 @@ +# (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.ir import ( + FindNodes, Assignment, StatementFunction, SubstituteExpressions +) +from loki.expression import LokiIdentityMapper +from loki.types import BasicType + + +__all__ = ['InlineSubstitutionMapper'] + + +class InlineSubstitutionMapper(LokiIdentityMapper): + """ + An expression mapper that defines symbolic substitution for inlining. + """ + + def map_algebraic_leaf(self, expr, *args, **kwargs): + raise NotImplementedError + + def map_scalar(self, expr, *args, **kwargs): + parent = self.rec(expr.parent, *args, **kwargs) if expr.parent is not None else None + + scope = kwargs.get('scope') or expr.scope + # We're re-scoping an imported symbol + if expr.scope != scope: + return expr.clone(scope=scope, type=expr.type.clone(), parent=parent) + return expr.clone(parent=parent) + + map_deferred_type_symbol = map_scalar + + def map_array(self, expr, *args, **kwargs): + if expr.dimensions: + dimensions = self.rec(expr.dimensions, *args, **kwargs) + else: + dimensions = None + parent = self.rec(expr.parent, *args, **kwargs) if expr.parent is not None else None + + scope = kwargs.get('scope') or expr.scope + # We're re-scoping an imported symbol + if expr.scope != scope: + return expr.clone(scope=scope, type=expr.type.clone(), parent=parent, dimensions=dimensions) + return expr.clone(parent=parent, dimensions=dimensions) + + def map_procedure_symbol(self, expr, *args, **kwargs): + parent = self.rec(expr.parent, *args, **kwargs) if expr.parent is not None else None + + scope = kwargs.get('scope') or expr.scope + # We're re-scoping an imported symbol + if expr.scope != scope: + return expr.clone(scope=scope, type=expr.type.clone(), parent=parent) + return expr.clone(parent=parent) + + def map_inline_call(self, expr, *args, **kwargs): + if expr.procedure_type is None or expr.procedure_type is BasicType.DEFERRED: + # Unkonw inline call, potentially an intrinsic + # We still need to recurse and ensure re-scoping + return super().map_inline_call(expr, *args, **kwargs) + + # if it is an inline call to a Statement Function + if isinstance(expr.routine, StatementFunction): + function = expr.routine + # Substitute all arguments through the elemental body + arg_map = dict(expr.arg_iter()) + fbody = SubstituteExpressions(arg_map).visit(function.rhs) + return fbody + + function = expr.procedure_type.procedure + v_result = [v for v in function.variables if v == function.name][0] + + # Substitute all arguments through the elemental body + arg_map = dict(expr.arg_iter()) + fbody = SubstituteExpressions(arg_map).visit(function.body) + + # Extract the RHS of the final result variable assignment + stmts = [s for s in FindNodes(Assignment).visit(fbody) if s.lhs == v_result] + assert len(stmts) == 1 + rhs = self.rec(stmts[0].rhs, *args, **kwargs) + return rhs diff --git a/loki/transformations/inline/procedures.py b/loki/transformations/inline/procedures.py new file mode 100644 index 000000000..f5a74764b --- /dev/null +++ b/loki/transformations/inline/procedures.py @@ -0,0 +1,395 @@ +# (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 collections import defaultdict, ChainMap + +from loki.ir import ( + Import, Comment, VariableDeclaration, CallStatement, Transformer, + FindNodes, FindVariables, FindInlineCalls, SubstituteExpressions, + pragmas_attached, is_loki_pragma, Interface, Pragma +) +from loki.expression import symbols as sym +from loki.types import BasicType +from loki.tools import as_tuple, CaseInsensitiveDict +from loki.logging import error +from loki.subroutine import Subroutine + +from loki.transformations.sanitise import transform_sequence_association_append_map +from loki.transformations.utilities import ( + single_variable_declaration, recursive_expression_map_update +) + + +__all__ = [ + 'inline_internal_procedures', 'inline_member_procedures', + 'inline_marked_subroutines', + 'resolve_sequence_association_for_inlined_calls' +] + + +def resolve_sequence_association_for_inlined_calls(routine, inline_internals, inline_marked): + """ + Resolve sequence association in calls to all member procedures (if ``inline_internals = True``) + or in calls to procedures that have been marked with an inline pragma (if ``inline_marked = True``). + If both ``inline_internals`` and ``inline_marked`` are ``False``, no processing is done. + """ + call_map = {} + with pragmas_attached(routine, node_type=CallStatement): + for call in FindNodes(CallStatement).visit(routine.body): + condition = ( + (inline_marked and is_loki_pragma(call.pragma, starts_with='inline')) or + (inline_internals and call.routine in routine.routines) + ) + if condition: + if call.routine == BasicType.DEFERRED: + # NOTE: Throwing error here instead of continuing, because the user has explicitly + # asked sequence assoc to happen with inlining, so source for routine should be + # found in calls to be inlined. + raise ValueError( + f"Cannot resolve sequence association for call to ``{call.name}`` " + + f"to be inlined in routine ``{routine.name}``, because " + + f"the ``CallStatement`` referring to ``{call.name}`` does not contain " + + "the source code of the procedure. " + + "If running in batch processing mode, please recheck Scheduler configuration." + ) + transform_sequence_association_append_map(call_map, call) + if call_map: + routine.body = Transformer(call_map).visit(routine.body) + + +def map_call_to_procedure_body(call, caller, callee=None): + """ + Resolve arguments of a call and map to the called procedure body. + + Parameters + ---------- + call : :any:`CallStatment` or :any:`InlineCall` + Call object that defines the argument mapping + caller : :any:`Subroutine` + Procedure (scope) into which the callee's body gets mapped + callee : :any:`Subroutine`, optional + Procedure (scope) called. Provide if it differs from + call.routine. + """ + + 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)) + + # Get callee from the procedure type + callee = callee or call.routine + if callee is BasicType.DEFERRED: + error( + '[Loki::TransformInline] Need procedure definition to resolve ' + f'call to {call.name} from {caller}' + ) + raise RuntimeError('Procedure definition not found! ') + + argmap = {} + callee_vars = FindVariables().visit(callee.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( + f'[Loki::TransformInline] Cannot resolve procedure call to {call.name}' + ) + arg_vars = tuple(v for v in callee_vars if v.name == arg.name) + argmap.update((v, _map_unbound_dims(v, qualified_value)) for v in arg_vars) + else: + argmap[arg] = val + + # Deal with PRESENT check for optional arguments + present_checks = tuple( + check for check in FindInlineCalls().visit(callee.body) if check.function == 'PRESENT' + ) + present_map = { + check: sym.Literal('.true.') if check.arguments[0] in [arg.name for arg in call.arg_map] + else sym.Literal('.false.') + for check in present_checks + } + argmap.update(present_map) + + # 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 + callee_body = SubstituteExpressions(argmap, rebuild_scopes=True).visit( + callee.body.body, scope=caller + ) + + # Remove 'loki routine' pragmas + callee_body = Transformer( + {pragma: None for pragma in FindNodes(Pragma).visit(callee_body) + if is_loki_pragma(pragma, starts_with='routine')} + ).visit(callee_body) + + # Inline substituted body within a pair of marker comments + comment = Comment(f'! [Loki] inlined child subroutine: {callee.name}') + c_line = Comment('! =========================================') + return (comment, c_line) + as_tuple(callee_body) + (c_line, ) + + +def inline_subroutine_calls(routine, calls, callee, allowed_aliases=None): + """ + Inline a set of call to an individual :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 + calls : tuple or list of :any:`CallStatement` + callee : :any:`Subroutine` + The called target subroutine to be inlined in the parent + allowed_aliases : tuple or list of str or :any:`Expression`, optional + List of variables that will not be renamed in the parent scope, even + if they alias with a local declaration. + """ + allowed_aliases = as_tuple(allowed_aliases) + + # Ensure we process sets of calls to the same callee + assert all(call.routine == callee for call in calls) + assert isinstance(callee, Subroutine) + + # Prevent shadowing of callee's variables by renaming them a priori + parent_variables = routine.variable_map + duplicates = tuple( + v for v in callee.variables + if v.name in parent_variables and v.name.lower() not in callee._dummies + ) + # Filter out allowed aliases to prevent suffixing + duplicates = tuple(v for v in duplicates if v.symbol not in allowed_aliases) + shadow_mapper = SubstituteExpressions( + {v: v.clone(name=f'{callee.name}_{v.name}') for v in duplicates} + ) + callee.spec = shadow_mapper.visit(callee.spec) + + var_map = {} + duplicate_names = {dl.name.lower() for dl in duplicates} + for v in FindVariables(unique=False).visit(callee.body): + if v.name.lower() in duplicate_names: + var_map[v] = v.clone(name=f'{callee.name}_{v.name}') + var_map = recursive_expression_map_update(var_map) + callee.body = SubstituteExpressions(var_map).visit(callee.body) + + # Separate allowed aliases from other variables to ensure clean hoisting + if allowed_aliases: + single_variable_declaration(callee, variables=allowed_aliases) + + # Get local variable declarations and hoist them + decls = FindNodes(VariableDeclaration).visit(callee.spec) + decls = tuple(d for d in decls if all(s.name.lower() not in callee._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)) + # Rescope the declaration symbols + decls = tuple(d.clone(symbols=tuple(s.clone(scope=routine) for s in d.symbols)) for d in decls) + + # Find and apply symbol remappings for array size expressions + symbol_map = dict(ChainMap(*[call.arg_map for call in calls])) + decls = SubstituteExpressions(symbol_map).visit(decls) + + routine.spec.append(decls) + + # Resolve the call by mapping arguments into the called procedure's body + call_map = { + call: map_call_to_procedure_body(call, caller=routine) for call in calls + } + + # Replace calls to child procedure with the child's body + routine.body = Transformer(call_map).visit(routine.body) + + # We need this to ensure that symbols, as well as nested scopes + # are correctly attached to each other (eg. nested associates). + routine.rescope_symbols() + + +def inline_internal_procedures(routine, allowed_aliases=None): + """ + Inline internal subroutines contained in an individual :any:`Subroutine`. + + Please note that internal functions are not yet supported! + + Parameters + ---------- + routine : :any:`Subroutine` + The subroutine in which to inline all member routines + allowed_aliases : tuple or list of str or :any:`Expression`, optional + List of variables that will not be renamed in the parent scope, even + if they alias with a local declaration. + """ + + from loki.transformations.inline import inline_functions # pylint: disable=cyclic-import,import-outside-toplevel + + # Run through all members and invoke individual inlining transforms + for child in routine.members: + if child.is_function: + inline_functions(routine, functions=(child,)) + else: + calls = tuple( + call for call in FindNodes(CallStatement).visit(routine.body) + if call.routine == child + ) + inline_subroutine_calls(routine, calls, child, allowed_aliases=allowed_aliases) + + # Can't use transformer to replace subroutine/function, so strip it manually + contains_body = tuple(n for n in routine.contains.body if not n == child) + routine.contains._update(body=contains_body) + + +inline_member_procedures = inline_internal_procedures + + +def inline_marked_subroutines(routine, allowed_aliases=None, adjust_imports=True): + """ + Inline :any:`Subroutine` objects guided by pragma annotations. + + When encountering :any:`CallStatement` objects that are marked with a + ``!$loki inline`` pragma, this utility will attempt to replace the call + with the body of the called procedure and remap all passed arguments + into the calling procedures scope. + + Please note that this utility requires :any:`CallStatement` objects + to be "enriched" with external type information. + + Parameters + ---------- + routine : :any:`Subroutine` + The subroutine in which to look for pragma-marked procedures to inline + allowed_aliases : tuple or list of str or :any:`Expression`, optional + List of variables that will not be renamed in the parent scope, even + if they alias with a local declaration. + adjust_imports : bool + Adjust imports by removing the symbol of the inlined routine or adding + imports needed by the imported routine (optional, default: True) + """ + + with pragmas_attached(routine, node_type=CallStatement): + + # Group the marked calls by callee routine + call_sets = defaultdict(list) + no_call_sets = defaultdict(list) + for call in FindNodes(CallStatement).visit(routine.body): + if call.routine == BasicType.DEFERRED: + continue + + if is_loki_pragma(call.pragma, starts_with='inline'): + call_sets[call.routine].append(call) + else: + no_call_sets[call.routine].append(call) + + # Trigger per-call inlining on collected sets + for callee, calls in call_sets.items(): + if callee: # Skip the unattached calls (collected under None) + inline_subroutine_calls( + routine, calls, callee, allowed_aliases=allowed_aliases + ) + + # Remove imported symbols that have become obsolete + if adjust_imports: + callees = tuple(callee.procedure_symbol for callee in call_sets.keys()) + not_inlined = tuple(callee.procedure_symbol for callee in no_call_sets.keys()) + + import_map = {} + for impt in FindNodes(Import).visit(routine.spec): + # Remove interface header imports + if any(f'{c.name.lower()}.intfb.h' == impt.module for c in callees): + import_map[impt] = None + + if any(s.name in callees for s in impt.symbols): + new_symbols = tuple( + s for s in impt.symbols if s.name not in callees or s.name in not_inlined + ) + # Remove import if no further symbols used, otherwise clone with new symbols + import_map[impt] = impt.clone(symbols=new_symbols) if new_symbols else None + + # Remove explicit interfaces of inlined routines + for intf in routine.interfaces: + if not intf.spec: + _body = tuple( + s.type.dtype.procedure for s in intf.symbols + if s.name not in callees or s.name in not_inlined + ) + if _body: + import_map[intf] = intf.clone(body=_body) + else: + import_map[intf] = None + + # Now move any callee imports we might need over to the caller + new_imports = set() + imported_module_map = CaseInsensitiveDict((im.module, im) for im in routine.imports) + for callee in call_sets.keys(): + for impt in callee.imports: + + # Add any callee module we do not yet know + if impt.module not in imported_module_map: + new_imports.add(impt) + + # If we're importing the same module, check for missing symbols + if m := imported_module_map.get(impt.module): + _m = import_map.get(m, m) + if not all(s in _m.symbols for s in impt.symbols): + new_symbols = tuple(s.rescope(routine) for s in impt.symbols) + import_map[m] = m.clone(symbols=tuple(set(_m.symbols + new_symbols))) + + # Finally, apply the import remapping + routine.spec = Transformer(import_map).visit(routine.spec) + + # Add missing explicit interfaces from inlined subroutines + new_intfs = [] + intf_symbols = routine.interface_symbols + for callee in call_sets.keys(): + for intf in callee.interfaces: + for s in intf.symbols: + if not s in intf_symbols: + new_intfs += [s.type.dtype.procedure,] + + if new_intfs: + routine.spec.append(Interface(body=as_tuple(new_intfs))) + + # Add Fortran imports to the top, and C-style interface headers at the bottom + c_imports = tuple(im for im in new_imports if im.c_import) + f_imports = tuple(im for im in new_imports if not im.c_import) + routine.spec.prepend(f_imports) + routine.spec.append(c_imports) From 23599f509a67459f654864760cd305d8174d1be1 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Wed, 11 Sep 2024 14:08:37 +0000 Subject: [PATCH 02/12] Transformations: Split tests and move into `transformations.inline` --- .../inline/tests/test_constants.py | 218 ++ .../inline/tests/test_functions.py | 241 +++ .../tests/test_inline_transformation.py | 477 +++++ .../inline/tests/test_procedures.py | 953 +++++++++ loki/transformations/tests/test_inline.py | 1838 ----------------- 5 files changed, 1889 insertions(+), 1838 deletions(-) create mode 100644 loki/transformations/inline/tests/test_constants.py create mode 100644 loki/transformations/inline/tests/test_functions.py create mode 100644 loki/transformations/inline/tests/test_inline_transformation.py create mode 100644 loki/transformations/inline/tests/test_procedures.py delete mode 100644 loki/transformations/tests/test_inline.py diff --git a/loki/transformations/inline/tests/test_constants.py b/loki/transformations/inline/tests/test_constants.py new file mode 100644 index 000000000..9ebf6babb --- /dev/null +++ b/loki/transformations/inline/tests/test_constants.py @@ -0,0 +1,218 @@ +# (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. + +import pytest + +from loki import Module, Subroutine +from loki.build import jit_compile_lib, Builder, Obj +from loki.frontend import available_frontends +from loki.ir import nodes as ir, FindNodes + +from loki.transformations.inline import inline_constant_parameters +from loki.transformations.utilities import replace_selected_kind + + +@pytest.fixture(name='builder') +def fixture_builder(tmp_path): + yield Builder(source_dirs=tmp_path, build_dir=tmp_path) + Obj.clear_cache() + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_transform_inline_constant_parameters(tmp_path, builder, frontend): + """ + Test correct inlining of constant parameters. + """ + fcode_module = """ +module parameters_mod + implicit none + integer, parameter :: a = 1 + integer, parameter :: b = -1 +contains + subroutine dummy + end subroutine dummy +end module parameters_mod +""" + + fcode = """ +module inline_const_param_mod + ! TODO: use parameters_mod, only: b + implicit none + integer, parameter :: c = 1+1 +contains + subroutine inline_const_param(v1, v2, v3) + use parameters_mod, only: a, b + integer, intent(in) :: v1 + integer, intent(out) :: v2, v3 + + v2 = v1 + b - a + v3 = c + end subroutine inline_const_param +end module inline_const_param_mod +""" + # Generate reference code, compile run and verify + param_module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) + module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + refname = f'ref_{module.name}_{ frontend}' + reference = jit_compile_lib([module, param_module], path=tmp_path, name=refname, builder=builder) + + v2, v3 = reference.inline_const_param_mod.inline_const_param(10) + assert v2 == 8 + assert v3 == 2 + (tmp_path/f'{module.name}.f90').unlink() + (tmp_path/f'{param_module.name}.f90').unlink() + + # Now transform with supplied elementals but without module + module = Module.from_source(fcode, definitions=param_module, frontend=frontend, xmods=[tmp_path]) + assert len(FindNodes(ir.Import).visit(module['inline_const_param'].spec)) == 1 + for routine in module.subroutines: + inline_constant_parameters(routine, external_only=True) + assert not FindNodes(ir.Import).visit(module['inline_const_param'].spec) + + # Hack: rename module to use a different filename in the build + module.name = f'{module.name}_' + obj = jit_compile_lib([module], path=tmp_path, name=f'{module.name}_{frontend}', builder=builder) + + v2, v3 = obj.inline_const_param_mod_.inline_const_param(10) + assert v2 == 8 + assert v3 == 2 + + (tmp_path/f'{module.name}.f90').unlink() + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_transform_inline_constant_parameters_kind(tmp_path, builder, frontend): + """ + Test correct inlining of constant parameters for kind symbols. + """ + fcode_module = """ +module kind_parameters_mod + implicit none + integer, parameter :: jprb = selected_real_kind(13, 300) +end module kind_parameters_mod +""" + + fcode = """ +module inline_const_param_kind_mod + implicit none +contains + subroutine inline_const_param_kind(v1) + use kind_parameters_mod, only: jprb + real(kind=jprb), intent(out) :: v1 + + v1 = real(2, kind=jprb) + 3. + end subroutine inline_const_param_kind +end module inline_const_param_kind_mod +""" + # Generate reference code, compile run and verify + param_module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) + module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + refname = f'ref_{module.name}_{frontend}' + reference = jit_compile_lib([module, param_module], path=tmp_path, name=refname, builder=builder) + + v1 = reference.inline_const_param_kind_mod.inline_const_param_kind() + assert v1 == 5. + (tmp_path/f'{module.name}.f90').unlink() + (tmp_path/f'{param_module.name}.f90').unlink() + + # Now transform with supplied elementals but without module + module = Module.from_source(fcode, definitions=param_module, frontend=frontend, xmods=[tmp_path]) + assert len(FindNodes(ir.Import).visit(module['inline_const_param_kind'].spec)) == 1 + for routine in module.subroutines: + inline_constant_parameters(routine, external_only=True) + assert not FindNodes(ir.Import).visit(module['inline_const_param_kind'].spec) + + # Hack: rename module to use a different filename in the build + module.name = f'{module.name}_' + obj = jit_compile_lib([module], path=tmp_path, name=f'{module.name}_{frontend}', builder=builder) + + v1 = obj.inline_const_param_kind_mod_.inline_const_param_kind() + assert v1 == 5. + + (tmp_path/f'{module.name}.f90').unlink() + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_transform_inline_constant_parameters_replace_kind(tmp_path, builder, frontend): + """ + Test correct inlining of constant parameters for kind symbols. + """ + fcode_module = """ +module replace_kind_parameters_mod + implicit none + integer, parameter :: jprb = selected_real_kind(13, 300) +end module replace_kind_parameters_mod +""" + + fcode = """ +module inline_param_repl_kind_mod + implicit none +contains + subroutine inline_param_repl_kind(v1) + use replace_kind_parameters_mod, only: jprb + real(kind=jprb), intent(out) :: v1 + real(kind=jprb) :: a = 3._JPRB + + v1 = 1._jprb + real(2, kind=jprb) + a + end subroutine inline_param_repl_kind +end module inline_param_repl_kind_mod +""" + # Generate reference code, compile run and verify + param_module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) + module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + refname = f'ref_{module.name}_{frontend}' + reference = jit_compile_lib([module, param_module], path=tmp_path, name=refname, builder=builder) + func = getattr(getattr(reference, module.name), module.subroutines[0].name) + + v1 = func() + assert v1 == 6. + (tmp_path/f'{module.name}.f90').unlink() + (tmp_path/f'{param_module.name}.f90').unlink() + + # Now transform with supplied elementals but without module + module = Module.from_source(fcode, definitions=param_module, frontend=frontend, xmods=[tmp_path]) + imports = FindNodes(ir.Import).visit(module.subroutines[0].spec) + assert len(imports) == 1 and imports[0].module.lower() == param_module.name.lower() + for routine in module.subroutines: + inline_constant_parameters(routine, external_only=True) + replace_selected_kind(routine) + imports = FindNodes(ir.Import).visit(module.subroutines[0].spec) + assert len(imports) == 1 and imports[0].module.lower() == 'iso_fortran_env' + + # Hack: rename module to use a different filename in the build + module.name = f'{module.name}_' + obj = jit_compile_lib([module], path=tmp_path, name=f'{module.name}_{frontend}', builder=builder) + + func = getattr(getattr(obj, module.name), module.subroutines[0].name) + v1 = func() + assert v1 == 6. + + (tmp_path/f'{module.name}.f90').unlink() + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_constant_replacement_internal(frontend): + """ + Test constant replacement for internally defined constants. + """ + fcode = """ +subroutine kernel(a, b) + integer, parameter :: par = 10 + integer, intent(inout) :: a, b + + a = b + par +end subroutine kernel + """ + routine = Subroutine.from_source(fcode, frontend=frontend) + inline_constant_parameters(routine=routine, external_only=False) + + assert len(routine.variables) == 2 + assert 'a' in routine.variables and 'b' in routine.variables + + stmts = FindNodes(ir.Assignment).visit(routine.body) + assert len(stmts) == 1 + assert stmts[0].rhs == 'b + 10' diff --git a/loki/transformations/inline/tests/test_functions.py b/loki/transformations/inline/tests/test_functions.py new file mode 100644 index 000000000..64b160407 --- /dev/null +++ b/loki/transformations/inline/tests/test_functions.py @@ -0,0 +1,241 @@ +# (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. + +import pytest + +from loki import Module, Subroutine +from loki.build import jit_compile_lib, Builder, Obj +from loki.frontend import available_frontends, OMNI, OFP +from loki.ir import ( + nodes as ir, FindNodes, FindVariables, FindInlineCalls +) + +from loki.transformations.inline import ( + inline_elemental_functions, inline_statement_functions +) + + +@pytest.fixture(name='builder') +def fixture_builder(tmp_path): + yield Builder(source_dirs=tmp_path, build_dir=tmp_path) + Obj.clear_cache() + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_transform_inline_elemental_functions(tmp_path, builder, frontend): + """ + Test correct inlining of elemental functions. + """ + fcode_module = """ +module multiply_mod + use iso_fortran_env, only: real64 + implicit none +contains + + elemental function multiply(a, b) + real(kind=real64) :: multiply + real(kind=real64), intent(in) :: a, b + real(kind=real64) :: temp + !$loki routine seq + + ! simulate multi-line function + temp = a * b + multiply = temp + end function multiply +end module multiply_mod +""" + + fcode = """ +subroutine transform_inline_elemental_functions(v1, v2, v3) + use iso_fortran_env, only: real64 + use multiply_mod, only: multiply + real(kind=real64), intent(in) :: v1 + real(kind=real64), intent(out) :: v2, v3 + + v2 = multiply(v1, 6._real64) + v3 = 600. + multiply(6._real64, 11._real64) +end subroutine transform_inline_elemental_functions +""" + + # Generate reference code, compile run and verify + module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) + routine = Subroutine.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + + refname = f'ref_{routine.name}_{frontend}' + reference = jit_compile_lib([module, routine], path=tmp_path, name=refname, builder=builder) + + v2, v3 = reference.transform_inline_elemental_functions(11.) + assert v2 == 66. + assert v3 == 666. + + (tmp_path/f'{module.name}.f90').unlink() + (tmp_path/f'{routine.name}.f90').unlink() + + # Now inline elemental functions + routine = Subroutine.from_source(fcode, definitions=module, frontend=frontend, xmods=[tmp_path]) + inline_elemental_functions(routine) + + # Make sure there are no more inline calls in the routine body + assert not FindInlineCalls().visit(routine.body) + + # Verify correct scope of inlined elements + assert all(v.scope is routine for v in FindVariables().visit(routine.body)) + + # Ensure the !$loki routine pragma has been removed + assert not FindNodes(ir.Pragma).visit(routine.body) + + # Hack: rename routine to use a different filename in the build + routine.name = f'{routine.name}_' + kernel = jit_compile_lib([routine], path=tmp_path, name=routine.name, builder=builder) + + v2, v3 = kernel.transform_inline_elemental_functions_(11.) + assert v2 == 66. + assert v3 == 666. + + builder.clean() + (tmp_path/f'{routine.name}.f90').unlink() + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_transform_inline_elemental_functions_extended(tmp_path, builder, frontend): + """ + Test correct inlining of elemental functions. + """ + fcode_module = """ +module multiply_extended_mod + use iso_fortran_env, only: real64 + implicit none +contains + + elemental function multiply(a, b) ! result (ret_mult) + ! real(kind=real64) :: ret_mult + real(kind=real64) :: multiply + real(kind=real64), intent(in) :: a, b + real(kind=real64) :: temp + + ! simulate multi-line function + temp = a * b + multiply = temp + ! ret_mult = temp + end function multiply + + elemental function multiply_single_line(a, b) + real(kind=real64) :: multiply_single_line + real(kind=real64), intent(in) :: a, b + real(kind=real64) :: temp + + multiply_single_line = a * b + end function multiply_single_line + + elemental function add(a, b) + real(kind=real64) :: add + real(kind=real64), intent(in) :: a, b + real(kind=real64) :: temp + + ! simulate multi-line function + temp = a + b + add = temp + end function add +end module multiply_extended_mod +""" + + fcode = """ +subroutine transform_inline_elemental_functions_extended(v1, v2, v3) + use iso_fortran_env, only: real64 + use multiply_extended_mod, only: multiply, multiply_single_line, add + real(kind=real64), intent(in) :: v1 + real(kind=real64), intent(out) :: v2, v3 + real(kind=real64), parameter :: param1 = 100. + + v2 = multiply(v1, 6._real64) + multiply_single_line(v1, 3._real64) + v3 = add(param1, 200._real64) + add(150._real64, 150._real64) + multiply(6._real64, 11._real64) +end subroutine transform_inline_elemental_functions_extended +""" + + # Generate reference code, compile run and verify + module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) + routine = Subroutine.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + + refname = f'ref_{routine.name}_{frontend}' + reference = jit_compile_lib([module, routine], path=tmp_path, name=refname, builder=builder) + + v2, v3 = reference.transform_inline_elemental_functions_extended(11.) + assert v2 == 99. + assert v3 == 666. + + (tmp_path/f'{module.name}.f90').unlink() + (tmp_path/f'{routine.name}.f90').unlink() + + # Now inline elemental functions + routine = Subroutine.from_source(fcode, definitions=module, frontend=frontend, xmods=[tmp_path]) + inline_elemental_functions(routine) + + + # Make sure there are no more inline calls in the routine body + assert not FindInlineCalls().visit(routine.body) + + # Verify correct scope of inlined elements + assert all(v.scope is routine for v in FindVariables().visit(routine.body)) + + # Hack: rename routine to use a different filename in the build + routine.name = f'{routine.name}_' + kernel = jit_compile_lib([routine], path=tmp_path, name=routine.name, builder=builder) + + v2, v3 = kernel.transform_inline_elemental_functions_extended_(11.) + assert v2 == 99. + assert v3 == 666. + + builder.clean() + (tmp_path/f'{routine.name}.f90').unlink() + + +@pytest.mark.parametrize('frontend', available_frontends( + skip={OFP: "OFP apparently has problems dealing with those Statement Functions", + OMNI: "OMNI automatically inlines Statement Functions"} +)) +@pytest.mark.parametrize('stmt_decls', (True, False)) +def test_inline_statement_functions(frontend, stmt_decls): + stmt_decls_code = """ + real :: PTARE + real :: FOEDELTA + FOEDELTA ( PTARE ) = PTARE + 1.0 + real :: FOEEW + FOEEW ( PTARE ) = PTARE + FOEDELTA(PTARE) + """.strip() + + fcode = f""" +subroutine stmt_func(arr, ret) + implicit none + real, intent(in) :: arr(:) + real, intent(inout) :: ret(:) + real :: ret2 + real, parameter :: rtt = 1.0 + {stmt_decls_code if stmt_decls else '#include "fcttre.func.h"'} + + ret = foeew(arr) + ret2 = foedelta(3.0) +end subroutine stmt_func + """.strip() + + routine = Subroutine.from_source(fcode, frontend=frontend) + if stmt_decls: + assert FindNodes(ir.StatementFunction).visit(routine.spec) + else: + assert not FindNodes(ir.StatementFunction).visit(routine.spec) + assert FindInlineCalls().visit(routine.body) + inline_statement_functions(routine) + + assert not FindNodes(ir.StatementFunction).visit(routine.spec) + if stmt_decls: + assert not FindInlineCalls().visit(routine.body) + assignments = FindNodes(ir.Assignment).visit(routine.body) + assert assignments[0].lhs == 'ret' + assert assignments[0].rhs == "arr + arr + 1.0" + assert assignments[1].lhs == 'ret2' + assert assignments[1].rhs == "3.0 + 1.0" + else: + assert FindInlineCalls().visit(routine.body) diff --git a/loki/transformations/inline/tests/test_inline_transformation.py b/loki/transformations/inline/tests/test_inline_transformation.py new file mode 100644 index 000000000..62da187f0 --- /dev/null +++ b/loki/transformations/inline/tests/test_inline_transformation.py @@ -0,0 +1,477 @@ +# (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. + +import pytest + +from loki import Module, Subroutine +from loki.frontend import available_frontends, OFP +from loki.ir import nodes as ir, FindNodes +from loki.batch import Scheduler, SchedulerConfig + +from loki.transformations.inline import InlineTransformation + + +@pytest.mark.parametrize('frontend', available_frontends( + (OFP, 'Prefix/elemental support not implemented')) +) +@pytest.mark.parametrize('pass_as_kwarg', (False, True)) +def test_inline_transformation(tmp_path, frontend, pass_as_kwarg): + """Test combining recursive inlining via :any:`InliningTransformation`.""" + + fcode_module = """ +module one_mod + real(kind=8), parameter :: one = 1.0 +end module one_mod +""" + + fcode_inner = """ +subroutine add_one_and_two(a) + use one_mod, only: one + implicit none + + real(kind=8), intent(inout) :: a + + a = a + one + + a = add_two(a) + +contains + elemental function add_two(x) + real(kind=8), intent(in) :: x + real(kind=8) :: add_two + + add_two = x + 2.0 + end function add_two +end subroutine add_one_and_two +""" + + fcode = f""" +subroutine test_inline_pragma(a, b) + implicit none + real(kind=8), intent(inout) :: a(3), b(3) + integer, parameter :: n = 3 + integer :: i + real :: stmt_arg + real :: some_stmt_func + some_stmt_func ( stmt_arg ) = stmt_arg + 3.1415 + +#include "add_one_and_two.intfb.h" + + do i=1, n + !$loki inline + call add_one_and_two({'a=' if pass_as_kwarg else ''}a(i)) + end do + + do i=1, n + !$loki inline + call add_one_and_two({'a=' if pass_as_kwarg else ''}b(i)) + end do + + a(1) = some_stmt_func({'stmt_arg=' if pass_as_kwarg else ''}a(2)) + +end subroutine test_inline_pragma +""" + module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) + inner = Subroutine.from_source(fcode_inner, definitions=module, frontend=frontend, xmods=[tmp_path]) + routine = Subroutine.from_source(fcode, frontend=frontend) + routine.enrich(inner) + + trafo = InlineTransformation( + inline_constants=True, external_only=True, inline_elementals=True, + inline_stmt_funcs=True + ) + + calls = FindNodes(ir.CallStatement).visit(routine.body) + assert len(calls) == 2 + assert all(c.routine == inner for c in calls) + + # Apply to the inner subroutine first to resolve parameter and calls + trafo.apply(inner) + + assigns = FindNodes(ir.Assignment).visit(inner.body) + assert len(assigns) == 3 + assert assigns[0].lhs == 'a' and assigns[0].rhs == 'a + 1.0' + assert assigns[1].lhs == 'result_add_two' and assigns[1].rhs == 'a + 2.0' + assert assigns[2].lhs == 'a' and assigns[2].rhs == 'result_add_two' + + # Apply to the outer routine, but with resolved body of the inner + trafo.apply(routine) + + calls = FindNodes(ir.CallStatement).visit(routine.body) + assert len(calls) == 0 + assigns = FindNodes(ir.Assignment).visit(routine.body) + assert len(assigns) == 7 + assert assigns[0].lhs == 'a(i)' and assigns[0].rhs == 'a(i) + 1.0' + assert assigns[1].lhs == 'result_add_two' and assigns[1].rhs == 'a(i) + 2.0' + assert assigns[2].lhs == 'a(i)' and assigns[2].rhs == 'result_add_two' + assert assigns[3].lhs == 'b(i)' and assigns[3].rhs == 'b(i) + 1.0' + assert assigns[4].lhs == 'result_add_two' and assigns[4].rhs == 'b(i) + 2.0' + assert assigns[5].lhs == 'b(i)' and assigns[5].rhs == 'result_add_two' + assert assigns[6].lhs == 'a(1)' and assigns[6].rhs == 'a(2) + 3.1415' + + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_transformation_local_seq_assoc(frontend, tmp_path): + fcode = """ +module somemod + implicit none + contains + + subroutine minusone_second(output, x) + real, intent(inout) :: output + real, intent(in) :: x(3) + output = x(2) - 1 + end subroutine minusone_second + + subroutine plusone(output, x) + real, intent(inout) :: output + real, intent(in) :: x + output = x + 1 + end subroutine plusone + + subroutine outer() + implicit none + real :: x(3, 3) + real :: y + x = 10.0 + + call inner(y, x(1, 1)) ! Sequence association tmp_path for member routine. + + !$loki inline + call plusone(y, x(3, 3)) ! Marked for inlining. + + call minusone_second(y, x(1, 3)) ! Standard call with sequence association (never processed). + + contains + + subroutine inner(output, x) + real, intent(inout) :: output + real, intent(in) :: x(3) + + output = x(2) + 2.0 + end subroutine inner + end subroutine outer + +end module somemod +""" + # Test case that nothing happens if `resolve_sequence_association=True` + # but inlining "marked" and "internals" is disabled. + module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + trafo = InlineTransformation( + inline_constants=True, external_only=True, inline_elementals=True, + inline_marked=False, inline_internals=False, resolve_sequence_association=True + ) + outer = module["outer"] + trafo.apply(outer) + callnames = [call.name for call in FindNodes(ir.CallStatement).visit(outer.body)] + assert 'plusone' in callnames + assert 'inner' in callnames + assert 'minusone_second' in callnames + + # Test case that only marked processed if + # `resolve_sequence_association=True` + # `inline_marked=True`, + # `inline_internals=False` + module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + trafo = InlineTransformation( + inline_constants=True, external_only=True, inline_elementals=True, + inline_marked=True, inline_internals=False, resolve_sequence_association=True + ) + outer = module["outer"] + trafo.apply(outer) + callnames = [call.name for call in FindNodes(ir.CallStatement).visit(outer.body)] + assert 'plusone' not in callnames + assert 'inner' in callnames + assert 'minusone_second' in callnames + + # Test case that a crash occurs if sequence association is not enabled even if it is needed. + module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + trafo = InlineTransformation( + inline_constants=True, external_only=True, inline_elementals=True, + inline_marked=True, inline_internals=True, resolve_sequence_association=False + ) + outer = module["outer"] + with pytest.raises(RuntimeError): + trafo.apply(outer) + callnames = [call.name for call in FindNodes(ir.CallStatement).visit(outer.body)] + + # Test case that sequence association is run and corresponding call inlined, avoiding crash. + module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + trafo = InlineTransformation( + inline_constants=True, external_only=True, inline_elementals=True, + inline_marked=False, inline_internals=True, resolve_sequence_association=True + ) + outer = module["outer"] + trafo.apply(outer) + callnames = [call.name for call in FindNodes(ir.CallStatement).visit(outer.body)] + assert 'plusone' in callnames + assert 'inner' not in callnames + assert 'minusone_second' in callnames + + # Test case that everything is enabled. + module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + trafo = InlineTransformation( + inline_constants=True, external_only=True, inline_elementals=True, + inline_marked=True, inline_internals=True, resolve_sequence_association=True + ) + outer = module["outer"] + trafo.apply(outer) + callnames = [call.name for call in FindNodes(ir.CallStatement).visit(outer.body)] + assert 'plusone' not in callnames + assert 'inner' not in callnames + assert 'minusone_second' in callnames + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_transformation_local_seq_assoc_crash_marked_no_seq_assoc(frontend, tmp_path): + # Test case that a crash occurs if marked routine with sequence association is + # attempted to inline without sequence association enabled. + fcode = """ +module somemod + implicit none + contains + + subroutine inner(output, x) + real, intent(inout) :: output + real, intent(in) :: x(3) + + output = x(2) + 2.0 + end subroutine inner + + subroutine outer() + real :: x(3, 3) + real :: y + x = 10.0 + + !$loki inline + call inner(y, x(1, 1)) ! Sequence association tmp_path for marked routine. + end subroutine outer + +end module somemod +""" + module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + trafo = InlineTransformation( + inline_constants=True, external_only=True, inline_elementals=True, + inline_marked=True, inline_internals=False, resolve_sequence_association=False + ) + outer = module["outer"] + with pytest.raises(RuntimeError): + trafo.apply(outer) + + # Test case that crash is avoided by activating sequence association. + module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + trafo = InlineTransformation( + inline_constants=True, external_only=True, inline_elementals=True, + inline_marked=True, inline_internals=False, resolve_sequence_association=True + ) + outer = module["outer"] + trafo.apply(outer) + assert len(FindNodes(ir.CallStatement).visit(outer.body)) == 0 + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_transformation_local_seq_assoc_crash_value_err_no_source(frontend, tmp_path): + # Testing that ValueError is thrown if sequence association is requested with inlining, + # but source code behind call is missing (not enough type information). + fcode = """ +module somemod + implicit none + contains + + subroutine outer() + real :: x(3, 3) + real :: y + x = 10.0 + + !$loki inline + call inner(y, x(1, 1)) ! Sequence association tmp_path for marked routine. + end subroutine outer + +end module somemod +""" + module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + trafo = InlineTransformation( + inline_constants=True, external_only=True, inline_elementals=True, + inline_marked=True, inline_internals=False, resolve_sequence_association=True + ) + outer = module["outer"] + with pytest.raises(ValueError): + trafo.apply(outer) + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_transformation_adjust_imports(frontend, tmp_path): + fcode_module = """ +module bnds_module + integer :: m + integer :: n + integer :: l +end module bnds_module + """ + + fcode_another = """ +module another_module + integer :: x +end module another_module + """ + + fcode_outer = """ +subroutine test_inline_outer(a, b) + use bnds_module, only: n + use test_inline_mod, only: test_inline_inner + use test_inline_another_mod, only: test_inline_another_inner + implicit none + + real(kind=8), intent(inout) :: a(n), b(n) + + !$loki inline + call test_inline_another_inner() + !$loki inline + call test_inline_inner(a, b) +end subroutine test_inline_outer + """ + + fcode_inner = """ +module test_inline_mod + implicit none + contains + +subroutine test_inline_inner(a, b) + use BNDS_module, only: n, m + use another_module, only: x + + real(kind=8), intent(inout) :: a(n), b(n) + real(kind=8) :: tmp(m) + integer :: i + + tmp(1:m) = x + do i=1, n + a(i) = b(i) + sum(tmp) + end do +end subroutine test_inline_inner +end module test_inline_mod + """ + + fcode_another_inner = """ +module test_inline_another_mod + implicit none + contains + +subroutine test_inline_another_inner() + use BNDS_module, only: n, m, l + +end subroutine test_inline_another_inner +end module test_inline_another_mod + """ + + _ = Module.from_source(fcode_another, frontend=frontend, xmods=[tmp_path]) + _ = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) + inner = Module.from_source(fcode_inner, frontend=frontend, xmods=[tmp_path]) + another_inner = Module.from_source(fcode_another_inner, frontend=frontend, xmods=[tmp_path]) + outer = Subroutine.from_source( + fcode_outer, definitions=(inner, another_inner), frontend=frontend, xmods=[tmp_path] + ) + + trafo = InlineTransformation( + inline_elementals=False, inline_marked=True, adjust_imports=True + ) + trafo.apply(outer) + + # Check that the inlining has happened + assign = FindNodes(ir.Assignment).visit(outer.body) + assert len(assign) == 2 + assert assign[0].lhs == 'tmp(1:m)' + assert assign[0].rhs == 'x' + assert assign[1].lhs == 'a(i)' + assert assign[1].rhs == 'b(i) + sum(tmp)' + + # Now check that the right modules have been moved, + # and the import of the call has been removed + imports = FindNodes(ir.Import).visit(outer.spec) + assert len(imports) == 2 + assert imports[0].module == 'another_module' + assert imports[0].symbols == ('x',) + assert imports[1].module == 'bnds_module' + assert all(_ in imports[1].symbols for _ in ['l', 'm', 'n']) + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_transformation_intermediate(tmp_path, frontend): + fcode_outermost = """ +module outermost_mod +implicit none +contains +subroutine outermost() +use intermediate_mod, only: intermediate + +!$loki inline +call intermediate() + +end subroutine outermost +end module outermost_mod +""" + + fcode_intermediate = """ +module intermediate_mod +implicit none +contains +subroutine intermediate() +use innermost_mod, only: innermost + +call innermost() + +end subroutine intermediate +end module intermediate_mod +""" + + fcode_innermost = """ +module innermost_mod +implicit none +contains +subroutine innermost() + +end subroutine innermost +end module innermost_mod +""" + + (tmp_path/'outermost_mod.F90').write_text(fcode_outermost) + (tmp_path/'intermediate_mod.F90').write_text(fcode_intermediate) + (tmp_path/'innermost_mod.F90').write_text(fcode_innermost) + + config = { + 'default': { + 'mode': 'idem', + 'role': 'kernel', + 'expand': True, + 'strict': True + }, + 'routines': { + 'outermost': {'role': 'kernel'} + } + } + + scheduler = Scheduler( + paths=[tmp_path], config=SchedulerConfig.from_dict(config), + frontend=frontend, xmods=[tmp_path] + ) + + def _get_successors(item): + return scheduler.sgraph.successors(scheduler[item]) + + # check graph edges before transformation + assert len(scheduler.items) == 3 + assert len(_get_successors('outermost_mod#outermost')) == 1 + assert scheduler['intermediate_mod#intermediate'] in _get_successors('outermost_mod#outermost') + assert len(_get_successors('intermediate_mod#intermediate')) == 1 + assert scheduler['innermost_mod#innermost'] in _get_successors('intermediate_mod#intermediate') + + scheduler.process( transformation=InlineTransformation() ) + + # check graph edges were updated correctly + assert len(scheduler.items) == 2 + assert len(_get_successors('outermost_mod#outermost')) == 1 + assert scheduler['innermost_mod#innermost'] in _get_successors('outermost_mod#outermost') diff --git a/loki/transformations/inline/tests/test_procedures.py b/loki/transformations/inline/tests/test_procedures.py new file mode 100644 index 000000000..e09b2bf17 --- /dev/null +++ b/loki/transformations/inline/tests/test_procedures.py @@ -0,0 +1,953 @@ +# (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. + +import pytest +import numpy as np + +from loki import Module, Subroutine +from loki.build import jit_compile +from loki.expression import symbols as sym +from loki.frontend import available_frontends, OMNI +from loki.ir import ( + nodes as ir, FindNodes, FindVariables, FindInlineCalls +) +from loki.types import BasicType, DerivedType + +from loki.transformations.inline import ( + inline_member_procedures, inline_marked_subroutines +) +from loki.transformations.sanitise import ResolveAssociatesTransformer + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_member_routines(tmp_path, frontend): + """ + Test inlining of member subroutines. + """ + fcode = """ +subroutine member_routines(a, b) + real(kind=8), intent(inout) :: a(3), b(3) + integer :: i + + do i=1, size(a) + call add_one(a(i)) + end do + + call add_to_a(b) + + do i=1, size(a) + call add_one(a(i)) + end do + + contains + + subroutine add_one(a) + real(kind=8), intent(inout) :: a + a = a + 1 + end subroutine + + subroutine add_to_a(b) + real(kind=8), intent(inout) :: b(:) + integer :: n + + n = size(a) + do i = 1, n + a(i) = a(i) + b(i) + end do + end subroutine +end subroutine member_routines + """ + routine = Subroutine.from_source(fcode, frontend=frontend) + + filepath = tmp_path/(f'ref_transform_inline_member_routines_{frontend}.f90') + reference = jit_compile(routine, filepath=filepath, objname='member_routines') + + a = np.array([1., 2., 3.], order='F') + b = np.array([3., 3., 3.], order='F') + reference(a, b) + + assert (a == [6., 7., 8.]).all() + assert (b == [3., 3., 3.]).all() + + # Now inline the member routines and check again + inline_member_procedures(routine=routine) + + assert not routine.members + assert not FindNodes(ir.CallStatement).visit(routine.body) + assert len(FindNodes(ir.Loop).visit(routine.body)) == 3 + assert 'n' in routine.variables + + # An verify compiled behaviour + filepath = tmp_path/(f'transform_inline_member_routines_{frontend}.f90') + function = jit_compile(routine, filepath=filepath, objname='member_routines') + + a = np.array([1., 2., 3.], order='F') + b = np.array([3., 3., 3.], order='F') + function(a, b) + + assert (a == [6., 7., 8.]).all() + assert (b == [3., 3., 3.]).all() + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_member_functions(tmp_path, frontend): + """ + Test inlining of member subroutines. + """ + fcode = """ +subroutine member_functions(a, b, c) + implicit none + real(kind=8), intent(inout) :: a(3), b(3), c(3) + integer :: i + + do i=1, size(a) + a(i) = add_one(a(i)) + end do + + c = add_to_a(b, 3) + + do i=1, size(a) + a(i) = add_one(a(i)) + end do + + contains + + function add_one(a) + real(kind=8) :: a + real(kind=8) :: add_one + add_one = a + 1 + end function + + function add_to_a(b, n) + integer, intent(in) :: n + real(kind=8), intent(in) :: b(n) + real(kind=8) :: add_to_a(n) + + do i = 1, n + add_to_a(i) = a(i) + b(i) + end do + end function +end subroutine member_functions + """ + routine = Subroutine.from_source(fcode, frontend=frontend) + + filepath = tmp_path/(f'ref_transform_inline_member_functions_{frontend}.f90') + reference = jit_compile(routine, filepath=filepath, objname='member_functions') + + a = np.array([1., 2., 3.], order='F') + b = np.array([3., 3., 3.], order='F') + c = np.array([0., 0., 0.], order='F') + reference(a, b, c) + + assert (a == [3., 4., 5.]).all() + assert (b == [3., 3., 3.]).all() + assert (c == [5., 6., 7.]).all() + + # Now inline the member routines and check again + inline_member_procedures(routine=routine) + + assert not routine.members + assert not FindNodes(ir.CallStatement).visit(routine.body) + assert len(FindNodes(ir.Loop).visit(routine.body)) == 3 + + # An verify compiled behaviour + filepath = tmp_path/(f'transform_inline_member_functions_{frontend}.f90') + function = jit_compile(routine, filepath=filepath, objname='member_functions') + + a = np.array([1., 2., 3.], order='F') + b = np.array([3., 3., 3.], order='F') + c = np.array([0., 0., 0.], order='F') + function(a, b, c) + + assert (a == [3., 4., 5.]).all() + assert (b == [3., 3., 3.]).all() + assert (c == [5., 6., 7.]).all() + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_member_routines_arg_dimensions(frontend): + """ + Test inlining of member subroutines when sub-arrays of rank less + than the formal argument are passed. + """ + fcode = """ +subroutine member_routines_arg_dimensions(matrix, tensor) + real(kind=8), intent(inout) :: matrix(3, 3), tensor(3, 3, 4) + integer :: i + do i=1, 3 + call add_one(3, matrix(1:3,i), tensor(:,i,:)) + end do + contains + subroutine add_one(n, a, b) + integer, intent(in) :: n + real(kind=8), intent(inout) :: a(3), b(3,1:n) + integer :: j + do j=1, n + a(j) = a(j) + 1 + b(j,:) = 66.6 + end do + end subroutine +end subroutine member_routines_arg_dimensions + """ + routine = Subroutine.from_source(fcode, frontend=frontend) + + # Ensure initial member arguments + assert len(routine.routines) == 1 + assert routine.routines[0].name == 'add_one' + assert len(routine.routines[0].arguments) == 3 + assert routine.routines[0].arguments[0].name == 'n' + assert routine.routines[0].arguments[1].name == 'a' + assert routine.routines[0].arguments[2].name == 'b' + + # Now inline the member routines and check again + inline_member_procedures(routine=routine) + + # Ensure member has been inlined and arguments adapated + assert len(routine.routines) == 0 + assert len([v for v in FindVariables().visit(routine.body) if v.name == 'a']) == 0 + assert len([v for v in FindVariables().visit(routine.body) if v.name == 'b']) == 0 + assert len([v for v in FindVariables().visit(routine.spec) if v.name == 'a']) == 0 + assert len([v for v in FindVariables().visit(routine.spec) if v.name == 'b']) == 0 + assigns = FindNodes(ir.Assignment).visit(routine.body) + assert len(assigns) == 2 + assert assigns[0].lhs == 'matrix(j, i)' and assigns[0].rhs =='matrix(j, i) + 1' + assert assigns[1].lhs == 'tensor(j, i, :)' + + # Ensure the `n` in the inner loop bound has been substituted too + loops = FindNodes(ir.Loop).visit(routine.body) + assert len(loops) == 2 + assert loops[0].bounds == '1:3' + assert loops[1].bounds == '1:3' + + +@pytest.mark.parametrize('frontend', available_frontends(xfail=[(OMNI, 'No header information in test')])) +def test_inline_member_routines_derived_type_member(frontend): + """ + Test inlining of member subroutines when the member routine + handles arrays that are derived type components and thus might + have the DEFERRED type. + """ + fcode = """ +subroutine outer(x, a) + real, intent(inout) :: x + type(my_type), intent(in) :: a + + ! Pass derived type arrays as arguments + call inner(a%b(:), a%c, a%k, a%n) + +contains + subroutine inner(y, z, k, n) + integer, intent(in) :: k, n + real, intent(inout) :: y(n), z(:,:) + integer :: j + + do j=1, n + x = x + y(j) + ! Use derived-type variable as index + ! to test for nested substitution + y(j) = z(k,j) + end do + end subroutine inner +end subroutine outer + """ + routine = Subroutine.from_source(fcode, frontend=frontend) + + assert routine.variable_map['x'].type.dtype == BasicType.REAL + assert isinstance(routine.variable_map['a'].type.dtype, DerivedType) + call = FindNodes(ir.CallStatement).visit(routine.body)[0] + assert isinstance(call.arguments[0], sym.Array) + assert isinstance(call.arguments[1], sym.DeferredTypeSymbol) + assert isinstance(call.arguments[2], sym.DeferredTypeSymbol) + + # Now inline the member routines and check again + inline_member_procedures(routine=routine) + + assigns = FindNodes(ir.Assignment).visit(routine.body) + assert len(assigns) == 2 + assert assigns[0].rhs =='x + a%b(j)' + assert assigns[1].lhs == 'a%b(j)' and assigns[1].rhs == 'a%c(a%k, j)' + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_member_routines_variable_shadowing(frontend): + """ + Test inlining of member subroutines when variable allocations + in child routine shadow different allocations in the parent. + """ + fcode = """ +subroutine outer() + real :: x = 3 ! 'x' is real in outer. + real :: y + + y = 1.0 + call inner(y) + x = x + y + +contains + subroutine inner(y) + real, intent(inout) :: Y + real :: x(3) ! 'x' is array in inner. + x = [1, 2, 3] + y = y + sum(x) + end subroutine inner +end subroutine outer + """ + routine = Subroutine.from_source(fcode, frontend=frontend) + + # Check outer and inner 'x' + assert routine.variable_map['x'] == 'x' + assert isinstance(routine.variable_map['x'], sym.Scalar) + assert routine.variable_map['x'].type.initial == 3 + + assert routine['inner'].variable_map['x'] in ['x(3)', 'x(1:3)'] + assert isinstance(routine['inner'].variable_map['x'], sym.Array) + assert routine['inner'].variable_map['x'].type.shape == (3,) + + inline_member_procedures(routine=routine) + + # Check outer has not changed + assert routine.variable_map['x'] == 'x' + assert isinstance(routine.variable_map['x'], sym.Scalar) + assert routine.variable_map['x'].type.initial == 3 + + # Check inner 'x' was moved correctly + assert routine.variable_map['inner_x'] in ['inner_x(3)', 'inner_x(1:3)'] + assert isinstance(routine.variable_map['inner_x'], sym.Array) + assert routine.variable_map['inner_x'].type.shape == (3,) + + # Check inner 'y' was substituted, not renamed! + assign = FindNodes(ir.Assignment).visit(routine.body) + assert routine.variable_map['y'] == 'y' + assert assign[2].lhs == 'y' and assign[2].rhs == 'y + sum(inner_x)' + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_internal_routines_aliasing_declaration(frontend): + """ + Test declaration splitting when inlining internal procedures. + """ + fcode = """ +subroutine outer() + integer :: z + integer :: jlon + z = 0 + jlon = 0 + + call inner(z) + + jlon = z + 4 +contains + subroutine inner(z) + integer, intent(inout) :: z + integer :: jlon, jg ! These two need to get separated + jlon = 1 + jg = 2 + z = jlon + jg + end subroutine inner +end subroutine outer + """ + routine = Subroutine.from_source(fcode, frontend=frontend) + + # Check outer and inner variables + assert len(routine.variable_map) == 2 + assert 'z' in routine.variable_map + assert 'jlon' in routine.variable_map + + assert len(routine['inner'].variable_map) == 3 + assert 'z' in routine['inner'].variable_map + assert 'jlon' in routine['inner'].variable_map + assert 'jg' in routine['inner'].variable_map + + inline_member_procedures(routine, allowed_aliases=('jlon',)) + + assert len(routine.variable_map) == 3 + assert 'z' in routine.variable_map + assert 'jlon' in routine.variable_map + assert 'jg' in routine.variable_map + + assigns = FindNodes(ir.Assignment).visit(routine.body) + assert len(assigns) == 6 + assert assigns[2].lhs == 'jlon' and assigns[2].rhs == '1' + assert assigns[3].lhs == 'jg' and assigns[3].rhs == '2' + assert assigns[4].lhs == 'z' and assigns[4].rhs == 'jlon + jg' + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_member_routines_indexing_of_shadowed_array(frontend): + """ + Test special case of inlining of member subroutines when inlined routine contains + shadowed array and array indices. + In particular, this test checks that also the variables indexing + the array in the inlined result get renamed correctly. + """ + fcode = """ + subroutine outer(klon) + integer :: jg, jlon + integer :: arr(3, 3) + + jg = 70000 + call inner2() + + contains + + subroutine inner2() + integer :: jlon, jg + integer :: arr(3, 3) + do jg=1,3 + do jlon=1,3 + arr(jlon, jg) = 11 + end do + end do + end subroutine inner2 + + end subroutine outer + """ + routine = Subroutine.from_source(fcode, frontend=frontend) + inline_member_procedures(routine) + innerloop = FindNodes(ir.Loop).visit(routine.body)[1] + innerloopvars = FindVariables().visit(innerloop) + assert 'inner2_arr(inner2_jlon,inner2_jg)' in innerloopvars + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_member_routines_sequence_assoc(frontend): + """ + Test inlining of member subroutines in the presence of sequence + associations. As this is not supported, we check for the + appropriate error. + """ + fcode = """ +subroutine member_routines_sequence_assoc(vector) + real(kind=8), intent(inout) :: vector(6) + integer :: i + + i = 2 + call inner(3, vector(i)) + + contains + subroutine inner(n, a) + integer, intent(in) :: n + real(kind=8), intent(inout) :: a(3) + integer :: j + do j=1, n + a(j) = a(j) + 1 + end do + end subroutine +end subroutine member_routines_sequence_assoc + """ + routine = Subroutine.from_source(fcode, frontend=frontend) + + # Expect to fail tmp_path due to use of sequence association + with pytest.raises(RuntimeError): + inline_member_procedures(routine=routine) + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_member_routines_with_associate(frontend): + """ + Ensure that internal routines with :any:`Associate` constructs get + inlined as expected. + """ + fcode = """ +subroutine acraneb_transt(klon, klev, kidia, kfdia, ktdia) + implicit none + + integer(kind=4), intent(in) :: klon, klev, kidia, kfdia, ktdia + integer(kind=4) :: jlon, jlev + + real(kind=8) :: zq1(klon) + real(kind=8) :: zq2(klon, klev) + + call delta_t(zq1) + + do jlev = ktdia, klev + call delta_t(zq2(1:klon,jlev)) + + enddo + +contains + +subroutine delta_t(pq) + implicit none + + real(kind=8), intent(in) :: pq(klon) + real(kind=8) :: x, z + + associate(zz => z) + + do jlon = 1,klon + x = x + pq(jlon) + enddo + end associate +end subroutine + +end subroutine acraneb_transt + """ + + routine = Subroutine.from_source(fcode, frontend=frontend) + + inline_member_procedures(routine=routine) + + assert not routine.members + loops = FindNodes(ir.Loop).visit(routine.body) + assert len(loops) == 3 + + assigns = FindNodes(ir.Assignment).visit(routine.body) + assert len(assigns) == 2 + assert assigns[0].rhs == 'x + zq1(jlon)' + assert assigns[1].rhs == 'x + zq2(jlon, jlev)' + + assocs = FindNodes(ir.Associate).visit(routine.body) + assert len(assocs) == 2 + + +@pytest.mark.parametrize('frontend', available_frontends( + xfail=[(OMNI, 'OMNI does not handle missing type definitions')] +)) +def test_inline_member_routines_with_optionals(frontend): + """ + Ensure that internal routines with optional arguments get + inlined as expected (esp. present instrinsics are correctly + evaluated for all variables types) + """ + fcode = """ +subroutine test_inline(klon, ydxfu, ydmf_phys_out) + + use yomxfu , only : txfu + use mf_phys_type_mod , only : mf_phys_out_type + + implicit none + + integer(kind=4), intent(in) :: klon + type(txfu) ,intent(inout) :: ydxfu + type(mf_phys_out_type) ,intent(in) :: ydmf_phys_out + + call member_rout (ydxfu%visicld, pvmin=ydmf_phys_out%visicld, psmax=1.0_8) + + contains + + subroutine member_rout (x, pvmin, pvmax, psmin, psmax) + + real(kind=8) ,intent(inout) :: x(1:klon) + real(kind=8) ,intent(in) ,optional :: pvmin(1:klon) + real(kind=8) ,intent(in) ,optional :: pvmax(1:klon) + real(kind=8) ,intent(in) ,optional :: psmin + real(kind=8) ,intent(in) ,optional :: psmax + + if (present (psmin)) x = psmin + if (present (psmax)) x = psmax + if (present (pvmin)) x = minval(pvmin(:)) + if (present (pvmax)) x = maxval(pvmax(:)) + + end subroutine member_rout + +end subroutine test_inline + """ + + routine = Subroutine.from_source(fcode, frontend=frontend) + + inline_member_procedures(routine=routine) + + assert not routine.members + + conds = FindNodes(ir.Conditional).visit(routine.body) + assert len(conds) == 4 + assert conds[0].condition == 'False' + assert conds[1].condition == 'True' + assert conds[2].condition == 'True' + assert conds[3].condition == 'False' + + +@pytest.mark.parametrize('frontend', available_frontends()) +@pytest.mark.parametrize('adjust_imports', [True, False]) +def test_inline_marked_subroutines(frontend, adjust_imports, tmp_path): + """ Test subroutine inlining via marker pragmas. """ + + fcode_driver = """ +subroutine test_pragma_inline(a, b) + use util_mod, only: add_one, add_a_to_b + implicit none + + real(kind=8), intent(inout) :: a(3), b(3) + integer, parameter :: n = 3 + integer :: i + + do i=1, n + !$loki inline + call add_one(a(i)) + end do + + !$loki inline + call add_a_to_b(a(:), b(:), 3) + + do i=1, n + call add_one(b(i)) + end do + +end subroutine test_pragma_inline + """ + + fcode_module = """ +module util_mod +implicit none + +contains + subroutine add_one(a) + interface + subroutine do_something() + end subroutine do_something + end interface + real(kind=8), intent(inout) :: a + a = a + 1 + end subroutine add_one + + subroutine add_a_to_b(a, b, n) + interface + subroutine do_something_else() + end subroutine do_something_else + end interface + real(kind=8), intent(inout) :: a(:), b(:) + integer, intent(in) :: n + integer :: i + + do i = 1, n + a(i) = a(i) + b(i) + end do + end subroutine add_a_to_b +end module util_mod +""" + module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) + driver = Subroutine.from_source(fcode_driver, frontend=frontend, xmods=[tmp_path]) + driver.enrich(module) + + calls = FindNodes(ir.CallStatement).visit(driver.body) + assert calls[0].routine == module['add_one'] + assert calls[1].routine == module['add_a_to_b'] + assert calls[2].routine == module['add_one'] + + inline_marked_subroutines( + routine=driver, allowed_aliases=('I',), adjust_imports=adjust_imports + ) + + # Check inlined loops and assignments + assert len(FindNodes(ir.Loop).visit(driver.body)) == 3 + assign = FindNodes(ir.Assignment).visit(driver.body) + assert len(assign) == 2 + assert assign[0].lhs == 'a(i)' and assign[0].rhs == 'a(i) + 1' + assert assign[1].lhs == 'a(i)' and assign[1].rhs == 'a(i) + b(i)' + + # Check that the last call is left untouched + calls = FindNodes(ir.CallStatement).visit(driver.body) + assert len(calls) == 1 + assert calls[0].routine.name == 'add_one' + assert calls[0].arguments == ('b(i)',) + + imports = FindNodes(ir.Import).visit(driver.spec) + assert len(imports) == 1 + if adjust_imports: + assert imports[0].symbols == ('add_one',) + else: + assert imports[0].symbols == ('add_one', 'add_a_to_b') + + if adjust_imports: + # check that explicit interfaces were imported + intfs = driver.interfaces + assert len(intfs) == 1 + assert all(isinstance(s, sym.ProcedureSymbol) for s in driver.interface_symbols) + assert 'do_something' in driver.interface_symbols + assert 'do_something_else' in driver.interface_symbols + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_marked_subroutines_with_interfaces(frontend, tmp_path): + """ Test inlining of subroutines with explicit interfaces via marker pragmas. """ + + fcode_driver = """ +subroutine test_pragma_inline(a, b) + implicit none + + interface + subroutine add_a_to_b(a, b, n) + real(kind=8), intent(inout) :: a(:), b(:) + integer, intent(in) :: n + end subroutine add_a_to_b + subroutine add_one(a) + real(kind=8), intent(inout) :: a + end subroutine add_one + end interface + + interface + subroutine add_two(a) + real(kind=8), intent(inout) :: a + end subroutine add_two + end interface + + real(kind=8), intent(inout) :: a(3), b(3) + integer, parameter :: n = 3 + integer :: i + + do i=1, n + !$loki inline + call add_one(a(i)) + end do + + !$loki inline + call add_a_to_b(a(:), b(:), 3) + + do i=1, n + call add_one(b(i)) + !$loki inline + call add_two(b(i)) + end do + +end subroutine test_pragma_inline + """ + + fcode_module = """ +module util_mod +implicit none + +contains + subroutine add_one(a) + real(kind=8), intent(inout) :: a + a = a + 1 + end subroutine add_one + + subroutine add_two(a) + real(kind=8), intent(inout) :: a + a = a + 2 + end subroutine add_two + + subroutine add_a_to_b(a, b, n) + real(kind=8), intent(inout) :: a(:), b(:) + integer, intent(in) :: n + integer :: i + + do i = 1, n + a(i) = a(i) + b(i) + end do + end subroutine add_a_to_b +end module util_mod +""" + + module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) + driver = Subroutine.from_source(fcode_driver, frontend=frontend, xmods=[tmp_path]) + driver.enrich(module.subroutines) + + calls = FindNodes(ir.CallStatement).visit(driver.body) + assert calls[0].routine == module['add_one'] + assert calls[1].routine == module['add_a_to_b'] + assert calls[2].routine == module['add_one'] + assert calls[3].routine == module['add_two'] + + inline_marked_subroutines(routine=driver, allowed_aliases=('I',)) + + # Check inlined loops and assignments + assert len(FindNodes(ir.Loop).visit(driver.body)) == 3 + assign = FindNodes(ir.Assignment).visit(driver.body) + assert len(assign) == 3 + assert assign[0].lhs == 'a(i)' and assign[0].rhs == 'a(i) + 1' + assert assign[1].lhs == 'a(i)' and assign[1].rhs == 'a(i) + b(i)' + assert assign[2].lhs == 'b(i)' and assign[2].rhs == 'b(i) + 2' + + # Check that the last call is left untouched + calls = FindNodes(ir.CallStatement).visit(driver.body) + assert len(calls) == 1 + assert calls[0].routine.name == 'add_one' + assert calls[0].arguments == ('b(i)',) + + intfs = FindNodes(ir.Interface).visit(driver.spec) + assert len(intfs) == 1 + assert intfs[0].symbols == ('add_one',) + + +@pytest.mark.parametrize('frontend', available_frontends()) +@pytest.mark.parametrize('adjust_imports', [True, False]) +def test_inline_marked_routine_with_optionals(frontend, adjust_imports, tmp_path): + """ Test subroutine inlining via marker pragmas with omitted optionals. """ + + fcode_driver = """ +subroutine test_pragma_inline_optionals(a, b) + use util_mod, only: add_one + implicit none + + real(kind=8), intent(inout) :: a(3), b(3) + integer, parameter :: n = 3 + integer :: i + + do i=1, n + !$loki inline + call add_one(a(i), two=2.0) + end do + + do i=1, n + !$loki inline + call add_one(b(i)) + end do + +end subroutine test_pragma_inline_optionals + """ + + fcode_module = """ +module util_mod +implicit none + +contains + subroutine add_one(a, two) + real(kind=8), intent(inout) :: a + real(kind=8), optional, intent(inout) :: two + a = a + 1 + + if (present(two)) then + a = a + two + end if + end subroutine add_one +end module util_mod +""" + module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) + driver = Subroutine.from_source(fcode_driver, frontend=frontend, xmods=[tmp_path]) + driver.enrich(module) + + calls = FindNodes(ir.CallStatement).visit(driver.body) + assert calls[0].routine == module['add_one'] + assert calls[1].routine == module['add_one'] + + inline_marked_subroutines(routine=driver, adjust_imports=adjust_imports) + + # Check inlined loops and assignments + assert len(FindNodes(ir.Loop).visit(driver.body)) == 2 + assign = FindNodes(ir.Assignment).visit(driver.body) + assert len(assign) == 4 + assert assign[0].lhs == 'a(i)' and assign[0].rhs == 'a(i) + 1' + assert assign[1].lhs == 'a(i)' and assign[1].rhs == 'a(i) + 2.0' + assert assign[2].lhs == 'b(i)' and assign[2].rhs == 'b(i) + 1' + # TODO: This is a problem, since it's not declared anymore + assert assign[3].lhs == 'b(i)' and assign[3].rhs == 'b(i) + two' + + # Check that the PRESENT checks have been resolved + assert len(FindNodes(ir.CallStatement).visit(driver.body)) == 0 + assert len(FindInlineCalls().visit(driver.body)) == 0 + checks = FindNodes(ir.Conditional).visit(driver.body) + assert len(checks) == 2 + assert checks[0].condition == 'True' + assert checks[1].condition == 'False' + + imports = FindNodes(ir.Import).visit(driver.spec) + assert len(imports) == 0 if adjust_imports else 1 + + +@pytest.mark.parametrize('frontend', available_frontends( + xfail=[(OMNI, 'OMNI has no sense of humour!')]) +) +def test_inline_marked_subroutines_with_associates(frontend): + """ Test subroutine inlining via marker pragmas with nested associates. """ + + fcode_outer = """ +subroutine test_pragma_inline_associates(never) + use peter_pan, only: neverland + implicit none + type(neverland), intent(inout) :: never + + associate(going=>never%going_to) + + associate(up=>give_you%up) + + !$loki inline + call dave(going, up) + + end associate + + end associate +end subroutine test_pragma_inline_associates + """ + + fcode_inner = """ +subroutine dave(going) + use your_imagination, only: astley + implicit none + type(astley), intent(inout) :: going + + associate(give_you=>going%give_you) + + associate(up=>give_you%up) + + call rick_is(up) + + end associate + + end associate +end subroutine dave + """ + + outer = Subroutine.from_source(fcode_outer, frontend=frontend) + inner = Subroutine.from_source(fcode_inner, frontend=frontend) + outer.enrich(inner) + + assert FindNodes(ir.CallStatement).visit(outer.body)[0].routine == inner + + inline_marked_subroutines(routine=outer, adjust_imports=True) + + # Ensure that all associates are perfectly nested afterwards + assocs = FindNodes(ir.Associate).visit(outer.body) + assert len(assocs) == 4 + assert assocs[1].parent == assocs[0] + assert assocs[2].parent == assocs[1] + assert assocs[3].parent == assocs[2] + + # And, because we can... + outer.body = ResolveAssociatesTransformer().visit(outer.body) + call = FindNodes(ir.CallStatement).visit(outer.body)[0] + assert call.name == 'rick_is' + assert call.arguments == ('never%going_to%give_you%up',) + # Q. E. D. + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_inline_marked_subroutines_declarations(frontend, tmp_path): + """Test symbol propagation to hoisted declaration when inlining.""" + fcode = """ +module inline_declarations + implicit none + + type bounds + integer :: start, end + end type bounds + + contains + + subroutine outer(a, bnds) + real(kind=8), intent(inout) :: a(bnds%end) + type(bounds), intent(in) :: bnds + real(kind=8) :: b(bnds%end) + + b(bnds%start:bnds%end) = a(bnds%start:bnds%end) + 42.0 + + !$loki inline + call inner(a, dims=bnds) + end subroutine outer + + subroutine inner(c, dims) + real(kind=8), intent(inout) :: c(dims%end) + type(bounds), intent(in) :: dims + real(kind=8) :: d(dims%end) + + d(dims%start:dims%end) = c(dims%start:dims%end) - 66.6 + c(dims%start) = sum(d) + end subroutine inner +end module inline_declarations +""" + module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) + outer = module['outer'] + + inline_marked_subroutines(routine=outer, adjust_imports=True) + + # Check that all declarations are using the ``bnds`` symbol + assert outer.symbols[0] == 'a(bnds%end)' + assert outer.symbols[2] == 'b(bnds%end)' + assert outer.symbols[3] == 'd(bnds%end)' + assert all( + a.shape == ('bnds%end',) for a in outer.symbols if isinstance(a, sym.Array) + ) diff --git a/loki/transformations/tests/test_inline.py b/loki/transformations/tests/test_inline.py deleted file mode 100644 index 4bf37cfd1..000000000 --- a/loki/transformations/tests/test_inline.py +++ /dev/null @@ -1,1838 +0,0 @@ -# (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. - -import pytest -import numpy as np - -from loki import ( - Module, Subroutine, FindVariables, BasicType, DerivedType, - FindInlineCalls -) -from loki.build import jit_compile, jit_compile_lib, Builder, Obj -from loki.expression import symbols as sym -from loki.frontend import available_frontends, OMNI, OFP -from loki.ir import nodes as ir, FindNodes -from loki.batch import Scheduler, SchedulerConfig - -from loki.transformations.inline import ( - inline_elemental_functions, inline_constant_parameters, - inline_member_procedures, inline_marked_subroutines, - inline_statement_functions, InlineTransformation, -) -from loki.transformations.sanitise import ResolveAssociatesTransformer -from loki.transformations.utilities import replace_selected_kind - -# pylint: disable=too-many-lines - - -@pytest.fixture(name='builder') -def fixture_builder(tmp_path): - yield Builder(source_dirs=tmp_path, build_dir=tmp_path) - Obj.clear_cache() - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_transform_inline_elemental_functions(tmp_path, builder, frontend): - """ - Test correct inlining of elemental functions. - """ - fcode_module = """ -module multiply_mod - use iso_fortran_env, only: real64 - implicit none -contains - - elemental function multiply(a, b) - real(kind=real64) :: multiply - real(kind=real64), intent(in) :: a, b - real(kind=real64) :: temp - !$loki routine seq - - ! simulate multi-line function - temp = a * b - multiply = temp - end function multiply -end module multiply_mod -""" - - fcode = """ -subroutine transform_inline_elemental_functions(v1, v2, v3) - use iso_fortran_env, only: real64 - use multiply_mod, only: multiply - real(kind=real64), intent(in) :: v1 - real(kind=real64), intent(out) :: v2, v3 - - v2 = multiply(v1, 6._real64) - v3 = 600. + multiply(6._real64, 11._real64) -end subroutine transform_inline_elemental_functions -""" - - # Generate reference code, compile run and verify - module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) - routine = Subroutine.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - - refname = f'ref_{routine.name}_{frontend}' - reference = jit_compile_lib([module, routine], path=tmp_path, name=refname, builder=builder) - - v2, v3 = reference.transform_inline_elemental_functions(11.) - assert v2 == 66. - assert v3 == 666. - - (tmp_path/f'{module.name}.f90').unlink() - (tmp_path/f'{routine.name}.f90').unlink() - - # Now inline elemental functions - routine = Subroutine.from_source(fcode, definitions=module, frontend=frontend, xmods=[tmp_path]) - inline_elemental_functions(routine) - - # Make sure there are no more inline calls in the routine body - assert not FindInlineCalls().visit(routine.body) - - # Verify correct scope of inlined elements - assert all(v.scope is routine for v in FindVariables().visit(routine.body)) - - # Ensure the !$loki routine pragma has been removed - assert not FindNodes(ir.Pragma).visit(routine.body) - - # Hack: rename routine to use a different filename in the build - routine.name = f'{routine.name}_' - kernel = jit_compile_lib([routine], path=tmp_path, name=routine.name, builder=builder) - - v2, v3 = kernel.transform_inline_elemental_functions_(11.) - assert v2 == 66. - assert v3 == 666. - - builder.clean() - (tmp_path/f'{routine.name}.f90').unlink() - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_transform_inline_elemental_functions_extended(tmp_path, builder, frontend): - """ - Test correct inlining of elemental functions. - """ - fcode_module = """ -module multiply_extended_mod - use iso_fortran_env, only: real64 - implicit none -contains - - elemental function multiply(a, b) ! result (ret_mult) - ! real(kind=real64) :: ret_mult - real(kind=real64) :: multiply - real(kind=real64), intent(in) :: a, b - real(kind=real64) :: temp - - ! simulate multi-line function - temp = a * b - multiply = temp - ! ret_mult = temp - end function multiply - - elemental function multiply_single_line(a, b) - real(kind=real64) :: multiply_single_line - real(kind=real64), intent(in) :: a, b - real(kind=real64) :: temp - - multiply_single_line = a * b - end function multiply_single_line - - elemental function add(a, b) - real(kind=real64) :: add - real(kind=real64), intent(in) :: a, b - real(kind=real64) :: temp - - ! simulate multi-line function - temp = a + b - add = temp - end function add -end module multiply_extended_mod -""" - - fcode = """ -subroutine transform_inline_elemental_functions_extended(v1, v2, v3) - use iso_fortran_env, only: real64 - use multiply_extended_mod, only: multiply, multiply_single_line, add - real(kind=real64), intent(in) :: v1 - real(kind=real64), intent(out) :: v2, v3 - real(kind=real64), parameter :: param1 = 100. - - v2 = multiply(v1, 6._real64) + multiply_single_line(v1, 3._real64) - v3 = add(param1, 200._real64) + add(150._real64, 150._real64) + multiply(6._real64, 11._real64) -end subroutine transform_inline_elemental_functions_extended -""" - - # Generate reference code, compile run and verify - module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) - routine = Subroutine.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - - refname = f'ref_{routine.name}_{frontend}' - reference = jit_compile_lib([module, routine], path=tmp_path, name=refname, builder=builder) - - v2, v3 = reference.transform_inline_elemental_functions_extended(11.) - assert v2 == 99. - assert v3 == 666. - - (tmp_path/f'{module.name}.f90').unlink() - (tmp_path/f'{routine.name}.f90').unlink() - - # Now inline elemental functions - routine = Subroutine.from_source(fcode, definitions=module, frontend=frontend, xmods=[tmp_path]) - inline_elemental_functions(routine) - - - # Make sure there are no more inline calls in the routine body - assert not FindInlineCalls().visit(routine.body) - - # Verify correct scope of inlined elements - assert all(v.scope is routine for v in FindVariables().visit(routine.body)) - - # Hack: rename routine to use a different filename in the build - routine.name = f'{routine.name}_' - kernel = jit_compile_lib([routine], path=tmp_path, name=routine.name, builder=builder) - - v2, v3 = kernel.transform_inline_elemental_functions_extended_(11.) - assert v2 == 99. - assert v3 == 666. - - builder.clean() - (tmp_path/f'{routine.name}.f90').unlink() - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_transform_inline_constant_parameters(tmp_path, builder, frontend): - """ - Test correct inlining of constant parameters. - """ - fcode_module = """ -module parameters_mod - implicit none - integer, parameter :: a = 1 - integer, parameter :: b = -1 -contains - subroutine dummy - end subroutine dummy -end module parameters_mod -""" - - fcode = """ -module inline_const_param_mod - ! TODO: use parameters_mod, only: b - implicit none - integer, parameter :: c = 1+1 -contains - subroutine inline_const_param(v1, v2, v3) - use parameters_mod, only: a, b - integer, intent(in) :: v1 - integer, intent(out) :: v2, v3 - - v2 = v1 + b - a - v3 = c - end subroutine inline_const_param -end module inline_const_param_mod -""" - # Generate reference code, compile run and verify - param_module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) - module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - refname = f'ref_{module.name}_{ frontend}' - reference = jit_compile_lib([module, param_module], path=tmp_path, name=refname, builder=builder) - - v2, v3 = reference.inline_const_param_mod.inline_const_param(10) - assert v2 == 8 - assert v3 == 2 - (tmp_path/f'{module.name}.f90').unlink() - (tmp_path/f'{param_module.name}.f90').unlink() - - # Now transform with supplied elementals but without module - module = Module.from_source(fcode, definitions=param_module, frontend=frontend, xmods=[tmp_path]) - assert len(FindNodes(ir.Import).visit(module['inline_const_param'].spec)) == 1 - for routine in module.subroutines: - inline_constant_parameters(routine, external_only=True) - assert not FindNodes(ir.Import).visit(module['inline_const_param'].spec) - - # Hack: rename module to use a different filename in the build - module.name = f'{module.name}_' - obj = jit_compile_lib([module], path=tmp_path, name=f'{module.name}_{frontend}', builder=builder) - - v2, v3 = obj.inline_const_param_mod_.inline_const_param(10) - assert v2 == 8 - assert v3 == 2 - - (tmp_path/f'{module.name}.f90').unlink() - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_transform_inline_constant_parameters_kind(tmp_path, builder, frontend): - """ - Test correct inlining of constant parameters for kind symbols. - """ - fcode_module = """ -module kind_parameters_mod - implicit none - integer, parameter :: jprb = selected_real_kind(13, 300) -end module kind_parameters_mod -""" - - fcode = """ -module inline_const_param_kind_mod - implicit none -contains - subroutine inline_const_param_kind(v1) - use kind_parameters_mod, only: jprb - real(kind=jprb), intent(out) :: v1 - - v1 = real(2, kind=jprb) + 3. - end subroutine inline_const_param_kind -end module inline_const_param_kind_mod -""" - # Generate reference code, compile run and verify - param_module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) - module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - refname = f'ref_{module.name}_{frontend}' - reference = jit_compile_lib([module, param_module], path=tmp_path, name=refname, builder=builder) - - v1 = reference.inline_const_param_kind_mod.inline_const_param_kind() - assert v1 == 5. - (tmp_path/f'{module.name}.f90').unlink() - (tmp_path/f'{param_module.name}.f90').unlink() - - # Now transform with supplied elementals but without module - module = Module.from_source(fcode, definitions=param_module, frontend=frontend, xmods=[tmp_path]) - assert len(FindNodes(ir.Import).visit(module['inline_const_param_kind'].spec)) == 1 - for routine in module.subroutines: - inline_constant_parameters(routine, external_only=True) - assert not FindNodes(ir.Import).visit(module['inline_const_param_kind'].spec) - - # Hack: rename module to use a different filename in the build - module.name = f'{module.name}_' - obj = jit_compile_lib([module], path=tmp_path, name=f'{module.name}_{frontend}', builder=builder) - - v1 = obj.inline_const_param_kind_mod_.inline_const_param_kind() - assert v1 == 5. - - (tmp_path/f'{module.name}.f90').unlink() - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_transform_inline_constant_parameters_replace_kind(tmp_path, builder, frontend): - """ - Test correct inlining of constant parameters for kind symbols. - """ - fcode_module = """ -module replace_kind_parameters_mod - implicit none - integer, parameter :: jprb = selected_real_kind(13, 300) -end module replace_kind_parameters_mod -""" - - fcode = """ -module inline_param_repl_kind_mod - implicit none -contains - subroutine inline_param_repl_kind(v1) - use replace_kind_parameters_mod, only: jprb - real(kind=jprb), intent(out) :: v1 - real(kind=jprb) :: a = 3._JPRB - - v1 = 1._jprb + real(2, kind=jprb) + a - end subroutine inline_param_repl_kind -end module inline_param_repl_kind_mod -""" - # Generate reference code, compile run and verify - param_module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) - module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - refname = f'ref_{module.name}_{frontend}' - reference = jit_compile_lib([module, param_module], path=tmp_path, name=refname, builder=builder) - func = getattr(getattr(reference, module.name), module.subroutines[0].name) - - v1 = func() - assert v1 == 6. - (tmp_path/f'{module.name}.f90').unlink() - (tmp_path/f'{param_module.name}.f90').unlink() - - # Now transform with supplied elementals but without module - module = Module.from_source(fcode, definitions=param_module, frontend=frontend, xmods=[tmp_path]) - imports = FindNodes(ir.Import).visit(module.subroutines[0].spec) - assert len(imports) == 1 and imports[0].module.lower() == param_module.name.lower() - for routine in module.subroutines: - inline_constant_parameters(routine, external_only=True) - replace_selected_kind(routine) - imports = FindNodes(ir.Import).visit(module.subroutines[0].spec) - assert len(imports) == 1 and imports[0].module.lower() == 'iso_fortran_env' - - # Hack: rename module to use a different filename in the build - module.name = f'{module.name}_' - obj = jit_compile_lib([module], path=tmp_path, name=f'{module.name}_{frontend}', builder=builder) - - func = getattr(getattr(obj, module.name), module.subroutines[0].name) - v1 = func() - assert v1 == 6. - - (tmp_path/f'{module.name}.f90').unlink() - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_constant_replacement_internal(frontend): - """ - Test constant replacement for internally defined constants. - """ - fcode = """ -subroutine kernel(a, b) - integer, parameter :: par = 10 - integer, intent(inout) :: a, b - - a = b + par -end subroutine kernel - """ - routine = Subroutine.from_source(fcode, frontend=frontend) - inline_constant_parameters(routine=routine, external_only=False) - - assert len(routine.variables) == 2 - assert 'a' in routine.variables and 'b' in routine.variables - - stmts = FindNodes(ir.Assignment).visit(routine.body) - assert len(stmts) == 1 - assert stmts[0].rhs == 'b + 10' - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_member_routines(tmp_path, frontend): - """ - Test inlining of member subroutines. - """ - fcode = """ -subroutine member_routines(a, b) - real(kind=8), intent(inout) :: a(3), b(3) - integer :: i - - do i=1, size(a) - call add_one(a(i)) - end do - - call add_to_a(b) - - do i=1, size(a) - call add_one(a(i)) - end do - - contains - - subroutine add_one(a) - real(kind=8), intent(inout) :: a - a = a + 1 - end subroutine - - subroutine add_to_a(b) - real(kind=8), intent(inout) :: b(:) - integer :: n - - n = size(a) - do i = 1, n - a(i) = a(i) + b(i) - end do - end subroutine -end subroutine member_routines - """ - routine = Subroutine.from_source(fcode, frontend=frontend) - - filepath = tmp_path/(f'ref_transform_inline_member_routines_{frontend}.f90') - reference = jit_compile(routine, filepath=filepath, objname='member_routines') - - a = np.array([1., 2., 3.], order='F') - b = np.array([3., 3., 3.], order='F') - reference(a, b) - - assert (a == [6., 7., 8.]).all() - assert (b == [3., 3., 3.]).all() - - # Now inline the member routines and check again - inline_member_procedures(routine=routine) - - assert not routine.members - assert not FindNodes(ir.CallStatement).visit(routine.body) - assert len(FindNodes(ir.Loop).visit(routine.body)) == 3 - assert 'n' in routine.variables - - # An verify compiled behaviour - filepath = tmp_path/(f'transform_inline_member_routines_{frontend}.f90') - function = jit_compile(routine, filepath=filepath, objname='member_routines') - - a = np.array([1., 2., 3.], order='F') - b = np.array([3., 3., 3.], order='F') - function(a, b) - - assert (a == [6., 7., 8.]).all() - assert (b == [3., 3., 3.]).all() - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_member_functions(tmp_path, frontend): - """ - Test inlining of member subroutines. - """ - fcode = """ -subroutine member_functions(a, b, c) - implicit none - real(kind=8), intent(inout) :: a(3), b(3), c(3) - integer :: i - - do i=1, size(a) - a(i) = add_one(a(i)) - end do - - c = add_to_a(b, 3) - - do i=1, size(a) - a(i) = add_one(a(i)) - end do - - contains - - function add_one(a) - real(kind=8) :: a - real(kind=8) :: add_one - add_one = a + 1 - end function - - function add_to_a(b, n) - integer, intent(in) :: n - real(kind=8), intent(in) :: b(n) - real(kind=8) :: add_to_a(n) - - do i = 1, n - add_to_a(i) = a(i) + b(i) - end do - end function -end subroutine member_functions - """ - routine = Subroutine.from_source(fcode, frontend=frontend) - - filepath = tmp_path/(f'ref_transform_inline_member_functions_{frontend}.f90') - reference = jit_compile(routine, filepath=filepath, objname='member_functions') - - a = np.array([1., 2., 3.], order='F') - b = np.array([3., 3., 3.], order='F') - c = np.array([0., 0., 0.], order='F') - reference(a, b, c) - - assert (a == [3., 4., 5.]).all() - assert (b == [3., 3., 3.]).all() - assert (c == [5., 6., 7.]).all() - - # Now inline the member routines and check again - inline_member_procedures(routine=routine) - - assert not routine.members - assert not FindNodes(ir.CallStatement).visit(routine.body) - assert len(FindNodes(ir.Loop).visit(routine.body)) == 3 - - # An verify compiled behaviour - filepath = tmp_path/(f'transform_inline_member_functions_{frontend}.f90') - function = jit_compile(routine, filepath=filepath, objname='member_functions') - - a = np.array([1., 2., 3.], order='F') - b = np.array([3., 3., 3.], order='F') - c = np.array([0., 0., 0.], order='F') - function(a, b, c) - - assert (a == [3., 4., 5.]).all() - assert (b == [3., 3., 3.]).all() - assert (c == [5., 6., 7.]).all() - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_member_routines_arg_dimensions(frontend): - """ - Test inlining of member subroutines when sub-arrays of rank less - than the formal argument are passed. - """ - fcode = """ -subroutine member_routines_arg_dimensions(matrix, tensor) - real(kind=8), intent(inout) :: matrix(3, 3), tensor(3, 3, 4) - integer :: i - do i=1, 3 - call add_one(3, matrix(1:3,i), tensor(:,i,:)) - end do - contains - subroutine add_one(n, a, b) - integer, intent(in) :: n - real(kind=8), intent(inout) :: a(3), b(3,1:n) - integer :: j - do j=1, n - a(j) = a(j) + 1 - b(j,:) = 66.6 - end do - end subroutine -end subroutine member_routines_arg_dimensions - """ - routine = Subroutine.from_source(fcode, frontend=frontend) - - # Ensure initial member arguments - assert len(routine.routines) == 1 - assert routine.routines[0].name == 'add_one' - assert len(routine.routines[0].arguments) == 3 - assert routine.routines[0].arguments[0].name == 'n' - assert routine.routines[0].arguments[1].name == 'a' - assert routine.routines[0].arguments[2].name == 'b' - - # Now inline the member routines and check again - inline_member_procedures(routine=routine) - - # Ensure member has been inlined and arguments adapated - assert len(routine.routines) == 0 - assert len([v for v in FindVariables().visit(routine.body) if v.name == 'a']) == 0 - assert len([v for v in FindVariables().visit(routine.body) if v.name == 'b']) == 0 - assert len([v for v in FindVariables().visit(routine.spec) if v.name == 'a']) == 0 - assert len([v for v in FindVariables().visit(routine.spec) if v.name == 'b']) == 0 - assigns = FindNodes(ir.Assignment).visit(routine.body) - assert len(assigns) == 2 - assert assigns[0].lhs == 'matrix(j, i)' and assigns[0].rhs =='matrix(j, i) + 1' - assert assigns[1].lhs == 'tensor(j, i, :)' - - # Ensure the `n` in the inner loop bound has been substituted too - loops = FindNodes(ir.Loop).visit(routine.body) - assert len(loops) == 2 - assert loops[0].bounds == '1:3' - assert loops[1].bounds == '1:3' - - -@pytest.mark.parametrize('frontend', available_frontends(xfail=[(OMNI, 'No header information in test')])) -def test_inline_member_routines_derived_type_member(frontend): - """ - Test inlining of member subroutines when the member routine - handles arrays that are derived type components and thus might - have the DEFERRED type. - """ - fcode = """ -subroutine outer(x, a) - real, intent(inout) :: x - type(my_type), intent(in) :: a - - ! Pass derived type arrays as arguments - call inner(a%b(:), a%c, a%k, a%n) - -contains - subroutine inner(y, z, k, n) - integer, intent(in) :: k, n - real, intent(inout) :: y(n), z(:,:) - integer :: j - - do j=1, n - x = x + y(j) - ! Use derived-type variable as index - ! to test for nested substitution - y(j) = z(k,j) - end do - end subroutine inner -end subroutine outer - """ - routine = Subroutine.from_source(fcode, frontend=frontend) - - assert routine.variable_map['x'].type.dtype == BasicType.REAL - assert isinstance(routine.variable_map['a'].type.dtype, DerivedType) - call = FindNodes(ir.CallStatement).visit(routine.body)[0] - assert isinstance(call.arguments[0], sym.Array) - assert isinstance(call.arguments[1], sym.DeferredTypeSymbol) - assert isinstance(call.arguments[2], sym.DeferredTypeSymbol) - - # Now inline the member routines and check again - inline_member_procedures(routine=routine) - - assigns = FindNodes(ir.Assignment).visit(routine.body) - assert len(assigns) == 2 - assert assigns[0].rhs =='x + a%b(j)' - assert assigns[1].lhs == 'a%b(j)' and assigns[1].rhs == 'a%c(a%k, j)' - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_member_routines_variable_shadowing(frontend): - """ - Test inlining of member subroutines when variable allocations - in child routine shadow different allocations in the parent. - """ - fcode = """ -subroutine outer() - real :: x = 3 ! 'x' is real in outer. - real :: y - - y = 1.0 - call inner(y) - x = x + y - -contains - subroutine inner(y) - real, intent(inout) :: Y - real :: x(3) ! 'x' is array in inner. - x = [1, 2, 3] - y = y + sum(x) - end subroutine inner -end subroutine outer - """ - routine = Subroutine.from_source(fcode, frontend=frontend) - - # Check outer and inner 'x' - assert routine.variable_map['x'] == 'x' - assert isinstance(routine.variable_map['x'], sym.Scalar) - assert routine.variable_map['x'].type.initial == 3 - - assert routine['inner'].variable_map['x'] in ['x(3)', 'x(1:3)'] - assert isinstance(routine['inner'].variable_map['x'], sym.Array) - assert routine['inner'].variable_map['x'].type.shape == (3,) - - inline_member_procedures(routine=routine) - - # Check outer has not changed - assert routine.variable_map['x'] == 'x' - assert isinstance(routine.variable_map['x'], sym.Scalar) - assert routine.variable_map['x'].type.initial == 3 - - # Check inner 'x' was moved correctly - assert routine.variable_map['inner_x'] in ['inner_x(3)', 'inner_x(1:3)'] - assert isinstance(routine.variable_map['inner_x'], sym.Array) - assert routine.variable_map['inner_x'].type.shape == (3,) - - # Check inner 'y' was substituted, not renamed! - assign = FindNodes(ir.Assignment).visit(routine.body) - assert routine.variable_map['y'] == 'y' - assert assign[2].lhs == 'y' and assign[2].rhs == 'y + sum(inner_x)' - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_internal_routines_aliasing_declaration(frontend): - """ - Test declaration splitting when inlining internal procedures. - """ - fcode = """ -subroutine outer() - integer :: z - integer :: jlon - z = 0 - jlon = 0 - - call inner(z) - - jlon = z + 4 -contains - subroutine inner(z) - integer, intent(inout) :: z - integer :: jlon, jg ! These two need to get separated - jlon = 1 - jg = 2 - z = jlon + jg - end subroutine inner -end subroutine outer - """ - routine = Subroutine.from_source(fcode, frontend=frontend) - - # Check outer and inner variables - assert len(routine.variable_map) == 2 - assert 'z' in routine.variable_map - assert 'jlon' in routine.variable_map - - assert len(routine['inner'].variable_map) == 3 - assert 'z' in routine['inner'].variable_map - assert 'jlon' in routine['inner'].variable_map - assert 'jg' in routine['inner'].variable_map - - inline_member_procedures(routine, allowed_aliases=('jlon',)) - - assert len(routine.variable_map) == 3 - assert 'z' in routine.variable_map - assert 'jlon' in routine.variable_map - assert 'jg' in routine.variable_map - - assigns = FindNodes(ir.Assignment).visit(routine.body) - assert len(assigns) == 6 - assert assigns[2].lhs == 'jlon' and assigns[2].rhs == '1' - assert assigns[3].lhs == 'jg' and assigns[3].rhs == '2' - assert assigns[4].lhs == 'z' and assigns[4].rhs == 'jlon + jg' - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_member_routines_indexing_of_shadowed_array(frontend): - """ - Test special case of inlining of member subroutines when inlined routine contains - shadowed array and array indices. - In particular, this test checks that also the variables indexing - the array in the inlined result get renamed correctly. - """ - fcode = """ - subroutine outer(klon) - integer :: jg, jlon - integer :: arr(3, 3) - - jg = 70000 - call inner2() - - contains - - subroutine inner2() - integer :: jlon, jg - integer :: arr(3, 3) - do jg=1,3 - do jlon=1,3 - arr(jlon, jg) = 11 - end do - end do - end subroutine inner2 - - end subroutine outer - """ - routine = Subroutine.from_source(fcode, frontend=frontend) - inline_member_procedures(routine) - innerloop = FindNodes(ir.Loop).visit(routine.body)[1] - innerloopvars = FindVariables().visit(innerloop) - assert 'inner2_arr(inner2_jlon,inner2_jg)' in innerloopvars - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_member_routines_sequence_assoc(frontend): - """ - Test inlining of member subroutines in the presence of sequence - associations. As this is not supported, we check for the - appropriate error. - """ - fcode = """ -subroutine member_routines_sequence_assoc(vector) - real(kind=8), intent(inout) :: vector(6) - integer :: i - - i = 2 - call inner(3, vector(i)) - - contains - subroutine inner(n, a) - integer, intent(in) :: n - real(kind=8), intent(inout) :: a(3) - integer :: j - do j=1, n - a(j) = a(j) + 1 - end do - end subroutine -end subroutine member_routines_sequence_assoc - """ - routine = Subroutine.from_source(fcode, frontend=frontend) - - # Expect to fail tmp_path due to use of sequence association - with pytest.raises(RuntimeError): - inline_member_procedures(routine=routine) - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_member_routines_with_associate(frontend): - """ - Ensure that internal routines with :any:`Associate` constructs get - inlined as expected. - """ - fcode = """ -subroutine acraneb_transt(klon, klev, kidia, kfdia, ktdia) - implicit none - - integer(kind=4), intent(in) :: klon, klev, kidia, kfdia, ktdia - integer(kind=4) :: jlon, jlev - - real(kind=8) :: zq1(klon) - real(kind=8) :: zq2(klon, klev) - - call delta_t(zq1) - - do jlev = ktdia, klev - call delta_t(zq2(1:klon,jlev)) - - enddo - -contains - -subroutine delta_t(pq) - implicit none - - real(kind=8), intent(in) :: pq(klon) - real(kind=8) :: x, z - - associate(zz => z) - - do jlon = 1,klon - x = x + pq(jlon) - enddo - end associate -end subroutine - -end subroutine acraneb_transt - """ - - routine = Subroutine.from_source(fcode, frontend=frontend) - - inline_member_procedures(routine=routine) - - assert not routine.members - loops = FindNodes(ir.Loop).visit(routine.body) - assert len(loops) == 3 - - assigns = FindNodes(ir.Assignment).visit(routine.body) - assert len(assigns) == 2 - assert assigns[0].rhs == 'x + zq1(jlon)' - assert assigns[1].rhs == 'x + zq2(jlon, jlev)' - - assocs = FindNodes(ir.Associate).visit(routine.body) - assert len(assocs) == 2 - - -@pytest.mark.parametrize('frontend', available_frontends( - xfail=[(OMNI, 'OMNI does not handle missing type definitions')] -)) -def test_inline_member_routines_with_optionals(frontend): - """ - Ensure that internal routines with optional arguments get - inlined as expected (esp. present instrinsics are correctly - evaluated for all variables types) - """ - fcode = """ -subroutine test_inline(klon, ydxfu, ydmf_phys_out) - - use yomxfu , only : txfu - use mf_phys_type_mod , only : mf_phys_out_type - - implicit none - - integer(kind=4), intent(in) :: klon - type(txfu) ,intent(inout) :: ydxfu - type(mf_phys_out_type) ,intent(in) :: ydmf_phys_out - - call member_rout (ydxfu%visicld, pvmin=ydmf_phys_out%visicld, psmax=1.0_8) - - contains - - subroutine member_rout (x, pvmin, pvmax, psmin, psmax) - - real(kind=8) ,intent(inout) :: x(1:klon) - real(kind=8) ,intent(in) ,optional :: pvmin(1:klon) - real(kind=8) ,intent(in) ,optional :: pvmax(1:klon) - real(kind=8) ,intent(in) ,optional :: psmin - real(kind=8) ,intent(in) ,optional :: psmax - - if (present (psmin)) x = psmin - if (present (psmax)) x = psmax - if (present (pvmin)) x = minval(pvmin(:)) - if (present (pvmax)) x = maxval(pvmax(:)) - - end subroutine member_rout - -end subroutine test_inline - """ - - routine = Subroutine.from_source(fcode, frontend=frontend) - - inline_member_procedures(routine=routine) - - assert not routine.members - - conds = FindNodes(ir.Conditional).visit(routine.body) - assert len(conds) == 4 - assert conds[0].condition == 'False' - assert conds[1].condition == 'True' - assert conds[2].condition == 'True' - assert conds[3].condition == 'False' - - -@pytest.mark.parametrize('frontend', available_frontends( - skip={OFP: "OFP apparently has problems dealing with those Statement Functions", - OMNI: "OMNI automatically inlines Statement Functions"} -)) -@pytest.mark.parametrize('stmt_decls', (True, False)) -def test_inline_statement_functions(frontend, stmt_decls): - stmt_decls_code = """ - real :: PTARE - real :: FOEDELTA - FOEDELTA ( PTARE ) = PTARE + 1.0 - real :: FOEEW - FOEEW ( PTARE ) = PTARE + FOEDELTA(PTARE) - """.strip() - - fcode = f""" -subroutine stmt_func(arr, ret) - implicit none - real, intent(in) :: arr(:) - real, intent(inout) :: ret(:) - real :: ret2 - real, parameter :: rtt = 1.0 - {stmt_decls_code if stmt_decls else '#include "fcttre.func.h"'} - - ret = foeew(arr) - ret2 = foedelta(3.0) -end subroutine stmt_func - """.strip() - - routine = Subroutine.from_source(fcode, frontend=frontend) - if stmt_decls: - assert FindNodes(ir.StatementFunction).visit(routine.spec) - else: - assert not FindNodes(ir.StatementFunction).visit(routine.spec) - assert FindInlineCalls().visit(routine.body) - inline_statement_functions(routine) - - assert not FindNodes(ir.StatementFunction).visit(routine.spec) - if stmt_decls: - assert not FindInlineCalls().visit(routine.body) - assignments = FindNodes(ir.Assignment).visit(routine.body) - assert assignments[0].lhs == 'ret' - assert assignments[0].rhs == "arr + arr + 1.0" - assert assignments[1].lhs == 'ret2' - assert assignments[1].rhs == "3.0 + 1.0" - else: - assert FindInlineCalls().visit(routine.body) - - -@pytest.mark.parametrize('frontend', available_frontends()) -@pytest.mark.parametrize('adjust_imports', [True, False]) -def test_inline_marked_subroutines(frontend, adjust_imports, tmp_path): - """ Test subroutine inlining via marker pragmas. """ - - fcode_driver = """ -subroutine test_pragma_inline(a, b) - use util_mod, only: add_one, add_a_to_b - implicit none - - real(kind=8), intent(inout) :: a(3), b(3) - integer, parameter :: n = 3 - integer :: i - - do i=1, n - !$loki inline - call add_one(a(i)) - end do - - !$loki inline - call add_a_to_b(a(:), b(:), 3) - - do i=1, n - call add_one(b(i)) - end do - -end subroutine test_pragma_inline - """ - - fcode_module = """ -module util_mod -implicit none - -contains - subroutine add_one(a) - interface - subroutine do_something() - end subroutine do_something - end interface - real(kind=8), intent(inout) :: a - a = a + 1 - end subroutine add_one - - subroutine add_a_to_b(a, b, n) - interface - subroutine do_something_else() - end subroutine do_something_else - end interface - real(kind=8), intent(inout) :: a(:), b(:) - integer, intent(in) :: n - integer :: i - - do i = 1, n - a(i) = a(i) + b(i) - end do - end subroutine add_a_to_b -end module util_mod -""" - module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) - driver = Subroutine.from_source(fcode_driver, frontend=frontend, xmods=[tmp_path]) - driver.enrich(module) - - calls = FindNodes(ir.CallStatement).visit(driver.body) - assert calls[0].routine == module['add_one'] - assert calls[1].routine == module['add_a_to_b'] - assert calls[2].routine == module['add_one'] - - inline_marked_subroutines( - routine=driver, allowed_aliases=('I',), adjust_imports=adjust_imports - ) - - # Check inlined loops and assignments - assert len(FindNodes(ir.Loop).visit(driver.body)) == 3 - assign = FindNodes(ir.Assignment).visit(driver.body) - assert len(assign) == 2 - assert assign[0].lhs == 'a(i)' and assign[0].rhs == 'a(i) + 1' - assert assign[1].lhs == 'a(i)' and assign[1].rhs == 'a(i) + b(i)' - - # Check that the last call is left untouched - calls = FindNodes(ir.CallStatement).visit(driver.body) - assert len(calls) == 1 - assert calls[0].routine.name == 'add_one' - assert calls[0].arguments == ('b(i)',) - - imports = FindNodes(ir.Import).visit(driver.spec) - assert len(imports) == 1 - if adjust_imports: - assert imports[0].symbols == ('add_one',) - else: - assert imports[0].symbols == ('add_one', 'add_a_to_b') - - if adjust_imports: - # check that explicit interfaces were imported - intfs = driver.interfaces - assert len(intfs) == 1 - assert all(isinstance(s, sym.ProcedureSymbol) for s in driver.interface_symbols) - assert 'do_something' in driver.interface_symbols - assert 'do_something_else' in driver.interface_symbols - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_marked_subroutines_with_interfaces(frontend, tmp_path): - """ Test inlining of subroutines with explicit interfaces via marker pragmas. """ - - fcode_driver = """ -subroutine test_pragma_inline(a, b) - implicit none - - interface - subroutine add_a_to_b(a, b, n) - real(kind=8), intent(inout) :: a(:), b(:) - integer, intent(in) :: n - end subroutine add_a_to_b - subroutine add_one(a) - real(kind=8), intent(inout) :: a - end subroutine add_one - end interface - - interface - subroutine add_two(a) - real(kind=8), intent(inout) :: a - end subroutine add_two - end interface - - real(kind=8), intent(inout) :: a(3), b(3) - integer, parameter :: n = 3 - integer :: i - - do i=1, n - !$loki inline - call add_one(a(i)) - end do - - !$loki inline - call add_a_to_b(a(:), b(:), 3) - - do i=1, n - call add_one(b(i)) - !$loki inline - call add_two(b(i)) - end do - -end subroutine test_pragma_inline - """ - - fcode_module = """ -module util_mod -implicit none - -contains - subroutine add_one(a) - real(kind=8), intent(inout) :: a - a = a + 1 - end subroutine add_one - - subroutine add_two(a) - real(kind=8), intent(inout) :: a - a = a + 2 - end subroutine add_two - - subroutine add_a_to_b(a, b, n) - real(kind=8), intent(inout) :: a(:), b(:) - integer, intent(in) :: n - integer :: i - - do i = 1, n - a(i) = a(i) + b(i) - end do - end subroutine add_a_to_b -end module util_mod -""" - - module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) - driver = Subroutine.from_source(fcode_driver, frontend=frontend, xmods=[tmp_path]) - driver.enrich(module.subroutines) - - calls = FindNodes(ir.CallStatement).visit(driver.body) - assert calls[0].routine == module['add_one'] - assert calls[1].routine == module['add_a_to_b'] - assert calls[2].routine == module['add_one'] - assert calls[3].routine == module['add_two'] - - inline_marked_subroutines(routine=driver, allowed_aliases=('I',)) - - # Check inlined loops and assignments - assert len(FindNodes(ir.Loop).visit(driver.body)) == 3 - assign = FindNodes(ir.Assignment).visit(driver.body) - assert len(assign) == 3 - assert assign[0].lhs == 'a(i)' and assign[0].rhs == 'a(i) + 1' - assert assign[1].lhs == 'a(i)' and assign[1].rhs == 'a(i) + b(i)' - assert assign[2].lhs == 'b(i)' and assign[2].rhs == 'b(i) + 2' - - # Check that the last call is left untouched - calls = FindNodes(ir.CallStatement).visit(driver.body) - assert len(calls) == 1 - assert calls[0].routine.name == 'add_one' - assert calls[0].arguments == ('b(i)',) - - intfs = FindNodes(ir.Interface).visit(driver.spec) - assert len(intfs) == 1 - assert intfs[0].symbols == ('add_one',) - - -@pytest.mark.parametrize('frontend', available_frontends()) -@pytest.mark.parametrize('adjust_imports', [True, False]) -def test_inline_marked_routine_with_optionals(frontend, adjust_imports, tmp_path): - """ Test subroutine inlining via marker pragmas with omitted optionals. """ - - fcode_driver = """ -subroutine test_pragma_inline_optionals(a, b) - use util_mod, only: add_one - implicit none - - real(kind=8), intent(inout) :: a(3), b(3) - integer, parameter :: n = 3 - integer :: i - - do i=1, n - !$loki inline - call add_one(a(i), two=2.0) - end do - - do i=1, n - !$loki inline - call add_one(b(i)) - end do - -end subroutine test_pragma_inline_optionals - """ - - fcode_module = """ -module util_mod -implicit none - -contains - subroutine add_one(a, two) - real(kind=8), intent(inout) :: a - real(kind=8), optional, intent(inout) :: two - a = a + 1 - - if (present(two)) then - a = a + two - end if - end subroutine add_one -end module util_mod -""" - module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) - driver = Subroutine.from_source(fcode_driver, frontend=frontend, xmods=[tmp_path]) - driver.enrich(module) - - calls = FindNodes(ir.CallStatement).visit(driver.body) - assert calls[0].routine == module['add_one'] - assert calls[1].routine == module['add_one'] - - inline_marked_subroutines(routine=driver, adjust_imports=adjust_imports) - - # Check inlined loops and assignments - assert len(FindNodes(ir.Loop).visit(driver.body)) == 2 - assign = FindNodes(ir.Assignment).visit(driver.body) - assert len(assign) == 4 - assert assign[0].lhs == 'a(i)' and assign[0].rhs == 'a(i) + 1' - assert assign[1].lhs == 'a(i)' and assign[1].rhs == 'a(i) + 2.0' - assert assign[2].lhs == 'b(i)' and assign[2].rhs == 'b(i) + 1' - # TODO: This is a problem, since it's not declared anymore - assert assign[3].lhs == 'b(i)' and assign[3].rhs == 'b(i) + two' - - # Check that the PRESENT checks have been resolved - assert len(FindNodes(ir.CallStatement).visit(driver.body)) == 0 - assert len(FindInlineCalls().visit(driver.body)) == 0 - checks = FindNodes(ir.Conditional).visit(driver.body) - assert len(checks) == 2 - assert checks[0].condition == 'True' - assert checks[1].condition == 'False' - - imports = FindNodes(ir.Import).visit(driver.spec) - assert len(imports) == 0 if adjust_imports else 1 - - -@pytest.mark.parametrize('frontend', available_frontends( - xfail=[(OMNI, 'OMNI has no sense of humour!')]) -) -def test_inline_marked_subroutines_with_associates(frontend): - """ Test subroutine inlining via marker pragmas with nested associates. """ - - fcode_outer = """ -subroutine test_pragma_inline_associates(never) - use peter_pan, only: neverland - implicit none - type(neverland), intent(inout) :: never - - associate(going=>never%going_to) - - associate(up=>give_you%up) - - !$loki inline - call dave(going, up) - - end associate - - end associate -end subroutine test_pragma_inline_associates - """ - - fcode_inner = """ -subroutine dave(going) - use your_imagination, only: astley - implicit none - type(astley), intent(inout) :: going - - associate(give_you=>going%give_you) - - associate(up=>give_you%up) - - call rick_is(up) - - end associate - - end associate -end subroutine dave - """ - - outer = Subroutine.from_source(fcode_outer, frontend=frontend) - inner = Subroutine.from_source(fcode_inner, frontend=frontend) - outer.enrich(inner) - - assert FindNodes(ir.CallStatement).visit(outer.body)[0].routine == inner - - inline_marked_subroutines(routine=outer, adjust_imports=True) - - # Ensure that all associates are perfectly nested afterwards - assocs = FindNodes(ir.Associate).visit(outer.body) - assert len(assocs) == 4 - assert assocs[1].parent == assocs[0] - assert assocs[2].parent == assocs[1] - assert assocs[3].parent == assocs[2] - - # And, because we can... - outer.body = ResolveAssociatesTransformer().visit(outer.body) - call = FindNodes(ir.CallStatement).visit(outer.body)[0] - assert call.name == 'rick_is' - assert call.arguments == ('never%going_to%give_you%up',) - # Q. E. D. - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_marked_subroutines_declarations(frontend, tmp_path): - """Test symbol propagation to hoisted declaration when inlining.""" - fcode = """ -module inline_declarations - implicit none - - type bounds - integer :: start, end - end type bounds - - contains - - subroutine outer(a, bnds) - real(kind=8), intent(inout) :: a(bnds%end) - type(bounds), intent(in) :: bnds - real(kind=8) :: b(bnds%end) - - b(bnds%start:bnds%end) = a(bnds%start:bnds%end) + 42.0 - - !$loki inline - call inner(a, dims=bnds) - end subroutine outer - - subroutine inner(c, dims) - real(kind=8), intent(inout) :: c(dims%end) - type(bounds), intent(in) :: dims - real(kind=8) :: d(dims%end) - - d(dims%start:dims%end) = c(dims%start:dims%end) - 66.6 - c(dims%start) = sum(d) - end subroutine inner -end module inline_declarations -""" - module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - outer = module['outer'] - - inline_marked_subroutines(routine=outer, adjust_imports=True) - - # Check that all declarations are using the ``bnds`` symbol - assert outer.symbols[0] == 'a(bnds%end)' - assert outer.symbols[2] == 'b(bnds%end)' - assert outer.symbols[3] == 'd(bnds%end)' - assert all( - a.shape == ('bnds%end',) for a in outer.symbols if isinstance(a, sym.Array) - ) - - -@pytest.mark.parametrize('frontend', available_frontends( - (OFP, 'Prefix/elemental support not implemented')) -) -@pytest.mark.parametrize('pass_as_kwarg', (False, True)) -def test_inline_transformation(tmp_path, frontend, pass_as_kwarg): - """Test combining recursive inlining via :any:`InliningTransformation`.""" - - fcode_module = """ -module one_mod - real(kind=8), parameter :: one = 1.0 -end module one_mod -""" - - fcode_inner = """ -subroutine add_one_and_two(a) - use one_mod, only: one - implicit none - - real(kind=8), intent(inout) :: a - - a = a + one - - a = add_two(a) - -contains - elemental function add_two(x) - real(kind=8), intent(in) :: x - real(kind=8) :: add_two - - add_two = x + 2.0 - end function add_two -end subroutine add_one_and_two -""" - - fcode = f""" -subroutine test_inline_pragma(a, b) - implicit none - real(kind=8), intent(inout) :: a(3), b(3) - integer, parameter :: n = 3 - integer :: i - real :: stmt_arg - real :: some_stmt_func - some_stmt_func ( stmt_arg ) = stmt_arg + 3.1415 - -#include "add_one_and_two.intfb.h" - - do i=1, n - !$loki inline - call add_one_and_two({'a=' if pass_as_kwarg else ''}a(i)) - end do - - do i=1, n - !$loki inline - call add_one_and_two({'a=' if pass_as_kwarg else ''}b(i)) - end do - - a(1) = some_stmt_func({'stmt_arg=' if pass_as_kwarg else ''}a(2)) - -end subroutine test_inline_pragma -""" - module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) - inner = Subroutine.from_source(fcode_inner, definitions=module, frontend=frontend, xmods=[tmp_path]) - routine = Subroutine.from_source(fcode, frontend=frontend) - routine.enrich(inner) - - trafo = InlineTransformation( - inline_constants=True, external_only=True, inline_elementals=True, - inline_stmt_funcs=True - ) - - calls = FindNodes(ir.CallStatement).visit(routine.body) - assert len(calls) == 2 - assert all(c.routine == inner for c in calls) - - # Apply to the inner subroutine first to resolve parameter and calls - trafo.apply(inner) - - assigns = FindNodes(ir.Assignment).visit(inner.body) - assert len(assigns) == 3 - assert assigns[0].lhs == 'a' and assigns[0].rhs == 'a + 1.0' - assert assigns[1].lhs == 'result_add_two' and assigns[1].rhs == 'a + 2.0' - assert assigns[2].lhs == 'a' and assigns[2].rhs == 'result_add_two' - - # Apply to the outer routine, but with resolved body of the inner - trafo.apply(routine) - - calls = FindNodes(ir.CallStatement).visit(routine.body) - assert len(calls) == 0 - assigns = FindNodes(ir.Assignment).visit(routine.body) - assert len(assigns) == 7 - assert assigns[0].lhs == 'a(i)' and assigns[0].rhs == 'a(i) + 1.0' - assert assigns[1].lhs == 'result_add_two' and assigns[1].rhs == 'a(i) + 2.0' - assert assigns[2].lhs == 'a(i)' and assigns[2].rhs == 'result_add_two' - assert assigns[3].lhs == 'b(i)' and assigns[3].rhs == 'b(i) + 1.0' - assert assigns[4].lhs == 'result_add_two' and assigns[4].rhs == 'b(i) + 2.0' - assert assigns[5].lhs == 'b(i)' and assigns[5].rhs == 'result_add_two' - assert assigns[6].lhs == 'a(1)' and assigns[6].rhs == 'a(2) + 3.1415' - - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_transformation_local_seq_assoc(frontend, tmp_path): - fcode = """ -module somemod - implicit none - contains - - subroutine minusone_second(output, x) - real, intent(inout) :: output - real, intent(in) :: x(3) - output = x(2) - 1 - end subroutine minusone_second - - subroutine plusone(output, x) - real, intent(inout) :: output - real, intent(in) :: x - output = x + 1 - end subroutine plusone - - subroutine outer() - implicit none - real :: x(3, 3) - real :: y - x = 10.0 - - call inner(y, x(1, 1)) ! Sequence association tmp_path for member routine. - - !$loki inline - call plusone(y, x(3, 3)) ! Marked for inlining. - - call minusone_second(y, x(1, 3)) ! Standard call with sequence association (never processed). - - contains - - subroutine inner(output, x) - real, intent(inout) :: output - real, intent(in) :: x(3) - - output = x(2) + 2.0 - end subroutine inner - end subroutine outer - -end module somemod -""" - # Test case that nothing happens if `resolve_sequence_association=True` - # but inlining "marked" and "internals" is disabled. - module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - trafo = InlineTransformation( - inline_constants=True, external_only=True, inline_elementals=True, - inline_marked=False, inline_internals=False, resolve_sequence_association=True - ) - outer = module["outer"] - trafo.apply(outer) - callnames = [call.name for call in FindNodes(ir.CallStatement).visit(outer.body)] - assert 'plusone' in callnames - assert 'inner' in callnames - assert 'minusone_second' in callnames - - # Test case that only marked processed if - # `resolve_sequence_association=True` - # `inline_marked=True`, - # `inline_internals=False` - module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - trafo = InlineTransformation( - inline_constants=True, external_only=True, inline_elementals=True, - inline_marked=True, inline_internals=False, resolve_sequence_association=True - ) - outer = module["outer"] - trafo.apply(outer) - callnames = [call.name for call in FindNodes(ir.CallStatement).visit(outer.body)] - assert 'plusone' not in callnames - assert 'inner' in callnames - assert 'minusone_second' in callnames - - # Test case that a crash occurs if sequence association is not enabled even if it is needed. - module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - trafo = InlineTransformation( - inline_constants=True, external_only=True, inline_elementals=True, - inline_marked=True, inline_internals=True, resolve_sequence_association=False - ) - outer = module["outer"] - with pytest.raises(RuntimeError): - trafo.apply(outer) - callnames = [call.name for call in FindNodes(ir.CallStatement).visit(outer.body)] - - # Test case that sequence association is run and corresponding call inlined, avoiding crash. - module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - trafo = InlineTransformation( - inline_constants=True, external_only=True, inline_elementals=True, - inline_marked=False, inline_internals=True, resolve_sequence_association=True - ) - outer = module["outer"] - trafo.apply(outer) - callnames = [call.name for call in FindNodes(ir.CallStatement).visit(outer.body)] - assert 'plusone' in callnames - assert 'inner' not in callnames - assert 'minusone_second' in callnames - - # Test case that everything is enabled. - module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - trafo = InlineTransformation( - inline_constants=True, external_only=True, inline_elementals=True, - inline_marked=True, inline_internals=True, resolve_sequence_association=True - ) - outer = module["outer"] - trafo.apply(outer) - callnames = [call.name for call in FindNodes(ir.CallStatement).visit(outer.body)] - assert 'plusone' not in callnames - assert 'inner' not in callnames - assert 'minusone_second' in callnames - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_transformation_local_seq_assoc_crash_marked_no_seq_assoc(frontend, tmp_path): - # Test case that a crash occurs if marked routine with sequence association is - # attempted to inline without sequence association enabled. - fcode = """ -module somemod - implicit none - contains - - subroutine inner(output, x) - real, intent(inout) :: output - real, intent(in) :: x(3) - - output = x(2) + 2.0 - end subroutine inner - - subroutine outer() - real :: x(3, 3) - real :: y - x = 10.0 - - !$loki inline - call inner(y, x(1, 1)) ! Sequence association tmp_path for marked routine. - end subroutine outer - -end module somemod -""" - module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - trafo = InlineTransformation( - inline_constants=True, external_only=True, inline_elementals=True, - inline_marked=True, inline_internals=False, resolve_sequence_association=False - ) - outer = module["outer"] - with pytest.raises(RuntimeError): - trafo.apply(outer) - - # Test case that crash is avoided by activating sequence association. - module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - trafo = InlineTransformation( - inline_constants=True, external_only=True, inline_elementals=True, - inline_marked=True, inline_internals=False, resolve_sequence_association=True - ) - outer = module["outer"] - trafo.apply(outer) - assert len(FindNodes(ir.CallStatement).visit(outer.body)) == 0 - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_transformation_local_seq_assoc_crash_value_err_no_source(frontend, tmp_path): - # Testing that ValueError is thrown if sequence association is requested with inlining, - # but source code behind call is missing (not enough type information). - fcode = """ -module somemod - implicit none - contains - - subroutine outer() - real :: x(3, 3) - real :: y - x = 10.0 - - !$loki inline - call inner(y, x(1, 1)) ! Sequence association tmp_path for marked routine. - end subroutine outer - -end module somemod -""" - module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - trafo = InlineTransformation( - inline_constants=True, external_only=True, inline_elementals=True, - inline_marked=True, inline_internals=False, resolve_sequence_association=True - ) - outer = module["outer"] - with pytest.raises(ValueError): - trafo.apply(outer) - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_transformation_adjust_imports(frontend, tmp_path): - fcode_module = """ -module bnds_module - integer :: m - integer :: n - integer :: l -end module bnds_module - """ - - fcode_another = """ -module another_module - integer :: x -end module another_module - """ - - fcode_outer = """ -subroutine test_inline_outer(a, b) - use bnds_module, only: n - use test_inline_mod, only: test_inline_inner - use test_inline_another_mod, only: test_inline_another_inner - implicit none - - real(kind=8), intent(inout) :: a(n), b(n) - - !$loki inline - call test_inline_another_inner() - !$loki inline - call test_inline_inner(a, b) -end subroutine test_inline_outer - """ - - fcode_inner = """ -module test_inline_mod - implicit none - contains - -subroutine test_inline_inner(a, b) - use BNDS_module, only: n, m - use another_module, only: x - - real(kind=8), intent(inout) :: a(n), b(n) - real(kind=8) :: tmp(m) - integer :: i - - tmp(1:m) = x - do i=1, n - a(i) = b(i) + sum(tmp) - end do -end subroutine test_inline_inner -end module test_inline_mod - """ - - fcode_another_inner = """ -module test_inline_another_mod - implicit none - contains - -subroutine test_inline_another_inner() - use BNDS_module, only: n, m, l - -end subroutine test_inline_another_inner -end module test_inline_another_mod - """ - - _ = Module.from_source(fcode_another, frontend=frontend, xmods=[tmp_path]) - _ = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) - inner = Module.from_source(fcode_inner, frontend=frontend, xmods=[tmp_path]) - another_inner = Module.from_source(fcode_another_inner, frontend=frontend, xmods=[tmp_path]) - outer = Subroutine.from_source( - fcode_outer, definitions=(inner, another_inner), frontend=frontend, xmods=[tmp_path] - ) - - trafo = InlineTransformation( - inline_elementals=False, inline_marked=True, adjust_imports=True - ) - trafo.apply(outer) - - # Check that the inlining has happened - assign = FindNodes(ir.Assignment).visit(outer.body) - assert len(assign) == 2 - assert assign[0].lhs == 'tmp(1:m)' - assert assign[0].rhs == 'x' - assert assign[1].lhs == 'a(i)' - assert assign[1].rhs == 'b(i) + sum(tmp)' - - # Now check that the right modules have been moved, - # and the import of the call has been removed - imports = FindNodes(ir.Import).visit(outer.spec) - assert len(imports) == 2 - assert imports[0].module == 'another_module' - assert imports[0].symbols == ('x',) - assert imports[1].module == 'bnds_module' - assert all(_ in imports[1].symbols for _ in ['l', 'm', 'n']) - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_inline_transformation_intermediate(tmp_path, frontend): - fcode_outermost = """ -module outermost_mod -implicit none -contains -subroutine outermost() -use intermediate_mod, only: intermediate - -!$loki inline -call intermediate() - -end subroutine outermost -end module outermost_mod -""" - - fcode_intermediate = """ -module intermediate_mod -implicit none -contains -subroutine intermediate() -use innermost_mod, only: innermost - -call innermost() - -end subroutine intermediate -end module intermediate_mod -""" - - fcode_innermost = """ -module innermost_mod -implicit none -contains -subroutine innermost() - -end subroutine innermost -end module innermost_mod -""" - - (tmp_path/'outermost_mod.F90').write_text(fcode_outermost) - (tmp_path/'intermediate_mod.F90').write_text(fcode_intermediate) - (tmp_path/'innermost_mod.F90').write_text(fcode_innermost) - - config = { - 'default': { - 'mode': 'idem', - 'role': 'kernel', - 'expand': True, - 'strict': True - }, - 'routines': { - 'outermost': {'role': 'kernel'} - } - } - - scheduler = Scheduler( - paths=[tmp_path], config=SchedulerConfig.from_dict(config), - frontend=frontend, xmods=[tmp_path] - ) - - def _get_successors(item): - return scheduler.sgraph.successors(scheduler[item]) - - # check graph edges before transformation - assert len(scheduler.items) == 3 - assert len(_get_successors('outermost_mod#outermost')) == 1 - assert scheduler['intermediate_mod#intermediate'] in _get_successors('outermost_mod#outermost') - assert len(_get_successors('intermediate_mod#intermediate')) == 1 - assert scheduler['innermost_mod#innermost'] in _get_successors('intermediate_mod#intermediate') - - scheduler.process( transformation=InlineTransformation() ) - - # check graph edges were updated correctly - assert len(scheduler.items) == 2 - assert len(_get_successors('outermost_mod#outermost')) == 1 - assert scheduler['innermost_mod#innermost'] in _get_successors('outermost_mod#outermost') From e666d8998cdf814686f136a7c4c77bcc92df2862 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Wed, 11 Sep 2024 14:33:39 +0000 Subject: [PATCH 03/12] Transformations: Move "extract" utilities to separate sub-package This also moves the test and renames the `extract_contained_` to `extract_internal_`, while keeping aliases to the old methods. --- loki/transformations/extract/__init__.py | 12 ++++ .../{extract.py => extract/internal.py} | 41 +++++++----- .../tests/test_extract_internal.py} | 62 +++++++++---------- 3 files changed, 68 insertions(+), 47 deletions(-) create mode 100644 loki/transformations/extract/__init__.py rename loki/transformations/{extract.py => extract/internal.py} (83%) rename loki/transformations/{tests/test_extract.py => extract/tests/test_extract_internal.py} (90%) diff --git a/loki/transformations/extract/__init__.py b/loki/transformations/extract/__init__.py new file mode 100644 index 000000000..9b82c36c9 --- /dev/null +++ b/loki/transformations/extract/__init__.py @@ -0,0 +1,12 @@ +# (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. +""" + +from loki.transformations.extract.internal import * # noqa diff --git a/loki/transformations/extract.py b/loki/transformations/extract/internal.py similarity index 83% rename from loki/transformations/extract.py rename to loki/transformations/extract/internal.py index f5c541fd0..54bf48649 100644 --- a/loki/transformations/extract.py +++ b/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,24 +82,25 @@ 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'). @@ -104,7 +108,7 @@ def extract_contained_procedure(procedure, name): # 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}, " @@ -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 diff --git a/loki/transformations/tests/test_extract.py b/loki/transformations/extract/tests/test_extract_internal.py similarity index 90% rename from loki/transformations/tests/test_extract.py rename to loki/transformations/extract/tests/test_extract_internal.py index 6d66a5857..6bbd1fd04 100644 --- a/loki/transformations/tests/test_extract.py +++ b/loki/transformations/extract/tests/test_extract_internal.py @@ -12,11 +12,11 @@ from loki.sourcefile import Sourcefile from loki.subroutine import Subroutine -from loki.transformations.extract import extract_contained_procedures +from loki.transformations.extract import extract_internal_procedures @pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_contained_procedures_basic_scalar(frontend): +def test_extract_internal_procedures_basic_scalar(frontend): """ Tests that a global scalar is correctly added as argument of `inner`. """ @@ -36,7 +36,7 @@ def test_extract_contained_procedures_basic_scalar(frontend): end subroutine outer """ src = Sourcefile.from_source(fcode, frontend=frontend) - routines = extract_contained_procedures(src.routines[0]) + routines = extract_internal_procedures(src.routines[0]) assert len(routines) == 1 assert routines[0].name == "inner" inner = routines[0] @@ -47,7 +47,7 @@ def test_extract_contained_procedures_basic_scalar(frontend): assert 'x' in (arg[0] for arg in call.kwarguments) @pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_contained_procedures_contains_emptied(frontend): +def test_extract_internal_procedures_contains_emptied(frontend): """ Tests that the contains section does not contain any functions or subroutines after processing. """ @@ -76,12 +76,12 @@ def test_extract_contained_procedures_contains_emptied(frontend): """ src = Sourcefile.from_source(fcode, frontend=frontend) outer = src.routines[0] - extract_contained_procedures(outer) + extract_internal_procedures(outer) # NOTE: Functions in Loki are also typed as Subroutines. assert not any(isinstance(r, Subroutine) for r in outer.contains.body) @pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_contained_procedures_basic_array(frontend): +def test_extract_internal_procedures_basic_array(frontend): """ Tests that a global array variable (and a scalar) is correctly added as argument of `inner`. """ @@ -104,7 +104,7 @@ def test_extract_contained_procedures_basic_array(frontend): end subroutine outer """ src = Sourcefile.from_source(fcode, frontend=frontend) - routines = extract_contained_procedures(src.routines[0]) + routines = extract_internal_procedures(src.routines[0]) assert len(routines) == 1 inner = routines[0] outer = src.routines[0] @@ -117,7 +117,7 @@ def test_extract_contained_procedures_basic_array(frontend): assert kwargdict['arr'] == 'arr' @pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_contained_procedures_existing_call_args(frontend): +def test_extract_internal_procedures_existing_call_args(frontend): """ Tests that variable resolution process works correctly when the parent contains a call to the extracted function that already has some calling arguments. @@ -149,7 +149,7 @@ def test_extract_contained_procedures_existing_call_args(frontend): """ src = Sourcefile.from_source(fcode, frontend=frontend) outer = src.routines[0] - extract_contained_procedures(outer) + extract_internal_procedures(outer) calls = FindNodes(CallStatement).visit(outer.body) for call in calls: @@ -176,7 +176,7 @@ def test_extract_contained_procedures_existing_call_args(frontend): assert kwargdict['y'] == 1 @pytest.mark.parametrize('frontend', available_frontends(xfail=[(OMNI, 'Parser fails on missing constants module')])) -def test_extract_contained_procedures_basic_import(frontend): +def test_extract_internal_procedures_basic_import(frontend): """ Tests that a global imported binding is correctly introduced to the contained subroutine. """ @@ -198,7 +198,7 @@ def test_extract_contained_procedures_basic_import(frontend): end subroutine outer """ src = Sourcefile.from_source(fcode, frontend=frontend) - routines = extract_contained_procedures(src.routines[0]) + routines = extract_internal_procedures(src.routines[0]) assert len(routines) == 1 inner = routines[0] assert "c2" in inner.import_map @@ -206,7 +206,7 @@ def test_extract_contained_procedures_basic_import(frontend): assert 'c2' not in inner.arguments @pytest.mark.parametrize('frontend', available_frontends(xfail=[(OMNI, 'Parser fails on missing type_mod module')])) -def test_extract_contained_procedures_recursive_definition(frontend): +def test_extract_internal_procedures_recursive_definition(frontend): """ Tests that whenever a global in the contained subroutine depends on another global variable, both are introduced as arguments, @@ -235,7 +235,7 @@ def test_extract_contained_procedures_recursive_definition(frontend): end subroutine outer """ src = Sourcefile.from_source(fcode, frontend=frontend) - routines = extract_contained_procedures(src.routines[0]) + routines = extract_internal_procedures(src.routines[0]) assert len(routines) == 1 outer = src.routines[0] inner = routines[0] @@ -274,7 +274,7 @@ def test_extract_contained_procedures_recursive_definition(frontend): assert klev.type.intent == "in" @pytest.mark.parametrize('frontend', available_frontends(xfail=[(OMNI, 'Parser fails on missing parkind1 module')])) -def test_extract_contained_procedures_recursive_definition_import(frontend): +def test_extract_internal_procedures_recursive_definition_import(frontend): """ Tests that whenever globals in the contained subroutine depend on imported bindings, the globals are introduced as arguments, and the imports are added to the contained subroutine. @@ -299,7 +299,7 @@ def test_extract_contained_procedures_recursive_definition_import(frontend): end subroutine outer """ src = Sourcefile.from_source(fcode, frontend=frontend) - routines = extract_contained_procedures(src.routines[0]) + routines = extract_internal_procedures(src.routines[0]) assert len(routines) == 1 outer = src.routines[0] inner = routines[0] @@ -324,7 +324,7 @@ def test_extract_contained_procedures_recursive_definition_import(frontend): assert len(symbols) == 2 @pytest.mark.parametrize('frontend', available_frontends(xfail=[(OMNI, 'Parser fails on missing parkind1 module')])) -def test_extract_contained_procedures_kind_resolution(frontend): +def test_extract_internal_procedures_kind_resolution(frontend): """ Tests that an unresolved kind parameter in inner scope is resolved from import in outer scope. """ @@ -342,12 +342,12 @@ def test_extract_contained_procedures_kind_resolution(frontend): end subroutine outer """ src = Sourcefile.from_source(fcode, frontend=frontend) - routines = extract_contained_procedures(src.routines[0]) + routines = extract_internal_procedures(src.routines[0]) inner = routines[0] assert "jpim" in inner.import_map @pytest.mark.parametrize('frontend', available_frontends(xfail=[(OMNI, 'Parser fails on missing stuff module')])) -def test_extract_contained_procedures_derived_type_resolution(frontend): +def test_extract_internal_procedures_derived_type_resolution(frontend): """ Tests that an unresolved derived type in inner scope is resolved from import in outer scope. """ @@ -365,12 +365,12 @@ def test_extract_contained_procedures_derived_type_resolution(frontend): end subroutine outer """ src = Sourcefile.from_source(fcode, frontend=frontend) - routines = extract_contained_procedures(src.routines[0]) + routines = extract_internal_procedures(src.routines[0]) inner = routines[0] assert "mytype" in inner.import_map @pytest.mark.parametrize('frontend', available_frontends(xfail=[(OMNI, 'Parser fails on missing types module')])) -def test_extract_contained_procedures_derived_type_field(frontend): +def test_extract_internal_procedures_derived_type_field(frontend): """ Test that when a derived type field, i.e 'a%b' is a global in the scope of the contained subroutine, the derived type itself, that is, 'a', is introduced as an the argument in the transformation. @@ -394,7 +394,7 @@ def test_extract_contained_procedures_derived_type_field(frontend): end subroutine outer """ src = Sourcefile.from_source(fcode, frontend=frontend) - routines = extract_contained_procedures(src.routines[0]) + routines = extract_internal_procedures(src.routines[0]) outer = src.routines[0] inner = routines[0] assert 'xtyp' in inner.arguments @@ -419,7 +419,7 @@ def test_extract_contained_procedures_derived_type_field(frontend): assert len(symbols) == 2 @pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_contained_procedures_intent(frontend): +def test_extract_internal_procedures_intent(frontend): """ This test is just to document the current behaviour: when a global is introduced as an argument to the extracted contained procedure, @@ -444,7 +444,7 @@ def test_extract_contained_procedures_intent(frontend): end subroutine outer """ src = Sourcefile.from_source(fcode, frontend=frontend) - routines = extract_contained_procedures(src.routines[0]) + routines = extract_internal_procedures(src.routines[0]) assert len(routines) == 1 outer = src.routines[0] inner = routines[0] @@ -458,7 +458,7 @@ def test_extract_contained_procedures_intent(frontend): assert outer.variable_map['p'].type.intent == "out" @pytest.mark.parametrize('frontend', available_frontends(xfail=[(OMNI, 'Parser fails on undefined symbols')])) -def test_extract_contained_procedures_undefined_in_parent(frontend): +def test_extract_internal_procedures_undefined_in_parent(frontend): """ This test is just to document current behaviour: an exception is raised if a global inside the contained procedure does not @@ -480,10 +480,10 @@ def test_extract_contained_procedures_undefined_in_parent(frontend): """ src = Sourcefile.from_source(fcode, frontend=frontend) with pytest.raises(RuntimeError): - extract_contained_procedures(src.routines[0]) + extract_internal_procedures(src.routines[0]) @pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_contained_procedures_multiple_contained_procedures(frontend): +def test_extract_internal_procedures_multiple_internal_procedures(frontend): """ Basic test to check that multiple contained procedures can also be handled. """ @@ -511,7 +511,7 @@ def test_extract_contained_procedures_multiple_contained_procedures(frontend): end subroutine outer """ src = Sourcefile.from_source(fcode, frontend=frontend) - routines = extract_contained_procedures(src.routines[0]) + routines = extract_internal_procedures(src.routines[0]) assert len(routines) == 2 assert routines[0].name == "inner1" assert routines[1].name == "inner2" @@ -527,7 +527,7 @@ def test_extract_contained_procedures_multiple_contained_procedures(frontend): assert 'gx' in (arg[0] for arg in call.kwarguments) @pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_contained_procedures_basic_scalar_function(frontend): +def test_extract_internal_procedures_basic_scalar_function(frontend): """ Basic test for scalars highlighting that the inner procedure may also be a function. """ @@ -548,7 +548,7 @@ def test_extract_contained_procedures_basic_scalar_function(frontend): end subroutine outer """ src = Sourcefile.from_source(fcode, frontend=frontend) - routines = extract_contained_procedures(src.routines[0]) + routines = extract_internal_procedures(src.routines[0]) assert len(routines) == 1 assert routines[0].name == "inner" inner = routines[0] @@ -561,7 +561,7 @@ def test_extract_contained_procedures_basic_scalar_function(frontend): @pytest.mark.parametrize( 'frontend', available_frontends(skip=(OFP, "ofp fails for unknown reason, likely frontend issue")) ) -def test_extract_contained_procedures_basic_scalar_function_both(frontend): +def test_extract_internal_procedures_basic_scalar_function_both(frontend): """ Basic test for scalars highlighting that the outer and inner procedure may be functions. """ @@ -584,7 +584,7 @@ def test_extract_contained_procedures_basic_scalar_function_both(frontend): end function outer """ src = Sourcefile.from_source(fcode, frontend=frontend) - routines = extract_contained_procedures(src.routines[0]) + routines = extract_internal_procedures(src.routines[0]) assert len(routines) == 1 assert routines[0].name == "inner" inner = routines[0] From 8d729184e17d843eaaaa5b92c02401e963b612ac Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Wed, 11 Sep 2024 15:24:18 +0000 Subject: [PATCH 04/12] Transformations: Rename `region_to_call` to `extract_marked_subroutines` --- .../tests/test_transform_region.py | 110 +++++++++--------- loki/transformations/transform_region.py | 16 +-- 2 files changed, 63 insertions(+), 63 deletions(-) diff --git a/loki/transformations/tests/test_transform_region.py b/loki/transformations/tests/test_transform_region.py index 4ae233174..8e88ef6fa 100644 --- a/loki/transformations/tests/test_transform_region.py +++ b/loki/transformations/tests/test_transform_region.py @@ -17,7 +17,7 @@ from loki.frontend import available_frontends from loki.transformations.transform_region import ( - region_hoist, region_to_call + region_hoist, extract_marked_subroutines ) @@ -346,23 +346,23 @@ def test_transform_region_hoist_promote(tmp_path, frontend): @pytest.mark.parametrize('frontend', available_frontends()) -def test_transform_region_to_call(tmp_path, frontend): +def test_extract_marked_subroutines(tmp_path, frontend): """ - A very simple region-to-call test case + A very simple :any:`extract_marked_subroutine` test case """ fcode = """ -subroutine reg_to_call(a, b, c) +subroutine test_extract(a, b, c) integer, intent(out) :: a, b, c a = 5 a = 1 -!$loki region-to-call in(a) out(b) +!$loki extract in(a) out(b) b = a -!$loki end region-to-call +!$loki end extract c = a + b -end subroutine reg_to_call +end subroutine test_extract """ routine = Subroutine.from_source(fcode, frontend=frontend) filepath = tmp_path/(f'{routine.name}_{frontend}.f90') @@ -376,8 +376,8 @@ def test_transform_region_to_call(tmp_path, frontend): assert len(FindNodes(CallStatement).visit(routine.body)) == 0 # Apply transformation - routines = region_to_call(routine) - assert len(routines) == 1 and routines[0].name == f'{routine.name}_region_to_call_0' + routines = extract_marked_subroutines(routine) + assert len(routines) == 1 and routines[0].name == f'{routine.name}_extracted_0' assert len(FindNodes(Assignment).visit(routine.body)) == 3 assert len(FindNodes(Assignment).visit(routines[0].body)) == 1 @@ -395,30 +395,30 @@ def test_transform_region_to_call(tmp_path, frontend): @pytest.mark.parametrize('frontend', available_frontends()) -def test_transform_region_to_call_multiple(tmp_path, frontend): +def test_extract_marked_subroutines_multiple(tmp_path, frontend): """ Test hoisting with multiple groups and multiple regions per group """ fcode = """ -subroutine reg_to_call_mult(a, b, c) +subroutine test_extract_mult(a, b, c) integer, intent(out) :: a, b, c a = 1 a = a + 1 a = a + 1 -!$loki region-to-call name(oiwjfklsf) inout(a) +!$loki extract name(oiwjfklsf) inout(a) a = a + 1 -!$loki end region-to-call +!$loki end extract a = a + 1 -!$loki region-to-call in(a) out(b) +!$loki extract in(a) out(b) b = a -!$loki end region-to-call +!$loki end extract -!$loki region-to-call in(a,b) out(c) +!$loki extract in(a,b) out(c) c = a + b -!$loki end region-to-call -end subroutine reg_to_call_mult +!$loki end extract +end subroutine test_extract_mult """ routine = Subroutine.from_source(fcode, frontend=frontend) filepath = tmp_path/(f'{routine.name}_{frontend}.f90') @@ -432,10 +432,10 @@ def test_transform_region_to_call_multiple(tmp_path, frontend): assert len(FindNodes(CallStatement).visit(routine.body)) == 0 # Apply transformation - routines = region_to_call(routine) + routines = extract_marked_subroutines(routine) assert len(routines) == 3 assert routines[0].name == 'oiwjfklsf' - assert all(routines[i].name == f'{routine.name}_region_to_call_{i}' for i in (1,2)) + assert all(routines[i].name == f'{routine.name}_extracted_{i}' for i in (1,2)) assert len(FindNodes(Assignment).visit(routine.body)) == 4 assert all(len(FindNodes(Assignment).visit(r.body)) == 1 for r in routines) @@ -453,32 +453,32 @@ def test_transform_region_to_call_multiple(tmp_path, frontend): @pytest.mark.parametrize('frontend', available_frontends()) -def test_transform_region_to_call_arguments(tmp_path, frontend): +def test_extract_marked_subroutines_arguments(tmp_path, frontend): """ Test hoisting with multiple groups and multiple regions per group and automatic derivation of arguments """ fcode = """ -subroutine reg_to_call_args(a, b, c) +subroutine test_extract_args(a, b, c) integer, intent(out) :: a, b, c a = 1 a = a + 1 a = a + 1 -!$loki region-to-call name(func_a) +!$loki extract name(func_a) a = a + 1 -!$loki end region-to-call +!$loki end extract a = a + 1 -!$loki region-to-call name(func_b) +!$loki extract name(func_b) b = a -!$loki end region-to-call +!$loki end extract ! partially override arguments -!$loki region-to-call name(func_c) inout(b) +!$loki extract name(func_c) inout(b) c = a + b -!$loki end region-to-call -end subroutine reg_to_call_args +!$loki end extract +end subroutine test_extract_args """ routine = Subroutine.from_source(fcode, frontend=frontend) filepath = tmp_path/(f'{routine.name}_{frontend}.f90') @@ -492,7 +492,7 @@ def test_transform_region_to_call_arguments(tmp_path, frontend): assert len(FindNodes(CallStatement).visit(routine.body)) == 0 # Apply transformation - routines = region_to_call(routine) + routines = extract_marked_subroutines(routine) assert len(routines) == 3 assert [r.name for r in routines] == ['func_a', 'func_b', 'func_c'] @@ -524,35 +524,35 @@ def test_transform_region_to_call_arguments(tmp_path, frontend): @pytest.mark.parametrize('frontend', available_frontends()) -def test_transform_region_to_call_arrays(tmp_path, frontend): +def test_extract_marked_subroutines_arrays(tmp_path, frontend): """ Test hoisting with array variables """ fcode = """ -subroutine reg_to_call_arr(a, b, n) +subroutine test_extract_arr(a, b, n) integer, intent(out) :: a(n), b(n) integer, intent(in) :: n integer :: j -!$loki region-to-call +!$loki extract do j=1,n a(j) = j end do -!$loki end region-to-call +!$loki end extract -!$loki region-to-call +!$loki extract do j=1,n b(j) = j end do -!$loki end region-to-call +!$loki end extract -!$loki region-to-call +!$loki extract do j=1,n-1 b(j) = b(j+1) - a(j) end do b(n) = 1 -!$loki end region-to-call -end subroutine reg_to_call_arr +!$loki end extract +end subroutine test_extract_arr """ routine = Subroutine.from_source(fcode, frontend=frontend) @@ -571,7 +571,7 @@ def test_transform_region_to_call_arrays(tmp_path, frontend): assert len(FindNodes(CallStatement).visit(routine.body)) == 0 # Apply transformation - routines = region_to_call(routine) + routines = extract_marked_subroutines(routine) assert len(FindNodes(Assignment).visit(routine.body)) == 0 assert len(FindNodes(CallStatement).visit(routine.body)) == 3 @@ -598,49 +598,49 @@ def test_transform_region_to_call_arrays(tmp_path, frontend): @pytest.mark.parametrize('frontend', available_frontends()) -def test_transform_region_to_call_imports(tmp_path, builder, frontend): +def test_extract_marked_subroutines_imports(tmp_path, builder, frontend): """ Test hoisting with correct treatment of imports """ fcode_module = """ -module region_to_call_mod +module extract_mod implicit none integer, parameter :: param = 1 integer :: arr1(10) integer :: arr2(10) -end module region_to_call_mod +end module extract_mod """.strip() fcode = """ -module reg_to_call_imps_mod +module test_extract_imps_mod implicit none contains - subroutine reg_to_call_imps(a, b) - use region_to_call_mod, only: param, arr1, arr2 + subroutine test_extract_imps(a, b) + use extract_mod, only: param, arr1, arr2 integer, intent(out) :: a(10), b(10) integer :: j -!$loki region-to-call +!$loki extract do j=1,10 a(j) = param end do -!$loki end region-to-call +!$loki end extract -!$loki region-to-call +!$loki extract do j=1,10 arr1(j) = j+1 end do -!$loki end region-to-call +!$loki end extract arr2(:) = arr1(:) -!$loki region-to-call +!$loki extract do j=1,10 b(j) = arr2(j) - a(j) end do -!$loki end region-to-call - end subroutine reg_to_call_imps -end module reg_to_call_imps_mod +!$loki end extract + end subroutine test_extract_imps +end module test_extract_imps_mod """ ext_module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) module = Module.from_source(fcode, frontend=frontend, definitions=ext_module, xmods=[tmp_path]) @@ -660,7 +660,7 @@ def test_transform_region_to_call_imports(tmp_path, builder, frontend): assert len(FindNodes(CallStatement).visit(module.subroutines[0].body)) == 0 # Apply transformation - routines = region_to_call(module.subroutines[0]) + routines = extract_marked_subroutines(module.subroutines[0]) assert len(FindNodes(Assignment).visit(module.subroutines[0].body)) == 1 assert len(FindNodes(CallStatement).visit(module.subroutines[0].body)) == 3 diff --git a/loki/transformations/transform_region.py b/loki/transformations/transform_region.py index 32a46a8b8..5b9c9c56f 100644 --- a/loki/transformations/transform_region.py +++ b/loki/transformations/transform_region.py @@ -28,7 +28,7 @@ ) -__all__ = ['region_hoist', 'region_to_call'] +__all__ = ['region_hoist', 'extract_marked_subroutines'] def region_hoist(routine): @@ -124,14 +124,14 @@ def region_hoist(routine): promote_nonmatching_variables(routine, promotion_vars_dims, promotion_vars_index) -def region_to_call(routine): +def extract_marked_subroutines(routine): """ - Convert regions annotated with ``!$loki region-to-call`` pragmas to subroutine calls. + Convert regions annotated with ``!$loki extract`` pragmas to subroutine calls. The pragma syntax for regions to convert to subroutines is - ``!$loki region-to-call [name(...)] [in(...)] [out(...)] [inout(...)]`` - and ``!$loki end region-to-call``. + ``!$loki extract [name(...)] [in(...)] [out(...)] [inout(...)]`` + and ``!$loki end extract``. 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. @@ -155,12 +155,12 @@ def region_to_call(routine): 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='region-to-call'): + if not is_loki_pragma(region.pragma, starts_with='extract'): continue # Name the external routine - parameters = get_pragma_parameters(region.pragma, starts_with='region-to-call') - name = parameters.get('name', f'{routine.name}_region_to_call_{counter}') + parameters = get_pragma_parameters(region.pragma, starts_with='extract') + name = parameters.get('name', f'{routine.name}_extracted_{counter}') counter += 1 # Create the external subroutine containing the routine's imports and the region's body From 2ed72882ddf4b24ac5737f6510ffbc82b84a4d73 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Thu, 12 Sep 2024 03:42:32 +0000 Subject: [PATCH 05/12] Transformations: Move `extract_marked_subroutine` into `transformations.extract` --- loki/transformations/extract/__init__.py | 1 + loki/transformations/extract/marked.py | 132 +++++++ .../extract/tests/test_extract_marked.py | 363 ++++++++++++++++++ .../tests/test_transform_region.py | 357 +---------------- loki/transformations/transform_region.py | 124 +----- 5 files changed, 503 insertions(+), 474 deletions(-) create mode 100644 loki/transformations/extract/marked.py create mode 100644 loki/transformations/extract/tests/test_extract_marked.py diff --git a/loki/transformations/extract/__init__.py b/loki/transformations/extract/__init__.py index 9b82c36c9..c6212cc33 100644 --- a/loki/transformations/extract/__init__.py +++ b/loki/transformations/extract/__init__.py @@ -10,3 +10,4 @@ """ from loki.transformations.extract.internal import * # noqa +from loki.transformations.extract.marked import * # noqa diff --git a/loki/transformations/extract/marked.py b/loki/transformations/extract/marked.py new file mode 100644 index 000000000..c323e9ab0 --- /dev/null +++ b/loki/transformations/extract/marked.py @@ -0,0 +1,132 @@ +# (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__ = ['extract_marked_subroutines'] + + +def extract_marked_subroutines(routine): + """ + Convert regions annotated with ``!$loki extract`` pragmas to subroutine calls. + + + The pragma syntax for regions to convert to subroutines is + ``!$loki extract [name(...)] [in(...)] [out(...)] [inout(...)]`` + and ``!$loki end extract``. + + 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. + + :param :class:``Subroutine`` routine: + the routine to modify. + + :return: 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='extract'): + continue + + # Name the external routine + parameters = get_pragma_parameters(region.pragma, starts_with='extract') + name = parameters.get('name', f'{routine.name}_extracted_{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 diff --git a/loki/transformations/extract/tests/test_extract_marked.py b/loki/transformations/extract/tests/test_extract_marked.py new file mode 100644 index 000000000..c9ef897fa --- /dev/null +++ b/loki/transformations/extract/tests/test_extract_marked.py @@ -0,0 +1,363 @@ +# (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. + +import pytest +import numpy as np + +from loki import Module, Subroutine +from loki.build import jit_compile, jit_compile_lib, Builder, Obj +from loki.frontend import available_frontends +from loki.ir import FindNodes, Section, Assignment, CallStatement, Intrinsic +from loki.tools import as_tuple + +from loki.transformations.extract.marked import extract_marked_subroutines + + +@pytest.fixture(scope='function', name='builder') +def fixture_builder(tmp_path): + yield Builder(source_dirs=tmp_path, build_dir=tmp_path/'build') + Obj.clear_cache() + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_extract_marked_subroutines(tmp_path, frontend): + """ + A very simple :any:`extract_marked_subroutine` test case + """ + fcode = """ +subroutine test_extract(a, b, c) + integer, intent(out) :: a, b, c + + a = 5 + a = 1 + +!$loki extract in(a) out(b) + b = a +!$loki end extract + + c = a + b +end subroutine test_extract +""" + routine = Subroutine.from_source(fcode, frontend=frontend) + filepath = tmp_path/(f'{routine.name}_{frontend}.f90') + function = jit_compile(routine, filepath=filepath, objname=routine.name) + + # Test the reference solution + a, b, c = function() + assert a == 1 and b == 1 and c == 2 + + assert len(FindNodes(Assignment).visit(routine.body)) == 4 + assert len(FindNodes(CallStatement).visit(routine.body)) == 0 + + # Apply transformation + routines = extract_marked_subroutines(routine) + assert len(routines) == 1 and routines[0].name == f'{routine.name}_extracted_0' + + assert len(FindNodes(Assignment).visit(routine.body)) == 3 + assert len(FindNodes(Assignment).visit(routines[0].body)) == 1 + assert len(FindNodes(CallStatement).visit(routine.body)) == 1 + + # Test transformation + contains = Section(body=as_tuple([Intrinsic('CONTAINS'), *routines, routine])) + module = Module(name=f'{routine.name}_mod', spec=None, contains=contains) + mod_filepath = tmp_path/(f'{module.name}_converted_{frontend}.f90') + mod = jit_compile(module, filepath=mod_filepath, objname=module.name) + mod_function = getattr(mod, routine.name) + + a, b, c = mod_function() + assert a == 1 and b == 1 and c == 2 + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_extract_marked_subroutines_multiple(tmp_path, frontend): + """ + Test hoisting with multiple groups and multiple regions per group + """ + fcode = """ +subroutine test_extract_mult(a, b, c) + integer, intent(out) :: a, b, c + + a = 1 + a = a + 1 + a = a + 1 +!$loki extract name(oiwjfklsf) inout(a) + a = a + 1 +!$loki end extract + a = a + 1 + +!$loki extract in(a) out(b) + b = a +!$loki end extract + +!$loki extract in(a,b) out(c) + c = a + b +!$loki end extract +end subroutine test_extract_mult +""" + routine = Subroutine.from_source(fcode, frontend=frontend) + filepath = tmp_path/(f'{routine.name}_{frontend}.f90') + function = jit_compile(routine, filepath=filepath, objname=routine.name) + + # Test the reference solution + a, b, c = function() + assert a == 5 and b == 5 and c == 10 + + assert len(FindNodes(Assignment).visit(routine.body)) == 7 + assert len(FindNodes(CallStatement).visit(routine.body)) == 0 + + # Apply transformation + routines = extract_marked_subroutines(routine) + assert len(routines) == 3 + assert routines[0].name == 'oiwjfklsf' + assert all(routines[i].name == f'{routine.name}_extracted_{i}' for i in (1,2)) + + assert len(FindNodes(Assignment).visit(routine.body)) == 4 + assert all(len(FindNodes(Assignment).visit(r.body)) == 1 for r in routines) + assert len(FindNodes(CallStatement).visit(routine.body)) == 3 + + # Test transformation + contains = Section(body=as_tuple([Intrinsic('CONTAINS'), *routines, routine])) + module = Module(name=f'{routine.name}_mod', spec=None, contains=contains) + mod_filepath = tmp_path/(f'{module.name}_converted_{frontend}.f90') + mod = jit_compile(module, filepath=mod_filepath, objname=module.name) + mod_function = getattr(mod, routine.name) + + a, b, c = mod_function() + assert a == 5 and b == 5 and c == 10 + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_extract_marked_subroutines_arguments(tmp_path, frontend): + """ + Test hoisting with multiple groups and multiple regions per group + and automatic derivation of arguments + """ + fcode = """ +subroutine test_extract_args(a, b, c) + integer, intent(out) :: a, b, c + + a = 1 + a = a + 1 + a = a + 1 +!$loki extract name(func_a) + a = a + 1 +!$loki end extract + a = a + 1 + +!$loki extract name(func_b) + b = a +!$loki end extract + +! partially override arguments +!$loki extract name(func_c) inout(b) + c = a + b +!$loki end extract +end subroutine test_extract_args +""" + routine = Subroutine.from_source(fcode, frontend=frontend) + filepath = tmp_path/(f'{routine.name}_{frontend}.f90') + function = jit_compile(routine, filepath=filepath, objname=routine.name) + + # Test the reference solution + a, b, c = function() + assert a == 5 and b == 5 and c == 10 + + assert len(FindNodes(Assignment).visit(routine.body)) == 7 + assert len(FindNodes(CallStatement).visit(routine.body)) == 0 + + # Apply transformation + routines = extract_marked_subroutines(routine) + assert len(routines) == 3 + assert [r.name for r in routines] == ['func_a', 'func_b', 'func_c'] + + assert len(routines[0].arguments) == 1 + assert routines[0].arguments[0] == 'a' and routines[0].arguments[0].type.intent == 'inout' + + assert {str(a) for a in routines[1].arguments} == {'a', 'b'} + assert routines[1].variable_map['a'].type.intent == 'in' + assert routines[1].variable_map['b'].type.intent == 'out' + + assert {str(a) for a in routines[2].arguments} == {'a', 'b', 'c'} + assert routines[2].variable_map['a'].type.intent == 'in' + assert routines[2].variable_map['b'].type.intent == 'inout' + assert routines[2].variable_map['c'].type.intent == 'out' + + assert len(FindNodes(Assignment).visit(routine.body)) == 4 + assert all(len(FindNodes(Assignment).visit(r.body)) == 1 for r in routines) + assert len(FindNodes(CallStatement).visit(routine.body)) == 3 + + # Test transformation + contains = Section(body=as_tuple([Intrinsic('CONTAINS'), *routines, routine])) + module = Module(name=f'{routine.name}_mod', spec=None, contains=contains) + mod_filepath = tmp_path/(f'{module.name}_converted_{frontend}.f90') + mod = jit_compile(module, filepath=mod_filepath, objname=module.name) + mod_function = getattr(mod, routine.name) + + a, b, c = mod_function() + assert a == 5 and b == 5 and c == 10 + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_extract_marked_subroutines_arrays(tmp_path, frontend): + """ + Test hoisting with array variables + """ + fcode = """ +subroutine test_extract_arr(a, b, n) + integer, intent(out) :: a(n), b(n) + integer, intent(in) :: n + integer :: j + +!$loki extract + do j=1,n + a(j) = j + end do +!$loki end extract + +!$loki extract + do j=1,n + b(j) = j + end do +!$loki end extract + +!$loki extract + do j=1,n-1 + b(j) = b(j+1) - a(j) + end do + b(n) = 1 +!$loki end extract +end subroutine test_extract_arr +""" + routine = Subroutine.from_source(fcode, frontend=frontend) + + filepath = tmp_path/(f'{routine.name}_{frontend}.f90') + function = jit_compile(routine, filepath=filepath, objname=routine.name) + + # Test the reference solution + n = 10 + a = np.zeros(shape=(n,), dtype=np.int32) + b = np.zeros(shape=(n,), dtype=np.int32) + function(a, b, n) + assert np.all(a == range(1,n+1)) + assert np.all(b == [1] * n) + + assert len(FindNodes(Assignment).visit(routine.body)) == 4 + assert len(FindNodes(CallStatement).visit(routine.body)) == 0 + + # Apply transformation + routines = extract_marked_subroutines(routine) + + assert len(FindNodes(Assignment).visit(routine.body)) == 0 + assert len(FindNodes(CallStatement).visit(routine.body)) == 3 + + assert len(routines) == 3 + + assert {(str(a), a.type.intent) for a in routines[0].arguments} == {('a(n)', 'out'), ('n', 'in')} + assert {(str(a), a.type.intent) for a in routines[1].arguments} == {('b(n)', 'out'), ('n', 'in')} + assert {(str(a), a.type.intent) for a in routines[2].arguments} == {('a(n)', 'in'), ('b(n)', 'inout'), ('n', 'in')} + assert routines[0].variable_map['a'].dimensions[0].scope is routines[0] + + # Test transformation + contains = Section(body=as_tuple([Intrinsic('CONTAINS'), *routines, routine])) + module = Module(name=f'{routine.name}_mod', spec=None, contains=contains) + mod_filepath = tmp_path/(f'{module.name}_converted_{frontend}.f90') + mod = jit_compile(module, filepath=mod_filepath, objname=module.name) + mod_function = getattr(mod, routine.name) + + a = np.zeros(shape=(n,), dtype=np.int32) + b = np.zeros(shape=(n,), dtype=np.int32) + mod_function(a, b, n) + assert np.all(a == range(1,n+1)) + assert np.all(b == [1] * n) + + +@pytest.mark.parametrize('frontend', available_frontends()) +def test_extract_marked_subroutines_imports(tmp_path, builder, frontend): + """ + Test hoisting with correct treatment of imports + """ + fcode_module = """ +module extract_mod + implicit none + integer, parameter :: param = 1 + integer :: arr1(10) + integer :: arr2(10) +end module extract_mod + """.strip() + + fcode = """ +module test_extract_imps_mod + implicit none +contains + subroutine test_extract_imps(a, b) + use extract_mod, only: param, arr1, arr2 + integer, intent(out) :: a(10), b(10) + integer :: j + +!$loki extract + do j=1,10 + a(j) = param + end do +!$loki end extract + +!$loki extract + do j=1,10 + arr1(j) = j+1 + end do +!$loki end extract + + arr2(:) = arr1(:) + +!$loki extract + do j=1,10 + b(j) = arr2(j) - a(j) + end do +!$loki end extract + end subroutine test_extract_imps +end module test_extract_imps_mod +""" + ext_module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) + module = Module.from_source(fcode, frontend=frontend, definitions=ext_module, xmods=[tmp_path]) + refname = f'ref_{module.name}_{frontend}' + reference = jit_compile_lib([module, ext_module], path=tmp_path, name=refname, builder=builder) + function = getattr(getattr(reference, module.name), module.subroutines[0].name) + + # Test the reference solution + a = np.zeros(shape=(10,), dtype=np.int32) + b = np.zeros(shape=(10,), dtype=np.int32) + function(a, b) + assert np.all(a == [1] * 10) + assert np.all(b == range(1,11)) + (tmp_path/f'{module.name}.f90').unlink() + + assert len(FindNodes(Assignment).visit(module.subroutines[0].body)) == 4 + assert len(FindNodes(CallStatement).visit(module.subroutines[0].body)) == 0 + + # Apply transformation + routines = extract_marked_subroutines(module.subroutines[0]) + + assert len(FindNodes(Assignment).visit(module.subroutines[0].body)) == 1 + assert len(FindNodes(CallStatement).visit(module.subroutines[0].body)) == 3 + + assert len(routines) == 3 + + assert {(str(a), a.type.intent) for a in routines[0].arguments} == {('a(10)', 'out')} + assert {(str(a), a.type.intent) for a in routines[1].arguments} == set() + assert {(str(a), a.type.intent) for a in routines[2].arguments} == {('a(10)', 'in'), ('b(10)', 'out')} + + # Insert created routines into module + module.contains.append(routines) + + obj = jit_compile_lib([module, ext_module], path=tmp_path, name=f'{module.name}_{frontend}', builder=builder) + mod_function = getattr(getattr(obj, module.name), module.subroutines[0].name) + + # Test transformation + a = np.zeros(shape=(10,), dtype=np.int32) + b = np.zeros(shape=(10,), dtype=np.int32) + mod_function(a, b) + assert np.all(a == [1] * 10) + assert np.all(b == range(1,11)) diff --git a/loki/transformations/tests/test_transform_region.py b/loki/transformations/tests/test_transform_region.py index 8e88ef6fa..d619dfa41 100644 --- a/loki/transformations/tests/test_transform_region.py +++ b/loki/transformations/tests/test_transform_region.py @@ -8,23 +8,12 @@ import pytest import numpy as np -from loki import ( - Module, Subroutine, Section, as_tuple, FindNodes, Loop, - Assignment, CallStatement, Intrinsic -) -from loki.build import jit_compile, jit_compile_lib, Builder, Obj +from loki import Subroutine, FindNodes, Loop +from loki.build import jit_compile from loki.expression import symbols as sym from loki.frontend import available_frontends -from loki.transformations.transform_region import ( - region_hoist, extract_marked_subroutines -) - - -@pytest.fixture(scope='function', name='builder') -def fixture_builder(tmp_path): - yield Builder(source_dirs=tmp_path, build_dir=tmp_path/'build') - Obj.clear_cache() +from loki.transformations.transform_region import region_hoist @pytest.mark.parametrize('frontend', available_frontends()) @@ -343,343 +332,3 @@ def test_transform_region_hoist_promote(tmp_path, frontend): hoisted_function(a=a, b=b, klon=klon, klev=klev) assert np.all(a == ref_a) assert np.all(b == ref_b) - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_marked_subroutines(tmp_path, frontend): - """ - A very simple :any:`extract_marked_subroutine` test case - """ - fcode = """ -subroutine test_extract(a, b, c) - integer, intent(out) :: a, b, c - - a = 5 - a = 1 - -!$loki extract in(a) out(b) - b = a -!$loki end extract - - c = a + b -end subroutine test_extract -""" - routine = Subroutine.from_source(fcode, frontend=frontend) - filepath = tmp_path/(f'{routine.name}_{frontend}.f90') - function = jit_compile(routine, filepath=filepath, objname=routine.name) - - # Test the reference solution - a, b, c = function() - assert a == 1 and b == 1 and c == 2 - - assert len(FindNodes(Assignment).visit(routine.body)) == 4 - assert len(FindNodes(CallStatement).visit(routine.body)) == 0 - - # Apply transformation - routines = extract_marked_subroutines(routine) - assert len(routines) == 1 and routines[0].name == f'{routine.name}_extracted_0' - - assert len(FindNodes(Assignment).visit(routine.body)) == 3 - assert len(FindNodes(Assignment).visit(routines[0].body)) == 1 - assert len(FindNodes(CallStatement).visit(routine.body)) == 1 - - # Test transformation - contains = Section(body=as_tuple([Intrinsic('CONTAINS'), *routines, routine])) - module = Module(name=f'{routine.name}_mod', spec=None, contains=contains) - mod_filepath = tmp_path/(f'{module.name}_converted_{frontend}.f90') - mod = jit_compile(module, filepath=mod_filepath, objname=module.name) - mod_function = getattr(mod, routine.name) - - a, b, c = mod_function() - assert a == 1 and b == 1 and c == 2 - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_marked_subroutines_multiple(tmp_path, frontend): - """ - Test hoisting with multiple groups and multiple regions per group - """ - fcode = """ -subroutine test_extract_mult(a, b, c) - integer, intent(out) :: a, b, c - - a = 1 - a = a + 1 - a = a + 1 -!$loki extract name(oiwjfklsf) inout(a) - a = a + 1 -!$loki end extract - a = a + 1 - -!$loki extract in(a) out(b) - b = a -!$loki end extract - -!$loki extract in(a,b) out(c) - c = a + b -!$loki end extract -end subroutine test_extract_mult -""" - routine = Subroutine.from_source(fcode, frontend=frontend) - filepath = tmp_path/(f'{routine.name}_{frontend}.f90') - function = jit_compile(routine, filepath=filepath, objname=routine.name) - - # Test the reference solution - a, b, c = function() - assert a == 5 and b == 5 and c == 10 - - assert len(FindNodes(Assignment).visit(routine.body)) == 7 - assert len(FindNodes(CallStatement).visit(routine.body)) == 0 - - # Apply transformation - routines = extract_marked_subroutines(routine) - assert len(routines) == 3 - assert routines[0].name == 'oiwjfklsf' - assert all(routines[i].name == f'{routine.name}_extracted_{i}' for i in (1,2)) - - assert len(FindNodes(Assignment).visit(routine.body)) == 4 - assert all(len(FindNodes(Assignment).visit(r.body)) == 1 for r in routines) - assert len(FindNodes(CallStatement).visit(routine.body)) == 3 - - # Test transformation - contains = Section(body=as_tuple([Intrinsic('CONTAINS'), *routines, routine])) - module = Module(name=f'{routine.name}_mod', spec=None, contains=contains) - mod_filepath = tmp_path/(f'{module.name}_converted_{frontend}.f90') - mod = jit_compile(module, filepath=mod_filepath, objname=module.name) - mod_function = getattr(mod, routine.name) - - a, b, c = mod_function() - assert a == 5 and b == 5 and c == 10 - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_marked_subroutines_arguments(tmp_path, frontend): - """ - Test hoisting with multiple groups and multiple regions per group - and automatic derivation of arguments - """ - fcode = """ -subroutine test_extract_args(a, b, c) - integer, intent(out) :: a, b, c - - a = 1 - a = a + 1 - a = a + 1 -!$loki extract name(func_a) - a = a + 1 -!$loki end extract - a = a + 1 - -!$loki extract name(func_b) - b = a -!$loki end extract - -! partially override arguments -!$loki extract name(func_c) inout(b) - c = a + b -!$loki end extract -end subroutine test_extract_args -""" - routine = Subroutine.from_source(fcode, frontend=frontend) - filepath = tmp_path/(f'{routine.name}_{frontend}.f90') - function = jit_compile(routine, filepath=filepath, objname=routine.name) - - # Test the reference solution - a, b, c = function() - assert a == 5 and b == 5 and c == 10 - - assert len(FindNodes(Assignment).visit(routine.body)) == 7 - assert len(FindNodes(CallStatement).visit(routine.body)) == 0 - - # Apply transformation - routines = extract_marked_subroutines(routine) - assert len(routines) == 3 - assert [r.name for r in routines] == ['func_a', 'func_b', 'func_c'] - - assert len(routines[0].arguments) == 1 - assert routines[0].arguments[0] == 'a' and routines[0].arguments[0].type.intent == 'inout' - - assert {str(a) for a in routines[1].arguments} == {'a', 'b'} - assert routines[1].variable_map['a'].type.intent == 'in' - assert routines[1].variable_map['b'].type.intent == 'out' - - assert {str(a) for a in routines[2].arguments} == {'a', 'b', 'c'} - assert routines[2].variable_map['a'].type.intent == 'in' - assert routines[2].variable_map['b'].type.intent == 'inout' - assert routines[2].variable_map['c'].type.intent == 'out' - - assert len(FindNodes(Assignment).visit(routine.body)) == 4 - assert all(len(FindNodes(Assignment).visit(r.body)) == 1 for r in routines) - assert len(FindNodes(CallStatement).visit(routine.body)) == 3 - - # Test transformation - contains = Section(body=as_tuple([Intrinsic('CONTAINS'), *routines, routine])) - module = Module(name=f'{routine.name}_mod', spec=None, contains=contains) - mod_filepath = tmp_path/(f'{module.name}_converted_{frontend}.f90') - mod = jit_compile(module, filepath=mod_filepath, objname=module.name) - mod_function = getattr(mod, routine.name) - - a, b, c = mod_function() - assert a == 5 and b == 5 and c == 10 - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_marked_subroutines_arrays(tmp_path, frontend): - """ - Test hoisting with array variables - """ - fcode = """ -subroutine test_extract_arr(a, b, n) - integer, intent(out) :: a(n), b(n) - integer, intent(in) :: n - integer :: j - -!$loki extract - do j=1,n - a(j) = j - end do -!$loki end extract - -!$loki extract - do j=1,n - b(j) = j - end do -!$loki end extract - -!$loki extract - do j=1,n-1 - b(j) = b(j+1) - a(j) - end do - b(n) = 1 -!$loki end extract -end subroutine test_extract_arr -""" - routine = Subroutine.from_source(fcode, frontend=frontend) - - filepath = tmp_path/(f'{routine.name}_{frontend}.f90') - function = jit_compile(routine, filepath=filepath, objname=routine.name) - - # Test the reference solution - n = 10 - a = np.zeros(shape=(n,), dtype=np.int32) - b = np.zeros(shape=(n,), dtype=np.int32) - function(a, b, n) - assert np.all(a == range(1,n+1)) - assert np.all(b == [1] * n) - - assert len(FindNodes(Assignment).visit(routine.body)) == 4 - assert len(FindNodes(CallStatement).visit(routine.body)) == 0 - - # Apply transformation - routines = extract_marked_subroutines(routine) - - assert len(FindNodes(Assignment).visit(routine.body)) == 0 - assert len(FindNodes(CallStatement).visit(routine.body)) == 3 - - assert len(routines) == 3 - - assert {(str(a), a.type.intent) for a in routines[0].arguments} == {('a(n)', 'out'), ('n', 'in')} - assert {(str(a), a.type.intent) for a in routines[1].arguments} == {('b(n)', 'out'), ('n', 'in')} - assert {(str(a), a.type.intent) for a in routines[2].arguments} == {('a(n)', 'in'), ('b(n)', 'inout'), ('n', 'in')} - assert routines[0].variable_map['a'].dimensions[0].scope is routines[0] - - # Test transformation - contains = Section(body=as_tuple([Intrinsic('CONTAINS'), *routines, routine])) - module = Module(name=f'{routine.name}_mod', spec=None, contains=contains) - mod_filepath = tmp_path/(f'{module.name}_converted_{frontend}.f90') - mod = jit_compile(module, filepath=mod_filepath, objname=module.name) - mod_function = getattr(mod, routine.name) - - a = np.zeros(shape=(n,), dtype=np.int32) - b = np.zeros(shape=(n,), dtype=np.int32) - mod_function(a, b, n) - assert np.all(a == range(1,n+1)) - assert np.all(b == [1] * n) - - -@pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_marked_subroutines_imports(tmp_path, builder, frontend): - """ - Test hoisting with correct treatment of imports - """ - fcode_module = """ -module extract_mod - implicit none - integer, parameter :: param = 1 - integer :: arr1(10) - integer :: arr2(10) -end module extract_mod - """.strip() - - fcode = """ -module test_extract_imps_mod - implicit none -contains - subroutine test_extract_imps(a, b) - use extract_mod, only: param, arr1, arr2 - integer, intent(out) :: a(10), b(10) - integer :: j - -!$loki extract - do j=1,10 - a(j) = param - end do -!$loki end extract - -!$loki extract - do j=1,10 - arr1(j) = j+1 - end do -!$loki end extract - - arr2(:) = arr1(:) - -!$loki extract - do j=1,10 - b(j) = arr2(j) - a(j) - end do -!$loki end extract - end subroutine test_extract_imps -end module test_extract_imps_mod -""" - ext_module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) - module = Module.from_source(fcode, frontend=frontend, definitions=ext_module, xmods=[tmp_path]) - refname = f'ref_{module.name}_{frontend}' - reference = jit_compile_lib([module, ext_module], path=tmp_path, name=refname, builder=builder) - function = getattr(getattr(reference, module.name), module.subroutines[0].name) - - # Test the reference solution - a = np.zeros(shape=(10,), dtype=np.int32) - b = np.zeros(shape=(10,), dtype=np.int32) - function(a, b) - assert np.all(a == [1] * 10) - assert np.all(b == range(1,11)) - (tmp_path/f'{module.name}.f90').unlink() - - assert len(FindNodes(Assignment).visit(module.subroutines[0].body)) == 4 - assert len(FindNodes(CallStatement).visit(module.subroutines[0].body)) == 0 - - # Apply transformation - routines = extract_marked_subroutines(module.subroutines[0]) - - assert len(FindNodes(Assignment).visit(module.subroutines[0].body)) == 1 - assert len(FindNodes(CallStatement).visit(module.subroutines[0].body)) == 3 - - assert len(routines) == 3 - - assert {(str(a), a.type.intent) for a in routines[0].arguments} == {('a(10)', 'out')} - assert {(str(a), a.type.intent) for a in routines[1].arguments} == set() - assert {(str(a), a.type.intent) for a in routines[2].arguments} == {('a(10)', 'in'), ('b(10)', 'out')} - - # Insert created routines into module - module.contains.append(routines) - - obj = jit_compile_lib([module, ext_module], path=tmp_path, name=f'{module.name}_{frontend}', builder=builder) - mod_function = getattr(getattr(obj, module.name), module.subroutines[0].name) - - # Test transformation - a = np.zeros(shape=(10,), dtype=np.int32) - b = np.zeros(shape=(10,), dtype=np.int32) - mod_function(a, b) - assert np.all(a == [1] * 10) - assert np.all(b == range(1,11)) diff --git a/loki/transformations/transform_region.py b/loki/transformations/transform_region.py index 5b9c9c56f..d6b93d6e8 100644 --- a/loki/transformations/transform_region.py +++ b/loki/transformations/transform_region.py @@ -11,24 +11,20 @@ """ from collections import defaultdict -from loki.analyse import dataflow_analysis_attached -from loki.expression import Variable from loki.ir import ( - CallStatement, Comment, Import, Loop, Pragma, PragmaRegion, - Section, FindNodes, FindScopes, FindVariables, MaskedTransformer, - NestedMaskedTransformer, Transformer, is_loki_pragma, + Comment, Loop, Pragma, PragmaRegion, FindNodes, FindScopes, + MaskedTransformer, NestedMaskedTransformer, 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, flatten, CaseInsensitiveDict +from loki.tools import as_tuple, flatten from loki.transformations.array_indexing import ( promotion_dimensions_from_loop_nest, promote_nonmatching_variables ) -__all__ = ['region_hoist', 'extract_marked_subroutines'] +__all__ = ['region_hoist'] def region_hoist(routine): @@ -122,115 +118,3 @@ def region_hoist(routine): num_targets = sum(1 for pragma in hoist_map if 'target' in get_pragma_parameters(pragma)) info('%s: hoisted %d region(s) in %d group(s)', routine.name, len(hoist_map) - num_targets, num_targets) promote_nonmatching_variables(routine, promotion_vars_dims, promotion_vars_index) - - -def extract_marked_subroutines(routine): - """ - Convert regions annotated with ``!$loki extract`` pragmas to subroutine calls. - - - The pragma syntax for regions to convert to subroutines is - ``!$loki extract [name(...)] [in(...)] [out(...)] [inout(...)]`` - and ``!$loki end extract``. - - 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. - - :param :class:``Subroutine`` routine: - the routine to modify. - - :return: 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='extract'): - continue - - # Name the external routine - parameters = get_pragma_parameters(region.pragma, starts_with='extract') - name = parameters.get('name', f'{routine.name}_extracted_{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 From f229bd1ff5d379e58473e96a3afd417b3ab97be3 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Thu, 12 Sep 2024 03:55:41 +0000 Subject: [PATCH 06/12] Transformations: Some linter fixes for transformation tests --- .../tests/test_hoist_variables.py | 2 +- .../transformations/tests/test_idempotence.py | 2 +- .../transformations/tests/test_parametrise.py | 4 +-- .../tests/test_raw_stack_allocator.py | 2 +- .../tests/test_transform_loop.py | 32 +++++++++---------- loki/transformations/tests/test_utilities.py | 24 +++++++------- 6 files changed, 32 insertions(+), 34 deletions(-) diff --git a/loki/transformations/tests/test_hoist_variables.py b/loki/transformations/tests/test_hoist_variables.py index 6e2d5a1fb..b3d1eff91 100644 --- a/loki/transformations/tests/test_hoist_variables.py +++ b/loki/transformations/tests/test_hoist_variables.py @@ -16,7 +16,7 @@ Scheduler, SchedulerConfig, is_iterable, FindInlineCalls ) from loki.build import jit_compile_lib, Builder -from loki.frontend import available_frontends, OMNI +from loki.frontend import available_frontends from loki.ir import nodes as ir, FindNodes from loki.transformations.hoist_variables import ( HoistVariablesAnalysis, HoistVariablesTransformation, diff --git a/loki/transformations/tests/test_idempotence.py b/loki/transformations/tests/test_idempotence.py index a5fb1e524..4b4c84cb0 100644 --- a/loki/transformations/tests/test_idempotence.py +++ b/loki/transformations/tests/test_idempotence.py @@ -15,7 +15,7 @@ @pytest.mark.parametrize('frontend', available_frontends()) -def test_transform_idempotence(frontend, tmp_path): +def test_transform_idempotence(frontend): """ Test the do-nothing equivalence of :any:`IdemTransformations` """ fcode_driver = """ diff --git a/loki/transformations/tests/test_parametrise.py b/loki/transformations/tests/test_parametrise.py index 484b5bcd3..0f4de1466 100644 --- a/loki/transformations/tests/test_parametrise.py +++ b/loki/transformations/tests/test_parametrise.py @@ -356,7 +356,7 @@ def stop_execution(**kwargs): abort_callbacks = (error_stop, stop_execution) - for i, abort_callback in enumerate(abort_callbacks): + for _, abort_callback in enumerate(abort_callbacks): scheduler = Scheduler( paths=[proj], config=config, seed_routines=['driver', 'another_driver'], frontend=frontend, xmods=[tmp_path] @@ -380,8 +380,6 @@ def test_parametrise_modified_callback_wrong_input(tmp_path, testdir, frontend, proj = testdir/'sources/projParametrise' dic2p = {'a': 12, 'b': 11} - a = dic2p['a'] - b = dic2p['b'] def only_warn(**kwargs): msg = kwargs.get("msg") diff --git a/loki/transformations/tests/test_raw_stack_allocator.py b/loki/transformations/tests/test_raw_stack_allocator.py index fc9ed0822..49e6f2346 100644 --- a/loki/transformations/tests/test_raw_stack_allocator.py +++ b/loki/transformations/tests/test_raw_stack_allocator.py @@ -8,7 +8,7 @@ import pytest from loki.backend import fgen -from loki.batch import Scheduler, SchedulerConfig, ProcedureItem +from loki.batch import Scheduler, SchedulerConfig from loki.dimension import Dimension from loki.expression import DeferredTypeSymbol, InlineCall, IntLiteral from loki.frontend import available_frontends, OMNI diff --git a/loki/transformations/tests/test_transform_loop.py b/loki/transformations/tests/test_transform_loop.py index a2209f4c3..52c7dc571 100644 --- a/loki/transformations/tests/test_transform_loop.py +++ b/loki/transformations/tests/test_transform_loop.py @@ -13,7 +13,7 @@ from loki import Subroutine from loki.build import jit_compile, clean_test from loki.expression import symbols as sym -from loki.frontend import available_frontends, OMNI +from loki.frontend import available_frontends from loki.ir import ( is_loki_pragma, pragmas_attached, FindNodes, Loop, Conditional, Assignment @@ -1631,7 +1631,7 @@ def test_transform_loop_unroll(tmp_path, frontend): # Test the reference solution s = np.zeros(1) function(s=s) - assert s == sum([x + 1 for x in range(1, 11)]) + assert s == sum(x + 1 for x in range(1, 11)) # Apply transformation assert len(FindNodes(Loop).visit(routine.body)) == 1 @@ -1644,7 +1644,7 @@ def test_transform_loop_unroll(tmp_path, frontend): # Test transformation s = np.zeros(1) unrolled_function(s=s) - assert s == sum([x + 1 for x in range(1, 11)]) + assert s == sum(x + 1 for x in range(1, 11)) clean_test(filepath) clean_test(unrolled_filepath) @@ -1673,7 +1673,7 @@ def test_transform_loop_unroll_step(tmp_path, frontend): # Test the reference solution s = np.zeros(1) function(s=s) - assert s == sum([x + 1 for x in range(1, 11, 2)]) + assert s == sum(x + 1 for x in range(1, 11, 2)) # Apply transformation assert len(FindNodes(Loop).visit(routine.body)) == 1 @@ -1686,7 +1686,7 @@ def test_transform_loop_unroll_step(tmp_path, frontend): # Test transformation s = np.zeros(1) unrolled_function(s=s) - assert s == sum([x + 1 for x in range(1, 11, 2)]) + assert s == sum(x + 1 for x in range(1, 11, 2)) clean_test(filepath) clean_test(unrolled_filepath) @@ -1717,7 +1717,7 @@ def test_transform_loop_unroll_non_literal_range(tmp_path, frontend): # Test the reference solution s = np.zeros(1) function(s=s) - assert s == sum([x + 1 for x in range(1, 11)]) + assert s == sum(x + 1 for x in range(1, 11)) # Apply transformation assert len(FindNodes(Loop).visit(routine.body)) == 1 @@ -1730,7 +1730,7 @@ def test_transform_loop_unroll_non_literal_range(tmp_path, frontend): # Test transformation s = np.zeros(1) unrolled_function(s=s) - assert s == sum([x + 1 for x in range(1, 11)]) + assert s == sum(x + 1 for x in range(1, 11)) clean_test(filepath) clean_test(unrolled_filepath) @@ -1762,7 +1762,7 @@ def test_transform_loop_unroll_nested(tmp_path, frontend): # Test the reference solution s = np.zeros(1) function(s=s) - assert s == sum([a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))]) + assert s == sum(a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))) # Apply transformation assert len(FindNodes(Loop).visit(routine.body)) == 2 @@ -1775,7 +1775,7 @@ def test_transform_loop_unroll_nested(tmp_path, frontend): # Test transformation s = np.zeros(1) unrolled_function(s=s) - assert s == sum([a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))]) + assert s == sum(a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))) clean_test(filepath) clean_test(unrolled_filepath) @@ -1807,7 +1807,7 @@ def test_transform_loop_unroll_nested_restricted_depth(tmp_path, frontend): # Test the reference solution s = np.zeros(1) function(s=s) - assert s == sum([a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))]) + assert s == sum(a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))) # Apply transformation assert len(FindNodes(Loop).visit(routine.body)) == 2 @@ -1820,7 +1820,7 @@ def test_transform_loop_unroll_nested_restricted_depth(tmp_path, frontend): # Test transformation s = np.zeros(1) unrolled_function(s=s) - assert s == sum([a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))]) + assert s == sum(a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))) clean_test(filepath) clean_test(unrolled_filepath) @@ -1854,7 +1854,7 @@ def test_transform_loop_unroll_nested_restricted_depth_unrollable(tmp_path, fron # Test the reference solution s = np.zeros(1) function(s=s) - assert s == sum([a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))]) + assert s == sum(a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))) # Apply transformation assert len(FindNodes(Loop).visit(routine.body)) == 2 @@ -1867,7 +1867,7 @@ def test_transform_loop_unroll_nested_restricted_depth_unrollable(tmp_path, fron # Test transformation s = np.zeros(1) unrolled_function(s=s) - assert s == sum([a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))]) + assert s == sum(a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))) clean_test(filepath) clean_test(unrolled_filepath) @@ -1915,7 +1915,7 @@ def test_transform_loop_unroll_nested_counters(tmp_path, frontend): # Test transformation s = np.zeros(1) unrolled_function(s=s) - assert s == sum([a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 11)) if b <= a]) + assert s == sum(a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 11)) if b <= a) clean_test(filepath) clean_test(unrolled_filepath) @@ -1953,7 +1953,7 @@ def test_transform_loop_unroll_nested_neighbours(tmp_path, frontend): # Test the reference solution s = np.zeros(1) function(s=s) - assert s == 2 * sum([a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))]) + assert s == 2 * sum(a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))) # Apply transformation assert len(FindNodes(Loop).visit(routine.body)) == 3 loop_unroll(routine) @@ -1965,7 +1965,7 @@ def test_transform_loop_unroll_nested_neighbours(tmp_path, frontend): # Test transformation s = np.zeros(1) unrolled_function(s=s) - assert s == 2 * sum([a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))]) + assert s == 2 * sum(a + b + 1 for (a, b) in itertools.product(range(1, 11), range(1, 6))) clean_test(filepath) clean_test(unrolled_filepath) diff --git a/loki/transformations/tests/test_utilities.py b/loki/transformations/tests/test_utilities.py index a9241a6d5..cbacd60d3 100644 --- a/loki/transformations/tests/test_utilities.py +++ b/loki/transformations/tests/test_utilities.py @@ -502,22 +502,22 @@ def test_transform_utilites_get_local_arrays(frontend, tmp_path): module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path], definitions=(global_mod,)) routine = module['test_get_local_arrays'] - locals = get_local_arrays(routine, routine.body, unique=True) - assert len(locals) == 1 - assert locals[0] == 'local(i)' + local_arrs = get_local_arrays(routine, routine.body, unique=True) + assert len(local_arrs) == 1 + assert local_arrs[0] == 'local(i)' - locals = get_local_arrays(routine, routine.body, unique=False) - assert len(locals) == 2 - assert all(l == 'local(i)' for l in locals) + local_arrs = get_local_arrays(routine, routine.body, unique=False) + assert len(local_arrs) == 2 + assert all(l == 'local(i)' for l in local_arrs) - locals = get_local_arrays(routine, routine.body.body[-1:], unique=False) - assert len(locals) == 1 - assert locals[0] == 'local(i)' + local_arrs = get_local_arrays(routine, routine.body.body[-1:], unique=False) + assert len(local_arrs) == 1 + assert local_arrs[0] == 'local(i)' # Test for component arrays on arguments in spec - locals = get_local_arrays(routine, routine.spec, unique=True) - assert len(locals) == 1 - assert locals[0] == 'local(n)' + local_arrs = get_local_arrays(routine, routine.spec, unique=True) + assert len(local_arrs) == 1 + assert local_arrs[0] == 'local(n)' @pytest.mark.parametrize('frontend', available_frontends()) From bf257a7676fb7105f246688f401a2e3bdd393500 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Thu, 12 Sep 2024 12:14:02 +0000 Subject: [PATCH 07/12] Transformation: Fix for internal routine extraction with literal kinds --- loki/transformations/extract/internal.py | 2 +- loki/transformations/extract/tests/test_extract_internal.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/loki/transformations/extract/internal.py b/loki/transformations/extract/internal.py index 54bf48649..55c34f5dd 100644 --- a/loki/transformations/extract/internal.py +++ b/loki/transformations/extract/internal.py @@ -156,7 +156,7 @@ def extract_internal_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 diff --git a/loki/transformations/extract/tests/test_extract_internal.py b/loki/transformations/extract/tests/test_extract_internal.py index 6bbd1fd04..832275d18 100644 --- a/loki/transformations/extract/tests/test_extract_internal.py +++ b/loki/transformations/extract/tests/test_extract_internal.py @@ -336,7 +336,7 @@ def test_extract_internal_procedures_kind_resolution(frontend): contains subroutine inner() integer(kind = jpim) :: y - integer :: z + integer(kind=8) :: z z = y end subroutine inner end subroutine outer From 63b75a381d19f80f4aa5a24eea8a7b1ac5c508c0 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Thu, 12 Sep 2024 18:42:16 +0000 Subject: [PATCH 08/12] Transformation: Add `ExtractTransformation` and tests --- loki/transformations/extract/__init__.py | 68 +++++++++ .../tests/test_extract_transformation.py | 133 ++++++++++++++++++ 2 files changed, 201 insertions(+) create mode 100644 loki/transformations/extract/tests/test_extract_transformation.py diff --git a/loki/transformations/extract/__init__.py b/loki/transformations/extract/__init__.py index c6212cc33..6cde05b07 100644 --- a/loki/transformations/extract/__init__.py +++ b/loki/transformations/extract/__init__.py @@ -7,7 +7,75 @@ """ 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.marked 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 + ---------- + inline_internals : bool + Extract internal procedure (see :any:`extract_internal_procedures`); + default: False. + inline_marked : bool + Extract :any:`Subroutine` objects marked by pragma annotations + (see :any:`extract_marked_subroutines`); default: True. + """ + def __init__(self, extract_internals=False, extract_marked=True): + self.extract_internals = extract_internals + self.extract_marked = extract_marked + + 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.extract_marked: + for routine in module.subroutines: + new_routines = extract_marked_subroutines(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.extract_marked: + for routine in sourcefile.subroutines: + new_routines = extract_marked_subroutines(routine) + sourcefile.ir.append(new_routines) diff --git a/loki/transformations/extract/tests/test_extract_transformation.py b/loki/transformations/extract/tests/test_extract_transformation.py new file mode 100644 index 000000000..8f3364842 --- /dev/null +++ b/loki/transformations/extract/tests/test_extract_transformation.py @@ -0,0 +1,133 @@ +# (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. + +import pytest + +from loki import Subroutine, Module, Sourcefile +from loki.frontend import available_frontends +from loki.ir import nodes as ir, FindNodes + +from loki.transformations.extract import ExtractTransformation + + +@pytest.mark.parametrize('frontend', available_frontends()) +@pytest.mark.parametrize('extract_marked', [False, True]) +@pytest.mark.parametrize('extract_internals', [False, True]) +def test_extract_transformation_module(extract_internals, extract_marked, frontend): + """ + Test basic subroutine extraction from marker pragmas in modules. + """ + fcode = """ +module test_extract_mod +implicit none +contains + +subroutine outer(n, a, b) + integer, intent(in) :: n + real(kind=8), intent(inout) :: a, b(n) + real(kind=8) :: x(n), y(n, n+1) + integer :: i, j + + x(:) = a + do i=1, n + y(i,:) = b(i) + end do + + !$loki extract name(test1) + do i=1, n + do j=1, n+1 + x(i) = x(i) + y(i, j) + end do + end do + !$loki end extract + + do i=1, n + call plus_one(x, i=i) + end do + +contains + subroutine plus_one(f, i) + real(kind=8), intent(inout) :: f(:) + integer, intent(in) :: i + + f(i) = f(i) + 1.0 + end subroutine plus_one +end subroutine outer +end module test_extract_mod +""" + module = Module.from_source(fcode, frontend=frontend) + + ExtractTransformation( + extract_internals=extract_internals, extract_marked=extract_marked + ).apply(module) + + routines = tuple(r for r in module.contains.body if isinstance(r, Subroutine)) + assert len(routines) == 1 + (1 if extract_internals else 0) + (1 if extract_marked else 0) + assert ('plus_one' in module) == extract_internals + assert ('test1' in module) == extract_marked + + outer = module['outer'] + assert len(FindNodes(ir.CallStatement).visit(outer.body)) == (2 if extract_marked else 1) + outer_internals = tuple(r for r in outer.contains.body if isinstance(r, Subroutine)) + assert len(outer_internals) == (0 if extract_internals else 1) + + +@pytest.mark.parametrize('frontend', available_frontends()) +@pytest.mark.parametrize('extract_marked', [False, True]) +@pytest.mark.parametrize('extract_internals', [False, True]) +def test_extract_transformation_sourcefile(extract_internals, extract_marked, frontend): + """ + Test internal procedure extraction from subroutines. + """ + fcode = """ +subroutine outer(n, a, b) + integer, intent(in) :: n + real(kind=8), intent(inout) :: a, b(n) + real(kind=8) :: x(n), y(n, n+1) + integer :: i, j + + x(:) = a + do i=1, n + y(i,:) = b(i) + end do + + !$loki extract name(test1) + do i=1, n + do j=1, n+1 + x(i) = x(i) + y(i, j) + end do + end do + !$loki end extract + + do i=1, n + call plus_one(x, i=i) + end do + +contains + subroutine plus_one(f, i) + real(kind=8), intent(inout) :: f(:) + integer, intent(in) :: i + + f(i) = f(i) + 1.0 + end subroutine plus_one +end subroutine outer +""" + sourcefile = Sourcefile.from_source(fcode, frontend=frontend) + + ExtractTransformation( + extract_internals=extract_internals, extract_marked=extract_marked + ).apply(sourcefile) + + routines = tuple(r for r in sourcefile.ir.body if isinstance(r, Subroutine)) + assert len(routines) == 1 + (1 if extract_internals else 0) + (1 if extract_marked else 0) + assert ('plus_one' in sourcefile) == extract_internals + assert ('test1' in sourcefile) == extract_marked + + outer = sourcefile['outer'] + assert len(FindNodes(ir.CallStatement).visit(outer.body)) == (2 if extract_marked else 1) + outer_internals = tuple(r for r in outer.contains.body if isinstance(r, Subroutine)) + assert len(outer_internals) == (0 if extract_internals else 1) From cc856941c3c5b76e6ce6dff3250e2925cf404092 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Tue, 8 Oct 2024 04:29:37 +0000 Subject: [PATCH 09/12] Transformations: Remove superfluous `enumerate` from tests --- loki/transformations/tests/test_parametrise.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/loki/transformations/tests/test_parametrise.py b/loki/transformations/tests/test_parametrise.py index 0f4de1466..d5ee63900 100644 --- a/loki/transformations/tests/test_parametrise.py +++ b/loki/transformations/tests/test_parametrise.py @@ -356,7 +356,7 @@ def stop_execution(**kwargs): abort_callbacks = (error_stop, stop_execution) - for _, abort_callback in enumerate(abort_callbacks): + for abort_callback in abort_callbacks: scheduler = Scheduler( paths=[proj], config=config, seed_routines=['driver', 'another_driver'], frontend=frontend, xmods=[tmp_path] From 4663da61aa5d2389be6c1756065626d2a81521d4 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Sun, 13 Oct 2024 18:49:46 +0000 Subject: [PATCH 10/12] Extracts: Rename `extract_marked_routines` to `outline_pragma_regions` This also changes the semantics of the associated pragma annotations to `!$loki outline`. --- loki/transformations/extract/__init__.py | 22 ++-- .../extract/{marked.py => outline.py} | 17 ++- .../tests/test_extract_transformation.py | 34 +++--- ...test_extract_marked.py => test_outline.py} | 110 +++++++++--------- 4 files changed, 91 insertions(+), 92 deletions(-) rename loki/transformations/extract/{marked.py => outline.py} (94%) rename loki/transformations/extract/tests/{test_extract_marked.py => test_outline.py} (84%) diff --git a/loki/transformations/extract/__init__.py b/loki/transformations/extract/__init__.py index 6cde05b07..fd2184c54 100644 --- a/loki/transformations/extract/__init__.py +++ b/loki/transformations/extract/__init__.py @@ -17,7 +17,7 @@ """ from loki.transformations.extract.internal import * # noqa -from loki.transformations.extract.marked import * # noqa +from loki.transformations.extract.outline import * # noqa from loki.batch import Transformation @@ -33,16 +33,16 @@ class ExtractTransformation(Transformation): Parameters ---------- - inline_internals : bool + extract_internals : bool Extract internal procedure (see :any:`extract_internal_procedures`); default: False. - inline_marked : bool - Extract :any:`Subroutine` objects marked by pragma annotations - (see :any:`extract_marked_subroutines`); default: True. + 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, extract_marked=True): + def __init__(self, extract_internals=False, outline_regions=True): self.extract_internals = extract_internals - self.extract_marked = extract_marked + self.outline_regions = outline_regions def transform_module(self, module, **kwargs): """ @@ -57,9 +57,9 @@ def transform_module(self, module, **kwargs): module.contains.append(new_routines) # Extract pragma-marked code regions into standalone subroutines - if self.extract_marked: + if self.outline_regions: for routine in module.subroutines: - new_routines = extract_marked_subroutines(routine) + new_routines = outline_pragma_regions(routine) module.contains.append(new_routines) def transform_file(self, sourcefile, **kwargs): @@ -75,7 +75,7 @@ def transform_file(self, sourcefile, **kwargs): sourcefile.ir.append(new_routines) # Extract pragma-marked code regions into standalone subroutines - if self.extract_marked: + if self.outline_regions: for routine in sourcefile.subroutines: - new_routines = extract_marked_subroutines(routine) + new_routines = outline_pragma_regions(routine) sourcefile.ir.append(new_routines) diff --git a/loki/transformations/extract/marked.py b/loki/transformations/extract/outline.py similarity index 94% rename from loki/transformations/extract/marked.py rename to loki/transformations/extract/outline.py index c323e9ab0..130f5de0d 100644 --- a/loki/transformations/extract/marked.py +++ b/loki/transformations/extract/outline.py @@ -17,17 +17,16 @@ from loki.tools import as_tuple, CaseInsensitiveDict -__all__ = ['extract_marked_subroutines'] +__all__ = ['outline_pragma_regions'] -def extract_marked_subroutines(routine): +def outline_pragma_regions(routine): """ - Convert regions annotated with ``!$loki extract`` pragmas to subroutine calls. - + Convert regions annotated with ``!$loki outline`` pragmas to subroutine calls. The pragma syntax for regions to convert to subroutines is - ``!$loki extract [name(...)] [in(...)] [out(...)] [inout(...)]`` - and ``!$loki end extract``. + ``!$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. @@ -51,12 +50,12 @@ def extract_marked_subroutines(routine): 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='extract'): + if not is_loki_pragma(region.pragma, starts_with='outline'): continue # Name the external routine - parameters = get_pragma_parameters(region.pragma, starts_with='extract') - name = parameters.get('name', f'{routine.name}_extracted_{counter}') + 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 diff --git a/loki/transformations/extract/tests/test_extract_transformation.py b/loki/transformations/extract/tests/test_extract_transformation.py index 8f3364842..94e99759d 100644 --- a/loki/transformations/extract/tests/test_extract_transformation.py +++ b/loki/transformations/extract/tests/test_extract_transformation.py @@ -15,9 +15,9 @@ @pytest.mark.parametrize('frontend', available_frontends()) -@pytest.mark.parametrize('extract_marked', [False, True]) +@pytest.mark.parametrize('outline_regions', [False, True]) @pytest.mark.parametrize('extract_internals', [False, True]) -def test_extract_transformation_module(extract_internals, extract_marked, frontend): +def test_extract_transformation_module(extract_internals, outline_regions, frontend): """ Test basic subroutine extraction from marker pragmas in modules. """ @@ -37,13 +37,13 @@ def test_extract_transformation_module(extract_internals, extract_marked, fronte y(i,:) = b(i) end do - !$loki extract name(test1) + !$loki outline name(test1) do i=1, n do j=1, n+1 x(i) = x(i) + y(i, j) end do end do - !$loki end extract + !$loki end outline do i=1, n call plus_one(x, i=i) @@ -62,26 +62,26 @@ def test_extract_transformation_module(extract_internals, extract_marked, fronte module = Module.from_source(fcode, frontend=frontend) ExtractTransformation( - extract_internals=extract_internals, extract_marked=extract_marked + extract_internals=extract_internals, outline_regions=outline_regions ).apply(module) routines = tuple(r for r in module.contains.body if isinstance(r, Subroutine)) - assert len(routines) == 1 + (1 if extract_internals else 0) + (1 if extract_marked else 0) + assert len(routines) == 1 + (1 if extract_internals else 0) + (1 if outline_regions else 0) assert ('plus_one' in module) == extract_internals - assert ('test1' in module) == extract_marked + assert ('test1' in module) == outline_regions outer = module['outer'] - assert len(FindNodes(ir.CallStatement).visit(outer.body)) == (2 if extract_marked else 1) + assert len(FindNodes(ir.CallStatement).visit(outer.body)) == (2 if outline_regions else 1) outer_internals = tuple(r for r in outer.contains.body if isinstance(r, Subroutine)) assert len(outer_internals) == (0 if extract_internals else 1) @pytest.mark.parametrize('frontend', available_frontends()) -@pytest.mark.parametrize('extract_marked', [False, True]) +@pytest.mark.parametrize('outline_regions', [False, True]) @pytest.mark.parametrize('extract_internals', [False, True]) -def test_extract_transformation_sourcefile(extract_internals, extract_marked, frontend): +def test_extract_transformation_sourcefile(extract_internals, outline_regions, frontend): """ - Test internal procedure extraction from subroutines. + Test internal procedure extraction and region outlining from subroutines. """ fcode = """ subroutine outer(n, a, b) @@ -95,13 +95,13 @@ def test_extract_transformation_sourcefile(extract_internals, extract_marked, fr y(i,:) = b(i) end do - !$loki extract name(test1) + !$loki outline name(test1) do i=1, n do j=1, n+1 x(i) = x(i) + y(i, j) end do end do - !$loki end extract + !$loki end outline do i=1, n call plus_one(x, i=i) @@ -119,15 +119,15 @@ def test_extract_transformation_sourcefile(extract_internals, extract_marked, fr sourcefile = Sourcefile.from_source(fcode, frontend=frontend) ExtractTransformation( - extract_internals=extract_internals, extract_marked=extract_marked + extract_internals=extract_internals, outline_regions=outline_regions ).apply(sourcefile) routines = tuple(r for r in sourcefile.ir.body if isinstance(r, Subroutine)) - assert len(routines) == 1 + (1 if extract_internals else 0) + (1 if extract_marked else 0) + assert len(routines) == 1 + (1 if extract_internals else 0) + (1 if outline_regions else 0) assert ('plus_one' in sourcefile) == extract_internals - assert ('test1' in sourcefile) == extract_marked + assert ('test1' in sourcefile) == outline_regions outer = sourcefile['outer'] - assert len(FindNodes(ir.CallStatement).visit(outer.body)) == (2 if extract_marked else 1) + assert len(FindNodes(ir.CallStatement).visit(outer.body)) == (2 if outline_regions else 1) outer_internals = tuple(r for r in outer.contains.body if isinstance(r, Subroutine)) assert len(outer_internals) == (0 if extract_internals else 1) diff --git a/loki/transformations/extract/tests/test_extract_marked.py b/loki/transformations/extract/tests/test_outline.py similarity index 84% rename from loki/transformations/extract/tests/test_extract_marked.py rename to loki/transformations/extract/tests/test_outline.py index c9ef897fa..a17ae7361 100644 --- a/loki/transformations/extract/tests/test_extract_marked.py +++ b/loki/transformations/extract/tests/test_outline.py @@ -14,7 +14,7 @@ from loki.ir import FindNodes, Section, Assignment, CallStatement, Intrinsic from loki.tools import as_tuple -from loki.transformations.extract.marked import extract_marked_subroutines +from loki.transformations.extract.outline import outline_pragma_regions @pytest.fixture(scope='function', name='builder') @@ -24,23 +24,23 @@ def fixture_builder(tmp_path): @pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_marked_subroutines(tmp_path, frontend): +def test_outline_pragma_regions(tmp_path, frontend): """ - A very simple :any:`extract_marked_subroutine` test case + A very simple :any:`outline_pragma_regions` test case """ fcode = """ -subroutine test_extract(a, b, c) +subroutine test_outline(a, b, c) integer, intent(out) :: a, b, c a = 5 a = 1 -!$loki extract in(a) out(b) +!$loki outline in(a) out(b) b = a -!$loki end extract +!$loki end outline c = a + b -end subroutine test_extract +end subroutine test_outline """ routine = Subroutine.from_source(fcode, frontend=frontend) filepath = tmp_path/(f'{routine.name}_{frontend}.f90') @@ -54,8 +54,8 @@ def test_extract_marked_subroutines(tmp_path, frontend): assert len(FindNodes(CallStatement).visit(routine.body)) == 0 # Apply transformation - routines = extract_marked_subroutines(routine) - assert len(routines) == 1 and routines[0].name == f'{routine.name}_extracted_0' + routines = outline_pragma_regions(routine) + assert len(routines) == 1 and routines[0].name == f'{routine.name}_outlined_0' assert len(FindNodes(Assignment).visit(routine.body)) == 3 assert len(FindNodes(Assignment).visit(routines[0].body)) == 1 @@ -73,30 +73,30 @@ def test_extract_marked_subroutines(tmp_path, frontend): @pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_marked_subroutines_multiple(tmp_path, frontend): +def test_outline_pragma_regions_multiple(tmp_path, frontend): """ Test hoisting with multiple groups and multiple regions per group """ fcode = """ -subroutine test_extract_mult(a, b, c) +subroutine test_outline_mult(a, b, c) integer, intent(out) :: a, b, c a = 1 a = a + 1 a = a + 1 -!$loki extract name(oiwjfklsf) inout(a) +!$loki outline name(oiwjfklsf) inout(a) a = a + 1 -!$loki end extract +!$loki end outline a = a + 1 -!$loki extract in(a) out(b) +!$loki outline in(a) out(b) b = a -!$loki end extract +!$loki end outline -!$loki extract in(a,b) out(c) +!$loki outline in(a,b) out(c) c = a + b -!$loki end extract -end subroutine test_extract_mult +!$loki end outline +end subroutine test_outline_mult """ routine = Subroutine.from_source(fcode, frontend=frontend) filepath = tmp_path/(f'{routine.name}_{frontend}.f90') @@ -110,10 +110,10 @@ def test_extract_marked_subroutines_multiple(tmp_path, frontend): assert len(FindNodes(CallStatement).visit(routine.body)) == 0 # Apply transformation - routines = extract_marked_subroutines(routine) + routines = outline_pragma_regions(routine) assert len(routines) == 3 assert routines[0].name == 'oiwjfklsf' - assert all(routines[i].name == f'{routine.name}_extracted_{i}' for i in (1,2)) + assert all(routines[i].name == f'{routine.name}_outlined_{i}' for i in (1,2)) assert len(FindNodes(Assignment).visit(routine.body)) == 4 assert all(len(FindNodes(Assignment).visit(r.body)) == 1 for r in routines) @@ -131,32 +131,32 @@ def test_extract_marked_subroutines_multiple(tmp_path, frontend): @pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_marked_subroutines_arguments(tmp_path, frontend): +def test_outline_pragma_regions_arguments(tmp_path, frontend): """ Test hoisting with multiple groups and multiple regions per group and automatic derivation of arguments """ fcode = """ -subroutine test_extract_args(a, b, c) +subroutine test_outline_args(a, b, c) integer, intent(out) :: a, b, c a = 1 a = a + 1 a = a + 1 -!$loki extract name(func_a) +!$loki outline name(func_a) a = a + 1 -!$loki end extract +!$loki end outline a = a + 1 -!$loki extract name(func_b) +!$loki outline name(func_b) b = a -!$loki end extract +!$loki end outline ! partially override arguments -!$loki extract name(func_c) inout(b) +!$loki outline name(func_c) inout(b) c = a + b -!$loki end extract -end subroutine test_extract_args +!$loki end outline +end subroutine test_outline_args """ routine = Subroutine.from_source(fcode, frontend=frontend) filepath = tmp_path/(f'{routine.name}_{frontend}.f90') @@ -170,7 +170,7 @@ def test_extract_marked_subroutines_arguments(tmp_path, frontend): assert len(FindNodes(CallStatement).visit(routine.body)) == 0 # Apply transformation - routines = extract_marked_subroutines(routine) + routines = outline_pragma_regions(routine) assert len(routines) == 3 assert [r.name for r in routines] == ['func_a', 'func_b', 'func_c'] @@ -202,35 +202,35 @@ def test_extract_marked_subroutines_arguments(tmp_path, frontend): @pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_marked_subroutines_arrays(tmp_path, frontend): +def test_outline_pragma_regions_arrays(tmp_path, frontend): """ Test hoisting with array variables """ fcode = """ -subroutine test_extract_arr(a, b, n) +subroutine test_outline_arr(a, b, n) integer, intent(out) :: a(n), b(n) integer, intent(in) :: n integer :: j -!$loki extract +!$loki outline do j=1,n a(j) = j end do -!$loki end extract +!$loki end outline -!$loki extract +!$loki outline do j=1,n b(j) = j end do -!$loki end extract +!$loki end outline -!$loki extract +!$loki outline do j=1,n-1 b(j) = b(j+1) - a(j) end do b(n) = 1 -!$loki end extract -end subroutine test_extract_arr +!$loki end outline +end subroutine test_outline_arr """ routine = Subroutine.from_source(fcode, frontend=frontend) @@ -249,7 +249,7 @@ def test_extract_marked_subroutines_arrays(tmp_path, frontend): assert len(FindNodes(CallStatement).visit(routine.body)) == 0 # Apply transformation - routines = extract_marked_subroutines(routine) + routines = outline_pragma_regions(routine) assert len(FindNodes(Assignment).visit(routine.body)) == 0 assert len(FindNodes(CallStatement).visit(routine.body)) == 3 @@ -276,49 +276,49 @@ def test_extract_marked_subroutines_arrays(tmp_path, frontend): @pytest.mark.parametrize('frontend', available_frontends()) -def test_extract_marked_subroutines_imports(tmp_path, builder, frontend): +def test_outline_pragma_regions_imports(tmp_path, builder, frontend): """ Test hoisting with correct treatment of imports """ fcode_module = """ -module extract_mod +module outline_mod implicit none integer, parameter :: param = 1 integer :: arr1(10) integer :: arr2(10) -end module extract_mod +end module outline_mod """.strip() fcode = """ -module test_extract_imps_mod +module test_outline_imps_mod implicit none contains - subroutine test_extract_imps(a, b) - use extract_mod, only: param, arr1, arr2 + subroutine test_outline_imps(a, b) + use outline_mod, only: param, arr1, arr2 integer, intent(out) :: a(10), b(10) integer :: j -!$loki extract +!$loki outline do j=1,10 a(j) = param end do -!$loki end extract +!$loki end outline -!$loki extract +!$loki outline do j=1,10 arr1(j) = j+1 end do -!$loki end extract +!$loki end outline arr2(:) = arr1(:) -!$loki extract +!$loki outline do j=1,10 b(j) = arr2(j) - a(j) end do -!$loki end extract - end subroutine test_extract_imps -end module test_extract_imps_mod +!$loki end outline + end subroutine test_outline_imps +end module test_outline_imps_mod """ ext_module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path]) module = Module.from_source(fcode, frontend=frontend, definitions=ext_module, xmods=[tmp_path]) @@ -338,7 +338,7 @@ def test_extract_marked_subroutines_imports(tmp_path, builder, frontend): assert len(FindNodes(CallStatement).visit(module.subroutines[0].body)) == 0 # Apply transformation - routines = extract_marked_subroutines(module.subroutines[0]) + routines = outline_pragma_regions(module.subroutines[0]) assert len(FindNodes(Assignment).visit(module.subroutines[0].body)) == 1 assert len(FindNodes(CallStatement).visit(module.subroutines[0].body)) == 3 From e74764368d5795cabd427590272bc19cef7fa698 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Mon, 14 Oct 2024 11:47:56 +0200 Subject: [PATCH 11/12] Transformations: Improved docstring for outline_pragma_regions Co-authored-by: Balthasar Reuter <6384870+reuterbal@users.noreply.github.com> --- loki/transformations/extract/outline.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/loki/transformations/extract/outline.py b/loki/transformations/extract/outline.py index 130f5de0d..b6082b82a 100644 --- a/loki/transformations/extract/outline.py +++ b/loki/transformations/extract/outline.py @@ -37,10 +37,15 @@ def outline_pragma_regions(routine): The pragma region in the original routine is replaced by a call to the new subroutine. - :param :class:``Subroutine`` routine: - the routine to modify. - - :return: the list of newly created subroutines. + 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 From 06774ae4f0f2da64d1298995c6272769f6fd0d2a Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Mon, 14 Oct 2024 10:19:39 +0000 Subject: [PATCH 12/12] Transformations: Simplify definition of retriever in _inline_functions --- loki/transformations/inline/functions.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/loki/transformations/inline/functions.py b/loki/transformations/inline/functions.py index ae1f8476d..f162a1de3 100644 --- a/loki/transformations/inline/functions.py +++ b/loki/transformations/inline/functions.py @@ -120,7 +120,10 @@ class FindInlineCallsSkipInlineCallParameters(ExpressionFinder): """ Find inline calls but skip/ignore parameters of inline calls. """ - retriever = ExpressionRetrieverSkipInlineCallParameters(lambda e: isinstance(e, sym.InlineCall)) + retriever = ExpressionRetrieverSkipInlineCallParameters( + query=lambda e: isinstance(e, sym.InlineCall), + inline_elementals_only=inline_elementals_only, functions=functions + ) # functions are provided, however functions is empty, thus early exit if functions is not None and not functions: @@ -133,12 +136,8 @@ class FindInlineCallsSkipInlineCallParameters(ExpressionFinder): # Find and filter inline calls and corresponding nodes function_calls = {} # Find inline calls but skip/ignore inline calls being parameters of other inline calls - # to ensure correct ordering of inlining. Those skipped/ignored inline calls will be handled - # in the next call to this function. - retriever = ExpressionRetrieverSkipInlineCallParameters(lambda e: isinstance(e, sym.InlineCall), - inline_elementals_only=inline_elementals_only, functions=functions) - # override retriever ... - FindInlineCallsSkipInlineCallParameters.retriever = retriever + # to ensure correct ordering of inlining. Those skipped/ignored inline calls will be handled + # in the next call to this function. for node, calls in FindInlineCallsSkipInlineCallParameters(with_ir_node=True).visit(routine.body): for call in calls: if call.procedure_type is BasicType.DEFERRED or isinstance(call.routine, StatementFunction):