diff --git a/psytran/loop.py b/psytran/loop.py index 0977a3b..f96dfb2 100644 --- a/psytran/loop.py +++ b/psytran/loop.py @@ -216,3 +216,35 @@ def is_parallelisable(loop): :rtype: :py:class:`bool` """ return loop.independent_iterations() + + +def get_perfectly_nested_loops(schedule): + """ + Finds the outermost of all perfectly nested loop structures within a + supplied schedule and returns them as a list. + + Does NOT nexcessarily return an outer loop, only the topmost loop of a + perfectly nested structure. + + :arg schedule: the schedule to search. + :type schedule: :py:class:`Schedule` + + :returns loops: all loops that are the outermost of a perfectly nested + structure. + :rtype loops: list[:py:class:`Loop`] + """ + loops = [] + + # Test all loops in schedule to see if they are perfectly nested + for loop in schedule.walk(nodes.Loop): + if is_perfectly_nested(loop): + loops.append(loop) + + # For each loop in the list check if loops below it are a descendent and + # remove them if they are + for top_loop in loops[:]: + for bottom_loop in get_descendents(top_loop): + if bottom_loop in loops: + loops.remove(bottom_loop) + + return loops diff --git a/test/test_loop.py b/test/test_loop.py index 13a5001..5021d45 100644 --- a/test/test_loop.py +++ b/test/test_loop.py @@ -20,6 +20,7 @@ is_parallelisable, is_perfectly_nested, is_simple_loop, + get_perfectly_nested_loops, ) perfectly_nested_loop = { @@ -257,3 +258,51 @@ def test_is_not_independent_triple_subloop(fortran_reader): assert is_perfectly_nested(loops[1]) assert not is_independent(loops[1]) assert is_independent(loops[2]) + + +def test_get_perfectly_nested_sibling_loops(fortran_reader): + """ + Test that :func:`get_perfectly_nested_loops` correctly identifies perfectly + nested looping structures even when they are siblings. + """ + schedule = get_schedule(fortran_reader, cs.double_loop_with_2_loops) + loops = get_perfectly_nested_loops(schedule) + assert len(loops) == 2 + assert loops[0].variable.name == "i" + assert loops[1].variable.name == "i" + assert loops[0] is not loops[1] + + +def test_get_perfectly_nested_loop_top_level(fortran_reader): + """ + Test that :func:`get_perfectly_nested_loops` correctly returns only the top + or outermost loop for a nest of loops where each qualifies as 'perfectly + nested'. + """ + schedule = get_schedule( + fortran_reader, cs.quadruple_loop_with_1_assignment + ) + loops = get_perfectly_nested_loops(schedule) + + # There are four loops in the nest but we only want to return the top level + assert len(loops) == 1 + # In this case, this should return an outer loop + assert is_outer_loop(loops[0]) + + +def test_get_perfectly_nested_loop_with_conditional(fortran_reader): + """ + Test that :func:`get_perfectly_nested_loops` correctly returns only the top + or outermost loop for a nest of loops where each qualifies as 'perfectly + nested'. + """ + schedule = get_schedule( + fortran_reader, cs.conditional_imperfectly_nested_triple_loop1 + ) + loops = get_perfectly_nested_loops(schedule) + + # There are four loops in the nest but we only want to return the top level + assert len(loops) == 1 + # In this case, this should return an outer loop + assert not is_outer_loop(loops[0]) + assert loops[0].variable.name == "j"