compiler: Patch halo exchange placement at stree construction time

devito/ir/clusters/algorithms.py

@@ -130,8 +130,61 @@ class Schedule(Queue):
 
     @timed_pass(name='schedule')
     def process(self, clusters):
+        clusters = self._schedule_adjacent(clusters)
+        clusters = self._constrain_nonadjacent(clusters)
+
+        return clusters
+
+    def _schedule_adjacent(self, clusters):
+        """
+        Schedule Clusters that are adjacent in program order.
+        """
         return self._process_fatd(clusters, 1)
 
+    def _constrain_nonadjacent(self, clusters):
+        """
+        Clusters that are non-adjacent in program order may require additional
+        constraints on their IterationSpaces.
+
+        For example, consider the following statements:
+
+            - `u[x, y] = f(w[x, y])`
+            - `f[x] = 3`
+            - `v[x, y] = g(u[x, y ± 1])`
+
+        The first and third statements are non-adjacent in program order, have
+        the same IterationSpace, and have a flow- or anti-dependence along `y`
+        due to `y ± 1`. Hence, those two statements will never be fusible, so
+        their IterationSpace must be constrained to prevent such a possibility
+        by any of the later passes (e.g., topofusion).
+        """
+        # A family collects the consecutive groups of Clusters with the same
+        # IterationSpace
+        key = lambda c: c.ispace
+        families = DefaultOrderedDict(list)
+        for ispace, group in groupby(clusters, key=key):
+            families[ispace].append(list(group))
+
+        mapper = {}
+        for family in families.values():
+            if len(family) < 2:
+                continue
+
+            candidates = flatten(family)
+            scheduled = self._schedule_adjacent(candidates)
+            for c, c1 in zip(candidates, scheduled, strict=True):
+                if c1.ispace.intervals == c.ispace.intervals:
+                    continue
+
+                # Only retain the new Interval constraints. Directions were
+                # inferred without the intervening Clusters.
+                ispace1 = IterationSpace(c1.ispace.intervals,
+                                         c.ispace.sub_iterators,
+                                         c.ispace.directions)
+                mapper[c] = c.rebuild(ispace=ispace1)
+
+        return [mapper.get(c, c) for c in clusters]
+
     def callback(self, clusters, prefix, backlog=None, known_break=None):
         if not prefix:
             return clusters
@@ -458,7 +511,7 @@ def callback(self, clusters, prefix, seen=None):
 
         # Construct a representation of the halo accesses
         processed = list(clusters)
-        for n, c in enumerate(clusters):
+        for c in clusters:
             if c.properties.is_sequential(d) or \
                c in seen:
                 continue
@@ -498,6 +551,7 @@ def callback(self, clusters, prefix, seen=None):
             # Insert `halo_touch` at the top of the IterationSpace within which
             # `c` is scheduled
             index = 0
+            n = processed.index(c)
             for i in reversed(range(n)):
                 if not processed[i].ispace.is_subset(c.ispace):
                     index = i + 1

devito/ir/stree/algorithms.py

@@ -192,24 +192,36 @@ def preprocess(clusters, options=None, **kwargs):
 
         else:
             dims = set(c.ispace.promote(lambda d: d.is_Block).itdims)
+            roots = {d.root for d in dims}
 
             found = []
             for c1 in list(queue):
-                distributed_aindices = c1.halo_scheme.distributed_aindices
-                h_indices = set().union(*[d._defines for d in c1.halo_scheme.loc_indices])
-
                 # Skip if the halo exchange would end up outside
                 # its iteration space
+                loc_indices = c1.halo_scheme.loc_indices
+                h_indices = set().union(*[d._defines for d in loc_indices])
                 if h_indices and not h_indices & dims:
                     continue
 
-                diff = dims - distributed_aindices
-                intersection = dims & distributed_aindices
+                # There must be at least one distributed Dimension
+                dist_aindices = c1.halo_scheme.distributed_aindices
+                if not (dims & dist_aindices):
+                    continue
+
+                # Ensure the guards are compatible
+                diff = dims - dist_aindices
+                if not all(c1.guards.get(d) == c.guards.get(d) for d in diff):
+                    continue
+
+                # Ensure we're inserting within a compatible IterationSpace
+                # E.g., if `dist_aindices` contains a SubDimension `yi`, we
+                # cannot proceed if `c` is over a different SubDimension `yi'`
+                if any(d.root in roots and d not in dims for d in dist_aindices):
+                    continue
 
-                if all(c1.guards.get(d) == c.guards.get(d) for d in diff) and \
-                   len(intersection) > 0:
-                    found.append(c1)
-                    queue.remove(c1)
+                # All good!
+                found.append(c1)
+                queue.remove(c1)
 
             syncs = normalize_syncs(*[c1.syncs for c1 in found])
             if syncs:

devito/ir/support/basic.py

@@ -379,6 +379,10 @@ def distance(self, other, logical=False):
                 # E.g., `self=R<f,[cy]>` and `self.itintervals=(y,)` => `sai=None`
                 pass
 
+            if skippable_interval(sai, self.ispace, sit) and \
+               skippable_interval(oai, other.ispace, oit):
+                continue
+
             # In some cases, the distance degenerates because `self` and
             # `other` never intersect, which essentially means there's no
             # dependence between them. In this case, we set the distance to a
@@ -406,8 +410,10 @@ def distance(self, other, logical=False):
 
                 # Case 3: `self` and `other` have some special form such that
                 # it's provable that they never intersect
-                if sai and sit == oit and disjoint_test(self[n], other[n], sai, sit):
-                    return Vector(S.ImaginaryUnit)
+                if sit == oit:
+                    dims = {sai, oai} - {None}
+                    if any(disjoint_test(self[n], other[n], d, sit) for d in dims):
+                        return Vector(S.ImaginaryUnit)
 
             # Compute the distance along the current IterationInterval
             if self.function._mem_shared:
@@ -426,10 +432,6 @@ def distance(self, other, logical=False):
                     ret.append(other[n] - self[n])
                 else:
                     ret.append(self[n] - other[n])
-            elif sai in self.ispace and oai in other.ispace:
-                # E.g., `self=R<f,[x, y]>`, `sai=time`,
-                #       `self.itintervals=(time, x, y)`, `n=0`
-                continue
             elif not sai and not oai:
                 if self[n] - other[n] == 0:
                     # E.g., `self=R<a,[4]>` and `other=W<a,[4]>`
@@ -1505,6 +1507,26 @@ def vinf(entries):
     return Vector(*(entries + [S.Infinity]))
 
 
+def skippable_interval(d, ispace, it):
+    """
+    Return True if the IterationInterval `it` can be skipped while matching
+    the access Dimension `d` to the IterationSpace `ispace`.
+
+    Constant accesses, represented by `d=None`, do not consume an
+    IterationInterval. A non-constant access can skip `it` only if `d`
+    occurs elsewhere in `ispace`. Otherwise the access is irregular and
+    must be handled conservatively.
+
+    Examples
+    --------
+    Given `W<a,[i]>` and `R<a,[4]>` over `(r, i)`, `r` is skippable
+    for both accesses, while `i` is not skippable for the write. By contrast,
+    `x` is not skippable for `a[y]` over `(x)` because `y` does not
+    occur in the IterationSpace.
+    """
+    return d is None or (d in ispace and not d._defines & it.dim._defines)
+
+
 def disjoint_test(e0, e1, d, it):
     """
     A rudimentary test to check if two accesses `e0` and `e1` along `d` within

tests/test_dle.py

@@ -477,16 +477,10 @@ def test_cache_blocking_imperfect_nest(blockinner):
     op1 = Operator(eqns, opt=('advanced', {'blockinner': blockinner}))
 
     # First, check the generated code
-    bns, _ = assert_blocking(op1, {'x0_blk0'})
+    bns, _ = assert_blocking(op1, {'x0_blk0', 'x1_blk0'})
     trees = retrieve_iteration_tree(bns['x0_blk0'])
-    assert len(trees) == 2
-    assert len(trees[0]) == len(trees[1])
-    assert all(i is j for i, j in zip(trees[0][:4], trees[1][:4], strict=True))
-    assert trees[0][4] is not trees[1][4]
+    assert len(trees) == 1
     assert trees[0].root.dim.is_Block
-    assert trees[1].root.dim.is_Block
-    assert op1.parameters[7] is trees[0][0].step
-    assert op1.parameters[10] is trees[0][1].step
 
     u.data[:] = 0.2
     v.data[:] = 1.5
@@ -633,7 +627,7 @@ def test_nthreads_generation(self):
          (False, False)),
         # two nests, each nest: outermost parallel, innermost sequential
         (['Eq(fc[x,y], fc[x,y+1] + fd[x-1,y])', 'Eq(fd[x-1,y+1], fd[x-1,y] + fc[x,y+1])'],
-         (True, False, False)),
+         (True, False, True, False)),
         # outermost sequential, innermost parallel w/ mixed dimensions
         (['Eq(fc[x+1,y], fc[x,y+1] + fc[x,y])', 'Eq(fc[x+1,y], 2. + fc[x,y+1])'],
          (False, True)),

tests/test_ir.py

@@ -25,7 +25,9 @@
 from devito.symbolics import DefFunction, FieldFromPointer
 from devito.tools import prod
 from devito.tools.data_structures import frozendict
-from devito.types import Array, Bundle, CriticalRegion, Jump, Scalar, Symbol
+from devito.types import (
+    Array, Bundle, CriticalRegion, CustomDimension, Jump, Scalar, Symbol
+)
 
 
 class TestVectorHierarchy:
@@ -912,6 +914,25 @@ def test_nodep(self, eqns):
         scope = Scope(eqns)
         assert len(scope.d_all) == 0
 
+    def test_nodep_with_nonaccessed_implicit_dim(self):
+        """
+        An outer implicit Dimension must not hide that an indexed write and a
+        constant read touch disjoint points of a CustomDimension.
+        """
+        r = Dimension(name='r')
+        i = CustomDimension(name='i', symbolic_min=0, symbolic_max=3)
+
+        a = Array(name='a', dimensions=(i,))
+        f = Function(name='f', dimensions=(r, i), shape=(3, 4))
+
+        eqns = [Eq(a[i], r, implicit_dims=(r, i)),
+                Eq(f, a[4], implicit_dims=(r, i))]
+        eqns = [LoweredEq(e) for e in eqns]
+
+        scope = Scope(eqns)
+        deps = [d for d in scope.d_all if d.function is a]
+        assert not deps
+
     @pytest.mark.parametrize('eqns', [
         ['Eq(a0[4], 1)', 'Eq(s, a0[4])'],
         ['Eq(a1[x+1, 4], 1)', 'Eq(s, a1[x, 4])'],

tests/test_mpi.py

@@ -3289,7 +3289,7 @@ def test_overriding_from_different_grid(self, mode):
         assert np.all(u3.data[0, 3:-3, 3:-3] == 1.)
 
     @pytest.mark.parallel(mode=4)
-    def test_fission_due_to_antidep(self, mode):
+    def test_fission_due_to_antidep_along_innermost_dim(self, mode):
         grid = Grid(shape=(16, 16, 64), dtype=np.float64)
 
         u = TimeFunction(name='u', grid=grid, space_order=4)
@@ -3305,8 +3305,8 @@ def test_fission_due_to_antidep(self, mode):
         # First, check the generated code
         assert_structure(op1, ['t',
                                't,x0_blk0,y0_blk0,x,y,z',
-                               't,x0_blk0,y0_blk0,x,y,z'],
-                         'tx0_blk0y0_blk0xyzz')
+                               't,x1_blk0,y1_blk0,x,y,z'],
+                         'tx0_blk0y0_blk0xyzx1_blk0y1_blk0xyz')
 
         def init(f, v=1):
             f.data[:] = np.indices(grid.shape).sum(axis=0) % (.004*v) + .01

tests/test_operator.py

@@ -1630,7 +1630,7 @@ def test_no_fission_as_illegal(self, exprs):
         (('Eq(tu[t,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x,y,z+2])',
           'Eq(tw[t,x,y,z], tv[t,x,y,z-1] + 1.)'),
-         '++++++++', ['txyz', 'txyz', 'txyz'], 'txyzxyzz'),
+         '++++++++++', ['txyz', 'txyz', 'txyz'], 'txyzxyzxyz'),
         # 8) WAR 1->2; WAW 1->3
         (('Eq(tu[t,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x+2,y,z])',
@@ -1640,7 +1640,7 @@ def test_no_fission_as_illegal(self, exprs):
         (('Eq(tu[t,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x,y,z-2])',
           'Eq(tw[t,x,y,z], tv[t,x,y+1,z] + 1.)'),
-         '+++++++++', ['txyz', 'txyz', 'txyz'], 'txyzxyzyz'),
+         '++++++++++', ['txyz', 'txyz', 'txyz'], 'txyzxyzxyz'),
         # 10) WAR 1->2; WAW 1->3
         (('Eq(tu[t-1,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x,y,z+2])',
@@ -2172,6 +2172,53 @@ def test_topofuse_w_numeric_dim(self):
 
         assert_structure(op, ['r,i', 'r'], 'r,i')
 
+    def test_topofuse_preserves_nonadjacent_fission(self):
+        """
+        A shallower IterationSpace may separate a stencil producer and consumer
+        during adjacent scheduling. Topofusion must preserve their fission.
+        """
+        x, y = dimensions('x y')
+
+        u = Function(name='u_nonadj', dimensions=(x, y), shape=(16, 16),
+                     space_order=1)
+        v = Function(name='v_nonadj', dimensions=(x, y), shape=(16, 16),
+                     space_order=1)
+        w = Function(name='w_nonadj', dimensions=(x,), shape=(16,))
+
+        eqns = [Eq(u, 1),
+                Eq(w, 2),
+                Eq(v, u[x, y - 1] + u[x, y + 1])]
+
+        op = Operator(eqns, opt=('topofuse', {'openmp': False}))
+
+        assert_structure(op, ['x,y', 'x', 'x,y'])
+
+    def test_nonadjacent_constraints_preserve_directions(self):
+        """
+        Constraints inferred from non-adjacent Clusters must not replace the
+        directions inferred while the intervening Clusters were still visible.
+        """
+        x, y = dimensions('x y')
+
+        u = Function(name='u_direction', dimensions=(x, y), shape=(16, 16),
+                     space_order=1)
+        p = Function(name='p_direction', dimensions=(x, y), shape=(16, 16),
+                     space_order=1)
+        q = Function(name='q_direction', dimensions=(x,), shape=(16,),
+                     space_order=1)
+
+        eqns = [Eq(u[x, y], 1),
+                Eq(q[x], p[x + 1, 0]),
+                Eq(u[x, y + 1], 2),
+                Eq(p[x, y], 3)]
+
+        op = Operator(eqns, opt=('topofuse', {'openmp': False}))
+
+        assert_structure(op, ['x,y', 'x', 'x,y'], 'x,y,y')
+
+        iterations = FindNodes(Iteration).visit(op)
+        assert iterations[0].direction is Backward
+
     @pytest.mark.parametrize('eqns, expected, exp_trees, exp_iters', [
         (['Eq(u[0, x], 1)',
             'Eq(u[1, x], u[0, x + h_x] + u[0, x - h_x] - 2*u[0, x])'],