SCC: Move vector-reduction region processing to SCCRevector

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):

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(<reduction clause>)``/
+        ``!$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)