compiler: Simplify HALO-write handling

Author: FabioLuporini

devito/ir/clusters/algorithms.py

@@ -637,11 +637,11 @@ def check_halo_writes(clusters):
     for c in clusters:
         dims = set()
         for f in c.scope.writes:
-            if not f.is_DiscreteFunction or f.grid is None:
+            if not f.is_DiscreteFunction:
                 continue
-            dist = f.grid.distributor
-            dims.update(d.root for d, n in zip(dist.dimensions, dist.topology,
-                                               strict=True) if n > 1)
+            for d, dec in zip(f.dimensions, f._decomposition, strict=True):
+                if dec is not None and len(dec) > 1:
+                    dims.add(d.root)
 
         key = lambda d: d.root in dims  # noqa: B023
         if dims and detect_halo_writes(c, key):

devito/ir/support/basic.py

@@ -505,67 +505,38 @@ def touched_nodomain(self, findex):
         if not self.affine(findex):
             return (False, False)
 
-        index = self[findex]
         d = self.aindices[findex]
-
-        if d is None:
-            index_min = index_max = index
-        else:
-            try:
-                m, M = self.intervals[d].offsets
-            except KeyError:
-                return (False, False)
-
-            coefficient = index.diff(d)
-            if coefficient.is_positive:
-                index_min = index.subs(d, d.symbolic_min + m)
-                index_max = index.subs(d, d.symbolic_max + M)
-            elif coefficient.is_negative:
-                index_min = index.subs(d, d.symbolic_max + M)
-                index_max = index.subs(d, d.symbolic_min + m)
-            elif coefficient.is_zero:
-                index_min = index_max = index
-            else:
+        limits = []
+        if d is not None:
+            i = self.intervals[d]
+            if i.is_Null:
                 return (False, False)
+            limits.append((d, d.symbolic_min + i.lower, d.symbolic_max + i.upper))
 
-        size_nodomain_left = self.function._size_nodomain[findex].left
-        domain_min = findex.symbolic_min + size_nodomain_left
-        domain_max = findex.symbolic_max + size_nodomain_left
+        # Runtime DOMAIN bounds may select any part of the allocated extent
+        for v in (findex.symbolic_min, findex.symbolic_max):
+            if v.is_Symbol:
+                limits.append((v, S.Zero, findex.symbolic_size - 1))
 
-        def bound(expr, maximize):
+        def outside(expr):
+            # A negative maximum distance proves the entire access is outside
             expr = sympy.expand(expr)
-
-            size = findex.symbolic_size
-            limits = [
-                (findex.symbolic_min, S.Zero, size - 1),
-                (findex.symbolic_max, S.Zero, size - 1)
-            ]
-
             for symbol, lower, upper in limits:
-                if not symbol.is_Symbol:
-                    continue
-
                 coefficient = expr.diff(symbol)
                 if coefficient.has(symbol):
-                    return None
-                elif coefficient.is_positive:
-                    value = upper if maximize else lower
-                elif coefficient.is_negative:
-                    value = lower if maximize else upper
-                elif coefficient.is_zero:
-                    continue
+                    return False
+                elif coefficient.is_nonnegative:
+                    expr = expr.subs(symbol, upper)
+                elif coefficient.is_nonpositive:
+                    expr = expr.subs(symbol, lower)
                 else:
-                    return None
-
-                expr = expr.subs(symbol, value)
-
-            return expr
+                    return False
 
-        left = bound(index_max - domain_min, True)
-        right = bound(index_min - domain_max, False)
+            return expr.is_negative is True
 
-        return (left is not None and left.is_negative is True,
-                right is not None and right.is_positive is True)
+        index = self[findex] - self.function._size_nodomain[findex].left
+        return (outside(index - findex.symbolic_min),
+                outside(findex.symbolic_max - index))
 
     def touched_halo(self, findex):
         """

devito/passes/clusters/buffering.py

@@ -9,8 +9,7 @@
 from devito.exceptions import CompilationError
 from devito.ir import (
     Backward, Cluster, Forward, GuardBound, GuardFactor, InitArray, Interval,
-    IntervalGroup, IterationSpace, Properties, Queue, Vector, detect_halo_writes,
-    lower_exprs, vmax, vmin
+    IterationSpace, Properties, Queue, Vector, detect_halo_writes, lower_exprs, vmax, vmin
 )
 from devito.logger import warning
 from devito.passes.clusters.utils import is_memcpy
@@ -496,61 +495,61 @@ def expand_halo_transfers(clusters, mapper):
     writes. For example, `usave` in `Eq(usave, u)` must eventually receive `u`'s
     populated HALO if a preceding `Eq` writes into `u`'s HALO.
     """
-    buffered = {f for f, _ in mapper}
-    if not buffered:
+    if not mapper:
         return clusters
 
+    # Get HALO writes along the buffered dimensions
     bdims = set()
     for b in mapper.values():
         bdims.update(d for d in b.dimensions if not isinstance(d, BufferDimension))
-    key = lambda d: d in bdims
 
     halo_writes = set()
     for c in clusters:
-        for w in detect_halo_writes(c, key):
+        for w in detect_halo_writes(c, bdims.__contains__):
             halo_writes.add(w.function)
+    if not halo_writes:
+        return clusters
 
+    # Expand the IterationSpace over the necessary amount of HALO; in doing so,
+    # check the expanded footprint of every access, including shifted reads.
+    # Writes must be pointwise so that the whole destination halo is filled
+    buffered = {f for f, _ in mapper}
     processed = []
     for c in clusters:
-        scope = c.scope
-        targets = set(scope.writes) & buffered
-        if c.is_wild or not targets or not halo_writes.intersection(scope.reads):
+        writes = c.scope.writes_tensor
+
+        if c.is_wild or \
+           writes.isdisjoint(buffered) or \
+           halo_writes.isdisjoint(c.scope.reads):
             processed.append(c)
             continue
 
-        if scope.writes_tensor != targets:
+        if not writes <= buffered:
             raise CompilationError(
                 "Cannot expand a mixed Cluster over the halo while buffering"
             )
 
         ispace = c.ispace
-        for f in targets:
+        for f in writes:
             ispace = _include_halo(ispace, f)
 
-        # Check the expanded footprint of every access, including shifted reads.
-        # Writes must be pointwise so that the whole destination halo is filled
-        for a in scope.accesses:
+        for a in c.scope.accesses:
             f = a.function
-            if not f.is_AbstractFunction:
-                continue
-
-            for d in f.dimensions:
-                if not key(d):
-                    continue
-                if d not in ispace.dimensions:
-                    raise CompilationError(
-                        f"Cannot expand access to `{f.name}` over the halo"
-                    )
 
+            for d in bdims.intersection(a.findices):
                 size = f._size_nodomain[d]
-                offset = simplify(a[d] - d)
-                if not is_integer(offset) or (a.is_write and offset != size.left):
+                offset = simplify(a[d] - d - size.left)
+
+                if d not in ispace.dimensions or \
+                   not is_integer(offset) or \
+                   (a.is_write and offset != 0):
                     raise CompilationError(
-                        f"Cannot expand non-pointwise access to `{f.name}` over the halo"
+                        f"Cannot expand access to `{f.name}` over the halo"
                     )
 
                 i = ispace[d]
-                if i.lower + offset < 0 or i.upper + offset > sum(size):
+                if i.lower + offset < -size.left or \
+                   i.upper + offset > size.right:
                     raise CompilationError(
                         f"Insufficient halo for `{f.name}` in buffered write"
                     )
@@ -562,10 +561,10 @@ def expand_halo_transfers(clusters, mapper):
 
 def _include_halo(ispace, f):
     """Extend `ispace` to include `f`'s HALO."""
-    ihalo = IntervalGroup([
+    ihalo = [
         Interval(i.dim, -f._size_halo[i.dim].left, f._size_halo[i.dim].right, i.stamp)
         for i in ispace if i.dim in f.dimensions
-    ])
+    ]
 
     return IterationSpace.union(ispace, IterationSpace(ihalo))
 

tests/test_buffering.py

@@ -94,18 +94,18 @@ def test_write_only_with_halo_source(forward):
     hx, hy = usave._size_halo.left[1:]
     for t in range(nt-1):
         assert np.all(usave.data[t] == t + forward)
-        for i in range(1, 5):
-            actual = usave.data_with_halo[t, hx:hx + grid.shape[0], hy - i]
-            assert np.all(actual == -(t + forward))
+        actual = usave.data_with_halo[t, hx:hx + grid.shape[0], hy-4:hy]
+        assert np.all(actual == -(t + forward))
 
 
 @pytest.mark.parametrize('space_order, shift', [(0, 0), (8, -1), (8, 1), (10, 1)])
+@switchconfig(autopadding=False)
 def test_write_only_with_halo_source_bounds(space_order, shift):
     grid = Grid(shape=(17, 17))
     y = grid.dimensions[-1]
 
     u = TimeFunction(name='u', grid=grid, space_order=8)
-    v = TimeFunction(name='v', grid=grid, space_order=space_order, padding=0)
+    v = TimeFunction(name='v', grid=grid, space_order=space_order)
     usave = TimeFunction(name='usave', grid=grid, space_order=8, save=5)
 
     k = CustomDimension(name='k', parent=y, symbolic_min=1,
@@ -116,7 +116,7 @@ def test_write_only_with_halo_source_bounds(space_order, shift):
             Eq(usave, u.forward + v.forward._subs(y, y + shift))]
 
     if space_order == 10:
-        # An extra halo point accommodates the shifted read
+        # A wider halo accommodates the shifted read
         v.data_with_halo[:] = 2
         op = Operator(eqns, opt='buffering', name='save_shifted_halo')
         op.apply(time_M=3)
@@ -128,7 +128,8 @@ def test_write_only_with_halo_source_bounds(space_order, shift):
             Operator(eqns, opt='buffering')
 
 
-def test_halo_transfers_non_time_dimension():
+@pytest.mark.parametrize('mixed', [False, True])
+def test_halo_transfers_non_time_dimension(mixed):
     s = Dimension(name='s')
     x = Dimension(name='x')
     u = Function(name='u', dimensions=(s, x), shape=(5, 17),
@@ -142,9 +143,18 @@ def test_halo_transfers_non_time_dimension():
 
     mirror = Cluster(lower_exprs(Eq(u[s+1, -k], -u[s+1, k])),
                      IterationSpace([Interval(s), Interval(k)]))
-    save = Cluster(lower_exprs(Eq(usave[s, x], u[s+1, x])),
-                   IterationSpace([Interval(s), Interval(x)]))
-    clusters = expand_halo_transfers([mirror, save], {(usave, save.guards): b})
+    eqns = [Eq(usave[s, x], u[s+1, x])]
+    if mixed:
+        eqns.append(Eq(u[s, x], 0))
+    save = Cluster(lower_exprs(eqns), IterationSpace([Interval(s), Interval(x)]))
+    mapper = {(usave, save.guards): b}
+
+    if mixed:
+        with pytest.raises(CompilationError, match='mixed Cluster'):
+            expand_halo_transfers([mirror, save], mapper)
+        return
+
+    clusters = expand_halo_transfers([mirror, save], mapper)
 
     assert clusters[0] is mirror
     assert clusters[1].ispace[x].offsets == (-4, 4)

tests/test_ir.py

@@ -154,64 +154,42 @@ def test_timedaccess_cached(self, fc, x, y):
 
         assert ta0 is ta1
 
-    def test_timedaccess_touched_nodomain(self):
+    @pytest.mark.parametrize('autopadding', [False, True])
+    def test_timedaccess_touched_nodomain(self, autopadding):
         grid = Grid(shape=(17, 17))
         x, y = grid.dimensions
 
-        f = Function(name='f', grid=grid, space_order=8)
+        with switchconfig(autopadding=autopadding):
+            f = Function(name='f', grid=grid, space_order=8)
         hx, hy = f._size_nodomain.left
 
         k = CustomDimension('k', parent=y, symbolic_min=1,
                             symbolic_max=4, symbolic_size=4)
         k0 = CustomDimension('k0', parent=y, symbolic_min=0,
                              symbolic_max=4, symbolic_size=5)
         yl = SubDimension.left('yl', y, thickness=4)
-
-        left = TimedAccess(
-            f.indexed[x + hx, hy - k], 'W', 0,
-            IterationSpace([Interval(x), Interval(k)])
-        )
-        right = TimedAccess(
-            f.indexed[x + hx, hy + y.symbolic_size - 1 + k], 'W', 0,
-            IterationSpace([Interval(x), Interval(k)])
-        )
-
-        depth = yl - y.symbolic_min + 1
-        left_sub = TimedAccess(
-            f.indexed[x + hx, hy + y.symbolic_min - depth], 'W', 0,
-            IterationSpace([Interval(x), Interval(yl)])
-        )
-        left_constant = TimedAccess(
-            f.indexed[x + hx, hy - 1], 'W', 0,
-            IterationSpace([Interval(x)])
-        )
-
-        domain = TimedAccess(
-            f.indexed[x + hx, y + hy], 'W', 0,
-            IterationSpace([Interval(x), Interval(y)])
-        )
-        straddling = TimedAccess(
-            f.indexed[x + hx, hy - k0], 'W', 0,
-            IterationSpace([Interval(x), Interval(k0)])
-        )
-        nonlinear = TimedAccess(
-            f.indexed[x + hx, hy - k**2], 'W', 0,
-            IterationSpace([Interval(x), Interval(k)])
-        )
-        shifted = TimedAccess(
-            f.indexed[x + hx, hy - k], 'W', 0,
-            IterationSpace([Interval(x), Interval(k, -1, 0)])
-        )
-
-        assert left.touched_nodomain(y) == (True, False)
-        assert right.touched_nodomain(y) == (False, True)
-        assert left_sub.touched_nodomain(y) == (True, False)
-        assert left_constant.touched_nodomain(y) == (True, False)
-
-        assert domain.touched_nodomain(y) == (False, False)
-        assert straddling.touched_nodomain(y) == (False, False)
-        assert nonlinear.touched_nodomain(y) == (False, False)
-        assert shifted.touched_nodomain(y) == (False, False)
+        yr = SubDimension.right('yr', y, thickness=4)
+        a = Scalar(name='a', is_const=True)
+
+        for index, interval, expected in [
+            (-k, Interval(k), (True, False)),
+            (y.symbolic_size - 1 + k, Interval(k), (False, True)),
+            (2*y.symbolic_min - yl - 1, Interval(yl), (True, False)),
+            (2*y.symbolic_max - yr + 1, Interval(yr), (False, True)),
+            (S.NegativeOne, None, (True, False)),
+            (y.symbolic_size, None, (False, True)),
+            (y, Interval(y), (False, False)),
+            (-k0, Interval(k0), (False, False)),
+            (-k**2, Interval(k), (False, False)),
+            (-k, Interval(k, -1, 0), (False, False)),
+            (-k, None, (False, False)),
+            (a*k, Interval(k), (False, False)),
+            (y.symbolic_min**2 - k, Interval(k), (False, False)),
+        ]:
+            intervals = [Interval(x)] + ([interval] if interval is not None else [])
+            access = TimedAccess(f.indexed[x + hx, hy + index], 'W', 0,
+                                 IterationSpace(intervals))
+            assert access.touched_nodomain(y) == expected
 
     def test_iteration_instance_arithmetic(self, x, y, ii_num, ii_literal):
         """