compiler: Refactor connected_components

devito/tools/data_structures.py

@@ -585,18 +585,25 @@ def topological_sort(self, choose_element=None):
 
     def connected_components(self, enumerated=False):
         """
-        Find all connected sub-graphs and return them as a list.
+        Find all weakly connected sub-graphs and return them as a list.
         """
-        groups = []
-
-        for n0 in self.graph:
-            found = {n0} | set(self.all_downstreams(n0))
-            for g in groups:
-                if g.intersection(found):
-                    g.update(found)
-                    break
-            else:
-                groups.append(found)
+        parent = {n: n for n in self.graph}
+
+        def find(n):
+            while parent[n] is not n:
+                parent[n] = parent[parent[n]]  # Path halving
+                n = parent[n]
+            return n
+
+        for n0, downstreams in self.graph.items():
+            for n1 in downstreams:
+                parent[find(n1)] = find(n0)
+
+        # Ordered by first appearance in `self.graph`
+        groups = OrderedDict()
+        for n in self.graph:
+            groups.setdefault(find(n), set()).add(n)
+        groups = list(groups.values())
 
         if enumerated:
             mapper = OrderedDict()

tests/test_tools.py

@@ -1,14 +1,15 @@
 import os
 import time
+from itertools import groupby
 
 import numpy as np
 import pytest
 from sympy.abc import a, b, c, d, e
 
 from devito import Eq, Operator, switchenv
 from devito.tools import (
-    CacheInstances, DefaultFrozenDict, UnboundedMultiTuple, UnboundTuple, ctypes_to_cstr,
-    filter_ordered, memoized_meth, toposort, transitive_closure
+    DAG, CacheInstances, DefaultFrozenDict, UnboundedMultiTuple, UnboundTuple,
+    ctypes_to_cstr, filter_ordered, memoized_meth, toposort, transitive_closure
 )
 from devito.types.basic import Symbol
 
@@ -28,6 +29,33 @@ def test_toposort(elements, expected):
         assert expected is None
 
 
+@pytest.mark.parametrize('edges, expected', [
+    # Downstreams overlapping two previously disjoint groups
+    ([(a, d), (b, e), (c, d), (c, e)], [{a, b, c, d, e}]),
+    # Disjoint chains
+    ([(a, b), (c, d)], [{a, b}, {c, d}, {e}]),
+    # Isolated nodes
+    ([], [{a}, {b}, {c}, {d}, {e}]),
+    # Suffix fan, as created by a `BREAK` hazard in `Fusion._build_dag`
+    ([(a, c), (b, c), (c, d), (c, e)], [{a, b, c, d, e}]),
+])
+def test_connected_components(edges, expected):
+    dag = DAG(nodes=[a, b, c, d, e], edges=edges)
+
+    # Components in order of first appearance
+    assert dag.connected_components() == tuple(expected)
+
+    # Same label within a component, distinct labels across components
+    mapper = dag.connected_components(enumerated=True)
+    assert all(len({mapper[n] for n in g}) == 1 for g in expected)
+    assert len(set(mapper.values())) == len(expected)
+
+    # The runs `Fusion` derives from the labels
+    key = lambda n: next(i for i, g in enumerate(expected) if n in g)
+    assert ([len(list(g)) for _, g in groupby(dag.nodes, key=mapper.get)] ==
+            [len(list(g)) for _, g in groupby(dag.nodes, key=key)])
+
+
 def test_sorting():
     key = lambda x: x