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.

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)

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,14 +116,20 @@ 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
         ----------
         routine : :any:`Subroutine`
             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)

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

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,16 +782,19 @@ 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
     assert pragmas[0].keyword == 'acc'
     assert pragmas[0].content == 'routine seq'
 
 
-
 @pytest.mark.parametrize('frontend', available_frontends())
 def test_scc_vector_reduction(frontend, horizontal, blocking):
     """

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)

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'])

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