Transformation utility to fix sequence association

cmake/loki_transform.cmake

@@ -25,6 +25,7 @@ include( loki_transform_helpers )
 #       [CPP]
 #       [FRONTEND <frontend>]
 #       [INLINE_MEMBERS]
+#       [RESOLVE_SEQUENCE_ASSOCIATION]
 #       [BUILDDIR <build-path>]
 #       [SOURCES <source1> [<source2> ...]]
 #       [HEADERS <header1> [<header2> ...]]
@@ -46,7 +47,7 @@ function( loki_transform )
 
     set( options
          CPP DATA_OFFLOAD REMOVE_OPENMP ASSUME_DEVICEPTR TRIM_VECTOR_SECTIONS GLOBAL_VAR_OFFLOAD
-         REMOVE_DERIVED_ARGS INLINE_MEMBERS DERIVE_ARGUMENT_ARRAY_SHAPE
+         REMOVE_DERIVED_ARGS INLINE_MEMBERS RESOLVE_SEQUENCE_ASSOCIATION DERIVE_ARGUMENT_ARRAY_SHAPE
     )
     set( oneValueArgs
          COMMAND MODE DIRECTIVE FRONTEND CONFIG BUILDDIR
@@ -193,7 +194,7 @@ endfunction()
 #       [DIRECTIVE <openacc|openmp|...>]
 #       [SOURCES <source1> [<source2> ...]]
 #       [HEADERS <header1> [<header2> ...]]
-#       [NO_PLAN_SOURCEDIR COPY_UNMODIFIED INLINE_MEMBERS]
+#       [NO_PLAN_SOURCEDIR COPY_UNMODIFIED INLINE_MEMBERS RESOLVE_SEQUENCE_ASSOCIATION]
 #   )
 #
 # Applies a Loki bulk transformation to the source files belonging to a particular
@@ -222,7 +223,7 @@ endfunction()
 
 function( loki_transform_target )
 
-    set( options NO_PLAN_SOURCEDIR COPY_UNMODIFIED CPP CPP_PLAN INLINE_MEMBERS )
+    set( options NO_PLAN_SOURCEDIR COPY_UNMODIFIED CPP CPP_PLAN INLINE_MEMBERS RESOLVE_SEQUENCE_ASSOCIATION )
     set( single_value_args TARGET COMMAND MODE DIRECTIVE FRONTEND CONFIG PLAN )
     set( multi_value_args SOURCES HEADERS )
 
@@ -291,6 +292,10 @@ function( loki_transform_target )
             list( APPEND _TRANSFORM_OPTIONS INLINE_MEMBERS )
         endif()
 
+        if( _PAR_RESOLVE_SEQUENCE_ASSOCIATION )
+            list( APPEND _TRANSFORM_OPTIONS RESOLVE_SEQUENCE_ASSOCIATION )
+        endif()
+
         loki_transform(
             COMMAND   ${_PAR_COMMAND}
             OUTPUT    ${LOKI_SOURCES_TO_APPEND}
@@ -384,7 +389,7 @@ or
 
     set( options
          CPP DATA_OFFLOAD REMOVE_OPENMP ASSUME_DEVICEPTR GLOBAL_VAR_OFFLOAD
-         TRIM_VECTOR_SECTIONS REMOVE_DERIVED_ARGS INLINE_MEMBERS
+         TRIM_VECTOR_SECTIONS REMOVE_DERIVED_ARGS INLINE_MEMBERS RESOLVE_SEQUENCE_ASSOCIATION
     )
     set( oneValueArgs
          MODE DIRECTIVE FRONTEND CONFIG PATH OUTPATH

cmake/loki_transform_helpers.cmake

@@ -112,6 +112,10 @@ macro( _loki_transform_parse_options )
         list( APPEND _ARGS --inline-members )
     endif()
 
+    if( _PAR_RESOLVE_SEQUENCE_ASSOCIATION )
+        list( APPEND _ARGS --resolve-sequence-association )
+    endif()
+
     if( _PAR_DERIVE_ARGUMENT_ARRAY_SHAPE )
         list( APPEND _ARGS --derive-argument-array-shape )
     endif()

loki/transform/__init__.py

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

loki/transform/transform_sequence_association.py

@@ -0,0 +1,258 @@
+# (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 pymbolic.primitives as pmbl
+
+from loki.expression import (
+    Sum, Product, IntLiteral, Array, RangeIndex,
+    SubstituteExpressions
+    )
+from loki.ir import CallStatement
+from loki.visitors import FindNodes, Transformer
+from loki.tools import as_tuple
+from loki.types import BasicType
+
+
+__all__ = [
+    'transform_sequence_association'
+]
+
+def check_if_scalar_syntax(arg, dummy):
+    """
+    Check if an array argument, arg,
+    is passed to an array dummy argument, dummy,
+    using scalar syntax. i.e. arg(1,1) -> d(m,n)
+
+    Parameters
+    ----------
+    arg:   variable
+    dummy: variable
+    """
+    if isinstance(arg, Array) and isinstance(dummy, Array):
+        if arg.dimensions:
+            if not any(isinstance(d, RangeIndex) for d in arg.dimensions):
+                return True
+    return False
+
+
+def single_sum(expr):
+    """
+    Return a Sum object of expr if expr is not an instance of pymbolic.primitives.Sum.
+    Otherwise return expr
+
+    Parameters
+    ----------
+    expr: any pymbolic expression
+    """
+
+    if isinstance(expr, pmbl.Sum):
+        return expr
+    return Sum((expr,))
+
+
+def product_value(expr):
+    """
+    If expr is an instance of pymbolic.primitives.Product, try to evaluate it
+    If it is possible, return the value as an int.
+    If it is not possible, try to simplify the the product and return as a Product
+    If it is not a pymbolic.primitives.Product , return expr
+
+    Note: Negative numbers and subtractions in Sums are represented as Product of
+          the integer -1 and the symbol. This complicates matters.
+    Note: Ensure that a Loki Product is returned, not a pymbolic Product
+
+    Parameters
+    ----------
+    expr: any pymbolic expression
+    """
+    if isinstance(expr, pmbl.Product):
+        m = 1
+        new_children = []
+        for c in expr.children:
+            if isinstance(c, IntLiteral):
+                m = m*c.value
+            elif isinstance(c, int):
+                m = m*c
+            else:
+                new_children += [c]
+        if m == 0:
+            return 0
+        if not new_children:
+            return m
+
+        if m > 1:
+            new_children = [IntLiteral(m)] + new_children
+        elif m == -1:
+            new_children = [-1] + new_children
+        elif m < -1:
+            new_children = [-1, IntLiteral(abs(m))] + new_children
+
+        return Product(as_tuple(new_children))
+
+    return expr
+
+
+def simplify_sum(expr):
+    """
+    If expr is an instance of pymbolic.primitives.Sum,
+    try to simplify it by evaluating any Products and adding up ints and IntLiterals.
+    If the sum can be reduced to a number, it returns an IntLiteral
+    If the Sum reduces to one expression, it returns that expression
+
+    Note: Ensure that a Loki Sum is returned, not a pymbolic Sum
+
+    Parameters
+    ----------
+    expr: any pymbolic expression
+    """
+
+    if isinstance(expr, pmbl.Sum):
+        n = 0
+        new_children = []
+        for c in expr.children:
+            c = product_value(c)
+            if isinstance(c, IntLiteral):
+                n += c.value
+            elif isinstance(c, int):
+                n += c
+            else:
+                new_children += [c]
+
+        if new_children:
+            if n > 0:
+                new_children += [IntLiteral(n)]
+            elif n < 0:
+                new_children += [Product((-1,IntLiteral(abs(n))))]
+
+            if len(new_children) > 1:
+                return Sum(as_tuple(new_children))
+            return new_children[0]
+        return IntLiteral(n)
+    return expr
+
+
+def construct_range_index(lower, length):
+    """
+    Construct a range index from lower to lower + length - 1
+
+    Parameters
+    ----------
+    lower : any pymbolic expression
+    length: any pymbolic expression
+    """
+
+    new_high = simplify_sum(single_sum(length) + lower - IntLiteral(1))
+
+    return RangeIndex((lower, new_high))
+
+
+def process_symbol(symbol, caller, call):
+    """
+    Map symbol in call.routine to the appropriate symbol in caller,
+    taking any parents into account
+
+    Parameters
+    ----------
+    symbol: Loki variable in call.routine
+    caller: Subroutine object containing call
+    call  : Call object
+    """
+
+    if isinstance(symbol, IntLiteral):
+        return symbol
+
+    if not symbol.parents:
+        if symbol in call.routine.arguments:
+            return call.arg_map[symbol]
+
+    elif symbol.parents[0] in call.routine.arguments:
+        return SubstituteExpressions(call.arg_map).visit(symbol.clone(scope=caller))
+
+    if call.routine in caller.members and symbol in caller.variables:
+        return symbol
+
+    raise RuntimeError('[Loki::transform_sequence_association] Unable to resolve argument dimension. Module variable?')
+
+
+def construct_length(xrange, caller, call):
+    """
+    Construct an expression for the length of xrange,
+    defined in call.routine, in caller.
+
+    Parameters
+    ----------
+    xrange: RangeIndex object defined in call.routine
+    caller: Subroutine object
+    call  : call contained in caller
+    """
+
+    new_start = process_symbol(xrange.start, caller, call)
+    new_stop  = process_symbol(xrange.stop, caller, call)
+
+    return single_sum(new_stop) - new_start + IntLiteral(1)
+
+
+def transform_sequence_association(routine):
+    """
+    Housekeeping routine to replace scalar syntax when passing arrays as arguments
+    For example, a call like
+
+    call myroutine(a(i,j))
+
+    where myroutine looks like
+
+    subroutine myroutine(a)
+        real :: a(5)
+    end subroutine myroutine
+
+    should be changed to
+
+    call myroutine(a(i:i+5,j)
+
+    Note: Using the __add__ and __mul__ functions of Sum and Product, respectively,
+          returns the pymbolic.primitives version of the objuect, not the loki.expressions version.
+          simplify_sum and product_value returns loki versions, so this is currently not an issue,
+          but this can cause unexpected behaviour
+
+    Parameters
+    ----------
+    routine : :any:`Subroutine`
+        The subroutine where calls will be changed
+    """
+
+    #List calls in routine, but make sure we have the called routine definition
+    calls = (c for c in FindNodes(CallStatement).visit(routine.body) if not c.procedure_type is BasicType.DEFERRED)
+    call_map = {}
+
+    for call in calls:
+
+        new_args = []
+
+        found_scalar = False
+        for dummy, arg in call.arg_map.items():
+            if check_if_scalar_syntax(arg, dummy):
+                found_scalar = True
+
+                new_dims = []
+                for s, lower in zip(dummy.shape, arg.dimensions):
+
+                    if isinstance(s, RangeIndex):
+                        new_dims += [construct_range_index(lower, construct_length(s, routine, call))]
+                    else:
+                        new_dims += [construct_range_index(lower, process_symbol(s, routine, call))]
+
+                if len(arg.dimensions) > len(dummy.shape):
+                    new_dims += arg.dimensions[len(dummy.shape):]
+                new_args += [arg.clone(dimensions=as_tuple(new_dims)),]
+            else:
+                new_args += [arg,]
+
+        if found_scalar:
+            call_map[call] = call.clone(arguments = as_tuple(new_args))
+
+    if call_map:
+        routine.body = Transformer(call_map).visit(routine.body)

scripts/loki_transform.py

@@ -108,12 +108,15 @@ def cli(debug):
               help="Remove derived-type arguments and replace with canonical arguments")
 @click.option('--inline-members/--no-inline-members', default=False,
               help='Inline member functions for SCC-class transformations.')
+@click.option('--resolve-sequence-association/--no-resolve-sequence-association', default=False,
+              help='Replace array arguments passed as scalars with arrays.')
 @click.option('--derive-argument-array-shape/--no-derive-argument-array-shape', default=False,
               help="Recursively derive explicit shape dimension for argument arrays")
 def convert(
         mode, config, build, source, header, cpp, directive, include, define, omni_include, xmod,
         data_offload, remove_openmp, assume_deviceptr, frontend, trim_vector_sections,
-        global_var_offload, remove_derived_args, inline_members, derive_argument_array_shape
+        global_var_offload, remove_derived_args, inline_members, resolve_sequence_association,
+        derive_argument_array_shape
 ):
     """
     Batch-processing mode for Fortran-to-Fortran transformations that
@@ -207,7 +210,8 @@ def convert(
     if mode in ['scc', 'scc-hoist', 'scc-stack']:
         # Apply the basic SCC transformation set
         scheduler.process( SCCBaseTransformation(
-            horizontal=horizontal, directive=directive, inline_members=inline_members
+            horizontal=horizontal, directive=directive,
+            inline_members=inline_members, resolve_sequence_association=resolve_sequence_association
         ))
         scheduler.process( SCCDevectorTransformation(
             horizontal=horizontal, trim_vector_sections=trim_vector_sections