compiler: Fix topofusion to honour the closest scope

devito/passes/clusters/fusion.py

@@ -1,4 +1,5 @@
 from collections import Counter, defaultdict
+from functools import cached_property
 from itertools import groupby
 
 from devito.finite_differences import IndexDerivative
@@ -8,7 +9,7 @@
 )
 from devito.symbolics import search
 from devito.tools import (
-    DAG, as_tuple, flatten, frozendict, memoized_func, memoized_meth, timed_pass
+    DAG, CacheInstances, as_tuple, flatten, frozendict, memoized_func, timed_pass
 )
 
 __all__ = ['fuse']
@@ -50,32 +51,92 @@ def _fusion_hazards(scope0, scope1, prefix):
     return NO_HAZARD
 
 
-class Key(tuple):
+class Keys(CacheInstances):
 
     """
-    A "fusion Key" for a Cluster (ClusterGroup) is a hashable tuple such that
-    two Clusters (ClusterGroups) are topo-fusible if and only if their Key is
-    identical.
-
-    A Key contains elements that can logically be split into two groups -- the
-    `strict` and the `weak` components of the Key. Two Clusters (ClusterGroups)
-    having same `strict` but different `weak` parts are, by definition, not
-    fusible; however, since at least their `strict` parts match, they can at
-    least be topologically reordered.
+    Provide different kind of keys for Clusters (ClusterGroups) to be used in
+    topological reordering and fusion.
     """
 
-    def __new__(cls, itintervals, guards, syncs, weak):
-        strict = [itintervals, guards, syncs]
-        obj = super().__new__(cls, strict + weak)
+    def __init__(self, c, fuse_tasks):
+        self.c = c
+        self.fuse_tasks = fuse_tasks
 
-        obj.itintervals = itintervals
-        obj.guards = guards
-        obj.syncs = syncs
+    @property
+    def first(self):
+        return self.c[0]
 
-        obj.strict = tuple(strict)
-        obj.weak = tuple(weak)
+    @property
+    def last(self):
+        return self.c[-1]
 
-        return obj
+    @cached_property
+    def itintervals(self):
+        return self.c.ispace.itintervals
+
+    @cached_property
+    def guards(self):
+        return self.c.guards if any(self.c.guards) else None
+
+    @cached_property
+    def syncs(self):
+        mapper = defaultdict(set)
+        for d, v in self.c.syncs.items():
+            for s in v:
+                if isinstance(s, PrefetchUpdate):
+                    continue
+                elif isinstance(s, WaitLock) and not self.fuse_tasks:
+                    # NOTE: A mix of Clusters w/ and w/o WaitLocks can safely
+                    # be fused, as in the worst case scenario the WaitLocks
+                    # get "hoisted" above the first Cluster in the sequence
+                    continue
+                elif isinstance(s, (InitArray, SyncArray, WaitLock, ReleaseLock)):
+                    mapper[d].add(type(s))
+                elif isinstance(s, WithLock) and self.fuse_tasks:
+                    # NOTE: Different WithLocks aren't fused unless the user
+                    # explicitly asks for it
+                    mapper[d].add(type(s))
+                else:
+                    mapper[d].add(s)
+            if d in mapper:
+                mapper[d] = frozenset(mapper[d])
+        return frozendict(mapper)
+
+    @cached_property
+    def strict(self):
+        return (self.itintervals, self.guards, self.syncs)
+
+    @cached_property
+    def weak(self):
+        c = self.c
+
+        # Clusters representing HaloTouches should get merged, if possible
+        weak = [c.is_halo_touch]
+
+        # If there are writes to thread-shared object, make it part of the key.
+        # This will promote fusion of non-adjacent Clusters writing to (some
+        # form of) shared memory, which in turn will minimize the number of
+        # necessary barriers. Same story for reads from thread-shared objects
+        weak.extend([
+            any(f._mem_shared for f in c.scope.writes),
+            any(f._mem_shared for f in c.scope.reads)
+        ])
+        weak.append(c.properties.is_core_init())
+
+        # Prefetchable Clusters should get merged, if possible
+        weak.append(c.is_glb_load_to_mem_shared)
+
+        # Promoting adjacency of IndexDerivatives will maximize their reuse
+        weak.append(any(search(c.exprs, IndexDerivative)))
+
+        # Promote adjacency of Clusters with same guard
+        weak.append(c.guards)
+
+        return tuple(weak)
+
+    @cached_property
+    def full(self):
+        return self.strict + self.weak
 
 
 class Fusion(Queue):
@@ -91,7 +152,7 @@ def __init__(self, toposort, options=None):
         options = options or {}
 
         self.toposort = toposort
-        self.fusetasks = options.get('fuse-tasks', False)
+        self.fuse_tasks = options.get('fuse-tasks', False)
 
         super().__init__()
 
@@ -111,8 +172,9 @@ def callback(self, cgroups, prefix):
             clusters = ClusterGroup(cgroups)
 
         # Fusion
+        key = lambda c: self._key(c).full
         processed = []
-        for _, group in groupby(clusters, key=self._key):
+        for _, group in groupby(clusters, key=key):
             g = list(group)
 
             for maybe_fusible in self._apply_heuristics(g):
@@ -134,60 +196,8 @@ def callback(self, cgroups, prefix):
         else:
             return [ClusterGroup(processed, prefix)]
 
-    @memoized_meth
     def _key(self, c):
-        itintervals = frozenset(c.ispace.itintervals)
-        guards = c.guards if any(c.guards) else None
-
-        # We allow fusing Clusters/ClusterGroups even in presence of WaitLocks and
-        # WithLocks, but not with any other SyncOps
-        mapper = defaultdict(set)
-        for d, v in c.syncs.items():
-            for s in v:
-                if isinstance(s, PrefetchUpdate):
-                    continue
-                elif isinstance(s, WaitLock) and not self.fusetasks:
-                    # NOTE: A mix of Clusters w/ and w/o WaitLocks can safely
-                    # be fused, as in the worst case scenario the WaitLocks
-                    # get "hoisted" above the first Cluster in the sequence
-                    continue
-                elif isinstance(s, (InitArray, SyncArray, WaitLock, ReleaseLock)):
-                    mapper[d].add(type(s))
-                elif isinstance(s, WithLock) and self.fusetasks:
-                    # NOTE: Different WithLocks aren't fused unless the user
-                    # explicitly asks for it
-                    mapper[d].add(type(s))
-                else:
-                    mapper[d].add(s)
-            if d in mapper:
-                mapper[d] = frozenset(mapper[d])
-        syncs = frozendict(mapper)
-
-        # Clusters representing HaloTouches should get merged, if possible
-        weak = [c.is_halo_touch]
-
-        # If there are writes to thread-shared object, make it part of the key.
-        # This will promote fusion of non-adjacent Clusters writing to (some
-        # form of) shared memory, which in turn will minimize the number of
-        # necessary barriers. Same story for reads from thread-shared objects
-        weak.extend([
-            any(f._mem_shared for f in c.scope.writes),
-            any(f._mem_shared for f in c.scope.reads)
-        ])
-        weak.append(c.properties.is_core_init())
-
-        # Prefetchable Clusters should get merged, if possible
-        weak.append(c.is_glb_load_to_mem_shared)
-
-        # Promoting adjacency of IndexDerivatives will maximize their reuse
-        weak.append(any(search(c.exprs, IndexDerivative)))
-
-        # Promote adjacency of Clusters with same guard
-        weak.append(c.guards)
-
-        key = Key(itintervals, guards, syncs, weak)
-
-        return key
+        return Keys(c, self.fuse_tasks)
 
     def _apply_heuristics(self, clusters):
         # We know at this point that `clusters` are potentially fusible since
@@ -232,6 +242,10 @@ def dump():
         return processed
 
     def _toposort(self, cgroups, prefix):
+        # If not enough ClusterGroups to do anything meaningful, don't waste time
+        if len(cgroups) <= 2:
+            return ClusterGroup(cgroups, prefix)
+
         # Are there any ClusterGroups that could potentially be topologically
         # reordered? If not, do not waste time
         counter = Counter(self._key(cg).strict for cg in cgroups)
@@ -248,13 +262,17 @@ def choose_element(queue, scheduled):
             k = self._key(scheduled[-1])
             m = {i: self._key(i) for i in queue}
 
-            # Process the `strict` part of the key
-            candidates = [i for i in queue if m[i].itintervals == k.itintervals]
+            # First of all, ensure we preserve the integrity of the current scope
+            candidates = [i for i in queue if k.last.ispace == m[i].first.ispace]
 
-            compatible = [i for i in candidates if m[i].guards == k.guards]
+            compatible = [i for i in candidates if k.last.guards == m[i].first.guards]
             candidates = compatible or candidates
 
-            compatible = [i for i in candidates if m[i].syncs == k.syncs]
+            # If the current scope is over, we maximize fusion
+            fusible = [i for i in queue if k.itintervals == m[i].itintervals]
+            candidates = candidates or fusible
+
+            compatible = [i for i in candidates if k.syncs == m[i].syncs]
             candidates = compatible or candidates
 
             # Process the `weak` part of the key

tests/test_dimension.py

@@ -2148,15 +2148,15 @@ def test_topofusion_w_subdims_conddims(self):
         assert exprs[1].write is g
 
         exprs = FindNodes(Expression).visit(bns['x1_blk0'])
-        assert len(exprs) == 1
-        assert exprs[0].write is h
-
-        exprs = FindNodes(Expression).visit(bns['x2_blk0'])
         assert len(exprs) == 3
         assert isinstance(exprs[0].expr.lhs, Temp)
         assert exprs[1].write is fsave
         assert exprs[2].write is gsave
 
+        exprs = FindNodes(Expression).visit(bns['x2_blk0'])
+        assert len(exprs) == 1
+        assert exprs[0].write is h
+
     def test_topofusion_w_subdims_conddims_v2(self):
         """
         Like `test_topofusion_w_subdims_conddims` but with more SubDomains,