diff --git a/devito/ir/support/guards.py b/devito/ir/support/guards.py index 909a5e59e1..f1aab4726e 100644 --- a/devito/ir/support/guards.py +++ b/devito/ir/support/guards.py @@ -183,8 +183,9 @@ def __new__(cls, d, index, direction, return cls._new(p0, p1, d, index, direction, d_min=d_min, d_max=d_max, **kwargs) @classmethod - def _new(cls, p0, p1, d, index, direction, d_min=None, d_max=None): - + def _new(cls, p0, p1, d, index, direction, d_min=None, d_max=None, **kwargs): + # `kwargs` may carry `evaluate`, e.g. from `negated`; the relation is + # never evaluated obj = super().__new__(cls, p0, p1, evaluate=False) obj.d = d diff --git a/devito/passes/clusters/buffering.py b/devito/passes/clusters/buffering.py index b1a0215bc2..eae3992562 100644 --- a/devito/passes/clusters/buffering.py +++ b/devito/passes/clusters/buffering.py @@ -1,6 +1,6 @@ from collections import defaultdict, namedtuple from functools import cached_property -from itertools import chain, groupby +from itertools import chain import numpy as np from sympy import And @@ -9,7 +9,7 @@ from devito.exceptions import CompilationError from devito.ir import ( - Backward, Cluster, Forward, GuardBound, GuardFactor, InitArray, Interval, + Backward, Cluster, Forward, GuardBound, GuardFactor, Guards, InitArray, Interval, IterationSpace, Properties, Queue, Vector, detect_halo_writes, lower_exprs, vmax, vmin ) from devito.logger import warning @@ -164,9 +164,15 @@ def callback(self, clusters, prefix): # that access it descriptors = defaultdict(list) for (f, guards), b in self.mapper.items(): - if f in bfmap: + if f not in bfmap: + continue + if guards is None: + # The Clusters sharing one buffer across their guards + other = {k for ff, k in self.mapper if ff == f and k is not None} + fc = [c for c in bfmap[f] if c.guards not in other] + else: fc = [c for c in bfmap[f] if c.guards == guards] - descriptors[b].append(BufferDescriptor(f, b, fc, guards)) + descriptors[b].append(BufferDescriptor(f, b, fc, guards)) # Are we inside the right `d`? descriptors = {b: [vi for vi in v if d in vi.itdims] @@ -176,6 +182,8 @@ def callback(self, clusters, prefix): if not descriptors: return clusters + _check_mutable_guards(descriptors, clusters) + # Then we create the initializing Clusters when necessary init = init_buffers(descriptors, self.options) @@ -196,15 +204,17 @@ def callback(self, clusters, prefix): # If a buffer is read but never written, then we need to add # an Eq to step through the next slot # E.g., `ub[0, x] = usave[time+2, x]` - for v in descriptors.flat_values(): - if not v.is_readonly: - continue - if c not in v.firstread: - continue - if c.guards.get(d) != v.guards.get(d): + # Likewise, a slot read before being written within the same + # iteration must be loaded from the buffered Function + # E.g., `ub[0, x] = usave[time, x]` before `ub[0, x] = ub[0, x] + 1` + steps = [(v, v.last_idx[c], c) for v in descriptors.flat_values() + if v.is_readonly and c in v.firstread] + steps.extend((v, i, c1) for v in descriptors.flat_values() + for i, c1 in v.loads.get(c, [])) + for v, idxf, cr in steps: + if v.guards is not None and c.guards.get(d) != v.guards.get(d): continue - idxf = v.last_idx[c] idxb = mds[(v.xd, idxf)] lhs = v.b.indexify()._subs(v.xd, idxb) @@ -216,13 +226,15 @@ def callback(self, clusters, prefix): ispace = v.step_to ispace = ispace.augment(subiters).augment(c.sub_iterators) + # The load precedes the first access to the slot, hence it drops + # the readers' guards reading a Function that may change before + # the reads it serves + guards = _stable_guards(cr.guards, clusters) if v.xd in ispace.itdims: - guards = c.guards.xandg(v.xd, GuardBound(0, v.first_idx.f)) - else: - guards = c.guards + guards = guards.xandg(v.xd, GuardBound(0, v.first_idx.f)) properties = c.properties.sequentialize(d) - if not _explicit_guard(c.guards.get(d, None), d): + if v.is_readonly and not _explicit_guard(c.guards.get(d, None), d): properties = properties.prefetchable(d) # `c` may be a HaloTouch Cluster, so with no vision of the `bdims` properties = properties.parallelize(v.bdims).affine(v.bdims) @@ -246,9 +258,19 @@ def callback(self, clusters, prefix): continue if c not in v.lastwrite: continue - if c.guards.get(d) != v.guards.get(d): + if v.guards is not None and c.guards.get(d) != v.guards.get(d): continue + # Only one slot per iteration is copied back, so writes at other + # indices (e.g., `f[time]` along with `f[time+1]`) would be lost + written = {w[v.dim] for c1 in v.clusters if c1.guards == c.guards + for w in c1.scope.writes.get(v.f, [])} + if len(written) > 1: + raise CompilationError( + f"Cannot buffer `{v.f.name}`, written at multiple indices " + f"along `{v.dim}` within the same iteration" + ) + idxf = v.last_idx[c] idxb = mds[(v.xd, idxf)] @@ -462,9 +484,21 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs): for f, clusters in bfmap.items(): dim = f.hdim - for k, g in groupby(clusters, key=lambda c: c.guards): - g = list(g) + # Clusters reading `f` after a write under other guards must see it, so + # they share one buffer with the writers + sguards = _shared_guards(f, dim, clusters) + if sguards: + g = [c for c in clusters if c.guards in sguards] + mapper[(f, None)] = _select_buffer( + f, dim, None, g, xds, async_degree, sregistry, callback + ) + groups[(f, None)] = g + + # NOTE: Clusters with the same guards may be non-consecutive + for k, g in as_mapper(clusters, lambda c: c.guards).items(): + if k in sguards: + continue if f in shared: mapper[(f, k)] = shared[f] groups[(f, k)] = g @@ -724,6 +758,63 @@ def firstread(self): mapper.setdefault(c.guards, c) return tuple(mapper.values()) + @cached_property + def loads(self): + """ + Map each Cluster to the indices of `f` to be loaded into the buffer right + before it, each paired with the Cluster reading it. An index is loaded if + read before being written in the same iteration, or ahead of all written + indices, as the corresponding buffer slot holds a value from a previous + iteration rather than the value stored in `f`. The load precedes the + first access to the index, as an earlier write under other guards may + not be performed. + """ + def writes(clusters): + return {a[self.dim] for c in clusters for a in c.scope.writes.get(self.f, [])} + + def reads(clusters): + return filter_ordered(a[self.dim] for c in clusters if not c.is_wild + for a in c.scope.reads.get(self.f, [])) + + # Slots read ahead of all written ones, hence never written by an earlier + # iteration. E.g., `f[time+1]` in `g[time] = f[time+1]; f[time] = ...` + windices = writes(self.clusters) + rindices = reads(self.clusters) + ahead = set() + if windices and rindices: + sign = 1 if self.is_forward_buffering else -1 + ahead = {i for i in rindices + if all(S(sign*(i - w)).is_positive for w in windices)} + + # A buffer shared across guards may hold, in a slot written by an earlier + # iteration under guards that did not hold, a value from an even earlier + # one. E.g., `f[time-1]` in `f[time] = ... if time < 3; g[time] = f[time-1]` + behind = set() + if self.guards is None: + behind = set(rindices) - windices - ahead + + mapper = defaultdict(list) + for n, c in enumerate(self.clusters): + if c.is_wild: + continue + + # Slots written or loaded earlier, unconditionally or under the same + # guards, unless these read a Function, which may change in between + earlier = [c1 for c1 in self.clusters[:n] + if not c1.guards or (c1.guards == c.guards and + not _reads_functions(c.guards))] + written = writes(earlier) + loaded = set(reads(earlier)) + + for i in reads([c]): + if i not in windices | ahead | behind or i in written | loaded: + continue + c0 = next(c1 for c1 in self.clusters[:n+1] + if i in writes([c1]) or i in reads([c1])) + mapper[c0].append((i, c)) + + return frozendict(mapper) + @cached_property def lastwrite(self): mapper = {} @@ -802,7 +893,8 @@ def written(self): def last_idx(self): """ The "last index" is the index that corresponds to the last slot that - must be read from the buffered Function. + must be read from the buffered Function or, if the buffer is written, + to the last written slot, that is the one to be copied back. Examples -------- @@ -810,16 +902,23 @@ def last_idx(self): * `time+1` in the case of `foo(u[time-1], u[time], u[time+1])` with a forward-propagating `time` Dimension; * `time-1` in the case of `foo(u[time-1], u[time], u[time+1])` - with a backwards-propagating `time` Dimension. + with a backwards-propagating `time` Dimension; + * `time` in the case of `g[time] = u[time+1]; u[time] = ...` + with a forward-propagating `time` Dimension. """ mapper = {} func = vmax if self.is_forward_buffering else vmin - for c in self.lastwrite + self.firstread: + for c in self.lastwrite or self.firstread: # Consider all Clusters sharing `c`'s guards, so the leading edge is # found even when `f` is accessed at different offsets across them # (e.g. `f` and `f.forward` in separate Eqs) group = [c1 for c1 in self.clusters if c1.guards == c.guards] - indices = extract_indices(self.f, self.dim, group) + if self.is_write: + # A slot read ahead of the written ones is not copied back + accesses = chain(*[c1.scope.writes.get(self.f, []) for c1 in group]) + indices = filter_ordered(i[self.dim] for i in accesses) + else: + indices = extract_indices(self.f, self.dim, group) idx = func(*[Vector(i) for i in indices])[0] mapper[c] = idx @@ -1045,7 +1144,7 @@ def _overwritten(v, c, idxf, descriptors, order, d): for c1 in v1.lastwrite: if order.get(c1, -1) <= order[c]: continue - if c1.guards.get(d) != v1.guards.get(d): + if v1.guards is not None and c1.guards.get(d) != v1.guards.get(d): continue # The later copy-back must write the whole slot written by `v` key = lambda i: i in v.bdims # noqa: B023 @@ -1069,10 +1168,12 @@ def _overwritten(v, c, idxf, descriptors, order, d): condition = And(*c1.guards.values(), coincide) # The condition is evaluated at the earlier copy-back, so nothing - # it reads, indices included, may be written in between + # it reads, indices included, may be written in between; and `v.f` + # may not be read in between, as the skipped slot is still unwritten reads = {i.function for i in retrieve_terminals(condition, deep=True)} between = [c2 for c2, n in order.items() if order[c] < n <= order[c1]] - if any(not reads.isdisjoint(c2.scope.writes) for c2 in between): + if any(not reads.isdisjoint(c2.scope.writes) or v.f in c2.scope.reads + for c2 in between): continue conditions.append(condition) @@ -1082,6 +1183,62 @@ def _overwritten(v, c, idxf, descriptors, order, d): return Or(*conditions) +def _reads_functions(guards): + """ + True if any of the guards reads a Function, indices included. + """ + return any(retrieve_functions(g, deep=True) for g in guards.values()) + + +def _stable_guards(guards, clusters): + """ + The guards, or conjuncts thereof, that read no Function written by any of + the given Clusters, hence holding throughout an iteration. + """ + written = {f for c in clusters for f in c.scope.writes} + mapper = {} + for d, g in guards.items(): + conjuncts = [i for i in (g.args if isinstance(g, And) else (g,)) + if not written & set(retrieve_functions(i, deep=True))] + if conjuncts: + mapper[d] = And(*conjuncts) + return Guards(mapper) + + +def _shared_guards(f, dim, clusters): + """ + The guards of the Clusters that must share a buffer for `f`, that is those + of Clusters reading `f` after a Cluster under other guards writes it, and + those of the latter. + """ + shared = set() + for n, c in enumerate(clusters): + if c.is_wild or not c.scope.reads.get(f): + continue + for c1 in clusters[:n]: + if c1.guards != c.guards and c1.scope.writes.get(f): + shared.update([c.guards, c1.guards]) + if not shared: + return shared + + # The slots of a shared buffer must be at a known distance from each other, + # e.g. not `time/factor` and a CondEq-shifted `time/factor + (time >= 4)`. + # Otherwise each set of guards keeps its own buffer + group = [c for c in clusters if c.guards in shared] + indices = extract_indices(f, dim, group) + if not all(is_integer(i - indices[0]) for i in indices): + return set() + + # Likewise if they iterate over different spatial Dimensions, e.g. different + # SubDomains, as a buffer has one IterationSpace + spaces = {frozenset(d for d in c.ispace.dimensions if not d._defines & dim._defines) + for c in group if not c.is_wild} + if len(spaces) > 1: + return set() + + return shared + + def _fits_buffer(f, clusters, b): """ True if the accesses to `f` in the given Clusters span no more slots than @@ -1191,7 +1348,7 @@ def make_buffer_subs(descriptors, mds, d): subs = defaultdict(dict) for b, v in descriptors.flat_items(): for c in v.clusters: - if c.guards.get(d) != v.guards.get(d): + if v.guards is not None and c.guards.get(d) != v.guards.get(d): continue accesses = c.scope[v.f] @@ -1223,6 +1380,29 @@ def _explicit_guard(g, d): return False +def _check_mutable_guards(descriptors, clusters): + """ + Raise if a buffered Function is written more than once under the same + guards, and these guards read, indices included, a Function written in + between. The copy-back follows the last of these writes only, hence it would + be skipped if the guards no longer held, thus losing the earlier writes. + """ + for v in descriptors.flat_values(): + writes = [c for c in v.clusters if c.scope.writes.get(v.f)] + for guards, group in as_mapper(writes, lambda c: c.guards).items(): + functions = set(retrieve_functions(list(guards.values()), deep=True)) + if len(group) < 2 or not functions: + continue + + start = clusters.index(group[0]) + end = clusters.index(group[-1]) + if any(functions & set(c.scope.writes) for c in clusters[start+1:end]): + raise CompilationError( + f"Unsupported `buffering` of `{v.f.name}`, written multiple " + "times under guards that read a Function modified in between" + ) + + class Descriptors(dict): def flat_items(self): diff --git a/tests/test_buffering.py b/tests/test_buffering.py index da6d160b13..e14251f548 100644 --- a/tests/test_buffering.py +++ b/tests/test_buffering.py @@ -1247,3 +1247,418 @@ def test_buffering_same_slot_skip(case): op.apply(time_m=0, time_M=3) assert np.all(f.data[:, 1, 1] == expected) + + +@pytest.mark.parametrize('eqns', ['g = f; f = T', 'f = f + T']) +def test_read_before_write(eqns): + """ + A slot of `f` read before being written in the same iteration must hold the + value stored in `f`, not the one left in the buffer by a previous iteration. + """ + grid = Grid(shape=(4, 4)) + time = grid.time_dim + nt = 3 + + f = TimeFunction(name='f', grid=grid, time_order=0, space_order=0, save=nt) + g = TimeFunction(name='g', grid=grid, time_order=0, space_order=0, save=nt) + T = TimeFunction(name='T', grid=grid, time_order=0, space_order=0) + f.data[:] = 10. + + if eqns == 'g = f; f = T': + eqs = [Eq(T, time), Eq(g, f), Eq(f, T)] + expected = [0., 1., 2.] + else: + eqs = [Eq(T, time), Eq(f, f + T)] + expected = [10., 11., 12.] + + op = Operator(eqs, opt='buffering', name='ReadBeforeWrite') + op.apply(time_m=0, time_M=nt-1) + + assert np.allclose(f.data[:, 1, 1], expected) + if eqns == 'g = f; f = T': + assert np.allclose(g.data[:, 1, 1], 10.) + + +@pytest.mark.parametrize('eqns', ['g = f.forward; f = h + 1', 'f = f.forward + h']) +def test_read_ahead(eqns): + """ + `h` imposes a forward-propagating `time`, so `f[time+1]` is read ahead of + the written `f[time]`. Its slot must be loaded from `f`, and only the + written slot copied back. + """ + grid = Grid(shape=(4, 4)) + nt = 4 + + f = TimeFunction(name='f', grid=grid, time_order=0, space_order=0, save=nt) + g = TimeFunction(name='g', grid=grid, time_order=0, space_order=0, save=nt) + h = TimeFunction(name='h', grid=grid, time_order=1, space_order=0) + f.data[:] = 10*np.arange(nt).reshape(nt, 1, 1) + + if eqns == 'g = f.forward; f = h + 1': + eqs = [Eq(h.forward, h + 1), Eq(g, f.forward), Eq(f, h + 1)] + expected_f = [1., 2., 3., 30.] + expected_g = [10., 20., 30., 0.] + else: + eqs = [Eq(h.forward, h + 1), Eq(f, f.forward + h)] + expected_f = [10., 21., 32., 30.] + expected_g = [0., 0., 0., 0.] + + op = Operator(eqs, opt='buffering', name='ReadAhead') + op.apply(time_m=0, time_M=nt-2) + + assert np.allclose(f.data[:, 1, 1], expected_f) + assert np.allclose(g.data[:, 1, 1], expected_g) + + +def test_non_consecutive_guards(): + """ + The Clusters reading `f` without guards are not consecutive, as a guarded + one sits in between. The buffer they share must span all of their accesses, + `f[time+1]` and `f[time]`, not just those of the last Cluster. + """ + grid = Grid(shape=(4, 4)) + time = grid.time_dim + nt = 4 + + f = TimeFunction(name='f', grid=grid, time_order=0, space_order=0, save=nt) + g = TimeFunction(name='g', grid=grid, time_order=0, space_order=0, save=nt) + h = TimeFunction(name='h', grid=grid, time_order=0, space_order=0, save=nt) + k = TimeFunction(name='k', grid=grid, time_order=0, space_order=0, save=nt) + f.data[:] = np.arange(nt).reshape(nt, 1, 1) + + c = ConditionalDimension(name='c', parent=time, condition=time >= 1, + relation='strict') + eqs = [Eq(g, f.forward), Eq(h, f, implicit_dims=c), Eq(k, f)] + + op = Operator(eqs, opt='buffering', name='NonConsecutiveGuards') + op.apply(time_m=0, time_M=nt-2) + + assert np.allclose(g.data[:, 1, 1], [1., 2., 3., 0.]) + assert np.allclose(h.data[:, 1, 1], [0., 1., 2., 0.]) + assert np.allclose(k.data[:, 1, 1], [0., 1., 2., 0.]) + + +@pytest.mark.parametrize('case', ['guarded-write-skipped']) +def test_read_write_different_guards(case): + """ + A read sees the earlier guarded write, or the initialized slot if that + write does not fire. + """ + grid = Grid(shape=(4, 4)) + time = grid.time_dim + nt = 4 + + f = TimeFunction(name='f', grid=grid, time_order=0, space_order=0, save=nt) + g = TimeFunction(name='g', grid=grid, time_order=0, space_order=0, save=nt) + f.data[:] = 10*np.arange(nt).reshape(nt, 1, 1) + + c = ConditionalDimension(name='c', parent=time, condition=time < 2, + relation='strict') + eqs = [Eq(f, time + 1, implicit_dims=c), Eq(g, f)] + + op = Operator(eqs, opt='buffering', name='DifferentGuards') + op.apply(time_m=0, time_M=nt-1) + + assert np.all(f.data[:, 1, 1] == [1, 2, 20, 30]) + assert np.all(g.data[:, 1, 1] == [1, 2, 20, 30]) + + +@pytest.mark.parametrize('case', ['read-then-write', 'write-then-read']) +def test_read_write_shifted_slot(case): + """ + The last timestep is saved into a slot shifted off the subsampled grid, so + at an unknown distance from the others. A read after this write, under other + guards, cannot share a buffer with it, so each keeps its own. + """ + grid = Grid(shape=(4, 4)) + time = grid.time_dim + + ct0 = ConditionalDimension(name='ct0', parent=time, factor=2, relation=Or) + ctend = ConditionalDimension(name='ctend', parent=time, + condition=CondEq(time, 3), relation='strict') + f = TimeFunction(name='f', grid=grid, time_order=0, space_order=0, + time_dim=ct0, save=3) + g = TimeFunction(name='g', grid=grid, time_order=0, space_order=0, save=4) + T = TimeFunction(name='T', grid=grid, time_order=0, space_order=0) + + if case == 'read-then-write': + cr = ConditionalDimension(name='cr', parent=time, condition=time >= 2, + relation='strict') + eqs = [Eq(T, time), Eq(f, T), Eq(g, f, implicit_dims=cr), + Eq(f, T + 1, implicit_dims=ctend)] + + op = Operator(eqs, opt='buffering', name='ShiftedSlot') + op.apply(time_m=0, time_M=3) + + assert np.all(f.data[:, 1, 1] == [0, 2, 4]) + assert np.all(g.data[:, 1, 1] == [0, 0, 2, 2]) + else: + f.data[:] = 10*np.arange(3).reshape(3, 1, 1) + 5 + eqs = [Eq(f, time, implicit_dims=ctend), Eq(g, f)] + + op = Operator(eqs, opt='buffering', name='ShiftedSlot') + op.apply(time_m=0, time_M=3) + + assert np.all(f.data[:, 1, 1] == [5, 15, 3]) + assert np.all(g.data[:, 1, 1] == [5, 0, 15, 0]) + + +def test_read_behind_guarded_write(): + """ + `f[time-1]` is read from the buffer shared with the write under `time < 3`. + When that write did not fire in the previous iteration, the slot must be + loaded from `f`. + """ + grid = Grid(shape=(4, 4)) + time = grid.time_dim + + f = TimeFunction(name='f', grid=grid, time_order=1, space_order=0, save=6) + g = TimeFunction(name='g', grid=grid, time_order=0, space_order=0, save=6) + T = TimeFunction(name='T', grid=grid, time_order=0, space_order=0) + f.data[:] = np.arange(6).reshape(6, 1, 1) + + cw = ConditionalDimension(name='cw', parent=time, condition=time < 3, + relation='strict') + cr = ConditionalDimension(name='cr', parent=time, condition=T[0, 0, 0] > 0.5, + relation='strict') + eqs = [Eq(f, time + 100, implicit_dims=cw), Eq(T, time - 1), + Eq(g, f.backward, implicit_dims=cr)] + + op = Operator(eqs, opt='buffering', name='ReadBehind') + op.apply(time_m=1, time_M=5) + + assert np.all(f.data[:, 1, 1] == [0, 101, 102, 3, 4, 5]) + assert np.all(g.data[:, 1, 1] == [0, 0, 101, 102, 3, 4]) + + +def test_read_write_different_subdomains(): + """ + A write over a SubDomain and a read over the whole domain, under other + guards, iterate over different spatial Dimensions, so each keeps its own + buffer. + """ + class Left(SubDomain): + name = 'left' + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', 2), y: y} + + grid = Grid(shape=(4, 4)) + left = Left(grid=grid) + time = grid.time_dim + + f = TimeFunction(name='f', grid=grid, time_order=0, space_order=0, save=4) + g = TimeFunction(name='g', grid=grid, time_order=0, space_order=0, save=4) + f.data[:] = 10*np.arange(4).reshape(4, 1, 1) + + cw = ConditionalDimension(name='cw', parent=time, condition=time < 2, + relation='strict') + eqs = [Eq(f, time + 1, implicit_dims=cw, subdomain=left), Eq(g, f)] + + op = Operator(eqs, opt='buffering', name='DifferentSubDomains') + op.apply(time_m=0, time_M=3) + + assert np.all(g.data[:, 1, 1] == [1, 2, 20, 30]) + assert np.all(g.data[:, 3, 1] == [0, 10, 20, 30]) + + +@pytest.mark.parametrize('case', ['reread', 'write-read', 'read-read']) +def test_guard_changes_between_accesses(case): + """ + Two accesses to `f` under the same guard, which reads `T`, updated in + between. The guard may not hold at the first access, so the slot must be + loaded for the second one. + """ + grid = Grid(shape=(2,)) + time = grid.time_dim + + f = TimeFunction(name='f', grid=grid, time_order=0, space_order=0, save=3) + g = TimeFunction(name='g', grid=grid, time_order=0, space_order=0, save=3) + h = TimeFunction(name='h', grid=grid, time_order=0, space_order=0, save=3) + T = TimeFunction(name='T', grid=grid, time_order=0, space_order=0) + f.data[:] = 10*np.arange(3).reshape(3, 1) + 5 + + c = ConditionalDimension(name='c', parent=time, condition=T[0, 0] > time, + relation='strict') + if case == 'reread': + eqs = [Eq(g, f, implicit_dims=c), Eq(T, time + 1), + Eq(f, f + 1, implicit_dims=c)] + expected = [[6, 16, 26], [0, 0, 0], [0, 0, 0]] + elif case == 'write-read': + eqs = [Eq(f, time + 100, implicit_dims=c), Eq(T, time + 1), + Eq(g, f, implicit_dims=c)] + expected = [[5, 15, 25], [5, 15, 25], [0, 0, 0]] + else: + eqs = [Eq(g, f, implicit_dims=c), Eq(T, time + 1), + Eq(h, f, implicit_dims=c)] + expected = [[5, 15, 25], [0, 0, 0], [5, 15, 25]] + + op = Operator(eqs, opt='buffering', name='GuardChanges') + op.apply(time_m=0, time_M=2) + + for i, v in zip([f, g, h], expected, strict=True): + assert np.all(i.data[:, 0] == v) + + +@pytest.mark.parametrize('case', ['stable', 'after', 'mutable', 'repeated', 'rise']) +def test_mutable_read_guards(case): + """ + `f` is written under `time < 2`, then read under guards reading `T`. The + slot is loaded before the write, as the write may not be performed, hence + the load must not depend on the value of `T` at that point. + """ + grid = Grid(shape=(4, 4)) + time = grid.time_dim + + f = TimeFunction(name='f', grid=grid, time_order=0, space_order=0, save=4) + g = TimeFunction(name='g', grid=grid, time_order=0, space_order=0, save=4) + h = TimeFunction(name='h', grid=grid, time_order=0, space_order=0, save=4) + T = TimeFunction(name='T', grid=grid, time_order=0, space_order=0) + + cw = ConditionalDimension(name='cw', parent=time, condition=time < 2, + relation='strict') + cr = ConditionalDimension(name='cr', parent=time, condition=T[0, 0, 0] > 0.5, + relation='strict') + + write = Eq(f, time + 1, implicit_dims=cw) + read = Eq(g, f, implicit_dims=cr) + expected_h = [0, 0, 0, 0] + if case == 'stable': + eqs = [Eq(T, 1), write, read] + expected_g = [1, 2, 20, 30] + elif case == 'after': + # The only update to T follows the last read + T.data[:] = 1 + eqs = [write, read, Eq(T, 0, implicit_dims=time), + Eq(f, time + 5, implicit_dims=cr)] + expected_g = [1, 0, 0, 0] + elif case == 'mutable': + eqs = [write, Eq(T, 1, implicit_dims=time), read] + expected_g = [1, 2, 20, 30] + elif case == 'repeated': + eqs = [write, read, Eq(T, 1, implicit_dims=time), + Eq(h, f, implicit_dims=cr)] + expected_g = [0, 2, 20, 30] + expected_h = [1, 2, 20, 30] + else: + # The guard turns true between the write and the read + eqs = [write, Eq(T, time - 1, implicit_dims=time), read] + expected_g = [0, 0, 20, 30] + + f.data[:] = 10*np.arange(4).reshape(4, 1, 1) + op = Operator(eqs, opt='buffering', name='MutableReadGuards') + op.apply(time_m=0, time_M=3) + + assert np.all(f.data[:, 1, 1] == [1, 2, 20, 30]) + assert np.all(g.data[:, 1, 1] == expected_g) + assert np.all(h.data[:, 1, 1] == expected_h) + + +@pytest.mark.parametrize('case', ['stable', 'mutable', 'indirect']) +def test_mutable_guards(case): + """ + Two writes to `f` under the same guards, which read `T`. The copy-back + follows the second write only, so `T` must not change in between, nor the + index it is read at. + """ + grid = Grid(shape=(4, 4)) + time = grid.time_dim + x, _ = grid.dimensions + + f = TimeFunction(name='f', grid=grid, time_order=0, space_order=0, save=2) + T = TimeFunction(name='T', grid=grid, time_order=0, space_order=0) + I = TimeFunction(name='I', grid=grid, time_order=0, space_order=0, + dtype=np.int32) + + if case == 'indirect': + condition = And(CondEq(time, 1), T[0, I[0, 0, 0], 0] < 1.5) + else: + condition = And(CondEq(time, 1), T[0, 0, 0] < 1.5) + c0 = ConditionalDimension(name='c0', parent=time, condition=condition, + relation='strict') + c1 = ConditionalDimension(name='c1', parent=time, condition=condition, + relation='strict') + + if case == 'stable': + eqs = [Eq(f, T, implicit_dims=c0), Eq(T, T + 1)] + + op = Operator(eqs, opt='buffering', name='MutableGuards') + op.apply(time_m=0, time_M=1) + + assert np.all(f.data[:, 1, 1] == [0, 1]) + else: + if case == 'mutable': + eqs = [Eq(f, T, implicit_dims=c0), Eq(T, T + 1), + Eq(f, T, implicit_dims=c1)] + else: + ci = ConditionalDimension(name='ci', parent=time, + condition=CondEq(time, 1), relation='strict') + eqs = [Eq(T, time + 2*x), Eq(f, T, implicit_dims=c0), + Eq(I, 1, implicit_dims=ci), Eq(f, T + 10, implicit_dims=c1)] + + with pytest.raises(CompilationError, match='modified in between'): + Operator(eqs, opt='buffering', name='MutableGuards') + + +@pytest.mark.parametrize('case', ['single', 'forward', 'readwrite', 'backward']) +def test_write_multi_index(case): + """ + `f` written at two time indices in the same iteration cannot be buffered, + as only one slot per iteration is copied back. + """ + grid = Grid(shape=(4, 4)) + time = grid.time_dim + + f = TimeFunction(name='f', grid=grid, time_order=1, space_order=0, save=5) + T = TimeFunction(name='T', grid=grid, time_order=0, space_order=0) + + if case == 'single': + eqns = [Eq(f.forward, f + T)] + + op = Operator([Eq(T, time)] + eqns, opt='buffering', name='WriteMultiIndex') + op.apply(time_m=0, time_M=3) + + assert np.all(f.data[:, 1, 1] == [0, 0, 1, 3, 6]) + else: + if case == 'forward': + eqns = [Eq(f.forward, T + 10), Eq(f, T)] + elif case == 'readwrite': + eqns = [Eq(f.forward, f + T + 10), Eq(f, T)] + else: + eqns = [Eq(f.backward, T + 10), Eq(f, T)] + + with pytest.raises(CompilationError, match='multiple indices'): + Operator([Eq(T, time)] + eqns, opt='buffering', name='WriteMultiIndex') + + +def test_slot_read_between_writes(): + """ + A subsampled save plus an explicit save at the last timestep write the + same slot at that timestep, so the first copy-back is normally skipped. + A read of the saved Function in between is served from its storage, so + the copy-back must not be skipped then. Both the read and the slot's + final value follow program order. + """ + nt, factor = 8, 2 + grid = Grid(shape=(4, 4)) + time = grid.time_dim + + tsub = ConditionalDimension(name='tsub', parent=time, factor=factor) + clast = ConditionalDimension(name='clast', parent=time, + condition=CondEq(time, nt-2)) + + u = TimeFunction(name='u', grid=grid, time_order=1, space_order=0) + usave = TimeFunction(name='usave', grid=grid, time_order=0, space_order=0, + save=nt//factor, time_dim=tsub) + g = Function(name='g', grid=grid) + + eqns = [Eq(u.forward, u + 1), Eq(usave, u), + Eq(g, usave, implicit_dims=clast), + Eq(usave, u + 100, implicit_dims=clast)] + + op = Operator(eqns, opt='buffering', name='SlotReadBetween') + op.apply(time_m=0, time_M=nt-2) + + assert np.all(usave.data[:, 1, 1] == [0, 2, 4, 106]) + assert np.all(g.data == 6) diff --git a/tests/test_ir.py b/tests/test_ir.py index 9a7b2e9fea..9f0c986815 100644 --- a/tests/test_ir.py +++ b/tests/test_ir.py @@ -1,6 +1,6 @@ import numpy as np import pytest -from sympy import S +from sympy import And, S from conftest import EVAL, skipif # noqa from devito import ( # noqa @@ -18,7 +18,9 @@ AFFINE, IRREGULAR, REGULAR, IterationInstance, Relation, Scope, TimedAccess, Vector, mocksym0, mocksym1 ) -from devito.ir.support.guards import GuardOverflow +from devito.ir.support.guards import ( + GuardBoundNextGt, GuardBoundNextLe, GuardOverflow +) from devito.ir.support.space import ( Any, Backward, Forward, Interval, IntervalGroup, IterationInterval, IterationSpace, NullInterval, null_ispace @@ -1731,3 +1733,24 @@ def test_guard_overflow(self): guard = GuardOverflow(freespace, size) assert ccode(guard) == 'freespace >= f_vec->size[0]*f_vec->size[1]' + + def test_guard_bound_next_negated(self): + """ + The negation of a GuardBoundNext is the opposite relation. SymPy + negates the relational arguments of an `And` to detect contradictions, + so this is also what makes an `And` of such guards constructible. + """ + grid = Grid(shape=(4, 4)) + time = grid.time_dim + + i = Scalar(name='i') + + guard = GuardBoundNextLe(time, i, Forward) + assert ccode(guard) == 'i <= time_M' + + negated = guard.negated + assert isinstance(negated, GuardBoundNextGt) + assert ccode(negated) == 'i > time_M' + + other = GuardBoundNextLe(time, i + 1, Forward) + assert set(And(guard, other).args) == {guard, other}