Repository: stfc/PSyclone

Issue:
Support parallelism clause on `ACC routine` directive.
![image](https://github.com/stfc/PSyclone/assets/17409212/49e941f5-3e49-4ace-9d7f-36c26076acd4)
where
![image](https://github.com/stfc/PSyclone/assets/17409212/35e39116-c40f-4b4a-b400-573d053c6f83)
This issue will allow `gang`, `worker` or `vector` to be specified when creating the directive.

Discussion:

Patch:
diff --git a/changelog b/changelog
index 21278fdbe..b7735c83d 100644
--- a/changelog
+++ b/changelog
@@ -131,6 +131,10 @@
 	47) PR #2625 for #2619. Adds a create function for OMPParallelDoDirective
 	to avoid issues from using its parent class' create method.
 
+	48) PR #2609 for #2407. The OpenACC routine directive and transformation
+	now accept a "parallelism" attribute to set the level of parallelism, which
+	is mandatory in recent OpenACC standards.
+
 release 2.5.0 14th of February 2024
 
 	1) PR #2199 for #2189. Fix bugs with missing maps in enter data
diff --git a/src/psyclone/psyir/nodes/acc_directives.py b/src/psyclone/psyir/nodes/acc_directives.py
index 0271564d8..ef5d9299c 100644
--- a/src/psyclone/psyir/nodes/acc_directives.py
+++ b/src/psyclone/psyir/nodes/acc_directives.py
@@ -144,12 +144,54 @@ class ACCStandaloneDirective(ACCDirective, StandaloneDirective,
 
 
 class ACCRoutineDirective(ACCStandaloneDirective):
-    ''' Class representing a "!$ACC routine" OpenACC directive in PSyIR. '''
+    '''
+    Class representing an "ACC routine" OpenACC directive in PSyIR.
+
+    :param str parallelism: the level of parallelism in the routine, one of
+        "gang", "worker", "vector", "seq".
+
+    '''
+    SUPPORTED_PARALLELISM = ["seq", "vector", "worker", "gang"]
+
+    def __init__(self, parallelism="seq", **kwargs):
+        self.parallelism = parallelism
+
+        super().__init__(self, **kwargs)
+
+    @property
+    def parallelism(self):
+        '''
+        :returns: the clause describing the level of parallelism within this
+                  routine (or a called one).
+        :rtype: str
+
+        '''
+        return self._parallelism
+
+    @parallelism.setter
+    def parallelism(self, value):
+        '''
+        :param str value: the new value for the level-of-parallelism within
+                          this routine (or a called one).
+
+        :raises TypeError: if `value` is not a str.
+        :raises ValueError: if `value` is not a recognised level of
+                            parallelism.
+        '''
+        if not isinstance(value, str):
+            raise TypeError(
+                f"Expected a str to specify the level of parallelism but got "
+                f"'{type(value).__name__}'")
+        if value.lower() not in self.SUPPORTED_PARALLELISM:
+            raise ValueError(
+                f"Expected one of {self.SUPPORTED_PARALLELISM} for the level "
+                f"of parallelism but got '{value}'")
+        self._parallelism = value.lower()
 
     def gen_code(self, parent):
-        '''Generate the fortran ACC Routine Directive and any associated code.
+        '''Generate the Fortran ACC Routine Directive and any associated code.
 
-        :param parent: the parent Node in the Schedule to which to add our \
+        :param parent: the parent Node in the Schedule to which to add our
                        content.
         :type parent: sub-class of :py:class:`psyclone.f2pygen.BaseGen`
         '''
@@ -157,7 +199,8 @@ class ACCRoutineDirective(ACCStandaloneDirective):
         self.validate_global_constraints()
 
         # Generate the code for this Directive
-        parent.add(DirectiveGen(parent, "acc", "begin", "routine", ""))
+        parent.add(DirectiveGen(parent, "acc", "begin", "routine",
+                                f"{self.parallelism}"))
 
     def begin_string(self):
         '''Returns the beginning statement of this directive, i.e.
@@ -168,7 +211,7 @@ class ACCRoutineDirective(ACCStandaloneDirective):
         :rtype: str
 
         '''
-        return "acc routine"
+        return f"acc routine {self.parallelism}"
 
 
 class ACCEnterDataDirective(ACCStandaloneDirective):
diff --git a/src/psyclone/transformations.py b/src/psyclone/transformations.py
index 55c7a4a64..2d999da7a 100644
--- a/src/psyclone/transformations.py
+++ b/src/psyclone/transformations.py
@@ -379,7 +379,9 @@ class MarkRoutineForGPUMixin:
         :raises TransformationError: if it is a kernel but without an
                                      associated PSyIR.
         :raises TransformationError: if any of the symbols in the kernel are
-                                     accessed via a module use statement.
+                                     accessed via a module use statement (and
+                                     are not compile-time constants).
+        :raises TransformationError: if the routine contains any CodeBlocks.
         :raises TransformationError: if the kernel contains any calls to other
                                      routines.
         '''
@@ -2526,14 +2528,6 @@ class ACCRoutineTrans(Transformation, MarkRoutineForGPUMixin):
     >>> rtrans.apply(kern)
 
     '''
-    @property
-    def name(self):
-        '''
-        :returns: the name of this transformation class.
-        :rtype: str
-        '''
-        return "ACCRoutineTrans"
-
     def apply(self, node, options=None):
         '''
         Add the '!$acc routine' OpenACC directive into the code of the
@@ -2541,10 +2535,15 @@ class ACCRoutineTrans(Transformation, MarkRoutineForGPUMixin):
         in the supplied Routine.
 
         :param node: the kernel call or routine implementation to transform.
-        :type node: :py:class:`psyclone.psyGen.Kern` or \
+        :type node: :py:class:`psyclone.psyGen.Kern` |
                     :py:class:`psyclone.psyir.nodes.Routine`
         :param options: a dictionary with options for transformations.
         :type options: Optional[Dict[str, Any]]
+        :param bool options["force"]: whether to allow routines with
+            CodeBlocks to run on the GPU.
+        :param str options["parallelism"]: the level of parallelism that the
+            target routine (or a callee) exposes. One of "seq" (the default),
+            "vector", "worker" or "gang".
 
         '''
         # Check that we can safely apply this transformation
@@ -2563,7 +2562,10 @@ class ACCRoutineTrans(Transformation, MarkRoutineForGPUMixin):
         for child in routine.children:
             if isinstance(child, ACCRoutineDirective):
                 return  # The routine is already marked with ACCRoutine
-        routine.children.insert(0, ACCRoutineDirective())
+
+        para = options.get("parallelism", "seq") if options else "seq"
+
+        routine.children.insert(0, ACCRoutineDirective(parallelism=para))
 
     def validate(self, node, options=None):
         '''
@@ -2581,16 +2583,27 @@ class ACCRoutineTrans(Transformation, MarkRoutineForGPUMixin):
         :raises TransformationError: if the node is not a kernel or a routine.
         :raises TransformationError: if the target is a built-in kernel.
         :raises TransformationError: if it is a kernel but without an
-                                     associated PSyIR.
+            associated PSyIR.
         :raises TransformationError: if any of the symbols in the kernel are
-                                     accessed via a module use statement.
+            accessed via a module use statement.
         :raises TransformationError: if the kernel contains any calls to other
-                                     routines.
+            routines.
+        :raises TransformationError: if the 'parallelism' option is supplied
+            but is not a recognised level of parallelism.
+
         '''
         super().validate(node, options)
 
         self.validate_it_can_run_on_gpu(node, options)
 
+        if options and "parallelism" in options:
+            para = options["parallelism"]
+            if para not in ACCRoutineDirective.SUPPORTED_PARALLELISM:
+                raise TransformationError(
+                    f"{self.name}: '{para}' is not a supported level of "
+                    f"parallelism. Should be one of "
+                    f"{ACCRoutineDirective.SUPPORTED_PARALLELISM}")
+
 
 class ACCDataTrans(RegionTrans):
     '''


Test patch:
diff --git a/src/psyclone/tests/domain/gocean/transformations/gocean1p0_transformations_test.py b/src/psyclone/tests/domain/gocean/transformations/gocean1p0_transformations_test.py
index ce8de067c..706848462 100644
--- a/src/psyclone/tests/domain/gocean/transformations/gocean1p0_transformations_test.py
+++ b/src/psyclone/tests/domain/gocean/transformations/gocean1p0_transformations_test.py
@@ -1500,7 +1500,7 @@ def test_accroutinetrans_with_kern(fortran_writer, monkeypatch):
     rtrans.apply(kern)
     # Check that there is a acc routine directive in the kernel
     code = fortran_writer(kern.get_kernel_schedule())
-    assert "!$acc routine\n" in code
+    assert "!$acc routine seq\n" in code
 
     # If the kernel schedule is not accessible, the transformation fails
     def raise_gen_error():
@@ -1526,7 +1526,7 @@ def test_accroutinetrans_with_routine(fortran_writer):
     rtrans.apply(routine)
     # Check that there is a acc routine directive in the routine
     code = fortran_writer(routine)
-    assert "!$acc routine\n" in code
+    assert "!$acc routine seq\n" in code
 
     # Even if applied multiple times the Directive is only there once
     previous_num_children = len(routine.children)
diff --git a/src/psyclone/tests/psyir/nodes/acc_directives_test.py b/src/psyclone/tests/psyir/nodes/acc_directives_test.py
index a6412915e..d1166a621 100644
--- a/src/psyclone/tests/psyir/nodes/acc_directives_test.py
+++ b/src/psyclone/tests/psyir/nodes/acc_directives_test.py
@@ -417,16 +417,41 @@ def test_acc_routine_directive_constructor_and_strings():
     ''' Test the ACCRoutineDirective constructor and its output
     strings.'''
     target = ACCRoutineDirective()
-    assert target.begin_string() == "acc routine"
+    # Defaults to sequential.
+    assert target.begin_string() == "acc routine seq"
     assert str(target) == "ACCRoutineDirective[]"
 
     temporary_module = ModuleGen("test")
     target.gen_code(temporary_module)
-    assert "!$acc routine\n" in str(temporary_module.root)
+    assert "!$acc routine seq\n" in str(temporary_module.root)
 
+    target2 = ACCRoutineDirective("VECTOR")
+    assert target2.parallelism == "vector"
+    assert target2.begin_string() == "acc routine vector"
+    target3 = ACCRoutineDirective("GANG")
+    assert target3.parallelism == "gang"
+    target4 = ACCRoutineDirective("WORKER")
+    assert target4.parallelism == "worker"
+
+
+def test_acc_routine_parallelism():
+    ''' Test the ACCRoutineDirective parallelism property. '''
+    target = ACCRoutineDirective()
+    assert target.parallelism == "seq"
+    target.parallelism = "vector"
+    assert target.parallelism == "vector"
+    with pytest.raises(TypeError) as err:
+        target.parallelism = 1
+    assert ("Expected a str to specify the level of parallelism but got 'int'"
+            in str(err.value))
+    with pytest.raises(ValueError) as err:
+        target.parallelism = "sequential"
+    assert ("Expected one of ['seq', 'vector', 'worker', 'gang'] for the level"
+            " of parallelism but got 'sequential'" in str(err.value))
 
 # Class ACCUpdateDirective
 
+
 def test_accupdatedirective_init():
     ''' Test the constructor of ACCUpdateDirective node. '''
 
diff --git a/src/psyclone/tests/psyir/transformations/acc_routine_trans_test.py b/src/psyclone/tests/psyir/transformations/acc_routine_trans_test.py
new file mode 100644
index 000000000..6fa16e1ba
--- /dev/null
+++ b/src/psyclone/tests/psyir/transformations/acc_routine_trans_test.py
@@ -0,0 +1,101 @@
+# -----------------------------------------------------------------------------
+# BSD 3-Clause License
+#
+# Copyright (c) 2024, Science and Technology Facilities Council.
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+# -----------------------------------------------------------------------------
+# Author: A. R. Porter, STFC Daresbury Lab
+
+''' Module containing tests for the ACCRoutineTrans transformation.
+
+Note that the implementation itself is still in the 'old' module
+psyclone.transformations. This is because it depends on the
+ `MarkRoutineForGPUMixin` class in that module.
+
+'''
+import pytest
+
+from psyclone.psyir.nodes import ACCRoutineDirective, Routine
+from psyclone.transformations import ACCRoutineTrans, TransformationError
+
+
+def test_accrt_validate_invalid_parallelism():
+    '''Check that the validate() method rejects an unrecognised "parallelism"
+    option.'''
+    artrans = ACCRoutineTrans()
+    with pytest.raises(TransformationError) as err:
+        artrans.validate(Routine("fake"), options={"parallelism": "wrong"})
+    assert ("'wrong' is not a supported level of parallelism. Should be one "
+            "of [" in str(err.value))
+
+
+def test_accrt_apply_calls_validate(monkeypatch):
+    '''Check that the apply() method calls the validate() method.'''
+
+    def fake_validate(_1, _2):
+        raise NotImplementedError("this one")
+    artrans = ACCRoutineTrans()
+    monkeypatch.setattr(artrans, "validate", fake_validate)
+    with pytest.raises(NotImplementedError) as err:
+        artrans.apply(Routine("fake"))
+    assert "this one" in str(err.value)
+
+
+def test_accrt_apply(fortran_reader):
+    '''Test that the apply() method works as expected.'''
+    psyir = fortran_reader.psyir_from_source('''
+module my_mod
+  implicit none
+contains
+  subroutine this_one(a)
+    real, dimension(:) :: a
+    a(:) = 0.0
+  end subroutine this_one
+end module my_mod''')
+    routine = psyir.walk(Routine)[0]
+    routine_copy = routine.copy()
+    assert routine.name == "this_one"
+    artrans = ACCRoutineTrans()
+    artrans.apply(routine)
+    directives = routine.walk(ACCRoutineDirective)
+    assert len(directives) == 1
+    # Default level of parallelism is "seq".
+    assert directives[0].begin_string() == "acc routine seq"
+    assert directives[0].parent is routine
+    # Applying it a second time should have no effect.
+    artrans.apply(routine)
+    directives2 = routine.walk(ACCRoutineDirective)
+    assert len(directives2) == 1
+    assert directives2[0] is directives[0]
+    assert directives2[0].parent is routine
+    # Check that we can specify the parallelism clause.
+    artrans.apply(routine_copy, options={"parallelism": "vector"})
+    directives3 = routine_copy.walk(ACCRoutineDirective)
+    assert directives3[0].begin_string() == "acc routine vector"
diff --git a/src/psyclone/tests/psyir/transformations/kernel_transformation_test.py b/src/psyclone/tests/psyir/transformations/kernel_transformation_test.py
index 1c71fa371..c6887947e 100644
--- a/src/psyclone/tests/psyir/transformations/kernel_transformation_test.py
+++ b/src/psyclone/tests/psyir/transformations/kernel_transformation_test.py
@@ -50,8 +50,7 @@ from psyclone.psyGen import Kern
 from psyclone.psyir.nodes import Routine, FileContainer, IntrinsicCall, Call
 from psyclone.psyir.symbols import DataSymbol, INTEGER_TYPE
 from psyclone.psyir.transformations import TransformationError
-from psyclone.transformations import (ACCRoutineTrans,
-                                      Dynamo0p3KernelConstTrans)
+from psyclone.transformations import ACCRoutineTrans, Dynamo0p3KernelConstTrans
 
 from psyclone.tests.gocean_build import GOceanBuild
 from psyclone.tests.lfric_build import LFRicBuild
@@ -270,10 +269,10 @@ def test_new_same_kern_single(kernel_outputdir, monkeypatch):
 # The following tests test the MarkRoutineForGPUMixin validation, for this
 # it uses the ACCRoutineTrans as instance of this Mixin.
 
-def test_accroutine_validate_wrong_node_type():
+def test_gpumixin_validate_wrong_node_type():
     '''
-    Test that the validate() method of ACCRoutineTrans rejects a node of the
-    wrong type.
+    Test that the MarkRoutineForGPUMixin.validate_it_can_run_on_gpu() method
+    rejects a node of the wrong type.
 
     '''
     rtrans = ACCRoutineTrans()
@@ -283,10 +282,10 @@ def test_accroutine_validate_wrong_node_type():
             "Routine but got 'FileContainer'" in str(err.value))
 
 
-def test_accroutine_validate_no_schedule(monkeypatch):
+def test_gpumixin_validate_no_schedule(monkeypatch):
     '''
-    Test that the validate() method of ACCRoutineTrans catches any errors
-    generated when attempting to get the PSyIR of a kernel.
+    Test that the MarkRoutineForGPUMixin.validate_it_can_run_on_gpu() method
+    catches any errors generated when attempting to get the PSyIR of a kernel.
 
     '''
     _, invoke = get_invoke("1_single_invoke.f90", api="lfric", idx=0)
@@ -307,11 +306,11 @@ def test_accroutine_validate_no_schedule(monkeypatch):
             "transform such a kernel." in str(err.value))
 
 
-def test_accroutinetrans_validate_no_import(fortran_reader):
+def test_gpumixin_validate_no_import(fortran_reader):
     '''
-    Test the validate() method of ACCRoutineTrans rejects a kernel that
-    accesses imported data unless that data is known to be a compile-
-    time constant.
+    Test the MarkRoutineForGPUMixin.validate_it_can_run_on_gpu() method
+    rejects a kernel that accesses imported data unless that data is known to
+    be a compile-time constant.
 
     '''
     code = '''\
@@ -340,10 +339,10 @@ end module my_mod'''
     rtrans.validate(routine)
 
 
-def test_accroutinetrans_validate_no_cblock(fortran_reader):
+def test_gpumixin_validate_no_cblock(fortran_reader):
     '''
-    Test the validate() method of ACCRoutineTrans rejects a kernel that
-    contains a CodeBlock.
+    Test the MarkRoutineForGPUMixin.validate_it_can_run_on_gpu() method rejects
+    a kernel that contains a CodeBlock.
 
     '''
     code = '''\
@@ -375,10 +374,10 @@ end module my_mod'''
             in str(err.value))
 
 
-def test_accroutinetrans_validate_no_call():
+def test_gpumixin_validate_no_call():
     '''
-    Test the validate() method of ACCRoutineTrans rejects a kernel that calls
-    another routine.
+    Test the MarkRoutineForGPUMixin.validate_it_can_run_on_gpu() method rejects
+    a kernel that calls another routine.
 
     '''
     psy, invoke = get_invoke("1.15_invoke_kern_with_call.f90", api="lfric",
@@ -462,7 +461,7 @@ def test_2kern_trans(kernel_outputdir):
     assert LFRicBuild(kernel_outputdir).code_compiles(psy)
 
 
-def test_builtin_no_trans():
+def test_gpumixin_builtin_no_trans():
     ''' Check that we reject attempts to transform built-in kernels. '''
     _, invoke = get_invoke("15.1.1_X_plus_Y_builtin.f90",
                            api="lfric", idx=0)