diff --git a/cmake/loki_transform.cmake b/cmake/loki_transform.cmake index dc14af58e..2c3cab406 100644 --- a/cmake/loki_transform.cmake +++ b/cmake/loki_transform.cmake @@ -25,6 +25,7 @@ include( loki_transform_helpers ) # [CPP] # [FRONTEND ] # [INLINE_MEMBERS] +# [RESOLVE_SEQUENCE_ASSOCIATION] # [BUILDDIR ] # [SOURCES [ ...]] # [HEADERS [ ...]] @@ -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 ] # [SOURCES [ ...]] # [HEADERS [ ...]] -# [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 index a3a5b344d..680ae0e72 100644 --- 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 index 56504cdb4..e3a51ed85 100644 --- 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 new file mode 100644 index 000000000..08f7b9d22 --- /dev/null +++ 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 index 084841429..6e459275f 100644 --- 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 new file mode 100644 index 000000000..bdd2c4978 --- /dev/null +++ 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 index 8294adcb6..b9f1be0ae 100644 --- 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 index 8cf608072..5bf2f9cb1 100644 --- 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)