#115 Add function that returns all perfectly nested loops

psytran/loop.py

@@ -216,3 +216,36 @@ 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 reversed(loops[:]):
+            if bottom_loop in get_descendents(top_loop):
+                loops.remove(bottom_loop)
+
+    # return list of perfectly nested loops
+    return loops

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"