compiler: Skip copy-backs overwritten by a later one to the same slot

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

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 {}
 

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

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