compiler: Simplify HALO-write handling

Author: FabioLuporini

devito/ir/clusters/algorithms.py

@@ -14,6 +14,7 @@
 from devito.ir.support import (
     Any, Backward, Forward, IterationSpace, Scope, detect_halo_writes, erange, pull_dims
 )
+from devito.logger import warning
 from devito.mpi.halo_scheme import HaloScheme, HaloTouch
 from devito.mpi.reduction_scheme import DistReduce
 from devito.symbolics import limits_mapper, retrieve_indexed, uxreplace, xreplace_indices
@@ -632,21 +633,22 @@ def _update(reductions):
 
 def check_halo_writes(clusters):
     """
-    Reject explicit HALO writes along Dimensions split across MPI ranks.
+    Warn about HALO writes along Dimensions not fixed to 1 in the Grid topology.
     """
     for c in clusters:
-        dims = set()
-        for f in c.scope.writes:
-            if not f.is_DiscreteFunction or f.grid is None:
-                continue
-            dist = f.grid.distributor
-            dims.update(d.root for d, n in zip(dist.dimensions, dist.topology,
-                                               strict=True) if n > 1)
-
-        key = lambda d: d.root in dims  # noqa: B023
-        if dims and detect_halo_writes(c, key):
-            raise CompilationError("Cannot write to the HALO along distributed "
-                                   "Dimensions")
+        try:
+            grid = c.grid
+        except ValueError:
+            grid = None
+
+        topology = {}
+        if grid is not None and grid.topology is not None:
+            topology = dict(zip(grid.dimensions, grid.topology, strict=True))
+
+        key = lambda d: d in c.dist_dimensions and topology.get(d) != 1  # noqa: B023
+        if detect_halo_writes(c, key):
+            warning("Writing to the HALO along potentially distributed Dimensions; "
+                    "set their Grid topology entries to 1")
 
 
 def normalize(clusters, sregistry=None, options=None, platform=None, **kwargs):

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))
 

devito/types/array.py

@@ -520,7 +520,7 @@ def initvalue(self):
               '_mem_rvalue', '__padding_dtype__', '_size_domain', '_size_halo',
               '_size_owned', '_size_padding', '_size_nopad', '_size_nodomain',
               '_offset_domain', '_offset_halo', '_offset_owned',
-              '_dist_dimensions', '_C_get_field', 'grid',
+              '_dist_dimensions', '_decomposition', '_C_get_field', 'grid',
               *AbstractFunction.__properties__):
         locals()[i] = property(lambda self, v=i: getattr(self.c0, v))
 

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)