compiler: Tweak caching in fusion pass, along with DAG construction

devito/ir/support/basic.py

@@ -339,7 +339,6 @@ def rebuild(self, **kwargs):
 
         return TimedAccess(access, mode, timestamp, ispace)
 
-    @memoized_meth
     def distance(self, other, logical=False):
         """
         Compute the distance from ``self`` to ``other``.
@@ -351,7 +350,19 @@ def distance(self, other, logical=False):
         logical : bool
             Compute a logical distance rather than true distance (i.e. ignoring
             degenerating indices created by size 1 buffers etc).
+
+        Notes
+        -----
+        The distance depends on the two accesses (including their Functions'
+        metadata), their IterationSpaces and `logical`, but never on the
+        timestamps or access modes. Hence it is cached on the former, and thus
+        shared by all TimedAccesses with the same geometry, such as the
+        timestamp-shifted copies created by `Scope.from_scopes`.
         """
+        return _cached_distance(self.access, self.ispace, other.access,
+                                other.ispace, logical)
+
+    def _distance(self, other, logical=False):
         if isinstance(self.access, ComponentAccess) and \
            isinstance(other.access, ComponentAccess) and \
            self.access.index != other.access.index:
@@ -1186,9 +1197,7 @@ def reads_smart_gen(self, f):
 
         if isinstance(f, (Function, Temp, TempArray, TBArray)):
             for i in self.getreads(f):
-                expand = erange if uses_subdims(i) else extrema
-
-                for j in expand(i.access):
+                for j in _expand_read(i.access, uses_subdims(i)):
                     yield TimedAccess(j, i.mode, i.timestamp, i.ispace)
 
         else:
@@ -1604,6 +1613,25 @@ def skippable_interval(d, ispace, it):
     return d is None or (d in ispace and not d._defines & it.dim._defines)
 
 
+@memoized_func(scope='build')
+def _expand_read(access, subdims):
+    """
+    The accesses a read may perform once its StencilDimensions are resolved;
+    see `Scope.reads_smart_gen`.
+    """
+    return erange(access) if subdims else extrema(access)
+
+
+@memoized_func(scope='build')
+def _cached_distance(access0, ispace0, access1, ispace1, logical):
+    """
+    The geometry-keyed cache behind `TimedAccess.distance`.
+    """
+    a0 = TimedAccess(access0, 'R', 0, ispace0)
+    a1 = TimedAccess(access1, 'R', 0, ispace1)
+    return a0._distance(a1, logical)
+
+
 def disjoint_subdims(a0, a1):
     """
     Determine whether two TimedAccesses touch disjoint SubDimension regions
@@ -1644,7 +1672,9 @@ def disjoint_subdims(a0, a1):
 
         f = a0.function.c0
         space_order = f.space_order if isinstance(f, Function) else 0
-        if disjoint_subdims_axis(e0, e1, d0, d1, it0, it1, space_order):
+        # The result doesn't depend on stamps; drop them from the cache key
+        if disjoint_subdims_axis(e0, e1, d0, d1, it0.reset(), it1.reset(),
+                                 space_order):
             return True
 
     return False
@@ -1748,10 +1778,8 @@ def disjoint_test(e0, e1, d, it):
     if any(not i.is_Number for i in [p00, p01, p10, p11]):
         return False
 
-    i0 = sympy.Interval(min(p00, p01), max(p00, p01))
-    i1 = sympy.Interval(min(p10, p11), max(p10, p11))
-
-    return not bool(i0.intersect(i1))
+    # Closed intervals, hence touching endpoints still overlap
+    return bool(max(p00, p01) < min(p10, p11) or max(p10, p11) < min(p00, p01))
 
 
 def degenerating_indices(i0, i1, function, logical=False):

devito/passes/clusters/fusion.py

@@ -7,6 +7,7 @@
 from devito.ir.support import (
     InitArray, PrefetchUpdate, ReleaseLock, Scope, SyncArray, WaitLock, WithLock
 )
+from devito.ir.support.basic import Dependence
 from devito.symbolics import search
 from devito.tools import (
     DAG, CacheInstances, as_tuple, flatten, frozendict, memoized_func, timed_pass
@@ -51,6 +52,44 @@ def _fusion_hazards(scope0, scope1, prefix):
     return NO_HAZARD
 
 
+def _maybe_break(cgroups, prefix):
+    """
+    True if any two of `cgroups` may induce a `BREAK` hazard, False if that's
+    provably impossible.
+
+    A `BREAK` requires an anti-dependence carried along `prefix`. All distinct
+    read/write access geometries that may occur in two different ClusterGroups
+    are checked, ignoring timestamps and `Scope.rules`, hence conservatively.
+    """
+    if not prefix:
+        return False
+
+    # Map each distinct access geometry to the ClusterGroups it occurs in
+    writes = defaultdict(dict)
+    for n, cg in enumerate(cgroups):
+        for f, v in cg.scope.writes.items():
+            for w in v:
+                writes[f].setdefault((w.access, w.ispace), (w, set()))[1].add(n)
+    reads = defaultdict(dict)
+    for n, cg in enumerate(cgroups):
+        for f in writes:
+            for r in cg.scope.reads_smart_gen(f):
+                reads[f].setdefault((r.access, r.ispace), (r, set()))[1].add(n)
+
+    for f, v in writes.items():
+        for r, rn in reads[f].values():
+            for w, wn in v.values():
+                if len(rn | wn) == 1:
+                    # Only ever within the same ClusterGroup
+                    continue
+
+                d = Dependence(r, w)
+                if not d.is_imaginary and d.cause & prefix:
+                    return True
+
+    return False
+
+
 class Keys(CacheInstances):
 
     """
@@ -297,10 +336,39 @@ def _build_dag(self, cgroups, prefix):
         """
         A DAG representing the data dependences across the ClusterGroups within
         a given scope.
+
+        If no `BREAK` hazard is possible, the DAG only encodes an ordering, so
+        just its transitive reduction is built.
         """
         prefix = frozenset(i.dim for i in as_tuple(prefix))
 
         dag = DAG(nodes=cgroups)
+
+        if not _maybe_break(cgroups, prefix):
+            # Visit the ClusterGroups backwards, tracking reachability as
+            # bitsets, so that already ordered pairs needn't be analyzed. E.g.,
+            # with edges `cg1 -> cg2` and `cg0 -> cg1`, `(cg0, cg2)` is skipped
+            reach = [0]*len(cgroups)
+            barrier = len(cgroups)
+            for n in reversed(range(len(cgroups))):
+                cg0 = cgroups[n]
+                if cg0.scope.has_barrier:
+                    barrier = n
+
+                r = 0
+                for n1 in range(n + 1, len(cgroups)):
+                    if r >> n1 & 1:
+                        continue
+
+                    # Fences have the same effect as any hazard
+                    cg1 = cgroups[n1]
+                    if n1 >= barrier or _fusion_hazards(cg0.scope, cg1.scope, prefix):
+                        dag.add_edge(cg0, cg1)
+                        r |= 1 << n1 | reach[n1]
+                reach[n] = r
+
+            return dag
+
         for n, cg0 in enumerate(cgroups):
             # Track whether there is any fence between `cg0` and the current `cg1`.
             fenced = cg0.scope.has_barrier