diff --git a/src/psyclone/psyir/nodes/dynamic_omp_task_directive.py b/src/psyclone/psyir/nodes/dynamic_omp_task_directive.py index 78286d66a7..382f16b124 100644 --- a/src/psyclone/psyir/nodes/dynamic_omp_task_directive.py +++ b/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 diff --git a/src/psyclone/psyir/nodes/omp_directives.py b/src/psyclone/psyir/nodes/omp_directives.py index 66ea04cb2c..4c69ab72bb 100644 --- a/src/psyclone/psyir/nodes/omp_directives.py +++ b/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() diff --git a/src/psyclone/psyir/transformations/maximal_omp_parallel_region_trans.py b/src/psyclone/psyir/transformations/maximal_omp_parallel_region_trans.py index a4d51a995d..17dc9ec8e6 100644 --- a/src/psyclone/psyir/transformations/maximal_omp_parallel_region_trans.py +++ b/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) diff --git a/src/psyclone/psyir/transformations/omp_loop_trans.py b/src/psyclone/psyir/transformations/omp_loop_trans.py index 694c6293be..0bb86e38b5 100644 --- a/src/psyclone/psyir/transformations/omp_loop_trans.py +++ b/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: diff --git a/src/psyclone/psyir/transformations/omp_parallel_loop_trans.py b/src/psyclone/psyir/transformations/omp_parallel_loop_trans.py index b1505b4c75..5cba11617b 100644 --- a/src/psyclone/psyir/transformations/omp_parallel_loop_trans.py +++ b/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,7 +106,7 @@ 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) @@ -113,6 +114,8 @@ def apply(self, node, options=None, **kwargs): # 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 diff --git a/src/psyclone/psyir/transformations/omp_parallel_trans.py b/src/psyclone/psyir/transformations/omp_parallel_trans.py index c6fe3843d7..4aa3610986 100644 --- a/src/psyclone/psyir/transformations/omp_parallel_trans.py +++ b/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"] diff --git a/src/psyclone/tests/domain/lfric/transformations/lfric_transformations_test.py b/src/psyclone/tests/domain/lfric/transformations/lfric_transformations_test.py index 1b959307b7..2a68ffa8d3 100644 --- a/src/psyclone/tests/domain/lfric/transformations/lfric_transformations_test.py +++ b/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]) diff --git a/src/psyclone/tests/nemo/transformations/openmp/openmp_test.py b/src/psyclone/tests/nemo/transformations/openmp/openmp_test.py index 93b197aa67..34bd158a5b 100644 --- a/src/psyclone/tests/nemo/transformations/openmp/openmp_test.py +++ b/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() diff --git a/src/psyclone/tests/psyir/backend/c_test.py b/src/psyclone/tests/psyir/backend/c_test.py index f3bc82ba36..fa4d325a6d 100644 --- a/src/psyclone/tests/psyir/backend/c_test.py +++ b/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) diff --git a/src/psyclone/tests/psyir/backend/fortran_test.py b/src/psyclone/tests/psyir/backend/fortran_test.py index a61d990672..7509bc83e1 100644 --- a/src/psyclone/tests/psyir/backend/fortran_test.py +++ b/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 diff --git a/src/psyclone/tests/psyir/nodes/omp_directives_test.py b/src/psyclone/tests/psyir/nodes/omp_directives_test.py index c113cf60d5..31ff33aa28 100644 --- a/src/psyclone/tests/psyir/nodes/omp_directives_test.py +++ b/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,6 +128,7 @@ 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 @@ -131,6 +136,7 @@ def test_ompparallel_lowering(fortran_reader, monkeypatch, caplog): 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' 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 4374ef8dd6..2085138eb9 100644 --- a/src/psyclone/tests/psyir/transformations/parallel_loop_trans_test.py +++ b/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 " diff --git a/src/psyclone/tests/psyir/transformations/transformations_test.py b/src/psyclone/tests/psyir/transformations/transformations_test.py index 3698f90949..125385be2a 100644 --- a/src/psyclone/tests/psyir/transformations/transformations_test.py +++ b/src/psyclone/tests/psyir/transformations/transformations_test.py @@ -608,6 +608,44 @@ def test_omploop_trans_new_options(sample_psyir): in str(excinfo.value)) +def test_omplooptrans_apply_teamsloop_collapse( + fortran_reader, fortran_writer +): + '''Test the behaviour of the OMPLoopTrans is as expected with + teams loop and collapse''' + code = """ + subroutine x() + integer :: i, j + integer, dimension(100, 100):: arr + do i = 1, 100 + do j = 1, 100 + arr(i, j) = 1 + end do + end do + end subroutine x""" + psyir = fortran_reader.psyir_from_source(code) + looptrans = OMPLoopTrans(omp_directive="teamsloop") + loops = psyir.walk(Loop) + looptrans.apply(loops[0], options={"collapse": True}) + out = fortran_writer(psyir) + correct = """subroutine x() + integer :: i + integer :: j + integer, dimension(100,100) :: arr + + !$omp teams loop collapse(2) default(shared) private(i,j) schedule(auto) + do i = 1, 100, 1 + do j = 1, 100, 1 + arr(i,j) = 1 + enddo + enddo + !$omp end teams loop + +end subroutine x +""" + assert correct == out + + def test_omplooptrans_apply_nowait(fortran_reader, fortran_writer): '''Test the behaviour of the OMPLoopTrans is as expected when we request nowait.''' diff --git a/src/psyclone/transformations.py b/src/psyclone/transformations.py index 72ed966b3c..5cfc45227d 100644 --- a/src/psyclone/transformations.py +++ b/src/psyclone/transformations.py @@ -177,6 +177,11 @@ def validate(self, node, options=None): local_options["force"] = True super().validate(node, options=local_options) + def apply(self, node, options=None): + super().apply(node, options=options) + node.ancestor(OMPParallelDirective, + include_self=True).generate_data_clauses() + class GOceanOMPParallelLoopTrans(OMPParallelLoopTrans): @@ -223,6 +228,9 @@ def apply(self, node, options=None): OMPParallelLoopTrans.apply(self, node) + node.ancestor(OMPParallelDirective, + include_self=True).generate_data_clauses() + class LFRicOMPLoopTrans(OMPLoopTrans):