Transformation utility to fix sequence association (#173)

* identified arrays passed as scalars * scalars sort of fixed * commit call * constructed new arguments * Might work now * Check that we have the called routine definition * some documentation * simplify TypedSymbol handling * some cleanup * Simplify (?) and add docstrings * fix style * more style * even more style * Ensure Loki versions of Sum and Product * fix another negative bug * add tests * going out in style * Add option for turning scalar fix on and off * moved option setting a bit * fix linter complaints * scalar_syntax -> sequence_association * fix -> resolve * inline and sequence association tests in single column transformation * cleanup * greatly simplify by using callers dimensions * cleanup * more cleanup
ecmwf-ifs · Nov 13, 2023 · f77fff1 · f77fff1
1 parent 9d27cab
commit f77fff1
Show file tree

Hide file tree

Showing 8 changed files with 270 additions and 8 deletions.
diff --git a/cmake/loki_transform.cmake b/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

diff --git a/cmake/loki_transform_helpers.cmake b/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()

diff --git a/loki/transform/__init__.py b/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
diff --git a/loki/transform/transform_sequence_association.py b/loki/transform/transform_sequence_association.py
@@ -0,0 +1,94 @@
+# (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.expression import Array, RangeIndex
+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 transform_sequence_association(routine):
+    """
+    Housekeeping routine to replace scalar syntax when passing arrays as arguments
+    For example, a call like
+
+    real :: a(m,n)
+
+    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:m,j)
+
+    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
+
+                n_dims = len(dummy.shape)
+                new_dims = []
+                for s, lower in zip(arg.shape[:n_dims], arg.dimensions[:n_dims]):
+
+                    if isinstance(s, RangeIndex):
+                        new_dims += [RangeIndex((lower, s.stop))]
+                    else:
+                        new_dims += [RangeIndex((lower, s))]
+
+                if len(arg.dimensions) > n_dims:
+                    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)
diff --git a/scripts/loki_transform.py b/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

diff --git a/tests/test_transform_sequence_association.py b/tests/test_transform_sequence_association.py
@@ -0,0 +1,73 @@
+# (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 conftest import available_frontends
+
+from loki.backend import fgen
+from loki.transform import transform_sequence_association
+from loki.module import Module
+from loki.ir import CallStatement
+from loki.visitors import FindNodes
+
+@pytest.mark.parametrize('frontend', available_frontends())
+def test_transform_scalar_notation(frontend):
+    fcode = """
+module mod_a
+    implicit none
+
+    type type_b
+        integer :: c
+        integer :: d
+    end type type_b
+
+    type type_a
+        type(type_b) :: b
+    end type type_a
+
+contains
+
+    subroutine main()
+
+        type(type_a) :: a
+        integer :: k, m, n
+
+        real    :: array(10,10)
+
+        call sub_x(array(1, 1), 1)
+        call sub_x(array(2, 2), 2)
+        call sub_x(array(m, 1), k)
+        call sub_x(array(m-1, 1), k-1)
+        call sub_x(array(a%b%c, 1), a%b%d)
+
+    contains
+
+        subroutine sub_x(array, k)
+    
+            integer, intent(in) :: k
+            real, intent(in)    :: array(k:n)
+    
+        end subroutine sub_x
+
+    end subroutine main
+
+end module mod_a
+    """.strip()
+
+    module = Module.from_source(fcode, frontend=frontend)
+    routine = module['main']
+
+    transform_sequence_association(routine)
+
+    calls = FindNodes(CallStatement).visit(routine.body)
+
+    assert fgen(calls[0]).lower() == 'call sub_x(array(1:10, 1), 1)'
+    assert fgen(calls[1]).lower() == 'call sub_x(array(2:10, 2), 2)'
+    assert fgen(calls[2]).lower() == 'call sub_x(array(m:10, 1), k)'
+    assert fgen(calls[3]).lower() == 'call sub_x(array(m - 1:10, 1), k - 1)'
+    assert fgen(calls[4]).lower() == 'call sub_x(array(a%b%c:10, 1), a%b%d)'
diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py
@@ -1721,3 +1721,79 @@ def test_single_column_coalesced_vector_section_trim_complex(frontend, horizonta
     else:
         assert assign in loop.body
         assert(len(FindNodes(Assignment).visit(loop.body)) == 4)
+
+
+@pytest.mark.parametrize('frontend', available_frontends())
+@pytest.mark.parametrize('inline_members', [False, True])
+@pytest.mark.parametrize('resolve_sequence_association', [False, True])
+def test_single_column_coalesced_inline_and_sequence_association(frontend, horizontal,
+                                                                 inline_members, resolve_sequence_association):
+    """
+    Test the combinations of routine inlining and sequence association
+    """
+
+    fcode_kernel = """
+    subroutine some_kernel(nlon, start, end)
+       implicit none
+
+       integer, intent(in) :: nlon, start, end
+       real, dimension(nlon) :: work
+
+       call contained_kernel(work(1))
+
+     contains
+
+       subroutine contained_kernel(work)
+          implicit none
+
+          real, dimension(nlon) :: work
+          integer :: jl
+
+          do jl = start, end
+             work(jl) = 1.
+          enddo
+
+       end subroutine contained_kernel
+    end subroutine some_kernel
+    """
+
+    routine = Subroutine.from_source(fcode_kernel, frontend=frontend)
+
+    scc_transform = SCCBaseTransformation(horizontal=horizontal,
+                                          inline_members=inline_members,
+                                          resolve_sequence_association=resolve_sequence_association)
+
+    #Not really doing anything for contained routines
+    if (not inline_members and not resolve_sequence_association):
+        scc_transform.apply(routine, role='kernel')
+
+        assert len(routine.members) == 1
+        assert not FindNodes(Loop).visit(routine.body)
+
+    #Should fail because it can't resolve sequence association
+    elif (inline_members and not resolve_sequence_association):
+        with pytest.raises(RuntimeError) as e_info:
+            scc_transform.apply(routine, role='kernel')
+        assert(e_info.exconly() ==
+               'RuntimeError: [Loki::TransformInline] Unable to resolve member subroutine call')
+
+    #Check that the call is properly modified
+    elif (not inline_members and resolve_sequence_association):
+        scc_transform.apply(routine, role='kernel')
+
+        assert len(routine.members) == 1
+        call = FindNodes(CallStatement).visit(routine.body)[0]
+        assert fgen(call).lower() == 'call contained_kernel(work(1:nlon))'
+
+    #Check that the contained subroutine has been inlined
+    else:
+        scc_transform.apply(routine, role='kernel')
+
+        assert len(routine.members) == 0
+
+        loop = FindNodes(Loop).visit(routine.body)[0]
+        assert loop.variable == 'jl'
+        assert loop.bounds == 'start:end'
+
+        assign = FindNodes(Assignment).visit(loop.body)[0]
+        assert fgen(assign).lower() == 'work(jl) = 1.'
diff --git a/transformations/transformations/single_column_coalesced.py b/transformations/transformations/single_column_coalesced.py
@@ -7,7 +7,7 @@
 
 import re
 from loki.expression import symbols as sym
-from loki.transform import resolve_associates, inline_member_procedures
+from loki.transform import resolve_associates, inline_member_procedures, transform_sequence_association
 from loki import (
     Transformation, FindNodes, Transformer, info,
     pragmas_attached, as_tuple, flatten, ir, FindExpressions,
@@ -40,13 +40,14 @@ class methods can be called directly.
         Enable full source-inlining of member subroutines; default: False.
     """
 
-    def __init__(self, horizontal, directive=None, inline_members=False):
+    def __init__(self, horizontal, directive=None, inline_members=False, resolve_sequence_association=False):
         self.horizontal = horizontal
 
         assert directive in [None, 'openacc']
         self.directive = directive
 
         self.inline_members = inline_members
+        self.resolve_sequence_association = resolve_sequence_association
 
     @classmethod
     def check_routine_pragmas(cls, routine, directive):
@@ -291,6 +292,10 @@ def process_kernel(self, routine):
         # Find the iteration index variable for the specified horizontal
         v_index = self.get_integer_variable(routine, name=self.horizontal.index)
 
+        # Transform arrays passed with scalar syntax to array syntax
+        if self.resolve_sequence_association:
+            transform_sequence_association(routine)
+
         # Perform full source-inlining for member subroutines if so requested
         if self.inline_members:
             inline_member_procedures(routine)