From 364b22eaf4e3256cb4b0d4c28de1817f7036266f Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Fri, 16 Feb 2024 16:15:23 +0100 Subject: [PATCH 01/12] Dimension: add separate bounds and size aliases --- loki/dimension.py | 14 +++++++++++++- loki/tests/test_dimension.py | 7 +++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/loki/dimension.py b/loki/dimension.py index e7c1f65ed..c45d94957 100644 --- a/loki/dimension.py +++ b/loki/dimension.py @@ -31,15 +31,27 @@ class Dimension: String representations of alternative size variables that are used to define arrays shapes of this dimension (eg. alternative names used in "driver" subroutines). + bounds_aliases : list or tuple of strings + String representations of alternative bounds variables that are + used to define loop ranges. """ - def __init__(self, name=None, index=None, bounds=None, size=None, aliases=None): + def __init__(self, name=None, index=None, bounds=None, size=None, aliases=None, + bounds_aliases=None): self.name = name self._index = index self._bounds = as_tuple(bounds) self._size = size self._aliases = as_tuple(aliases) + if bounds_aliases: + if len(bounds_aliases) != 2: + raise RuntimeError(f'Start and end both needed for horizontal bounds aliases in {self.name}') + if bounds_aliases[0].split('%')[0] != bounds_aliases[1].split('%')[0]: + raise RuntimeError(f'Inconsistent root name for horizontal bounds aliases in {self.name}') + + self._bounds_aliases = as_tuple(bounds_aliases) + def __repr__(self): """ Pretty-print dimension details """ name = f'<{self.name}>' if self.name else '' diff --git a/loki/tests/test_dimension.py b/loki/tests/test_dimension.py index 5d142b0dc..b71c17ad6 100644 --- a/loki/tests/test_dimension.py +++ b/loki/tests/test_dimension.py @@ -64,3 +64,10 @@ def test_dimension_index_range(frontend): assert FindNodes(Loop).visit(routine.body)[0].bounds == dim.range assert FindNodes(Loop).visit(routine.body)[0].bounds.lower == dim.bounds[0] assert FindNodes(Loop).visit(routine.body)[0].bounds.upper == dim.bounds[1] + + # Test the correct creation of horizontal dim with aliased bounds vars + _ = Dimension('test_dim_alias', bounds_aliases=('bnds%start', 'bnds%end')) + with pytest.raises(RuntimeError): + _ = Dimension('test_dim_alias', bounds_aliases=('bnds%start',)) + with pytest.raises(RuntimeError): + _ = Dimension('test_dim_alias', bounds_aliases=('bnds%start', 'some_other_bnds%end')) From dca406a330d6462095d462bee6ba303de96d9330 Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Fri, 16 Feb 2024 16:18:17 +0100 Subject: [PATCH 02/12] SCCBase: horizontal bounds vars no longer compulsory if aliases are given --- .../tests/test_single_column_coalesced.py | 43 ++++++++++++++++++- .../transformations/single_column_base.py | 24 ++++++++--- 2 files changed, 61 insertions(+), 6 deletions(-) diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index 7609572dd..c7e68e4d1 100644 --- a/transformations/tests/test_single_column_coalesced.py +++ b/transformations/tests/test_single_column_coalesced.py @@ -29,6 +29,10 @@ def fixture_horizontal(): return Dimension(name='horizontal', size='nlon', index='jl', bounds=('start', 'end'), aliases=('nproma',)) +@pytest.fixture(scope='module', name='horizontal_bounds_aliases') +def fixture_horizontal_bounds_aliases(): + return Dimension(name='horizontal_bounds_aliases', size='nlon', index='jl', bounds=('start', 'end'), + aliases=('nproma',), bounds_aliases=('bnds%start', 'bnds%end')) @pytest.fixture(scope='module', name='vertical') def fixture_vertical(): @@ -1696,11 +1700,19 @@ def test_single_column_coalesced_demotion_parameter(frontend, horizontal): @pytest.mark.parametrize('frontend', available_frontends()) -def test_scc_base_horizontal_bounds_checks(frontend, horizontal): +def test_scc_base_horizontal_bounds_checks(frontend, horizontal, horizontal_bounds_aliases): """ Test the SCCBaseTransformation checks for horizontal loop bounds. """ + fcode = """ + subroutine kernel(start, end, work) + real, intent(inout) :: work + integer, intent(in) :: start, end + + end subroutine kernel +""" + fcode_no_start = """ subroutine kernel(end, work) real, intent(inout) :: work @@ -1717,8 +1729,27 @@ def test_scc_base_horizontal_bounds_checks(frontend, horizontal): end subroutine kernel """ + fcode_alias = """ + module bnds_type_mod + implicit none + type bnds_type + integer :: start + integer :: end + end type bnds_type + end module bnds_type_mod + + subroutine kernel(bnds, work) + use bnds_type_mod, only : bnds_type + type(bnds_type), intent(in) :: bnds + real, intent(inout) :: work + + end subroutine kernel +""" + + routine = Subroutine.from_source(fcode, frontend=frontend) no_start = Subroutine.from_source(fcode_no_start, frontend=frontend) no_end = Subroutine.from_source(fcode_no_end, frontend=frontend) + alias = Sourcefile.from_source(fcode_alias, frontend=frontend) transform = SCCBaseTransformation(horizontal=horizontal) with pytest.raises(RuntimeError): @@ -1726,6 +1757,16 @@ def test_scc_base_horizontal_bounds_checks(frontend, horizontal): with pytest.raises(RuntimeError): transform.apply(no_end, role='kernel') + transform = SCCBaseTransformation(horizontal=horizontal_bounds_aliases) + transform.apply(alias, role='kernel') + + bounds = SCCBaseTransformation.check_horizontal_var(routine, horizontal_bounds_aliases) + assert bounds[0] == 'start' + assert bounds[1] == 'end' + + bounds = SCCBaseTransformation.check_horizontal_var(alias, horizontal_bounds_aliases) + assert bounds[0] == 'bnds%start' + assert bounds[1] == 'bnds%end' @pytest.mark.parametrize('frontend', available_frontends()) @pytest.mark.parametrize('trim_vector_sections', [False, True]) diff --git a/transformations/transformations/single_column_base.py b/transformations/transformations/single_column_base.py index 795d7d1f4..441c39129 100644 --- a/transformations/transformations/single_column_base.py +++ b/transformations/transformations/single_column_base.py @@ -93,10 +93,24 @@ def check_horizontal_var(cls, routine, horizontal): to define the horizontal data dimension and iteration space. """ - if horizontal.bounds[0] not in routine.variable_map: - raise RuntimeError(f'No horizontal start variable found in {routine.name}') - if horizontal.bounds[1] not in routine.variable_map: - raise RuntimeError(f'No horizontal end variable found in {routine.name}') + variables = routine.variables + bounds = as_tuple(horizontal.bounds[0]) + bounds += as_tuple(horizontal.bounds[1]) + try: + if bounds[0] not in variables: + raise RuntimeError(f'No horizontal start variable found in {routine.name}') + if bounds[1] not in variables: + raise RuntimeError(f'No horizontal end variable found in {routine.name}') + except RuntimeError as exc: + if horizontal._bounds_aliases: + bounds = as_tuple(horizontal._bounds_aliases[0]) + bounds += as_tuple(horizontal._bounds_aliases[1]) + if bounds[0].split('%', maxsplit=1)[0] not in variables: + raise RuntimeError(f'No horizontal start variable found in {routine.name}') from exc + if bounds[1].split('%', maxsplit=1)[0] not in variables: + raise RuntimeError(f'No horizontal end variable found in {routine.name}') from exc + + return bounds @classmethod def get_integer_variable(cls, routine, name): @@ -275,7 +289,7 @@ def process_kernel(self, routine): return # check for horizontal loop bounds in subroutine symbol table - self.check_horizontal_var(routine, self.horizontal) + bounds = self.check_horizontal_var(routine, self.horizontal) # Find the iteration index variable for the specified horizontal v_index = self.get_integer_variable(routine, name=self.horizontal.index) From d71de4997c97bda27f5c4f95774cc3b99528e839 Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Fri, 16 Feb 2024 16:20:19 +0100 Subject: [PATCH 03/12] SCCRevector: use horizontal bounds aliases if vars don't exist --- loki/tools/util.py | 21 +++- .../tests/test_single_column_coalesced.py | 101 ++++++++++++++++++ .../single_column_coalesced_vector.py | 11 +- 3 files changed, 129 insertions(+), 4 deletions(-) diff --git a/loki/tools/util.py b/loki/tools/util.py index a6da4dfc5..dde39d633 100644 --- a/loki/tools/util.py +++ b/loki/tools/util.py @@ -38,7 +38,7 @@ 'auto_post_mortem_debugger', 'set_excepthook', 'timeout', 'WeakrefProperty', 'group_by_class', 'replace_windowed', 'dict_override', 'stdchannel_redirected', - 'stdchannel_is_captured', 'graphviz_present' + 'stdchannel_is_captured', 'graphviz_present', 'resolve_typebound_var' ] @@ -794,3 +794,22 @@ def graphviz_present(): return False return True + + +def resolve_typebound_var(name, variable_map): + """ + A small convenience utility to resolve type-bound variables. + + Parameters + ---------- + name : str + The full name of the variable to be resolved, e.g., a%b%c%d. + variable_map : dict + A map of the variables defined in the current scope. + """ + + name_parts = name.split('%', maxsplit=1) + if (var := variable_map.get(name_parts[0], None)) and len(name_parts) > 1: + var = var.get_derived_type_member(name_parts[1]) + return var + diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index c7e68e4d1..7745057dc 100644 --- a/transformations/tests/test_single_column_coalesced.py +++ b/transformations/tests/test_single_column_coalesced.py @@ -129,6 +129,107 @@ def test_scc_revector_transformation(frontend, horizontal): assert kernel_calls[0].name == 'compute_column' +@pytest.mark.parametrize('frontend', available_frontends()) +def test_scc_revector_transformation_aliased_bounds(frontend, horizontal_bounds_aliases): + """ + Test removal of vector loops in kernel and re-insertion of a single + hoisted horizontal loop in the kernel with aliased loop bounds. + """ + + fcode_bnds_type_mod = """ + module bnds_type_mod + implicit none + type bnds_type + integer :: start + integer :: end + end type bnds_type + end module bnds_type_mod +""" + + fcode_driver = """ + SUBROUTINE column_driver(nlon, nz, q, t, nb) + USE bnds_type_mod, only : bnds_type + INTEGER, INTENT(IN) :: nlon, nz, nb ! Size of the horizontal and vertical + REAL, INTENT(INOUT) :: t(nlon,nz,nb) + REAL, INTENT(INOUT) :: q(nlon,nz,nb) + INTEGER :: b, start, end + TYPE(bnds_type) :: bnds + + bnds%start = 1 + bnds%end = nlon + do b=1, nb + call compute_column(bnds, nlon, nz, q(:,:,b), t(:,:,b)) + end do + END SUBROUTINE column_driver +""" + + fcode_kernel = """ + SUBROUTINE compute_column(bnds, nlon, nz, q, t) + USE bnds_type_mod, only : bnds_type + TYPE(bnds_type), INTENT(IN) :: bnds + INTEGER, INTENT(IN) :: nlon, nz ! Size of the horizontal and vertical + REAL, INTENT(INOUT) :: t(nlon,nz) + REAL, INTENT(INOUT) :: q(nlon,nz) + INTEGER :: jl, jk + REAL :: c + + c = 5.345 + DO jk = 2, nz + DO jl = bnds%start, bnds%end + t(jl, jk) = c * jk + q(jl, jk) = q(jl, jk-1) + t(jl, jk) * c + END DO + END DO + + ! The scaling is purposefully upper-cased + DO JL = BNDS%START, BNDS%END + Q(JL, NZ) = Q(JL, NZ) * C + END DO + END SUBROUTINE compute_column +""" + bnds_type_mod = Sourcefile.from_source(fcode_bnds_type_mod, frontend=frontend) + kernel = Sourcefile.from_source(fcode_kernel, frontend=frontend, + definitions=bnds_type_mod.definitions).subroutines[0] + driver = Sourcefile.from_source(fcode_driver, frontend=frontend, + definitions=bnds_type_mod.definitions).subroutines[0] + + # Ensure we have three loops in the kernel prior to transformation + kernel_loops = FindNodes(Loop).visit(kernel.body) + assert len(kernel_loops) == 3 + + scc_transform = (SCCDevectorTransformation(horizontal=horizontal_bounds_aliases),) + scc_transform += (SCCRevectorTransformation(horizontal=horizontal_bounds_aliases),) + for transform in scc_transform: + transform.apply(driver, role='driver') + transform.apply(kernel, role='kernel') + + # Ensure we have two nested loops in the kernel + # (the hoisted horizontal and the native vertical) + kernel_loops = FindNodes(Loop).visit(kernel.body) + assert len(kernel_loops) == 2 + assert kernel_loops[1] in FindNodes(Loop).visit(kernel_loops[0].body) + assert kernel_loops[0].variable == 'jl' + assert kernel_loops[0].bounds == 'bnds%start:bnds%end' + assert kernel_loops[1].variable == 'jk' + assert kernel_loops[1].bounds == '2:nz' + + # Ensure all expressions and array indices are unchanged + assigns = FindNodes(Assignment).visit(kernel.body) + assert fgen(assigns[1]).lower() == 't(jl, jk) = c*jk' + assert fgen(assigns[2]).lower() == 'q(jl, jk) = q(jl, jk - 1) + t(jl, jk)*c' + assert fgen(assigns[3]).lower() == 'q(jl, nz) = q(jl, nz)*c' + + # Ensure driver remains unaffected + driver_loops = FindNodes(Loop).visit(driver.body) + assert len(driver_loops) == 1 + assert driver_loops[0].variable == 'b' + assert driver_loops[0].bounds == '1:nb' + + kernel_calls = FindNodes(CallStatement).visit(driver_loops[0]) + assert len(kernel_calls) == 1 + assert kernel_calls[0].name == 'compute_column' + + @pytest.mark.parametrize('frontend', available_frontends()) def test_scc_base_resolve_vector_notation(frontend, horizontal): """ diff --git a/transformations/transformations/single_column_coalesced_vector.py b/transformations/transformations/single_column_coalesced_vector.py index 0f35085d7..d27ce606a 100644 --- a/transformations/transformations/single_column_coalesced_vector.py +++ b/transformations/transformations/single_column_coalesced_vector.py @@ -11,7 +11,8 @@ from loki import ( Transformation, FindNodes, ir, FindScopes, as_tuple, flatten, Transformer, NestedTransformer, FindVariables, demote_variables, is_dimension_constant, - is_loki_pragma, dataflow_analysis_attached, BasicType, pragmas_attached + is_loki_pragma, dataflow_analysis_attached, BasicType, pragmas_attached, + resolve_typebound_var ) from transformations.single_column_base import SCCBaseTransformation @@ -265,9 +266,13 @@ def wrap_vector_section(cls, section, routine, horizontal): The dimension specifying the horizontal vector dimension """ + variable_map = routine.variable_map + bounds = SCCBaseTransformation.check_horizontal_var(routine, horizontal) + + v_start = resolve_typebound_var(bounds[0], variable_map) + v_end = resolve_typebound_var(bounds[1], variable_map) + # Create a single loop around the horizontal from a given body - v_start = routine.variable_map[horizontal.bounds[0]] - v_end = routine.variable_map[horizontal.bounds[1]] index = SCCBaseTransformation.get_integer_variable(routine, horizontal.index) bounds = sym.LoopRange((v_start, v_end)) From 8aacaafed678d9b8e18df1eed91a00796ede2e8f Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Fri, 16 Feb 2024 16:20:52 +0100 Subject: [PATCH 04/12] SCCDemote: use horizontal size aliases to determine demotion candidates --- transformations/tests/test_single_column_coalesced.py | 5 +++-- .../transformations/single_column_coalesced_vector.py | 6 ++++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index 7745057dc..d9b5d5299 100644 --- a/transformations/tests/test_single_column_coalesced.py +++ b/transformations/tests/test_single_column_coalesced.py @@ -341,12 +341,13 @@ def test_scc_demote_transformation(frontend, horizontal): """ fcode_kernel = """ - SUBROUTINE compute_column(start, end, nlon, nz, q) + SUBROUTINE compute_column(start, end, nlon, nproma, nz, q) INTEGER, INTENT(IN) :: start, end ! Iteration indices INTEGER, INTENT(IN) :: nlon, nz ! Size of the horizontal and vertical + INTEGER, INTENT(IN) :: nproma ! Horizontal size alias REAL, INTENT(INOUT) :: q(nlon,nz) REAL :: t(nlon,nz) - REAL :: a(nlon) + REAL :: a(nproma) REAL :: b(nlon,psize) REAL :: unused(nlon) INTEGER, PARAMETER :: psize = 3 diff --git a/transformations/transformations/single_column_coalesced_vector.py b/transformations/transformations/single_column_coalesced_vector.py index d27ce606a..4f5a5d806 100644 --- a/transformations/transformations/single_column_coalesced_vector.py +++ b/transformations/transformations/single_column_coalesced_vector.py @@ -331,7 +331,8 @@ def _get_local_arrays(section): # Only demote local arrays with the horizontal as fast dimension arrays = [v for v in arrays if isinstance(v, sym.Array)] arrays = [v for v in arrays if v.name not in argument_names] - arrays = [v for v in arrays if v.shape and v.shape[0] == horizontal.size] + arrays = [v for v in arrays if v.shape and + v.shape[0] in [horizontal.size, *horizontal._aliases]] # Also demote arrays whose remaning dimensions are known constants arrays = [v for v in arrays if all(is_dimension_constant(d) for d in v.shape[1:])] @@ -401,4 +402,5 @@ def process_kernel(self, routine, demote_locals=True): if demote_locals: variables = tuple(v.name for v in to_demote) if variables: - demote_variables(routine, variable_names=variables, dimensions=self.horizontal.size) + demote_variables(routine, variable_names=variables, + dimensions=[self.horizontal.size, *self.horizontal._aliases]) From 9e6781e20357a54eb3dcb4f936b32985bf955839 Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Sat, 17 Feb 2024 18:07:16 +0100 Subject: [PATCH 05/12] SCCDemote: don't demote arrays marked explicitly for preservation --- .../tests/test_single_column_coalesced.py | 17 ++++++++++++++--- .../single_column_coalesced_vector.py | 10 ++++++++-- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index d9b5d5299..61e42ec9e 100644 --- a/transformations/tests/test_single_column_coalesced.py +++ b/transformations/tests/test_single_column_coalesced.py @@ -350,6 +350,7 @@ def test_scc_demote_transformation(frontend, horizontal): REAL :: a(nproma) REAL :: b(nlon,psize) REAL :: unused(nlon) + REAL :: d(nlon,psize) INTEGER, PARAMETER :: psize = 3 INTEGER :: jl, jk REAL :: c @@ -369,18 +370,24 @@ def test_scc_demote_transformation(frontend, horizontal): b(jl, 2) = Q(JL, 3) b(jl, 3) = a(jl) * (b(jl, 1) + b(jl, 2)) + d(jl, 1) = b(jl, 1) + d(jl, 2) = b(jl, 2) + d(jl, 3) = b(jl, 3) + Q(JL, NZ) = Q(JL, NZ) * C + b(jl, 3) END DO END SUBROUTINE compute_column """ - kernel = Subroutine.from_source(fcode_kernel, frontend=frontend) + kernel_source = Sourcefile.from_source(fcode_kernel, frontend=frontend) + kernel_item = ProcedureItem(name='#compute_column', source=kernel_source, config={'preserve_arrays': ['d',]}) + kernel = kernel_source.subroutines[0] # Must run SCCDevector first because demotion relies on knowledge # of vector sections scc_transform = (SCCDevectorTransformation(horizontal=horizontal),) scc_transform += (SCCDemoteTransformation(horizontal=horizontal),) for transform in scc_transform: - transform.apply(kernel, role='kernel') + transform.apply(kernel, role='kernel', item=kernel_item) # Ensure correct array variables shapes assert isinstance(kernel.variable_map['a'], Scalar) @@ -389,6 +396,7 @@ def test_scc_demote_transformation(frontend, horizontal): assert isinstance(kernel.variable_map['t'], Array) assert isinstance(kernel.variable_map['q'], Array) assert isinstance(kernel.variable_map['unused'], Scalar) + assert isinstance(kernel.variable_map['d'], Array) # Ensure that parameter-sized array b got demoted only assert kernel.variable_map['b'].shape == ((3,) if frontend is OMNI else ('psize',)) @@ -403,7 +411,10 @@ def test_scc_demote_transformation(frontend, horizontal): assert fgen(assigns[4]).lower() == 'b(1) = q(jl, 2)' assert fgen(assigns[5]).lower() == 'b(2) = q(jl, 3)' assert fgen(assigns[6]).lower() == 'b(3) = a*(b(1) + b(2))' - assert fgen(assigns[7]).lower() == 'q(jl, nz) = q(jl, nz)*c + b(3)' + assert fgen(assigns[7]).lower() == 'd(jl, 1) = b(1)' + assert fgen(assigns[8]).lower() == 'd(jl, 2) = b(2)' + assert fgen(assigns[9]).lower() == 'd(jl, 3) = b(3)' + assert fgen(assigns[10]).lower() == 'q(jl, nz) = q(jl, nz)*c + b(3)' @pytest.mark.parametrize('frontend', available_frontends()) diff --git a/transformations/transformations/single_column_coalesced_vector.py b/transformations/transformations/single_column_coalesced_vector.py index 4f5a5d806..ecb4394a1 100644 --- a/transformations/transformations/single_column_coalesced_vector.py +++ b/transformations/transformations/single_column_coalesced_vector.py @@ -376,11 +376,13 @@ def transform_subroutine(self, routine, **kwargs): if role == 'kernel': demote_locals = self.demote_local_arrays + preserve_arrays = [] if item: demote_locals = item.config.get('demote_locals', self.demote_local_arrays) - self.process_kernel(routine, demote_locals=demote_locals) + preserve_arrays = item.config.get('preserve_arrays', []) + self.process_kernel(routine, demote_locals=demote_locals, preserve_arrays=preserve_arrays) - def process_kernel(self, routine, demote_locals=True): + def process_kernel(self, routine, demote_locals=True, preserve_arrays=None): """ Applies the SCCDemote utilities to a "kernel" and demotes all suitable local arrays. @@ -398,6 +400,10 @@ def process_kernel(self, routine, demote_locals=True): # may carry buffered values between them, so that we may not demote those! to_demote = self.kernel_get_locals_to_demote(routine, sections, self.horizontal) + # Filter out arrays marked explicitly for preservation + if preserve_arrays: + to_demote = [v for v in to_demote if not v.name in preserve_arrays] + # Demote all private local variables that do not buffer values between sections if demote_locals: variables = tuple(v.name for v in to_demote) From 26f9752df84800a98f8ef8378a4d27dc29a0ce01 Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Fri, 1 Mar 2024 12:44:30 +0100 Subject: [PATCH 06/12] SCCBase: vector notation resolution now works for aliased bounds vars --- transformations/tests/test_single_column_coalesced.py | 2 +- transformations/transformations/single_column_base.py | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index 61e42ec9e..6f34366be 100644 --- a/transformations/tests/test_single_column_coalesced.py +++ b/transformations/tests/test_single_column_coalesced.py @@ -1862,7 +1862,7 @@ def test_scc_base_horizontal_bounds_checks(frontend, horizontal, horizontal_boun routine = Subroutine.from_source(fcode, frontend=frontend) no_start = Subroutine.from_source(fcode_no_start, frontend=frontend) no_end = Subroutine.from_source(fcode_no_end, frontend=frontend) - alias = Sourcefile.from_source(fcode_alias, frontend=frontend) + alias = Sourcefile.from_source(fcode_alias, frontend=frontend).subroutines[0] transform = SCCBaseTransformation(horizontal=horizontal) with pytest.raises(RuntimeError): diff --git a/transformations/transformations/single_column_base.py b/transformations/transformations/single_column_base.py index 441c39129..9453b0e11 100644 --- a/transformations/transformations/single_column_base.py +++ b/transformations/transformations/single_column_base.py @@ -11,6 +11,7 @@ ir, Transformation, FindNodes, Transformer, as_tuple, FindExpressions, SymbolAttributes, BasicType, SubstituteExpressions, + resolve_typebound_var ) @@ -180,9 +181,12 @@ def resolve_vector_dimension(cls, routine, loop_variable, bounds): bounds : tuple of :any:`Scalar` Tuple defining the iteration space of the inserted loops. """ + bounds_str = f'{bounds[0]}:{bounds[1]}' - bounds_v = (sym.Variable(name=bounds[0]), sym.Variable(name=bounds[1])) + variable_map = routine.variable_map + bounds_v = (resolve_typebound_var(bounds[0], variable_map), + resolve_typebound_var(bounds[1], variable_map)) mapper = {} for stmt in FindNodes(ir.Assignment).visit(routine.body): @@ -302,7 +306,7 @@ def process_kernel(self, routine): self.resolve_masked_stmts(routine, loop_variable=v_index) # Resolve vector notation, eg. VARIABLE(KIDIA:KFDIA) - self.resolve_vector_dimension(routine, loop_variable=v_index, bounds=self.horizontal.bounds) + self.resolve_vector_dimension(routine, loop_variable=v_index, bounds=bounds) def process_driver(self, routine): """ From d12fe4c1842f4e87f81ba65e2a39c350437bffc5 Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Fri, 15 Mar 2024 11:25:30 +0100 Subject: [PATCH 07/12] SCCAnnotate: check against size alias when marking arrays for privatization --- .../tests/test_single_column_coalesced.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index 6f34366be..4ca96fcb3 100644 --- a/transformations/tests/test_single_column_coalesced.py +++ b/transformations/tests/test_single_column_coalesced.py @@ -36,7 +36,7 @@ def fixture_horizontal_bounds_aliases(): @pytest.fixture(scope='module', name='vertical') def fixture_vertical(): - return Dimension(name='vertical', size='nz', index='jk') + return Dimension(name='vertical', size='nz', index='jk', aliases=('nlev',)) @pytest.fixture(scope='module', name='blocking') @@ -731,26 +731,30 @@ def test_scc_annotate_openacc(frontend, horizontal, blocking): """ fcode_driver = """ - SUBROUTINE column_driver(nlon, nz, q, nb) + SUBROUTINE column_driver(nlon, nproma, nlev, nz, q, nb) INTEGER, INTENT(IN) :: nlon, nz, nb ! Size of the horizontal and vertical + INTEGER, INTENT(IN) :: nproma, nlev ! Aliases of horizontal and vertical sizes REAL, INTENT(INOUT) :: q(nlon,nz,nb) INTEGER :: b, start, end start = 1 end = nlon do b=1, nb - call compute_column(start, end, nlon, nz, q(:,:,b)) + call compute_column(start, end, nlon, nproma, nz, q(:,:,b)) end do END SUBROUTINE column_driver """ fcode_kernel = """ - SUBROUTINE compute_column(start, end, nlon, nz, q) - INTEGER, INTENT(IN) :: start, end ! Iteration indices - INTEGER, INTENT(IN) :: nlon, nz ! Size of the horizontal and vertical + SUBROUTINE compute_column(start, end, nlon, nproma, nlev, nz, q) + INTEGER, INTENT(IN) :: start, end ! Iteration indices + INTEGER, INTENT(IN) :: nlon, nz ! Size of the horizontal and vertical + INTEGER, INTENT(IN) :: nproma, nlev ! Aliases of horizontal and vertical sizes REAL, INTENT(INOUT) :: q(nlon,nz) REAL :: t(nlon,nz) REAL :: a(nlon) + REAL :: d(nproma) + REAL :: e(nlev) REAL :: b(nlon,psize) INTEGER, PARAMETER :: psize = 3 INTEGER :: jl, jk From 17ab087c6172ffb68e9fbc89d03a01c23e69b880 Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Fri, 12 Apr 2024 09:59:53 +0200 Subject: [PATCH 08/12] PoolAllocator: switch to typebound var resolution convenience utility --- transformations/transformations/pool_allocator.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/transformations/transformations/pool_allocator.py b/transformations/transformations/pool_allocator.py index 294716a7e..c063bfccd 100644 --- a/transformations/transformations/pool_allocator.py +++ b/transformations/transformations/pool_allocator.py @@ -23,7 +23,7 @@ DetachScopesMapper, SymbolAttributes, BasicType, DerivedType, is_dimension_constant, recursive_expression_map_update, get_pragma_parameters, FindInlineCalls, Interface, - dataflow_analysis_attached + dataflow_analysis_attached, resolve_typebound_var ) __all__ = ['TemporariesPoolAllocatorTransformation'] @@ -435,11 +435,7 @@ def _get_stack_storage_and_size_var(self, routine, stack_size): ) variables_append += [stack_storage] - name_parts = self.block_dim.size.split('%', maxsplit=1) - block_size = routine.symbol_map[name_parts[0]] - if len(name_parts) > 1: - block_size = block_size.get_derived_type_member(name_parts[1]) - + block_size = resolve_typebound_var(self.block_dim.size, routine.symbol_map) stack_alloc = Allocation(variables=(stack_storage.clone(dimensions=( # pylint: disable=no-member stack_size_var, block_size)),)) stack_dealloc = Deallocation(variables=(stack_storage.clone(dimensions=None),)) # pylint: disable=no-member From c679b91958c3d4bd12faf6dee934e1299acabaa9 Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Fri, 19 Apr 2024 12:13:48 +0000 Subject: [PATCH 09/12] Appease pylint --- loki/tools/util.py | 1 - 1 file changed, 1 deletion(-) diff --git a/loki/tools/util.py b/loki/tools/util.py index dde39d633..a15dcb3de 100644 --- a/loki/tools/util.py +++ b/loki/tools/util.py @@ -812,4 +812,3 @@ def resolve_typebound_var(name, variable_map): if (var := variable_map.get(name_parts[0], None)) and len(name_parts) > 1: var = var.get_derived_type_member(name_parts[1]) return var - From af7def5d402dd92791717065f41c85e603903602 Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Tue, 23 Apr 2024 10:31:53 +0200 Subject: [PATCH 10/12] SCCBase: simplify horizontal loop bounds checking --- loki/dimension.py | 12 +++++++++ .../tests/test_single_column_coalesced.py | 4 +-- .../transformations/single_column_base.py | 27 +++++++------------ .../single_column_coalesced_vector.py | 2 +- 4 files changed, 25 insertions(+), 20 deletions(-) diff --git a/loki/dimension.py b/loki/dimension.py index c45d94957..a85a1a662 100644 --- a/loki/dimension.py +++ b/loki/dimension.py @@ -106,3 +106,15 @@ def size_expressions(self): if self._bounds: exprs += (f'{self._bounds[1]} - {self._bounds[0]} + 1', ) return exprs + + @property + def bounds_expressions(self): + """ + A list of all expression strings representing the bounds of a data space. + """ + + exprs = [(b,) for b in self.bounds] + if self._bounds_aliases: + exprs = [expr + (b,) for expr, b in zip(exprs, self._bounds_aliases)] + + return as_tuple(exprs) diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index 4ca96fcb3..0112df6af 100644 --- a/transformations/tests/test_single_column_coalesced.py +++ b/transformations/tests/test_single_column_coalesced.py @@ -1877,11 +1877,11 @@ def test_scc_base_horizontal_bounds_checks(frontend, horizontal, horizontal_boun transform = SCCBaseTransformation(horizontal=horizontal_bounds_aliases) transform.apply(alias, role='kernel') - bounds = SCCBaseTransformation.check_horizontal_var(routine, horizontal_bounds_aliases) + bounds = SCCBaseTransformation.get_horizontal_loop_bounds(routine, horizontal_bounds_aliases) assert bounds[0] == 'start' assert bounds[1] == 'end' - bounds = SCCBaseTransformation.check_horizontal_var(alias, horizontal_bounds_aliases) + bounds = SCCBaseTransformation.get_horizontal_loop_bounds(alias, horizontal_bounds_aliases) assert bounds[0] == 'bnds%start' assert bounds[1] == 'bnds%end' diff --git a/transformations/transformations/single_column_base.py b/transformations/transformations/single_column_base.py index 9453b0e11..56677b506 100644 --- a/transformations/transformations/single_column_base.py +++ b/transformations/transformations/single_column_base.py @@ -81,7 +81,7 @@ def check_routine_pragmas(cls, routine, directive): return False @classmethod - def check_horizontal_var(cls, routine, horizontal): + def get_horizontal_loop_bounds(cls, routine, horizontal): """ Check for horizontal loop bounds in a :any:`Subroutine`. @@ -94,22 +94,15 @@ def check_horizontal_var(cls, routine, horizontal): to define the horizontal data dimension and iteration space. """ + bounds = () variables = routine.variables - bounds = as_tuple(horizontal.bounds[0]) - bounds += as_tuple(horizontal.bounds[1]) - try: - if bounds[0] not in variables: - raise RuntimeError(f'No horizontal start variable found in {routine.name}') - if bounds[1] not in variables: - raise RuntimeError(f'No horizontal end variable found in {routine.name}') - except RuntimeError as exc: - if horizontal._bounds_aliases: - bounds = as_tuple(horizontal._bounds_aliases[0]) - bounds += as_tuple(horizontal._bounds_aliases[1]) - if bounds[0].split('%', maxsplit=1)[0] not in variables: - raise RuntimeError(f'No horizontal start variable found in {routine.name}') from exc - if bounds[1].split('%', maxsplit=1)[0] not in variables: - raise RuntimeError(f'No horizontal end variable found in {routine.name}') from exc + for name, _bounds in zip(['start', 'end'], horizontal.bounds_expressions): + for bound in _bounds: + if bound.split('%', maxsplit=1)[0] in variables: + bounds += (bound,) + break + else: + raise RuntimeError(f'No horizontol {name} variable matching {_bounds[0]} found in {routine.name}') return bounds @@ -293,7 +286,7 @@ def process_kernel(self, routine): return # check for horizontal loop bounds in subroutine symbol table - bounds = self.check_horizontal_var(routine, self.horizontal) + bounds = self.get_horizontal_loop_bounds(routine, self.horizontal) # Find the iteration index variable for the specified horizontal v_index = self.get_integer_variable(routine, name=self.horizontal.index) diff --git a/transformations/transformations/single_column_coalesced_vector.py b/transformations/transformations/single_column_coalesced_vector.py index ecb4394a1..8782cddc4 100644 --- a/transformations/transformations/single_column_coalesced_vector.py +++ b/transformations/transformations/single_column_coalesced_vector.py @@ -267,7 +267,7 @@ def wrap_vector_section(cls, section, routine, horizontal): """ variable_map = routine.variable_map - bounds = SCCBaseTransformation.check_horizontal_var(routine, horizontal) + bounds = SCCBaseTransformation.get_horizontal_loop_bounds(routine, horizontal) v_start = resolve_typebound_var(bounds[0], variable_map) v_end = resolve_typebound_var(bounds[1], variable_map) From 44127cc35116f97cb90f5d089da8047220abb9b4 Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Tue, 23 Apr 2024 11:09:50 +0200 Subject: [PATCH 11/12] Make typebound var resolution utility a method of ProgramUnit --- loki/expression/tests/test_expression.py | 2 ++ loki/program_unit.py | 20 +++++++++++++++++++ loki/tools/util.py | 20 +------------------ .../transformations/pool_allocator.py | 4 ++-- .../transformations/single_column_base.py | 7 +++---- .../single_column_coalesced_vector.py | 7 +++---- 6 files changed, 31 insertions(+), 29 deletions(-) diff --git a/loki/expression/tests/test_expression.py b/loki/expression/tests/test_expression.py index 82da2f60d..096f294c2 100644 --- a/loki/expression/tests/test_expression.py +++ b/loki/expression/tests/test_expression.py @@ -1647,6 +1647,8 @@ def test_typebound_resolution_type_info(frontend): for var_name in var_tt_to_try: assert f'var_tt%{var_name}' not in sub.symbol_attrs + assert 'var_c%c_b%b_a%a' == sub.resolve_typebound_var('var_c%c_b%b_a%a') + # Create each derived type member and verify its type for var_name, dtype in var_c_to_try.items(): var = var_c.get_derived_type_member(var_name) diff --git a/loki/program_unit.py b/loki/program_unit.py index 11f797363..f73063020 100644 --- a/loki/program_unit.py +++ b/loki/program_unit.py @@ -798,3 +798,23 @@ def apply(self, op, **kwargs): """ # TODO: Should type-check for an `Operation` object here op.apply(self, **kwargs) + + def resolve_typebound_var(self, name, variable_map=None): + """ + A small convenience utility to resolve type-bound variables. + + Parameters + ---------- + name : str + The full name of the variable to be resolved, e.g., a%b%c%d. + variable_map : dict + A map of the variables defined in the current scope. + """ + + if not (_variable_map := variable_map): + _variable_map = self.variable_map + + name_parts = name.split('%', maxsplit=1) + if (var := _variable_map.get(name_parts[0], None)) and len(name_parts) > 1: + var = var.get_derived_type_member(name_parts[1]) + return var diff --git a/loki/tools/util.py b/loki/tools/util.py index a15dcb3de..a6da4dfc5 100644 --- a/loki/tools/util.py +++ b/loki/tools/util.py @@ -38,7 +38,7 @@ 'auto_post_mortem_debugger', 'set_excepthook', 'timeout', 'WeakrefProperty', 'group_by_class', 'replace_windowed', 'dict_override', 'stdchannel_redirected', - 'stdchannel_is_captured', 'graphviz_present', 'resolve_typebound_var' + 'stdchannel_is_captured', 'graphviz_present' ] @@ -794,21 +794,3 @@ def graphviz_present(): return False return True - - -def resolve_typebound_var(name, variable_map): - """ - A small convenience utility to resolve type-bound variables. - - Parameters - ---------- - name : str - The full name of the variable to be resolved, e.g., a%b%c%d. - variable_map : dict - A map of the variables defined in the current scope. - """ - - name_parts = name.split('%', maxsplit=1) - if (var := variable_map.get(name_parts[0], None)) and len(name_parts) > 1: - var = var.get_derived_type_member(name_parts[1]) - return var diff --git a/transformations/transformations/pool_allocator.py b/transformations/transformations/pool_allocator.py index c063bfccd..e99278366 100644 --- a/transformations/transformations/pool_allocator.py +++ b/transformations/transformations/pool_allocator.py @@ -23,7 +23,7 @@ DetachScopesMapper, SymbolAttributes, BasicType, DerivedType, is_dimension_constant, recursive_expression_map_update, get_pragma_parameters, FindInlineCalls, Interface, - dataflow_analysis_attached, resolve_typebound_var + dataflow_analysis_attached ) __all__ = ['TemporariesPoolAllocatorTransformation'] @@ -435,7 +435,7 @@ def _get_stack_storage_and_size_var(self, routine, stack_size): ) variables_append += [stack_storage] - block_size = resolve_typebound_var(self.block_dim.size, routine.symbol_map) + block_size = routine.resolve_typebound_var(self.block_dim.size, routine.symbol_map) stack_alloc = Allocation(variables=(stack_storage.clone(dimensions=( # pylint: disable=no-member stack_size_var, block_size)),)) stack_dealloc = Deallocation(variables=(stack_storage.clone(dimensions=None),)) # pylint: disable=no-member diff --git a/transformations/transformations/single_column_base.py b/transformations/transformations/single_column_base.py index 56677b506..5901d6944 100644 --- a/transformations/transformations/single_column_base.py +++ b/transformations/transformations/single_column_base.py @@ -10,8 +10,7 @@ from loki import ( ir, Transformation, FindNodes, Transformer, as_tuple, FindExpressions, - SymbolAttributes, BasicType, SubstituteExpressions, - resolve_typebound_var + SymbolAttributes, BasicType, SubstituteExpressions ) @@ -178,8 +177,8 @@ def resolve_vector_dimension(cls, routine, loop_variable, bounds): bounds_str = f'{bounds[0]}:{bounds[1]}' variable_map = routine.variable_map - bounds_v = (resolve_typebound_var(bounds[0], variable_map), - resolve_typebound_var(bounds[1], variable_map)) + bounds_v = (routine.resolve_typebound_var(bounds[0], variable_map), + routine.resolve_typebound_var(bounds[1], variable_map)) mapper = {} for stmt in FindNodes(ir.Assignment).visit(routine.body): diff --git a/transformations/transformations/single_column_coalesced_vector.py b/transformations/transformations/single_column_coalesced_vector.py index 8782cddc4..effae644b 100644 --- a/transformations/transformations/single_column_coalesced_vector.py +++ b/transformations/transformations/single_column_coalesced_vector.py @@ -11,8 +11,7 @@ from loki import ( Transformation, FindNodes, ir, FindScopes, as_tuple, flatten, Transformer, NestedTransformer, FindVariables, demote_variables, is_dimension_constant, - is_loki_pragma, dataflow_analysis_attached, BasicType, pragmas_attached, - resolve_typebound_var + is_loki_pragma, dataflow_analysis_attached, BasicType, pragmas_attached ) from transformations.single_column_base import SCCBaseTransformation @@ -269,8 +268,8 @@ def wrap_vector_section(cls, section, routine, horizontal): variable_map = routine.variable_map bounds = SCCBaseTransformation.get_horizontal_loop_bounds(routine, horizontal) - v_start = resolve_typebound_var(bounds[0], variable_map) - v_end = resolve_typebound_var(bounds[1], variable_map) + v_start = routine.resolve_typebound_var(bounds[0], variable_map) + v_end = routine.resolve_typebound_var(bounds[1], variable_map) # Create a single loop around the horizontal from a given body index = SCCBaseTransformation.get_integer_variable(routine, horizontal.index) From aec79bfd0e2c7b3155842cee7b954b047ad6236c Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Tue, 23 Apr 2024 11:15:36 +0200 Subject: [PATCH 12/12] SCCDemote: add documentation for array demotion options --- .../transformations/single_column_coalesced_vector.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/transformations/transformations/single_column_coalesced_vector.py b/transformations/transformations/single_column_coalesced_vector.py index effae644b..3a36ea296 100644 --- a/transformations/transformations/single_column_coalesced_vector.py +++ b/transformations/transformations/single_column_coalesced_vector.py @@ -304,6 +304,11 @@ class SCCDemoteTransformation(Transformation): A set of utilities to determine which local arrays can be safely demoted in a :any:`Subroutine` as part of a transformation pass. + Unless the option `demote_local_arrays` is set to `False`, this transformation will demote + local arrays that do not buffer values between vector loops. Specific arrays in individual + routines can also be marked for preservation by assigning them to the `preserve_arrays` list + in the :any:`SchedulerConfig`. + Parameters ---------- horizontal : :any:`Dimension`