diff --git a/devito/tools/data_structures.py b/devito/tools/data_structures.py index ced3fd2d9b..bde8e338c4 100644 --- a/devito/tools/data_structures.py +++ b/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() diff --git a/tests/test_tools.py b/tests/test_tools.py index e5b4617e41..49e24e3900 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -1,5 +1,6 @@ import os import time +from itertools import groupby import numpy as np import pytest @@ -7,8 +8,8 @@ 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