diff --git a/devito/ir/clusters/algorithms.py b/devito/ir/clusters/algorithms.py index 374536fa62..f594d35203 100644 --- a/devito/ir/clusters/algorithms.py +++ b/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']: diff --git a/devito/ir/support/basic.py b/devito/ir/support/basic.py index cdf2ab51e8..518e1d0f44 100644 --- a/devito/ir/support/basic.py +++ b/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 diff --git a/devito/ir/support/utils.py b/devito/ir/support/utils.py index 67844287a9..1882d6267f 100644 --- a/devito/ir/support/utils.py +++ b/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` diff --git a/devito/passes/clusters/buffering.py b/devito/passes/clusters/buffering.py index 9c627284ec..02fe2144f3 100644 --- a/devito/passes/clusters/buffering.py +++ b/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 diff --git a/devito/types/array.py b/devito/types/array.py index ea80bed19b..f47cb0ebac 100644 --- a/devito/types/array.py +++ b/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)) diff --git a/tests/test_buffering.py b/tests/test_buffering.py index eb2ac804ab..e61c50c5fc 100644 --- a/tests/test_buffering.py +++ b/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)) diff --git a/tests/test_ir.py b/tests/test_ir.py index 50f7e8d5d8..1757fe29a7 100644 --- a/tests/test_ir.py +++ b/tests/test_ir.py @@ -4,8 +4,8 @@ from conftest import EVAL, skipif # noqa from devito import ( # noqa - Constant, Dimension, Eq, Function, Grid, Inc, Operator, SubDimension, TimeFunction, - switchconfig + Constant, Dimension, Eq, Function, Grid, Inc, Operator, SubDimension, TimeDimension, + TimeFunction, switchconfig ) from devito.finite_differences.differentiable import IndexSum, LocalSum from devito.ir.cgen import ccode @@ -23,6 +23,7 @@ Any, Backward, Forward, Interval, IntervalGroup, IterationInterval, IterationSpace, NullInterval, null_ispace ) +from devito.ir.support.utils import detect_halo_writes from devito.symbolics import DefFunction, FieldFromPointer, uxreplace from devito.tools import prod from devito.tools.data_structures import frozendict @@ -153,6 +154,43 @@ def test_timedaccess_cached(self, fc, x, y): assert ta0 is ta1 + @pytest.mark.parametrize('autopadding', [False, True]) + def test_timedaccess_touched_nodomain(self, autopadding): + grid = Grid(shape=(17, 17)) + x, y = grid.dimensions + + 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) + 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): """ Test arithmetic operations involving objects of type IterationInstance. @@ -1289,6 +1327,36 @@ def test_reduction_dimensions(self): class TestCluster: + @pytest.mark.parametrize('dimtype', [Dimension, TimeDimension]) + def test_detect_halo_writes(self, dimtype): + x = Dimension(name='x') + y = dimtype(name='y') + f = Function(name='f', dimensions=(x, y), shape=(17, 17), + halo=((4, 4), (4, 4))) + hx, hy = f._size_nodomain.left + k = CustomDimension(name='k', parent=y, symbolic_min=1, + symbolic_max=4, symbolic_size=4) + ispace = IterationSpace([Interval(x), Interval(y), Interval(k)]) + + halo = f.indexed[x + hx, hy - k] + domain = f.indexed[x + hx, y + hy] + nonlinear = f.indexed[x + hx, hy - k**2] + c = Cluster([Eq(halo, 1), Eq(domain, 2), Eq(nonlinear, 3), + Eq(Symbol(name='r'), 4)], ispace) + + # A halo write does not imply all writes to the same Function are halo-only + writes = detect_halo_writes(c, key=lambda d: d is y) + assert {w.access for w in writes} == {halo} + assert not detect_halo_writes(c, key=lambda d: d is x) + assert not detect_halo_writes(c, key=lambda d: False) + + # No dimension type, including TimeDimension, is special to this query + assert detect_halo_writes(c, key=lambda d: True) == writes + + wild = Cluster(Eq(Symbol(name='r'), CriticalRegion(True)), ispace) + assert wild.is_wild + assert not detect_halo_writes(wild, key=lambda d: True) + def test_from_clusters_mixed_dtypes(self): grid = Grid(shape=(4,)) x, = grid.dimensions diff --git a/tests/test_mpi.py b/tests/test_mpi.py index 5b78305670..6081fa627e 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -13,6 +13,8 @@ ) from devito.arch.compiler import OneapiCompiler from devito.data import LEFT, RIGHT +from devito.ir import Cluster, Interval, IterationSpace +from devito.ir.clusters.algorithms import check_halo_writes from devito.ir.iet import ( Call, Conditional, FindNodes, FindSymbols, Iteration, retrieve_iteration_tree ) @@ -21,6 +23,7 @@ from devito.mpi.distributed import CustomTopology from devito.mpi.routines import ComputeCall, HaloUpdateCall, HaloUpdateList, MPICall from devito.tools import Bunch +from devito.types import Bundle from devito.types.dimension import ModuloDimension from examples.seismic.acoustic import acoustic_setup @@ -1107,6 +1110,84 @@ def check_halo_exchanges(op, exp0, exp1): return calls, tloop +class TestHaloWrites: + + @pytest.mark.parallel(mode=[1, 2]) + @pytest.mark.parametrize('axis', [0, 1]) + @pytest.mark.parametrize('side', [LEFT, RIGHT]) + @pytest.mark.parametrize('with_grid', [False, True]) + @pytest.mark.parametrize('topology', [None, ('*', 1)]) + def test_check_halo_writes(self, axis, side, with_grid, topology, mode, caplog): + grid = Grid(shape=(16, 16), topology=topology) + d = grid.dimensions[axis] + k = CustomDimension(name='k', parent=d, symbolic_min=1, + symbolic_max=2, symbolic_size=2) + if with_grid: + kwargs = {'grid': grid} + else: + kwargs = {'dimensions': grid.dimensions, 'shape': grid.shape_local, + 'distributor': grid.distributor} + f = Function(name='f', space_order=2, **kwargs) + g = Function(name='g', grid=grid) + index = -k if side is LEFT else d.symbolic_size - 1 + k + eqns = [Eq(f, 1), Eq(f._subs(d, index), 0), Eq(g, f.dx)] + + op = Operator(eqns, name='halo_writes') + expected = not with_grid or topology is None or topology[axis] != 1 + assert ('HALO along potentially distributed' in caplog.text) == expected + + # Grid-backed derivatives still require normal halo exchanges, + # even on a single rank + assert bool(FindNodes(HaloUpdateCall).visit(op)) == with_grid + + @switchconfig(mpi=False) + @pytest.mark.parametrize('axis', [0, 1]) + @pytest.mark.parametrize('topology', [None, ('*', 1), (1, 1), (2, 1)]) + def test_check_halo_writes_serial(self, axis, topology, caplog): + grid = Grid(shape=(16, 16), topology=topology) + f = Function(name='f', grid=grid) + d = grid.dimensions[axis] + + Operator(Eq(f, 1), name='domain_writes') + assert 'HALO along potentially distributed' not in caplog.text + + # The warning uses the requested topology, not the serial decomposition + Operator(Eq(f._subs(d, -1), 0), name='halo_writes_serial') + expected = topology is None or topology[axis] != 1 + assert ('HALO along potentially distributed' in caplog.text) == expected + + @pytest.mark.parallel(mode=[1, 2]) + @pytest.mark.parametrize('axis', [0, 1]) + def test_check_halo_writes_bundle(self, axis, mode, caplog): + grid = Grid(shape=(16, 16), topology=('*', 1)) + f = Function(name='f', grid=grid) + g = Function(name='g', grid=grid) + fg = Bundle(name='fg', components=(f, g)) + assert fg._decomposition is f._decomposition + + x, y = grid.dimensions + hx, hy = fg._size_nodomain.left + index = fg.indexed[x + hx, y + hy]._subs(grid.dimensions[axis], -1) + c = Cluster([Eq(index, 0)], IterationSpace([Interval(x), Interval(y)])) + + check_halo_writes([c]) + assert ('HALO along potentially distributed' in caplog.text) == (axis == 0) + + @switchconfig(mpi=False) + @pytest.mark.parametrize('with_halo', [False, True]) + def test_check_halo_writes_multiple_grids(self, with_halo, caplog): + grid = Grid(shape=(16, 16)) + grid1 = Grid(shape=(16, 16), dimensions=grid.dimensions) + x, _ = grid.dimensions + f = Function(name='f', grid=grid) + g = Function(name='g', grid=grid1) + + eqn = Eq(f._subs(x, -1), g._subs(x, 1)) if with_halo else Eq(f, g) + Operator(eqn, name='mixed_grid_halo_writes') + + assert ('HALO along potentially distributed' in caplog.text) == with_halo + + class TestCodeGeneration: @pytest.mark.parallel(mode=1)