From 291ad457155bc7cafd82300c6fbb1b9195d35cb1 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Fri, 9 Aug 2024 10:45:43 +0000 Subject: [PATCH] SCC: Let SCCRevector mark routine and SCCAnnotate translates This also renames and refactors the `check_routine_pragmas` utility, which now only needs to check for genuine `!$loki routine seq` annotations. --- .../block_index_transformations.py | 6 ++-- .../transformations/single_column/annotate.py | 21 ++++++----- loki/transformations/single_column/base.py | 4 +-- .../single_column/tests/test_scc.py | 14 +++++--- loki/transformations/single_column/vector.py | 11 ++++-- loki/transformations/tests/test_utilities.py | 18 +++++----- loki/transformations/utilities.py | 36 ++++--------------- 7 files changed, 50 insertions(+), 60 deletions(-) diff --git a/loki/transformations/block_index_transformations.py b/loki/transformations/block_index_transformations.py index 722b844e9..59a0d8c47 100644 --- a/loki/transformations/block_index_transformations.py +++ b/loki/transformations/block_index_transformations.py @@ -14,7 +14,7 @@ from loki.transformations.sanitise import resolve_associates from loki.transformations.utilities import ( recursive_expression_map_update, get_integer_variable, - get_loop_bounds, check_routine_pragmas + get_loop_bounds, check_routine_sequential ) from loki.transformations.single_column.base import SCCBaseTransformation @@ -238,8 +238,8 @@ def process_kernel(self, routine, item, successors, targets, exclude_arrays): v_index = get_integer_variable(routine, name=self.horizontal.index) SCCBaseTransformation.resolve_masked_stmts(routine, loop_variable=v_index) - # Bail if routine is marked as sequential or routine has already been processed - if check_routine_pragmas(routine, directive=None): + # Bail if routine is marked as sequential + if check_routine_sequential(routine): return bounds = get_loop_bounds(routine, self.horizontal) diff --git a/loki/transformations/single_column/annotate.py b/loki/transformations/single_column/annotate.py index 7cbbb47fd..46b6fd468 100644 --- a/loki/transformations/single_column/annotate.py +++ b/loki/transformations/single_column/annotate.py @@ -18,7 +18,7 @@ from loki.types import DerivedType from loki.transformations.utilities import ( - find_driver_loops, get_local_arrays, check_routine_pragmas + find_driver_loops, get_local_arrays ) @@ -116,7 +116,8 @@ def kernel_annotate_sequential_loops_openacc(cls, routine): @classmethod def kernel_annotate_subroutine_present_openacc(cls, routine): """ - Insert ``!$acc data present`` annotations around the body of a subroutine. + Insert ``!$acc routine seq/vector`` directives and wrap + subroutine body in ``!$acc data present`` directives. Parameters ---------- @@ -124,6 +125,11 @@ def kernel_annotate_subroutine_present_openacc(cls, routine): The subroutine to which annotations will be added """ + # Update `!$loki routine seq/vector` pragmas with `!$acc` + for pragma in FindNodes(ir.Pragma).visit(routine.ir): + if is_loki_pragma(pragma, starts_with='routine'): + pragma._update(keyword='acc') + # Get the names of all array and derived type arguments args = [a for a in routine.arguments if isinstance(a, sym.Array)] args += [a for a in routine.arguments if isinstance(a.type.dtype, DerivedType)] @@ -146,9 +152,6 @@ def insert_annotations(cls, routine, horizontal): # to ensure device-resident data is used for array and struct arguments. cls.kernel_annotate_subroutine_present_openacc(routine) - # Mark routine as `!$acc routine vector` to make it device-callable - routine.spec.append(ir.Pragma(keyword='acc', content='routine vector')) - def transform_subroutine(self, routine, **kwargs): """ Apply SCCAnnotate utilities to a :any:`Subroutine`. @@ -180,9 +183,11 @@ def process_kernel(self, routine): Subroutine to apply this transformation to. """ - # Bail if routine is marked as sequential - if check_routine_pragmas(routine, self.directive): - return + # Bail if this routine has been processed before + for p in FindNodes(ir.Pragma).visit(routine.ir): + # Check if `!$acc routine` has already been added + if p.keyword.lower() == 'acc' and 'routine' in p.content.lower(): + return if self.directive == 'openacc': self.insert_annotations(routine, self.horizontal) diff --git a/loki/transformations/single_column/base.py b/loki/transformations/single_column/base.py index 5c2f17c92..a08a60a93 100644 --- a/loki/transformations/single_column/base.py +++ b/loki/transformations/single_column/base.py @@ -14,7 +14,7 @@ from loki.transformations.sanitise import resolve_associates from loki.transformations.utilities import ( - get_integer_variable, get_loop_bounds, check_routine_pragmas + get_integer_variable, get_loop_bounds, check_routine_sequential ) @@ -146,7 +146,7 @@ def process_kernel(self, routine): """ # Bail if routine is marked as sequential or routine has already been processed - if check_routine_pragmas(routine, self.directive): + if check_routine_sequential(routine): return # check for horizontal loop bounds in subroutine symbol table diff --git a/loki/transformations/single_column/tests/test_scc.py b/loki/transformations/single_column/tests/test_scc.py index 0059a27b7..52febbabf 100644 --- a/loki/transformations/single_column/tests/test_scc.py +++ b/loki/transformations/single_column/tests/test_scc.py @@ -753,9 +753,10 @@ def test_scc_multiple_acc_pragmas(frontend, horizontal, blocking): @pytest.mark.parametrize('frontend', available_frontends()) -def test_scc_base_routine_seq_pragma(frontend, horizontal): +def test_scc_annotate_routine_seq_pragma(frontend, horizontal, blocking): """ - Test that `!$loki routine seq` pragmas are replaced correctly by `!$acc routine seq` pragmas. + Test that `!$loki routine seq` pragmas are replaced correctly by + `!$acc routine seq` pragmas. """ fcode = """ @@ -781,8 +782,12 @@ def test_scc_base_routine_seq_pragma(frontend, horizontal): assert pragmas[0].keyword == 'loki' assert pragmas[0].content == 'routine seq' - transformation = SCCBaseTransformation(horizontal=horizontal, directive='openacc') - transformation.transform_subroutine(routine, role='kernel', targets=['some_kernel',]) + transformation = SCCAnnotateTransformation( + horizontal=horizontal, directive='openacc', block_dim=blocking + ) + transformation.transform_subroutine( + routine, role='kernel', targets=['some_kernel',] + ) pragmas = FindNodes(Pragma).visit(routine.spec) assert len(pragmas) == 1 @@ -790,7 +795,6 @@ def test_scc_base_routine_seq_pragma(frontend, horizontal): assert pragmas[0].content == 'routine seq' - @pytest.mark.parametrize('frontend', available_frontends()) def test_scc_vector_reduction(frontend, horizontal, blocking): """ diff --git a/loki/transformations/single_column/vector.py b/loki/transformations/single_column/vector.py index 805be9045..6fb147ef4 100644 --- a/loki/transformations/single_column/vector.py +++ b/loki/transformations/single_column/vector.py @@ -24,7 +24,7 @@ from loki.transformations.array_indexing import demote_variables from loki.transformations.utilities import ( get_integer_variable, get_loop_bounds, find_driver_loops, - get_local_arrays, check_routine_pragmas + get_local_arrays, check_routine_sequential ) @@ -99,7 +99,7 @@ def extract_vector_sections(cls, section, horizontal): # check if calls have been enriched if not call.routine is BasicType.DEFERRED: # check if called routine is marked as sequential - if check_routine_pragmas(routine=call.routine, directive=None): + if check_routine_sequential(routine=call.routine): continue if call in section: @@ -414,6 +414,10 @@ def transform_subroutine(self, routine, **kwargs): targets = kwargs.get('targets', ()) if role == 'kernel': + # Skip if kernel is marked as `!$loki routine seq` + if check_routine_sequential(routine): + return + # Revector all marked vector sections within the kernel body routine.body = self.revector_section(routine, routine.body) @@ -424,6 +428,9 @@ def transform_subroutine(self, routine, **kwargs): # Mark sequential loops inside vector sections self.mark_seq_loops(routine.body) + # Mark subroutine as vector parallel for later annotation + routine.spec.append(ir.Pragma(keyword='loki', content='routine vector')) + if role == 'driver': with pragmas_attached(routine, ir.Loop): driver_loops = find_driver_loops(routine=routine, targets=targets) diff --git a/loki/transformations/tests/test_utilities.py b/loki/transformations/tests/test_utilities.py index 7dc3158e8..7f35b4799 100644 --- a/loki/transformations/tests/test_utilities.py +++ b/loki/transformations/tests/test_utilities.py @@ -19,7 +19,7 @@ single_variable_declaration, recursive_expression_map_update, convert_to_lower_case, replace_intrinsics, rename_variables, get_integer_variable, get_loop_bounds, is_driver_loop, - find_driver_loops, get_local_arrays, check_routine_pragmas + find_driver_loops, get_local_arrays, check_routine_sequential ) @@ -511,11 +511,11 @@ def test_transform_utilites_get_local_arrays(frontend, tmp_path): @pytest.mark.parametrize('frontend', available_frontends()) -def test_transform_utilites_check_routine_pragmas(frontend, tmp_path): - """ Test :any:`check_routine_pragmas` utility. """ +def test_transform_utilites_check_routine_sequential(frontend, tmp_path): + """ Test :any:`check_routine_sequential` utility. """ fcode = """ -module test_check_routine_pragmas_mod +module test_check_routine_sequential_mod implicit none contains @@ -537,12 +537,10 @@ def test_transform_utilites_check_routine_pragmas(frontend, tmp_path): i = i + 1 end subroutine test_acc_vec -end module test_check_routine_pragmas_mod +end module test_check_routine_sequential_mod """ module = Module.from_source(fcode, frontend=frontend, xmods=[tmp_path]) - # TODO: This utility needs some serious clean-up, so we're just testing - # the bare basics here and promise to do better next time ;) - assert check_routine_pragmas(module['test_acc_seq'], directive=None) - assert check_routine_pragmas(module['test_loki_seq'], directive=None) - assert check_routine_pragmas(module['test_acc_vec'], directive='openacc') + assert not check_routine_sequential(module['test_acc_seq']) + assert check_routine_sequential(module['test_loki_seq']) + assert not check_routine_sequential(module['test_acc_vec']) diff --git a/loki/transformations/utilities.py b/loki/transformations/utilities.py index 6c98e5435..04489f71f 100644 --- a/loki/transformations/utilities.py +++ b/loki/transformations/utilities.py @@ -32,7 +32,7 @@ 'sanitise_imports', 'replace_selected_kind', 'single_variable_declaration', 'recursive_expression_map_update', 'get_integer_variable', 'get_loop_bounds', 'find_driver_loops', - 'get_local_arrays', 'check_routine_pragmas' + 'get_local_arrays', 'check_routine_sequential' ] @@ -650,41 +650,17 @@ def get_local_arrays(routine, section, unique=True): return arrays -def check_routine_pragmas(routine, directive): +def check_routine_sequential(routine): """ - Check if routine is marked as sequential or has already been processed. + Check if routine is marked as "sequential". Parameters ---------- routine : :any:`Subroutine` Subroutine to perform checks on. - directive: string or None - Directives flavour to use for parallelism annotations; either - ``'openacc'`` or ``None``. - """ - - pragmas = FindNodes(ir.Pragma).visit(routine.ir) - routine_pragmas = [p for p in pragmas if p.keyword.lower() in ['loki', 'acc']] - routine_pragmas = [p for p in routine_pragmas if 'routine' in p.content.lower()] - - seq_pragmas = [r for r in routine_pragmas if 'seq' in r.content.lower()] - if seq_pragmas: - loki_seq_pragmas = [r for r in routine_pragmas if 'loki' == r.keyword.lower()] - if loki_seq_pragmas: - if directive == 'openacc': - # Mark routine as acc seq - mapper = {seq_pragmas[0]: None} - routine.spec = Transformer(mapper).visit(routine.spec) - routine.body = Transformer(mapper).visit(routine.body) - - # Append the acc pragma to routine.spec, regardless of where the corresponding - # loki pragma is found - routine.spec.append(ir.Pragma(keyword='acc', content='routine seq')) - return True - - vec_pragmas = [r for r in routine_pragmas if 'vector' in r.content.lower()] - if vec_pragmas: - if directive == 'openacc': + """ + for pragma in FindNodes(ir.Pragma).visit(routine.ir): + if is_loki_pragma(pragma, starts_with='routine seq'): return True return False