diff --git a/devito/passes/clusters/fusion.py b/devito/passes/clusters/fusion.py index 994b4c9b65..1459a32dd9 100644 --- a/devito/passes/clusters/fusion.py +++ b/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 diff --git a/tests/test_dimension.py b/tests/test_dimension.py index c8de790431..96bec12935 100644 --- a/tests/test_dimension.py +++ b/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,