diff --git a/devito/passes/clusters/buffering.py b/devito/passes/clusters/buffering.py index e8904558d8..60f2deb7b2 100644 --- a/devito/passes/clusters/buffering.py +++ b/devito/passes/clusters/buffering.py @@ -3,8 +3,9 @@ from itertools import chain, groupby import numpy as np +from sympy import And from sympy import Eq as SympyEq -from sympy import Ne, S, simplify +from sympy import Ne, Not, Or, S, simplify, true from devito.exceptions import CompilationError from devito.ir import ( @@ -13,7 +14,7 @@ ) from devito.logger import warning from devito.passes.clusters.utils import is_memcpy -from devito.symbolics import IntDiv, retrieve_functions, uxreplace +from devito.symbolics import CondEq, IntDiv, retrieve_functions, uxreplace from devito.tools import ( Stamp, as_mapper, as_tuple, filter_ordered, flatten, frozendict, is_integer, timed_pass @@ -184,6 +185,10 @@ def callback(self, clusters, prefix): # E.g., `usave[time+1, x+1, y+1] -> ub0[t1, x+1, y+1]` subs = make_buffer_subs(descriptors, mds, d) + # Program order of the Clusters, to spot copy-backs into a slot that a + # later copy-back in the same iteration overwrites + order = {c: i for i, c in enumerate(clusters)} + processed = [] for c in clusters: # If a buffer is read but never written, then we need to add @@ -259,6 +264,15 @@ def callback(self, clusters, prefix): else: guards = c.guards + # Skip the copy-back when a later one in the same iteration + # writes the same slot, so that a single (possibly asynchronous) + # copy-back per slot is performed, carrying the last value + overwritten = _overwritten(v, c, idxf, descriptors, order, d) + if overwritten is true: + continue + elif overwritten is not None: + guards = guards.xandg(d, Not(overwritten)) + properties = c.properties.sequentialize(d) syncs = c.syncs @@ -985,6 +999,58 @@ def _underlying_buffer(clusters): return None +def _overwritten(v, c, idxf, descriptors, order, d): + """ + The condition under which the copy-back of `v` following the Cluster `c`, + into the slot `idxf` of `v.f`, is overwritten by a later copy-back into the + same slot within the same iteration of `d`. None if it is never overwritten. + + For example, given a subsampled save `Eq(usave, u)` and a save of the last + timestep `Eq(usave, u, implicit_dims=clast)`, both copy-backs target the same + slot whenever the last timestep falls onto the subsampled grid. Performing + both would be redundant and, with asynchronous copy-backs (e.g. to disk), + racy, as the two writes into the slot would be performed concurrently. + """ + conditions = [] + for v1 in descriptors.flat_values(): + if v1.f is not v.f or v1.is_readonly: + continue + # Copy-backs out of the same buffer are already serialized, since the + # buffer cannot be refilled before its previous copy-back is completed + if v1.b is v.b: + continue + + for c1 in v1.lastwrite: + if order.get(c1, -1) <= order[c]: + continue + if 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 + if c1.ispace.project(key) != c.ispace.project(key): + continue + # A guard over other Dimensions may restrict the later copy-back to + # a portion of the slot only + if any(not dd._defines & d._defines for dd in c1.guards): + continue + + shift = v1.last_idx[c1] - idxf + if shift.is_Number: + if shift != 0: + continue + coincide = true + else: + coincide = CondEq(shift, 0) + + # The later copy-back fires under all of its guards, which may be + # keyed on Dimensions derived from `d` (e.g., a layer Dimension) + conditions.append(And(*c1.guards.values(), coincide)) + + if not conditions: + return None + return Or(*conditions) + + def _fits_buffer(f, clusters, b): """ True if the accesses to `f` in the given Clusters span no more slots than diff --git a/devito/passes/clusters/cse.py b/devito/passes/clusters/cse.py index 87770d4039..9c899eb23e 100644 --- a/devito/passes/clusters/cse.py +++ b/devito/passes/clusters/cse.py @@ -14,7 +14,7 @@ from devito.finite_differences.differentiable import IndexSum from devito.ir import Cluster, Scope, cluster_pass from devito.symbolics import ( - DefFunction, Reserved, estimate_cost, q_leaf, q_terminal, search + DefFunction, IndexedPointer, Reserved, estimate_cost, q_leaf, q_terminal, search ) from devito.symbolics.manipulation import _uxreplace from devito.tools import DAG, as_list, as_tuple, extract_dtype, frozendict @@ -418,11 +418,17 @@ def _(expr): @_catch.register(Indexed) +@_catch.register(IndexedPointer) @_catch.register(Symbol) @_catch.register(Reserved) def _(expr): """ Handler for objects preventing CSE to propagate through their arguments. + + For IndexedPointer (e.g. `f_vec->sizes[(time / ct0f) + ...]`) this also + keeps integer index expressions out of a floating-point cluster's CSE + temporaries, which `cse_dtype` would otherwise declare as floating-point + and thus produce non-integral array subscripts. """ return {} diff --git a/tests/test_buffering.py b/tests/test_buffering.py index ef2a06c3fb..2a46203a85 100644 --- a/tests/test_buffering.py +++ b/tests/test_buffering.py @@ -1,6 +1,6 @@ import numpy as np import pytest -from sympy import Or +from sympy import Not, Or from conftest import skipif from devito import ( @@ -11,7 +11,8 @@ from devito.arch.archinfo import AppleArm from devito.exceptions import CompilationError from devito.ir import ( - Cluster, FindSymbols, Interval, IterationSpace, lower_exprs, retrieve_iteration_tree + Cluster, Conditional, FindNodes, FindSymbols, Interval, IterationSpace, lower_exprs, + retrieve_iteration_tree ) from devito.passes.clusters.buffering import BufferDimension, expand_halo_transfers from devito.types import Array @@ -1090,6 +1091,47 @@ def test_deferred_guard_no_fit(): assert np.allclose(f.data[:, 1, 1], expected) +@pytest.mark.parametrize("factor", [1, 2, 3]) +def test_buffering_same_slot_single_copyback(factor): + """ + A subsampled save and a CondEq save into the same Function land on the same + slot when the CondEq point is on the subsampled grid. Only one copy-back, + carrying the last value, must then be performed: the copy-back of the first + write is skipped whenever the second one overwrites its slot. Asynchronous + copy-backs (e.g., to disk) would otherwise write the slot concurrently. + """ + grid = Grid((4, 4)) + + nt = 5 + ntmod = (nt - 1) * factor + 1 + last = ntmod - 2 + on_grid = last % factor == 0 + nsave = nt - 1 if on_grid else nt + + ct0 = ConditionalDimension(name="ct0", parent=grid.time_dim, factor=factor, + relation=Or) + ctend = ConditionalDimension(name="ctend", parent=grid.time_dim, + condition=CondEq(grid.time_dim, last), + relation='strict') + f = TimeFunction(grid=grid, name='f', time_order=0, space_order=0, + time_dim=ct0, save=nsave) + T = TimeFunction(grid=grid, name='T', time_order=0, space_order=0) + + eqs = [Eq(T, grid.time_dim), Eq(f, T), Eq(f, T+1, implicit_dims=ctend)] + + op = Operator(eqs, opt='buffering') + + # The first copy-back is skipped when the second one targets the same slot + skips = [i for i in FindNodes(Conditional).visit(op) if i.condition.has(Not)] + assert len(skips) == 1 + + op.apply(time_m=0, time_M=last) + + expected = factor * np.arange(nsave) + expected[-1] = last + 1 + assert np.allclose(f.data[:, 1, 1], expected) + + @pytest.mark.parametrize("factor", [1, 2, 3]) def test_buffering_multi_cond(factor): grid = Grid((16, 16)) diff --git a/tests/test_cse.py b/tests/test_cse.py index 976593e9ff..6ab68f05a0 100644 --- a/tests/test_cse.py +++ b/tests/test_cse.py @@ -11,7 +11,7 @@ from devito.ir import Conditional, DummyEq, FindNodes, FindSymbols from devito.ir.support import generator from devito.passes.clusters.cse import CTemp, _cse -from devito.symbolics import DefFunction, indexify +from devito.symbolics import DefFunction, FieldFromPointer, IndexedPointer, indexify from devito.types import Array, Symbol, Temp @@ -154,6 +154,33 @@ def test_deffunction_not_captured(): assert processed == exprs +def test_indexedpointer_index_not_captured(): + """ + The index of an IndexedPointer (e.g. `f_vec->sizes[...]`) must not be + CSE-captured: in a floating-point cluster `cse_dtype` would bind the + integer index to a floating-point temporary, yielding a non-integral + array subscript that does not compile (nvcc, hipcc). + """ + time = Symbol(name='time', dtype=np.int32) + fac = Symbol(name='fac', dtype=np.int32) + vec = Symbol(name='f_vec', dtype=np.int32) + a = CTemp(name='a', dtype=np.float32) + b = CTemp(name='b', dtype=np.float32) + + index = time/fac + 2*time + sizes = IndexedPointer(FieldFromPointer('sizes', vec), index) + offs = IndexedPointer(FieldFromPointer('offs', vec), index) + exprs = [DummyEq(a, sizes + 1.), DummyEq(b, offs + sizes)] + + counter = generator() + make = lambda e: CTemp(name=f'r{counter()}', dtype=np.float32) + processed = _cse(exprs, make) + + for e in processed: + for i in e.find(IndexedPointer): + assert not any(isinstance(s, CTemp) for s in i.index.free_symbols) + + def test_w_conditionals(): grid = Grid(shape=(10, 10, 10)) x, _, _ = grid.dimensions