compiler: Fix buffering in presence of written HALO regions

devito/ir/clusters/algorithms.py

@@ -12,8 +12,9 @@
 from devito.ir.clusters.visitors import Queue, cluster_pass
 from devito.ir.equations import OpMax, OpMin, OpMinMax, identity_mapper
 from devito.ir.support import (
-    Any, Backward, Forward, IterationSpace, Scope, erange, pull_dims
+    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
@@ -486,6 +487,8 @@ def communications(clusters):
     clusters = HaloComms().process(clusters)
     clusters = reduction_comms(clusters)
 
+    check_halo_writes(clusters)
+
     return clusters
 
 
@@ -628,6 +631,26 @@ def _update(reductions):
     return processed
 
 
+def check_halo_writes(clusters):
+    """
+    Warn about HALO writes along Dimensions not fixed to 1 in the Grid topology.
+    """
+    for c in clusters:
+        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):
     clusters = normalize_nested_indexeds(clusters, sregistry)
     if options['mapify-reduce']:

devito/ir/support/basic.py

@@ -492,6 +492,52 @@ def distance(self, other, logical=False):
 
         return Vector(*ret)
 
+    def touched_nodomain(self, findex):
+        """
+        Return a boolean 2-tuple, one entry for each ``findex`` DataSide. True
+        means that the entire access lies outside the DOMAIN along that
+        DataSide.
+
+        If containment outside the DOMAIN cannot be proven, return False for
+        that DataSide. Unlike ``touched_halo``, this is a containment query and
+        applies irrespective of whether ``findex`` is distributed.
+        """
+        if not self.affine(findex):
+            return (False, False)
+
+        d = self.aindices[findex]
+        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))
+
+        # 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 outside(expr):
+            # A negative maximum distance proves the entire access is outside
+            expr = sympy.expand(expr)
+            for symbol, lower, upper in limits:
+                coefficient = expr.diff(symbol)
+                if coefficient.has(symbol):
+                    return False
+                elif coefficient.is_nonnegative:
+                    expr = expr.subs(symbol, upper)
+                elif coefficient.is_nonpositive:
+                    expr = expr.subs(symbol, lower)
+                else:
+                    return False
+
+            return expr.is_negative 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):
         """
         Return a boolean 2-tuple, one entry for each ``findex`` DataSide. True

devito/ir/support/utils.py

@@ -15,6 +15,7 @@
     'Stencil',
     'bounded',
     'detect_accesses',
+    'detect_halo_writes',
     'erange',
     'extrema',
     'maximum',
@@ -217,6 +218,24 @@ def detect_accesses(exprs):
     return mapper
 
 
+def detect_halo_writes(c, key):
+    """
+    Return the write accesses in `c` proven entirely outside DOMAIN along at
+    least one Dimension selected by `key`. Wild Clusters are ignored.
+    """
+    writes = set()
+    if c.is_wild:
+        return writes
+
+    for w in c.scope.writes_gen():
+        for d in w.findices:
+            if key(d) and any(w.touched_nodomain(d)):
+                writes.add(w)
+                break
+
+    return writes
+
+
 def pull_dims(exprs, flag=True):
     """
     Extract all Dimensions from one or more expressions. If `flag=True`

devito/passes/clusters/buffering.py

@@ -9,7 +9,7 @@
 from devito.exceptions import CompilationError
 from devito.ir import (
     Backward, Cluster, Forward, GuardBound, GuardFactor, InitArray, Interval,
-    IntervalGroup, IterationSpace, Properties, Queue, Vector, 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
@@ -118,6 +118,10 @@ def key(f):
     # First we generate all the necessary buffers
     mapper = generate_buffers(clusters, key, sregistry, options)
 
+    # Take into account writes into the HALO regions so that the buffered
+    # Functions can be populated accordingly
+    clusters = expand_halo_transfers(clusters, mapper)
+
     # Then we inject them into the Clusters. This involves creating the
     # initializing Clusters, and replacing the buffered Functions with the buffers
     clusters = InjectBuffers(mapper, sregistry, options).process(clusters)
@@ -485,6 +489,86 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs):
     return mapper
 
 
+def expand_halo_transfers(clusters, mapper):
+    """
+    Include the halo in buffered writes reading Functions with explicit HALO
+    writes. For example, `usave` in `Eq(usave, u)` must eventually receive `u`'s
+    populated HALO if a preceding `Eq` writes into `u`'s HALO.
+    """
+    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))
+
+    halo_writes = set()
+    for c in clusters:
+        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:
+        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 not writes <= buffered:
+            raise CompilationError(
+                "Cannot expand a mixed Cluster over the halo while buffering"
+            )
+
+        ispace = c.ispace
+        for f in writes:
+            ispace = _include_halo(ispace, f)
+
+        for a in c.scope.accesses:
+            f = a.function
+
+            for d in bdims.intersection(a.findices):
+                size = f._size_nodomain[d]
+                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 access to `{f.name}` over the halo"
+                    )
+
+                i = ispace[d]
+                if i.lower + offset < -size.left or \
+                   i.upper + offset > size.right:
+                    raise CompilationError(
+                        f"Insufficient halo for `{f.name}` in buffered write"
+                    )
+
+        processed.append(c.rebuild(ispace=ispace))
+
+    return processed
+
+
+def _include_halo(ispace, f):
+    """Extend `ispace` to include `f`'s HALO."""
+    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))
+
+
 def map_buffered_functions(clusters, key):
     """
     Map each candidate Function to the Clusters that access it.
@@ -641,12 +725,7 @@ def write_to(self):
         ispace = ispace.promote(lambda d: d.is_AbstractSub, mode='total')
 
         # Analogous to the above, we need to include the halo region as well
-        ihalo = IntervalGroup([
-            Interval(i.dim, -h.left, h.right, i.stamp)
-            for i, h in zip(ispace, self.b._size_halo, strict=False)
-        ])
-
-        ispace = IterationSpace.union(ispace, IterationSpace(ihalo))
+        ispace = _include_halo(ispace, self.b)
 
         return ispace
 

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

@@ -4,12 +4,17 @@
 
 from conftest import skipif
 from devito import (
-    CondEq, ConditionalDimension, Constant, Dimension, Eq, Function, Grid, Operator,
-    SparseTimeFunction, SubDimension, SubDomain, TimeFunction, configuration, switchconfig
+    CondEq, ConditionalDimension, Constant, CustomDimension, Dimension, Eq, Function,
+    Grid, Operator, SparseTimeFunction, SubDimension, SubDomain, TimeFunction,
+    configuration, switchconfig
 )
 from devito.arch.archinfo import AppleArm
 from devito.exceptions import CompilationError
-from devito.ir import FindSymbols, retrieve_iteration_tree
+from devito.ir import (
+    Cluster, FindSymbols, Interval, IterationSpace, lower_exprs, retrieve_iteration_tree
+)
+from devito.passes.clusters.buffering import BufferDimension, expand_halo_transfers
+from devito.types import Array
 
 
 def test_read_write():
@@ -64,6 +69,100 @@ def test_write_only():
     assert np.all(v.data == v1.data)
 
 
+@pytest.mark.parametrize('forward', [False, True])
+def test_write_only_with_halo_source(forward):
+    """
+    A buffered save of a Function with a populated halo must preserve that halo.
+    """
+    nt = 5
+    grid = Grid(shape=(17, 17))
+    y = grid.dimensions[-1]
+
+    u = TimeFunction(name='u', grid=grid, space_order=8)
+    usave = TimeFunction(name='usave', grid=grid, space_order=8, save=nt)
+
+    k = CustomDimension(name='k', parent=y, symbolic_min=1,
+                        symbolic_max=4, symbolic_size=4)
+
+    eqns = [Eq(u.forward, u + 1),
+            Eq(u.forward._subs(y, -k), -u.forward._subs(y, k)),
+            Eq(usave, u.forward if forward else u)]
+
+    op = Operator(eqns, opt='buffering', name='save_halo')
+    op.apply(time_M=nt-2)
+
+    hx, hy = usave._size_halo.left[1:]
+    for t in range(nt-1):
+        assert np.all(usave.data[t] == 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)
+    usave = TimeFunction(name='usave', grid=grid, space_order=8, save=5)
+
+    k = CustomDimension(name='k', parent=y, symbolic_min=1,
+                        symbolic_max=4, symbolic_size=4)
+
+    eqns = [Eq(u.forward, u + 1),
+            Eq(u.forward._subs(y, -k), -u.forward._subs(y, k)),
+            Eq(usave, u.forward + v.forward._subs(y, y + shift))]
+
+    if space_order == 10:
+        # 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)
+        assert np.all(usave.data[3] == 6)
+        hx, hy = usave._size_halo.left[1:]
+        assert np.all(usave.data_with_halo[3, hx:hx + grid.shape[0], hy-4:hy] == -2)
+    else:
+        with pytest.raises(CompilationError, match='Insufficient halo for `v`'):
+            Operator(eqns, opt='buffering')
+
+
+@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),
+                 halo=((0, 0), (4, 4)))
+    usave = Function(name='usave', dimensions=(s, x), shape=(5, 17),
+                     halo=u.halo)
+    db = BufferDimension('db', 0, 0, 1, s)
+    b = Array(name='b', dimensions=(db, x), halo=usave.halo)
+    k = CustomDimension(name='k', parent=x, symbolic_min=1,
+                        symbolic_max=4, symbolic_size=4)
+
+    mirror = Cluster(lower_exprs(Eq(u[s+1, -k], -u[s+1, k])),
+                     IterationSpace([Interval(s), Interval(k)]))
+    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)
+    # The streaming axis is not part of the halo footprint, even with a shifted read
+    assert clusters[1].ispace[s] == save.ispace[s]
+    assert clusters[1].exprs[0].args == save.exprs[0].args
+
+
 def test_read_only():
     nt = 10
     grid = Grid(shape=(2, 2))