(#Towards 3420) Move OMPParallelTrans to create the data sharing attributes at transformation time

src/psyclone/psyir/nodes/dynamic_omp_task_directive.py

@@ -2102,6 +2102,28 @@ def _compute_clauses(self):
             out_clause,
         )
 
+    def generate_data_and_depend_clauses(self):
+        """
+        Adds the data and depend clauses needed for this task directive.
+        """
+        # Create the clauses
+        (
+            private_clause,
+            firstprivate_clause,
+            shared_clause,
+            in_clause,
+            out_clause,
+        ) = self._compute_clauses()
+
+        # Replace the children with the new children
+        old_children = self.pop_all_children()
+        self.addchild(old_children[0])
+        self.addchild(private_clause)
+        self.addchild(firstprivate_clause)
+        self.addchild(shared_clause)
+        self.addchild(in_clause)
+        self.addchild(out_clause)
+
     def lower_to_language_level(self):
         """
         Lowers the structure of the PSyIR tree inside the Directive
@@ -2137,22 +2159,7 @@ def lower_to_language_level(self):
                 )
 
         # Create the clauses
-        (
-            private_clause,
-            firstprivate_clause,
-            shared_clause,
-            in_clause,
-            out_clause,
-        ) = self._compute_clauses()
-
-        # Replace the children with the new children
-        old_children = self.pop_all_children()
-        self.addchild(old_children[0])
-        self.addchild(private_clause)
-        self.addchild(firstprivate_clause)
-        self.addchild(shared_clause)
-        self.addchild(in_clause)
-        self.addchild(out_clause)
+        self.generate_data_and_depend_clauses()
         super().lower_to_language_level()
 
         # Replace this node with an OMPTaskDirective

src/psyclone/psyir/nodes/omp_directives.py

@@ -1282,6 +1282,27 @@ def private_clause(self):
         '''
         return self.children[2]
 
+    def generate_data_clauses(self):
+        '''
+        Adds the private and firstprivate clauses to this OMPParallelDirective
+        according to the contained nodes.
+        '''
+        reprod_red_call_list = self.reductions(reprod=True)
+        private, fprivate, need_sync = self.infer_sharing_attributes()
+        # Order the clause variables alphabetically (to facilitate comparisons)
+        if reprod_red_call_list:
+            table = self.ancestor(Routine).symbol_table
+            thread_idx = table.find_or_create_tag(
+                "omp_thread_index", root_name="th_idx",
+                symbol_type=DataSymbol, datatype=INTEGER_TYPE)
+            private.add(thread_idx)
+        private_clause = OMPPrivateClause.create(
+                            sorted(private, key=lambda x: x.name))
+        fprivate_clause = OMPFirstprivateClause.create(
+                            sorted(fprivate, key=lambda x: x.name))
+        self.children[2].replace_with(private_clause)
+        self.children[3].replace_with(fprivate_clause)
+
     def lower_to_language_level(self) -> Node:
         '''
         In-place construction of clauses as PSyIR constructs.
@@ -1356,15 +1377,7 @@ def lower_to_language_level(self) -> Node:
 
         # Create data sharing clauses (before lowering, so the CodedKern
         # semantics are taken into account)
-        private, fprivate, need_sync = self.infer_sharing_attributes()
-
-        # Order the clause variables alphabetically (to facilitate comparisons)
-        if reprod_red_call_list:
-            private.add(thread_idx)
-        private_clause = OMPPrivateClause.create(
-                            sorted(private, key=lambda x: x.name))
-        fprivate_clause = OMPFirstprivateClause.create(
-                            sorted(fprivate, key=lambda x: x.name))
+        _, _, need_sync = self.infer_sharing_attributes()
 
         # Check all of the need_sync nodes are synchronized in children.
         # unless it has reduction_kernels which are handled separately
@@ -1387,9 +1400,6 @@ def lower_to_language_level(self) -> Node:
                         " or the code includes the necessary "
                         "synchronisations.", type(self).__name__, sym.name)
 
-        self.children[2].replace_with(private_clause)
-        self.children[3].replace_with(fprivate_clause)
-
         # Continue lowering children (but not the clauses we just added)
         for child in self.children[:2]:
             child.lower_to_language_level()

src/psyclone/psyir/transformations/maximal_omp_parallel_region_trans.py

@@ -88,6 +88,9 @@ class MaximalOMPParallelRegionTrans(MaximalRegionTrans):
     def apply(self, nodes: Union[Node, Schedule, list[Node]], **kwargs):
         '''Applies the transformation to the nodes provided.
 
+        In addition to its own keyword argument options, it also accepts any
+        options valid for :py:class:`OMPParallelTrans`
+
         :param nodes: can be a single node, a schedule or a list of nodes.
         '''
         super().apply(nodes, **kwargs)

src/psyclone/psyir/transformations/omp_loop_trans.py

@@ -342,7 +342,9 @@ def apply(self, node, options=None,
         parent = node.parent
         position = node.position
         super().apply(node, local_options, reduction_ops=red_ops, **kwargs)
-
+        if self.omp_directive not in ["do", "loop"]:
+            node.ancestor((OMPParallelDoDirective),
+                          include_self=True).generate_data_clauses()
         # Add reduction clauses to the newly introduced directive
         directive = parent.children[position]
         for (op, ref) in self.inferred_reduction_clauses:

src/psyclone/psyir/transformations/omp_parallel_loop_trans.py

@@ -41,12 +41,13 @@
 
 
 from psyclone.psyir.nodes import (
-    OMPParallelDoDirective, OMPReductionClause)
+    OMPParallelDoDirective, OMPReductionClause, OMPParallelDirective)
 from psyclone.psyir.nodes.omp_directives import MAP_REDUCTION_OP_TO_OMP
 from psyclone.psyir.transformations.omp_loop_trans import OMPLoopTrans
 
 
 class OMPParallelLoopTrans(OMPLoopTrans):
+
     ''' Adds an OpenMP PARALLEL DO directive to a loop.
 
         For example:
@@ -105,14 +106,16 @@ def apply(self, node, options=None, **kwargs):
         directive = OMPParallelDoDirective(children=[node.detach()],
                                            omp_schedule=self.omp_schedule)
 
-        # add any inferred reduction clauses to the newly introduced directive
+        # Add any inferred reduction clauses to the newly introduced directive
         for (op, ref) in self.inferred_reduction_clauses:
             clause = OMPReductionClause(MAP_REDUCTION_OP_TO_OMP[op])
             clause.addchild(ref)
             directive.addchild(clause)
 
         # add the OpenMP loop directive as a child of the node's parent
         node_parent.addchild(directive, index=node_position)
+        node.ancestor(OMPParallelDirective,
+                      include_self=True).generate_data_clauses()
 
 
 # For Sphinx AutoAPI documentation generation

src/psyclone/psyir/transformations/omp_parallel_trans.py

@@ -137,7 +137,11 @@ def apply(self, nodes: list[Node], options=None, **kwargs):
         :param nodes: list of Nodes to put within parallel region.
         '''
         # TODO #2668: Remove options.
-        super().apply(nodes, options, **kwargs)
+        super().apply(nodes, options=options, **kwargs)
+
+        node_list = self.get_node_list(nodes)
+        node_list[0].ancestor(OMPParallelDirective,
+                              include_self=True).generate_data_clauses()
 
 
 __all__ = ["OMPParallelTrans"]

src/psyclone/tests/domain/lfric/transformations/lfric_transformations_test.py

@@ -3876,6 +3876,8 @@ def test_reprod_view(monkeypatch, annexed, dist_mem):
             7*indent + "0: " + call + " x_innerproduct_y(asum,f1,f2)\n" +
             2*indent + ompdefault + "[default=DefaultClauseTypes.SHARED]\n" +
             2*indent + ompprivate + "[]\n" +
+            3*indent + "Reference[name:'df']\n" +
+            3*indent + "Reference[name:'th_idx']\n" +
             2*indent + ompfprivate + "[]\n" +
             indent + "1: " + gsum + "[scalar='asum']\n" +
             indent + "2: " + ompparallel + "[]\n" +
@@ -3892,6 +3894,7 @@ def test_reprod_view(monkeypatch, annexed, dist_mem):
             7*indent + "0: " + call + " inc_a_times_x(asum,f1)\n" +
             2*indent + ompdefault + "[default=DefaultClauseTypes.SHARED]\n" +
             2*indent + ompprivate + "[]\n" +
+            3*indent + "Reference[name:'df']\n" +
             2*indent + ompfprivate + "[]\n" +
             indent + "3: " + ompparallel + "[]\n" +
             2*indent + sched + "[]\n" +
@@ -3907,6 +3910,8 @@ def test_reprod_view(monkeypatch, annexed, dist_mem):
             7*indent + "0: " + call + " sum_x(bsum,f2)\n" +
             2*indent + ompdefault + "[default=DefaultClauseTypes.SHARED]\n" +
             2*indent + ompprivate + "[]\n" +
+            3*indent + "Reference[name:'df']\n" +
+            3*indent + "Reference[name:'th_idx']\n" +
             2*indent + ompfprivate + "[]\n" +
             indent + "4: " + gsum + "[scalar='bsum']\n")
         if not annexed:
@@ -3928,6 +3933,8 @@ def test_reprod_view(monkeypatch, annexed, dist_mem):
             7*indent + "0: " + call + " x_innerproduct_y(asum,f1,f2)\n" +
             2*indent + ompdefault + "[default=DefaultClauseTypes.SHARED]\n" +
             2*indent + ompprivate + "[]\n" +
+            3*indent + "Reference[name:'df']\n" +
+            3*indent + "Reference[name:'th_idx']\n" +
             2*indent + ompfprivate + "[]\n" +
             indent + "1: " + ompparallel + "[]\n" +
             2*indent + sched + "[]\n" +
@@ -3943,6 +3950,7 @@ def test_reprod_view(monkeypatch, annexed, dist_mem):
             7*indent + "0: " + call + " inc_a_times_x(asum,f1)\n" +
             2*indent + ompdefault + "[default=DefaultClauseTypes.SHARED]\n" +
             2*indent + ompprivate + "[]\n" +
+            3*indent + "Reference[name:'df']\n" +
             2*indent + ompfprivate + "[]\n" +
             indent + "2: " + ompparallel + "[]\n" +
             2*indent + sched + "[]\n" +
@@ -3958,13 +3966,15 @@ def test_reprod_view(monkeypatch, annexed, dist_mem):
             7*indent + "0: " + call + " sum_x(bsum,f2)\n" +
             2*indent + ompdefault + "[default=DefaultClauseTypes.SHARED]\n" +
             2*indent + ompprivate + "[]\n" +
+            3*indent + "Reference[name:'df']\n" +
+            3*indent + "Reference[name:'th_idx']\n" +
             2*indent + ompfprivate + "[]\n")
     if expected not in result:
         print("Expected ...")
         print(expected)
         print("Found ...")
         print(result)
-        assert 0
+    assert expected in result
 
 
 @pytest.mark.parametrize("reprod", [True, False])

src/psyclone/tests/nemo/transformations/openmp/openmp_test.py

@@ -89,6 +89,7 @@ def test_omp_explicit_gen(fortran_reader, fortran_writer):
         "  !$omp end parallel do\n"
         "\n"
         "end program explicit_do")
+    print(fortran_writer(psyir).lower())
     assert expected in fortran_writer(psyir).lower()
     # Check that calling gen a second time gives the same code
     assert expected in fortran_writer(psyir).lower()

src/psyclone/tests/psyir/backend/c_test.py

@@ -571,6 +571,7 @@ def test_cw_directive_with_clause(fortran_reader):
     master = OMPMasterDirective(children=[directive])
     parallel = OMPParallelDirective.create(children=[master])
     schedule.addchild(parallel, 0)
+    parallel.generate_data_clauses()
     # Add a barrier to cover the StandaloneDirective visitor
     parallel.children[0].addchild(OMPBarrierDirective(), 0)
 

src/psyclone/tests/psyir/backend/fortran_test.py

@@ -2125,6 +2125,7 @@ def test_fw_directive_with_clause(fortran_reader, fortran_writer):
     master = OMPMasterDirective(children=[directive])
     parallel = OMPParallelDirective.create(children=[master])
     schedule.addchild(parallel, 0)
+    parallel.generate_data_clauses()
     assert '''!$omp parallel default(shared) private(i)
   !$omp master
   !$omp taskloop num_tasks(32) nogroup

src/psyclone/tests/psyir/nodes/omp_directives_test.py

@@ -80,8 +80,11 @@
 TEST_LOGGER_INF = "psyclone.psyir.tools.reduction_inference"
 
 
-def test_ompparallel_lowering(fortran_reader, monkeypatch, caplog):
-    ''' Check that lowering an OMP Parallel region leaves it with the
+def test_ompparallel_generate_data_clauses_and_lowering(
+    fortran_reader, monkeypatch, caplog
+):
+    ''' Check that lowering an OMP Parallel region and calling
+    generate_data_clauses (as the transformations do) leaves it with the
     appropriate begin_string and clauses for the backend to generate
     the right code'''
     code = '''
@@ -106,6 +109,7 @@ def test_ompparallel_lowering(fortran_reader, monkeypatch, caplog):
     ptrans.apply(loops[0].parent.parent)
     assert isinstance(tree.children[0].children[0], OMPParallelDirective)
     pdir = tree.children[0].children[0]
+    pdir.generate_data_clauses()
     pdir.lower_to_language_level()
     assert pdir.begin_string() == "omp parallel"
     assert len(pdir.children) == 4
@@ -124,13 +128,15 @@ def test_ompparallel_lowering(fortran_reader, monkeypatch, caplog):
     # Add loop
     pdir.children[0].addchild(new_loop)
 
+    pdir.generate_data_clauses()
     pdir.lower_to_language_level()
     assert pdir.children[2] is not priv_clause
 
     # Monkeypatch a case with private and firstprivate clauses
     monkeypatch.setattr(pdir, "infer_sharing_attributes",
                         lambda: ({Symbol("a")}, {Symbol("b")}, None))
 
+    pdir.generate_data_clauses()
     pdir.lower_to_language_level()
     assert isinstance(pdir.children[2], OMPPrivateClause)
     assert len(pdir.children[2].children) == 1
@@ -169,9 +175,11 @@ def test_ompparallel_lowering(fortran_reader, monkeypatch, caplog):
             "\n" in caplog.text)
 
 
-def test_omp_parallel_do_lowering(fortran_reader, monkeypatch, caplog):
-    ''' Check that lowering an OMP Parallel Do leaves it with the
-    appropriate begin_string and clauses for the backend to generate
+def test_omp_parallel_do_generate_data_clauses(
+    fortran_reader, monkeypatch, caplog
+):
+    ''' Check that generate_data_clauses on an OMP Parallel Do leaves it with
+    the appropriate begin_string and clauses for the backend to generate
     the right code'''
 
     code = '''
@@ -212,6 +220,8 @@ def test_omp_parallel_do_lowering(fortran_reader, monkeypatch, caplog):
     # Change the schedule
     pdir._omp_schedule = "dynamic"
 
+    pdir.generate_data_clauses()
+    # Schedule is only updated in lowering.
     pdir.lower_to_language_level()
     assert pdir.children[2] is not priv_clause
     assert isinstance(pdir.children[2], OMPPrivateClause)
@@ -224,7 +234,7 @@ def test_omp_parallel_do_lowering(fortran_reader, monkeypatch, caplog):
     monkeypatch.setattr(pdir, "infer_sharing_attributes",
                         lambda: ({Symbol("a")}, {Symbol("b")}, None))
 
-    pdir.lower_to_language_level()
+    pdir.generate_data_clauses()
     assert isinstance(pdir.children[2], OMPPrivateClause)
     assert len(pdir.children[2].children) == 1
     assert pdir.children[2].children[0].name == 'a'

src/psyclone/tests/psyir/transformations/parallel_loop_trans_test.py

@@ -64,6 +64,14 @@ def _directive(self, children, collapse=None):
         '''
         return OMPParallelDoDirective(children=children, collapse=collapse)
 
+    def apply(self, node, options=None, **kwargs):
+        '''Applies the ParaTrans to the provided node.'''
+        super().apply(node, options=options, **kwargs)
+        # Since we create OMPParallelDoDirectives we have to call
+        # the relevant method.
+        node.ancestor(OMPParallelDoDirective,
+                      include_self=True).generate_data_clauses()
+
 
 CODE = '''
 subroutine my_sub()
@@ -149,6 +157,7 @@ def test_paralooptrans_validate_pure_calls(fortran_reader, fortran_writer):
     loop = psyir.walk(Loop)[0]
     trans.apply(loop)
     code = fortran_writer(psyir)
+    print(code)
     assert "!$omp parallel do default(shared) private(i,j,k)" in code
 
 
@@ -665,6 +674,7 @@ def test_paralooptrans_with_array_privatisation(fortran_reader,
                          logger="psyclone.psyir.transformations"):
         trans.apply(loop, {"privatise_arrays": True,
                            "force_private": ["ztmp2", "symbol_not_in_loop"]})
+    print(fortran_writer(psyir))
     assert ("!$omp parallel do default(shared) private(ji,jj,ztmp) "
             "firstprivate(ztmp2)" in fortran_writer(psyir))
     assert ("ParaTrans has been provided with the 'symbol_not_in_loop' symbol "