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

devito/ir/support/basic.py

@@ -279,6 +279,17 @@ def itintervals(self):
     def is_local(self):
         return self.function.is_Symbol
 
+    @cached_property
+    def expanded(self):
+        """
+        The accesses `self` may perform once its StencilDimensions are resolved;
+        see `Scope.reads_smart_gen`.
+        """
+        if any(d.is_Sub for d in self.ispace.dimensions):
+            return erange(self.access)
+        else:
+            return extrema(self.access)
+
     @cached_property
     def is_regular(self):
         if not super().is_regular:
@@ -339,7 +350,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 +361,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:
@@ -1182,13 +1204,9 @@ def reads_smart_gen(self, f):
         be found. For example, a DiscreteFunction would never appear among
         the iteration symbols.
         """
-        uses_subdims = lambda i: any(d.is_Sub for d in i.ispace.dimensions)
-
         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 i.expanded:
                     yield TimedAccess(j, i.mode, i.timestamp, i.ispace)
 
         else:
@@ -1604,6 +1622,16 @@ 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 _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

tests/test_ir.py

@@ -7,7 +7,9 @@
     Constant, Dimension, Eq, Function, Grid, Inc, Operator, SubDimension, TimeDimension,
     TimeFunction, switchconfig
 )
-from devito.finite_differences.differentiable import IndexSum, LocalSum
+from devito.finite_differences.differentiable import (
+    IndexDerivative, IndexSum, LocalSum, Weights
+)
 from devito.ir.cgen import ccode
 from devito.ir.clusters import Cluster, ClusterGroup
 from devito.ir.equations import LoweredEq
@@ -16,20 +18,21 @@
 from devito.ir.stree import stree_build
 from devito.ir.support.basic import (
     AFFINE, IRREGULAR, REGULAR, IterationInstance, Relation, Scope, TimedAccess, Vector,
-    mocksym0, mocksym1
+    disjoint_subdims_axis, disjoint_test, mocksym0, mocksym1
 )
 from devito.ir.support.guards import GuardOverflow
 from devito.ir.support.space import (
     Any, Backward, Forward, Interval, IntervalGroup, IterationInterval, IterationSpace,
     NullInterval, null_ispace
 )
 from devito.ir.support.utils import detect_halo_writes
+from devito.passes.clusters.fusion import Fusion, _maybe_break
 from devito.symbolics import DefFunction, FieldFromPointer, uxreplace
-from devito.tools import prod
+from devito.tools import Stamp, prod
 from devito.tools.data_structures import frozendict
 from devito.types import (
-    Array, BlockDimension, Bundle, ConditionalDimension, CriticalRegion, CustomDimension,
-    Jump, Scalar, StencilDimension, Symbol
+    Array, BlockDimension, Bundle, ComponentAccess, ConditionalDimension, CriticalRegion,
+    CustomDimension, Jump, Scalar, StencilDimension, Symbol
 )
 
 
@@ -677,6 +680,113 @@ def test_subdimension_stencil_gap_function_order(self, orders):
             expected = S.ImaginaryUnit if order == 8 else S.Infinity
             assert a.distance(b) == b.distance(a) == (expected,)
 
+    def test_distance_cache(self):
+        grid = Grid(shape=(8, 8))
+        x, y = grid.dimensions
+        f = Function(name='f', grid=grid)
+        g = Function(name='g', grid=grid)
+        ispace = IterationSpace([Interval(x), Interval(y)])
+        a = TimedAccess(f[x + 1, y], 'W', 0, ispace)
+        b = TimedAccess(f[x, y], 'R', 1, ispace)
+
+        d = a.distance(b)
+        assert d == (1, 0)
+
+        # Timestamps and access modes are not part of the geometry, hence a cache hit
+        assert a.rebuild(timestamp=5).distance(b.rebuild(mode='W')) is d
+
+        # Anything else is
+        rev = IterationSpace([Interval(x), Interval(y)], directions={x: Backward})
+        shifted = IterationSpace([Interval(x, 0, 1), Interval(y)])
+        c0 = a.rebuild(access=ComponentAccess(f[x, y], 0))
+        c1 = a.rebuild(access=ComponentAccess(f[x, y], 1))
+        pairs = [
+            (b, a, False),
+            (a.rebuild(ispace=rev), b.rebuild(ispace=rev), False),
+            (a.rebuild(ispace=shifted), b, False),
+            (a.rebuild(access=g[x + 1, y]), b.rebuild(access=g[x, y]), False),
+            (a, b, True),
+            (c0, c1, False),
+        ]
+        for i, j, logical in pairs:
+            assert i.distance(j, logical=logical) is not d
+
+        assert b.distance(a) == (-1, 0)
+        assert c0.distance(c1) == (S.ImaginaryUnit,)
+
+        # Cleared at the end of each Operator build
+        Operator(Eq(f, 1))
+        assert a.distance(b) == d
+        assert a.distance(b) is not d
+
+    def test_expand_read(self):
+        grid = Grid(shape=(8,))
+        x, = grid.dimensions
+        xl = SubDimension.left('xl', x, 4)
+        f = Function(name='f', grid=grid)
+        h = StencilDimension('h', 0, 2)
+        w = Weights(name='w', dimensions=h, initvalue=[1., 2., 3.])
+
+        plain = IterationSpace([Interval(x)])
+        sub = IterationSpace([Interval(xl)])
+
+        def expanded(access, ispace):
+            return tuple(TimedAccess(access, 'R', 0, ispace).expanded)
+
+        # Unbounded StencilDimension: every point over SubDimensions, else the extrema
+        assert expanded(f[x + h], sub) == (f[x], f[x + 1], f[x + 2])
+        assert expanded(f[x + h], plain) == (f[x], f[x + 2])
+
+        assert expanded(f[x], sub) == (f[x],)
+        assert expanded(f[x], plain) == (f[x], f[x])
+
+        # A StencilDimension bound by an IndexDerivative isn't expanded
+        bound = IndexDerivative(f[x + h]*w, {x: h})
+        assert expanded(bound, sub) == (bound,)
+
+    def test_disjoint_subdims_cache(self):
+        grid = Grid(shape=(24, 24))
+        x, y = grid.dimensions
+        xl = SubDimension.left('xl', x, 4, separated=True)
+        xm = SubDimension.middle('xm', x, 4, 4, separated=True)
+        xm = xm._rebuild(thickness=(xl.ltkn, xm.rtkn))
+        f = Function(name='f', grid=grid)
+
+        def accesses(stamp):
+            il, im = Interval(xl, 0, 0, stamp), Interval(xm, 0, 0, stamp)
+            left = TimedAccess(f[xl, y], 'W', 0, IterationSpace([il, Interval(y)]))
+            middle = TimedAccess(f[xm, y], 'R', 1, IterationSpace([im, Interval(y)]))
+            return il, left, middle
+
+        il0, left0, middle0 = accesses(Stamp())
+        il1, left1, middle1 = accesses(Stamp())
+        assert il0 != il1 and il0.reset() == il1.reset()
+
+        disjoint_subdims_axis.clear()
+
+        d0 = left0.distance(middle0)
+        assert d0 == (S.ImaginaryUnit,)
+        assert len(disjoint_subdims_axis.cache) == 1
+
+        # Stamps aren't part of the per-axis cache key, unlike the distance one
+        d1 = left1.distance(middle1)
+        assert d1 == d0
+        assert d1 is not d0
+        assert len(disjoint_subdims_axis.cache) == 1
+
+    @pytest.mark.parametrize('offset,expected', [(4, True), (6, False), (8, False)])
+    @pytest.mark.parametrize('swap', [False, True])
+    def test_disjoint_test(self, offset, expected, swap):
+        grid = Grid(shape=(16,))
+        x, = grid.dimensions
+        xl = SubDimension.left('xl', x, 4)
+
+        # `12 - xl` spans [9, 12]; `xl + offset` spans [offset, offset + 3]
+        e0, e1 = 12 - xl, xl + offset
+        if swap:
+            e0, e1 = e1, e0
+        assert disjoint_test(e0, e1, xl, Interval(xl)) is expected
+
 
 class TestSpace:
 
@@ -1714,6 +1824,50 @@ def test_local_sums(self):
         assert not cluster.rebuild(exprs=[Eq(f[x], 0)]).local_sums
 
 
+class TestFusionDAG:
+
+    @pytest.fixture
+    def setup(self):
+        d = Dimension(name='d')
+        x = Dimension(name='x')
+        ispace = IterationSpace([Interval(d), Interval(x)])
+        f, g, h = [Function(name=n, dimensions=(d, x), shape=(4, 4)) for n in 'fgh']
+
+        def cgroups(*exprs):
+            return [ClusterGroup(Cluster(e, ispace), ispace) for e in exprs]
+
+        return d, x, f, g, h, cgroups
+
+    def test_break_stencil_reads(self):
+        x = Dimension(name='x')
+        s = StencilDimension('s', 0, 1)
+        ispace = IterationSpace([Interval(x)])
+        f, g = [Function(name=n, dimensions=(x,), shape=(4,)) for n in 'fg']
+
+        # The read `f[x + s]` expands to `f[x]` and `f[x + 1]`, the latter
+        # being ahead of the write along `x`
+        cgs = [ClusterGroup(Cluster(e, ispace), ispace)
+               for e in (Eq(g[x], f[x + s]), Eq(f[x + s], 1))]
+        assert _maybe_break(cgs, frozenset({x}))
+
+    def test_transitive_reduction(self, setup):
+        d, x, f, g, h, cgroups = setup
+
+        # `h` depends on both `f` and `g`, but `(cg0, cg2)` is implied by
+        # `cg0 -> cg1 -> cg2`
+        cgs = cgroups(Eq(f[d, x], 1), Eq(g[d, x], f[d, x]),
+                      Eq(h[d, x], f[d, x] + g[d, x]))
+        dag = Fusion(toposort=True)._build_dag(cgs, [Interval(d)])
+        assert set(dag.edges) == {(cgs[0], cgs[1]), (cgs[1], cgs[2])}
+        assert dag.all_downstreams(cgs[0]) == [cgs[1], cgs[2]]
+
+        # A fence orders everything around it, even absent data dependences
+        cgs = cgroups(Eq(f[d, x], 1), Eq(Symbol(name='s'), CriticalRegion(True)),
+                      Eq(g[d, x], 1))
+        dag = Fusion(toposort=True)._build_dag(cgs, [Interval(d)])
+        assert set(dag.edges) == {(cgs[0], cgs[1]), (cgs[1], cgs[2])}
+
+
 class TestGuards:
 
     def test_guard_overflow(self):