From 2f097148e22e0c5244be262f1fcbbdcb99425cf6 Mon Sep 17 00:00:00 2001 From: mloubout Date: Wed, 14 Oct 2026 20:28:42 -0400 Subject: [PATCH] compiler: Serialize copy-backs into the same Function from different buffers --- devito/passes/clusters/asynchrony.py | 41 ++++++++++++++++++++++++ tests/test_gpu_common.py | 47 ++++++++++++++++++++++++++-- 2 files changed, 86 insertions(+), 2 deletions(-) diff --git a/devito/passes/clusters/asynchrony.py b/devito/passes/clusters/asynchrony.py index a64ccf6122..2c2c72949f 100644 --- a/devito/passes/clusters/asynchrony.py +++ b/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): diff --git a/tests/test_gpu_common.py b/tests/test_gpu_common.py index 3efda2820f..144f117907 100644 --- a/tests/test_gpu_common.py +++ b/tests/test_gpu_common.py @@ -3,10 +3,11 @@ 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 @@ -14,7 +15,8 @@ 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}),