compiler: Serialize copy-backs into the same Function from different buffers

devito/passes/clusters/asynchrony.py

@@ -86,6 +86,7 @@ def callback(self, clusters, prefix):
 
         locks = {}
         syncs = defaultdict(lambda: defaultdict(OrderedSet))
+        copybacks = []
         for c0 in clusters:
             d = self.key0(c0)
             if d is not dim:
@@ -102,6 +103,11 @@ def callback(self, clusters, prefix):
             protected = self._schedule_waitlocks(c0, d, clusters, locks, syncs)
             self._schedule_withlocks(c0, d, protected, locks, syncs)
 
+            if wraps_memcpy(c0):
+                copybacks.append((c0, d))
+
+        self._serialize_copybacks(copybacks, clusters, syncs)
+
         processed = [c.rebuild(syncs={**c.syncs, **syncs[c]}) for c in clusters]
 
         return processed
@@ -194,6 +200,41 @@ def _schedule_withlocks(self, c0, d, protected, locks, syncs):
                     WithLock(lock[i], target, i, function, findex, d, gid=gid)
                 ])
 
+    def _serialize_copybacks(self, copybacks, clusters, syncs):
+        """
+        Order the copy-backs into the same Function from different buffers.
+
+        Each buffer has its own lock and thus its own asynchronous task, so
+        two such copy-backs may otherwise run concurrently and race on the
+        Function's state (e.g., the next free offset of a compressed
+        Function). Before a buffer is refilled, wait for the pending
+        copy-backs of the other buffers into the same Function.
+        """
+        withlocks = {}
+        for c0, d in copybacks:
+            withlocks[c0] = [s for s in syncs[c0][d] if isinstance(s, WithLock)]
+
+        for c0, d in copybacks:
+            function = c0.exprs[0].lhs.function
+            targets = {s.target for s in withlocks[c0]}
+            gids = {s.gid for s in withlocks[c0]}
+
+            waits = [WaitLock(s.handle, s.target)
+                     for c1, _ in copybacks if c1 is not c0
+                     for s in withlocks[c1]
+                     if s.function is function and
+                     s.target not in targets and
+                     # Tasks of the same group run in one thread, in order
+                     (s.gid is None or s.gid not in gids)]
+            if not waits:
+                continue
+
+            # The refill of `c0`'s buffers precedes `c0`'s own task
+            for c1 in clusters:
+                if targets & set(c1.scope.writes):
+                    c2 = in_critical_region(c1, clusters) or c1
+                    syncs[c2][d].update(waits)
+
 
 @timed_pass(name='memcpy_prefetch')
 def memcpy_prefetch(clusters, key0, sregistry):

tests/test_gpu_common.py

@@ -3,18 +3,20 @@
 import pytest
 import scipy.sparse
 import sympy
+from sympy import Or
 
 from conftest import assert_structure, skipif
 from devito import (
-    Buffer, ConditionalDimension, Constant, Dimension, Eq, Function, Grid, Inc,
+    Buffer, CondEq, ConditionalDimension, Constant, Dimension, Eq, Function, Grid, Inc,
     MatrixSparseTimeFunction, Operator, SparseTimeFunction, SubDimension, SubDomain,
     SubDomainSet, TensorTimeFunction, TimeFunction, assign, configuration, exp,
     switchconfig, switchenv
 )
 from devito.arch import Cpu64, Device, get_cpu_info, get_gpu_info
 from devito.exceptions import InvalidArgument
 from devito.ir import (
-    Conditional, Expression, FindNodes, FindSymbols, Section, retrieve_iteration_tree
+    BusyWait, Conditional, Expression, FindNodes, FindSymbols, Section,
+    retrieve_iteration_tree
 )
 from devito.passes.iet.languages.openmp import OmpIteration
 from devito.types import DeviceID, DeviceRM, Lock, NPThreads, PThreadArray, Symbol
@@ -1314,6 +1316,47 @@ def test_save_multi_output(self):
         assert all(np.all(usave.data[i] == 2*i + 1) for i in range(usave.save))
         assert all(np.all(vsave.data[i] == 2*i + 1) for i in range(vsave.save))
 
+    def test_save_two_buffers_ordered(self):
+        """
+        A subsampled save and a CondEq save into the same Function use two
+        buffers, hence two locks and two asynchronous tasks. The tasks must
+        not run concurrently, as they update the same Function: each buffer
+        is refilled only once the other buffer's pending copy-back is done.
+        """
+        grid = Grid(shape=(4, 4))
+        time_dim = grid.time_dim
+
+        factor = 2
+        last = 7
+        ct0 = ConditionalDimension(name='ct0', parent=time_dim, factor=factor,
+                                   relation=Or)
+        ctend = ConditionalDimension(name='ctend', parent=time_dim,
+                                     condition=CondEq(time_dim, last),
+                                     relation='strict')
+
+        f = TimeFunction(name='f', grid=grid, time_order=0, space_order=0,
+                         time_dim=ct0, save=last//factor + 2)
+        T = TimeFunction(name='T', grid=grid, time_order=0, space_order=0)
+
+        eqns = [Eq(T, time_dim),
+                Eq(f, T),
+                Eq(f, T + 1, implicit_dims=ctend)]
+
+        op = Operator(eqns, opt=('buffering', 'tasking', 'orchestrate'))
+
+        # Check generated code
+        locks = {i.name for i in FindSymbols().visit(op) if isinstance(i, Lock)}
+        assert len(locks) == 2
+        waits = FindNodes(BusyWait).visit(op.body)
+        assert len(waits) == 2
+        for w in waits:
+            assert {i.name for i in FindSymbols().visit(w)} >= locks
+
+        op.apply(time_m=0, time_M=last)
+
+        expected = np.append(factor*np.arange(last//factor + 1), last + 1)
+        assert np.all(f.data[:, 1, 1] == expected)
+
     @pytest.mark.parametrize('opt', [
         ('buffering', 'tasking', 'orchestrate'),
         ('buffering', 'tasking', 'orchestrate', {'linearize': True}),