Repository: stfc/PSyclone (python)

Issue:
[LFRic] Dependence analysis issues
Transforming Gravity Wave for GPU using `parallel` and `loop` directives has revealed a couple of problems:

 1. ParallelLoopTrans does not take into account whether the loop is to be run in serial (e.g. ACC LOOP seq);
 2. `DynLoop.independent_iterations` says that all loops over dofs can be parallelised whereas some Builtins perform reductions that will require special treatment.

Pull request description:
(Closes #2360) Fixes for LFRic dependency analysis
A couple of small fixes identified when doing OpenACC for LFRic.

Patch:
diff --git a/changelog b/changelog
index f0305d8fa..2e0f2e543 100644
--- a/changelog
+++ b/changelog
@@ -34,6 +34,9 @@
 	12) PR #2321 for #2319. Refactors DynStencil into a dataclass named
 	LFRicArgStencil.
 
+	13) PR #2361 for #2360. Fix LFRic Loop.independent_iterations with domain
+	information and allow sequential directives on loops with dependencies.
+
 release 2.4.0 29th of September 2023
 
 	1) PR #1758 for #1741. Splits the PSyData read functionality into a
diff --git a/src/psyclone/dynamo0p3.py b/src/psyclone/dynamo0p3.py
index dafd2492a..b2f7fa9e0 100644
--- a/src/psyclone/dynamo0p3.py
+++ b/src/psyclone/dynamo0p3.py
@@ -7320,7 +7320,7 @@ class DynLoop(PSyLoop):
 
         '''
         # pylint: disable=import-outside-toplevel
-        from psyclone.psyir.tools import DependencyTools
+        from psyclone.psyir.tools import DependencyTools, DTCode
         if not dep_tools:
             dtools = DependencyTools()
         else:
@@ -7350,11 +7350,22 @@ class DynLoop(PSyLoop):
 
         # The generic DA says that this loop cannot be parallelised. However,
         # we use domain-specific information to qualify this.
-        if self.loop_type in ["colour", "dof"]:
-            # This loop is either over cells of a single colour or DoFs.
+        if self.loop_type == "colour":
+            # This loop is either over cells of a single colour.
             # According to LFRic rules this is safe to parallelise.
             return True
 
+        if self.loop_type == "dof":
+            # The generic DA can't see the PSyIR of this Builtin (because it
+            # hasn't been lowered to language level) so we use
+            # domain-specific knowledge about its properties.
+            if self.kernel.is_reduction:
+                dtools._add_message(
+                    f"Builtin '{self.kernel.name}' performs a reduction",
+                    DTCode.WARN_SCALAR_REDUCTION)
+                return False
+            return True
+
         if self.loop_type == "":
             # We can parallelise a non-coloured loop if it only updates
             # quantities on discontinuous function spaces. If an LFRic kernel
@@ -7362,7 +7373,12 @@ class DynLoop(PSyLoop):
             # have at least one argument with GH_INC access. Therefore, we
             # can simply check whether or not it has such an argument in order
             # to infer the continuity of the space.
-            return not self.has_inc_arg()
+            if self.has_inc_arg():
+                dtools._add_message(
+                    f"Kernel '{self.kernel.name}' performs an INC update",
+                    DTCode.ERROR_WRITE_WRITE_RACE)
+                return False
+            return True
 
         raise InternalError(f"independent_iterations: loop of type "
                             f"'{self.loop_type}' is not supported.")
diff --git a/src/psyclone/psyir/transformations/parallel_loop_trans.py b/src/psyclone/psyir/transformations/parallel_loop_trans.py
index 49b810df9..26916c741 100644
--- a/src/psyclone/psyir/transformations/parallel_loop_trans.py
+++ b/src/psyclone/psyir/transformations/parallel_loop_trans.py
@@ -87,10 +87,11 @@ class ParallelLoopTrans(LoopTrans, metaclass=abc.ABCMeta):
                         This transform supports "collapse", which is the\
                         number of nested loops to collapse.
         :type options: Optional[Dict[str, Any]]
-        :param int options["collapse"]: number of nested loops to collapse \
+        :param int options["collapse"]: number of nested loops to collapse
                                         or None.
-        :param bool options["force"]: whether to force parallelisation of the \
+        :param bool options["force"]: whether to force parallelisation of the
                 target loop (i.e. ignore any dependence analysis).
+        :param bool options["sequential"]: whether this is a sequential loop.
 
         :raises TransformationError: if the \
                 :py:class:`psyclone.psyir.nodes.Loop` loop iterates over \
@@ -106,16 +107,18 @@ class ParallelLoopTrans(LoopTrans, metaclass=abc.ABCMeta):
         # unsupported nodes.
         super().validate(node, options=options)
 
-        # Check we are not a sequential loop
-        if isinstance(node, PSyLoop) and node.loop_type == 'colours':
-            raise TransformationError(f"Error in {self.name} transformation. "
-                                      f"The target loop is over colours and "
-                                      f"must be computed serially.")
-
         if not options:
             options = {}
         collapse = options.get("collapse", None)
         ignore_dep_analysis = options.get("force", False)
+        sequential = options.get("sequential", False)
+
+        # Check we are not a sequential loop
+        if (not sequential and isinstance(node, PSyLoop) and
+                node.loop_type == 'colours'):
+            raise TransformationError(f"Error in {self.name} transformation. "
+                                      f"The target loop is over colours and "
+                                      f"must be computed serially.")
 
         # If 'collapse' is specified, check that it is an int and that the
         # loop nest has at least that number of loops in it
@@ -141,7 +144,7 @@ class ParallelLoopTrans(LoopTrans, metaclass=abc.ABCMeta):
                     f"containing only {loop_count} loops")
 
         # Check that there are no loop-carried dependencies
-        if ignore_dep_analysis:
+        if sequential or ignore_dep_analysis:
             return
 
         dep_tools = DependencyTools()
diff --git a/src/psyclone/transformations.py b/src/psyclone/transformations.py
index 91f5b5778..2cc35294a 100644
--- a/src/psyclone/transformations.py
+++ b/src/psyclone/transformations.py
@@ -638,10 +638,10 @@ class DynamoOMPParallelLoopTrans(OMPParallelLoopTrans):
         validity checks. Actual transformation is done by the
         :py:class:`base class <OMPParallelLoopTrans>`.
 
-        :param str omp_directive: choose which OpenMP loop directive to use. \
+        :param str omp_directive: choose which OpenMP loop directive to use.
             Defaults to "do".
-        :param str omp_schedule: the OpenMP schedule to use. Must be one of \
-            'runtime', 'static', 'dynamic', 'guided' or 'auto'. Defaults to \
+        :param str omp_schedule: the OpenMP schedule to use. Must be one of
+            'runtime', 'static', 'dynamic', 'guided' or 'auto'. Defaults to
             'static'.
 
     '''
@@ -681,8 +681,14 @@ class DynamoOMPParallelLoopTrans(OMPParallelLoopTrans):
                 raise TransformationError(
                     f"Error in {self.name} transformation. The kernel has an "
                     f"argument with INC access. Colouring is required.")
-
-        super().validate(node, options=options)
+        # As this is a domain-specific loop, we don't perform general
+        # dependence analysis because it is too conservative and doesn't
+        # account for the special steps taken for such a loop at code-
+        # generation time (e.g. the way we ensure variables are given the
+        # correct sharing attributes).
+        local_options = options.copy() if options else {}
+        local_options["force"] = True
+        super().validate(node, options=local_options)
 
 
 class GOceanOMPParallelLoopTrans(OMPParallelLoopTrans):


Test patch:
diff --git a/src/psyclone/tests/domain/lfric/lfric_loop_test.py b/src/psyclone/tests/domain/lfric/lfric_loop_test.py
index bc66858fb..7adf17156 100644
--- a/src/psyclone/tests/domain/lfric/lfric_loop_test.py
+++ b/src/psyclone/tests/domain/lfric/lfric_loop_test.py
@@ -54,7 +54,6 @@ from psyclone.parse.algorithm import parse
 from psyclone.psyGen import PSyFactory
 from psyclone.psyir.nodes import (ArrayReference, Call, Literal, Reference,
                                   Schedule, ScopingNode)
-from psyclone.psyir.symbols import DataSymbol, INTEGER_TYPE
 from psyclone.psyir.tools import DependencyTools
 from psyclone.psyir.tools.dependency_tools import Message, DTCode
 from psyclone.tests.lfric_build import LFRicBuild
@@ -1003,7 +1002,7 @@ end module testkern_mod
 
 
 def test_loop_independent_iterations(monkeypatch, dist_mem):
-    '''Tests for the independent_iterations method.'''
+    '''Tests for the independent_iterations() method.'''
     # A 'null' loop cannot be parallelised (because there's nothing to
     # parallelise).
     loop = DynLoop(loop_type="null")
@@ -1017,39 +1016,25 @@ def test_loop_independent_iterations(monkeypatch, dist_mem):
     schedule = psy.invokes.invoke_list[0].schedule
     loop = schedule.walk(DynLoop)[0]
     assert not loop.independent_iterations()
+    dtools = DependencyTools()
+    # Check that we get the expected message back from the DA
+    loop.independent_iterations(dep_tools=dtools)
+    msgs = dtools.get_all_messages()
+    assert msgs[0].code == DTCode.ERROR_WRITE_WRITE_RACE
     # Colour the loop.
     trans = Dynamo0p3ColourTrans()
     trans.apply(loop)
     loops = schedule.walk(DynLoop)
     assert not loops[0].independent_iterations()
     assert loops[1].independent_iterations()
-    # Loop over dofs.
-    _, invoke_info = parse(os.path.join(BASE_PATH,
-                                        "15.1.6_aX_plus_bY_builtin.f90"),
-                           api=TEST_API)
-    psy = PSyFactory(TEST_API, distributed_memory=dist_mem).create(invoke_info)
-    schedule = psy.invokes.invoke_list[0].schedule
-    loop = schedule.walk(DynLoop)[0]
-    assert loop.independent_iterations()
-    # Test that we can get hold of any DA messages if we supply our own
-    # instance of DependencyTools.
-    dtools = DependencyTools()
-    loop.independent_iterations(dep_tools=dtools)
-    assert len(dtools.get_all_messages()) == 1
     # Unsupported/unknown loop type
-    monkeypatch.setattr(loop, "_loop_type", "broken")
+    monkeypatch.setattr(loops[1], "_loop_type", "broken")
     with pytest.raises(InternalError) as err:
-        loop.independent_iterations()
+        loops[1].independent_iterations()
     assert "loop of type 'broken' is not supported" in str(err.value)
     monkeypatch.undo()
-    # Test that DA warnings about a scalar variable are ignored if the loop is
-    # over DoFs.
-    monkeypatch.setattr(DependencyTools, "get_all_messages",
-                        lambda _1: [Message("just a test",
-                                            DTCode.WARN_SCALAR_REDUCTION)])
-    assert loop.independent_iterations()
     # Test when the DA returns True. Since this currently never happens for an
-    # LFRic kernel we use monkeypatch.
+    # LFRic kernel that operates on cell columns we use monkeypatch.
     monkeypatch.setattr(DependencyTools, "can_loop_be_parallelised",
                         lambda _1, _2, test_all_variables=False,
                         signatures_to_ignore=[]: True)
@@ -1062,3 +1047,41 @@ def test_loop_independent_iterations(monkeypatch, dist_mem):
         raise KeyError("This is just a test")
     monkeypatch.setattr(DependencyTools, "can_loop_be_parallelised", fake)
     assert loop.independent_iterations()
+
+
+def test_dof_loop_independent_iterations(monkeypatch, dist_mem):
+    '''
+    Test that independent_iterations() behaves as expected when a loop
+    is over dofs.
+    '''
+    # Loop over dofs containing a reduction.
+    _, invoke_info = parse(os.path.join(BASE_PATH,
+                                        "15.9.1_X_innerproduct_Y_builtin.f90"),
+                           api=TEST_API)
+    psy = PSyFactory(TEST_API, distributed_memory=dist_mem).create(invoke_info)
+    loop = psy.invokes.invoke_list[0].schedule.walk(DynLoop)[0]
+    assert loop.loop_type == "dof"
+    assert not loop.independent_iterations()
+    # Test that we can get hold of any DA messages if we supply our own
+    # instance of DependencyTools.
+    dtools = DependencyTools()
+    loop.independent_iterations(dep_tools=dtools)
+    msgs = dtools.get_all_messages()
+    assert len(msgs) == 2
+    msg_codes = [msg.code for msg in msgs]
+    assert DTCode.WARN_SCALAR_WRITTEN_ONCE in msg_codes
+    assert DTCode.WARN_SCALAR_REDUCTION in msg_codes
+    # Loop over dofs with no reduction.
+    _, invoke_info = parse(os.path.join(BASE_PATH,
+                                        "15.1.6_aX_plus_bY_builtin.f90"),
+                           api=TEST_API)
+    psy = PSyFactory(TEST_API, distributed_memory=dist_mem).create(invoke_info)
+    schedule = psy.invokes.invoke_list[0].schedule
+    loop = schedule.walk(DynLoop)[0]
+    assert loop.independent_iterations()
+    # Test that DA warnings about a scalar variable are ignored if the loop is
+    # over DoFs.
+    monkeypatch.setattr(DependencyTools, "get_all_messages",
+                        lambda _1: [Message("just a test",
+                                            DTCode.WARN_SCALAR_REDUCTION)])
+    assert loop.independent_iterations()
diff --git a/src/psyclone/tests/psyir/transformations/parallel_loop_trans_test.py b/src/psyclone/tests/psyir/transformations/parallel_loop_trans_test.py
index 2624c4bc4..d63264135 100644
--- a/src/psyclone/tests/psyir/transformations/parallel_loop_trans_test.py
+++ b/src/psyclone/tests/psyir/transformations/parallel_loop_trans_test.py
@@ -87,6 +87,22 @@ def test_paralooptrans_validate_force(fortran_reader):
     trans.validate(loop, {"force": True})
 
 
+def test_paralooptrans_validate_sequential(fortran_reader):
+    '''
+    Test that the 'sequential' option allows the validate check to succeed even
+    when the dependency analysis finds a possible loop-carried dependency.
+
+    '''
+    psyir = fortran_reader.psyir_from_source(CODE)
+    loop = psyir.walk(Loop)[0]
+    trans = ParaTrans()
+    with pytest.raises(TransformationError) as err:
+        trans.validate(loop)
+    assert "Dependency analysis failed with the following" in str(err.value)
+    # Set the 'sequential' option to True - no exception should be raised.
+    trans.validate(loop, {"sequential": True})
+
+
 def test_paralooptrans_validate_collapse(fortran_reader):
     '''
     Test the various validation checks on the 'collapse' option.
@@ -118,7 +134,9 @@ def test_paralooptrans_validate_colours(monkeypatch):
     '''
     Test that we raise an error if the user attempts to apply the
     transformation to a loop over colours (since any such
-    loop must be sequential).
+    loop must be sequential). If the user explicitly requests a 'sequential'
+    loop transformation (e.g. for "acc loop seq") then that should be
+    permitted.
 
     '''
     _, invoke = get_invoke("single_invoke_three_kernels.f90", "gocean1.0",
@@ -132,6 +150,8 @@ def test_paralooptrans_validate_colours(monkeypatch):
         trans.validate(child)
     assert ("The target loop is over colours and must be computed serially"
             in str(err.value))
+    # However, if we are requesting a sequential loop then all is fine.
+    trans.validate(child, options={"sequential": True})
 
 
 def test_paralooptrans_validate_ignore_written_once(fortran_reader):