diff --git a/loki/transformations/single_column/annotate.py b/loki/transformations/single_column/annotate.py index 2dcff09de..7cbbb47fd 100644 --- a/loki/transformations/single_column/annotate.py +++ b/loki/transformations/single_column/annotate.py @@ -5,15 +5,13 @@ # granted to it by virtue of its status as an intergovernmental organisation # nor does it submit to any jurisdiction. -import re - from loki.batch import Transformation from loki.expression import ( symbols as sym, FindVariables, is_dimension_constant ) from loki.ir import ( nodes as ir, FindNodes, Transformer, pragmas_attached, - pragma_regions_attached, is_loki_pragma, get_pragma_parameters + is_loki_pragma, get_pragma_parameters ) from loki.logging import info from loki.tools import as_tuple, flatten @@ -76,28 +74,19 @@ def kernel_annotate_vector_loops_openacc(cls, routine): f'{[a.name for a in private_arrays]}' ) - mapper = {} - with pragma_regions_attached(routine): - for region in FindNodes(ir.PragmaRegion).visit(routine.body): - if is_loki_pragma(region.pragma, starts_with='vector-reduction'): - if (reduction_clause := re.search(r'reduction\([\w:0-9 \t]+\)', region.pragma.content)): - - loops = FindNodes(ir.Loop).visit(region) - assert len(loops) == 1 - pragma = ir.Pragma(keyword='acc', content=f'loop vector {reduction_clause[0]}') - # Update loop and region in place to remove marker pragmas - loops[0]._update(pragma=(pragma,)) - region._update(pragma=None, pragma_post=None) - with pragmas_attached(routine, ir.Loop): for loop in FindNodes(ir.Loop).visit(routine.body): - if is_loki_pragma(loop.pragma, starts_with='loop vector'): - # Construct pragma and wrap entire body in vector loop - private_arrs = ', '.join(v.name for v in private_arrays) - pragma = () - private_clause = '' if not private_arrays else f' private({private_arrs})' - pragma = ir.Pragma(keyword='acc', content=f'loop vector{private_clause}') - loop._update(pragma=(pragma,)) + for pragma in as_tuple(loop.pragma): + if is_loki_pragma(pragma, starts_with='loop vector reduction'): + # Turn reduction pragmas into `!$acc` equivalent + pragma._update(keyword='acc') + continue + + if is_loki_pragma(pragma, starts_with='loop vector'): + # Turn general vector pragmas into `!$acc` and add private clause + private_arrs = ', '.join(v.name for v in private_arrays) + private_clause = '' if not private_arrays else f' private({private_arrs})' + pragma._update(keyword='acc', content=f'loop vector{private_clause}') @classmethod def kernel_annotate_sequential_loops_openacc(cls, routine): diff --git a/loki/transformations/single_column/vector.py b/loki/transformations/single_column/vector.py index c66ae7101..805be9045 100644 --- a/loki/transformations/single_column/vector.py +++ b/loki/transformations/single_column/vector.py @@ -5,6 +5,8 @@ # granted to it by virtue of its status as an intergovernmental organisation # nor does it submit to any jurisdiction. +import re + from more_itertools import split_at from loki.analyse import dataflow_analysis_attached @@ -14,7 +16,7 @@ ) from loki.ir import ( nodes as ir, FindNodes, FindScopes, Transformer, - NestedTransformer, is_loki_pragma, pragmas_attached + NestedTransformer, is_loki_pragma, pragmas_attached, pragma_regions_attached ) from loki.tools import as_tuple, flatten from loki.types import BasicType @@ -319,6 +321,31 @@ def revector_section(self, routine, section): } return Transformer(mapper).visit(section) + def mark_vector_reductions(self, routine, section): + """ + Mark vector-reduction loops in marked vector-reduction + regions. + + If a region explicitly marked with + ``!$loki vector-reduction()``/ + ``!$loki end vector-reduction`` is encountered, we replace + existing ``!$loki loop vector`` loop pragmas and add the + reduction keyword and clause. These will be turned into + OpenACC equivalents by :any:`SCCAnnotate`. + """ + with pragma_regions_attached(routine): + for region in FindNodes(ir.PragmaRegion).visit(section): + if is_loki_pragma(region.pragma, starts_with='vector-reduction'): + if (reduction_clause := re.search(r'reduction\([\w:0-9 \t]+\)', region.pragma.content)): + + loops = FindNodes(ir.Loop).visit(region) + assert len(loops) == 1 + pragma = ir.Pragma(keyword='loki', content=f'loop vector {reduction_clause[0]}') + # Update loop and region in place to remove marker pragmas + loops[0]._update(pragma=(pragma,)) + region._update(pragma=None, pragma_post=None) + + def mark_seq_loops(self, section): """ Mark interior sequential loops in a thread-parallel section @@ -390,8 +417,11 @@ def transform_subroutine(self, routine, **kwargs): # Revector all marked vector sections within the kernel body routine.body = self.revector_section(routine, routine.body) - # Mark sequential loops inside vector sections with pragmas_attached(routine, ir.Loop): + # Check for explicitly labelled vector-reduction regions + self.mark_vector_reductions(routine, routine.body) + + # Mark sequential loops inside vector sections self.mark_seq_loops(routine.body) if role == 'driver': @@ -402,6 +432,9 @@ def transform_subroutine(self, routine, **kwargs): # Revector all marked sections within the driver loop body loop._update(body=self.revector_section(routine, loop.body)) + # Check for explicitly labelled vector-reduction regions + self.mark_vector_reductions(routine, loop.body) + # Mark sequential loops inside vector sections self.mark_seq_loops(loop.body)