From eb0ecb9e1dd163c5ed14a024cb963847539e52b4 Mon Sep 17 00:00:00 2001 From: mloubout Date: Wed, 14 Oct 2026 21:54:09 -0400 Subject: [PATCH] compiler: Share a single buffer across all writes to a write-only Function --- devito/passes/clusters/buffering.py | 34 +++++++++++++++++++++++------ tests/test_buffering.py | 18 +++++++++------ 2 files changed, 38 insertions(+), 14 deletions(-) diff --git a/devito/passes/clusters/buffering.py b/devito/passes/clusters/buffering.py index 60f2deb7b2..18826dff2d 100644 --- a/devito/passes/clusters/buffering.py +++ b/devito/passes/clusters/buffering.py @@ -448,11 +448,26 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs): # buffer created for the regular `Eq(usave, u)` copy-back) extras = {} groups = {} + + # A Function that is only written to has a single buffer, shared by all of + # its writes regardless of their guards. Each slot is copied back right + # after being written, and the buffer is refilled only once the previous + # copy-back has completed, so all copy-backs into the Function are + # performed in program order + shared = {} + writeonly = {f for f, clusters in bfmap.items() + if all(f not in c.scope.reads for c in clusters)} + for f, clusters in bfmap.items(): dim = f.hdim for k, g in groupby(clusters, key=lambda c: c.guards): g = list(g) + if f in shared: + mapper[(f, k)] = shared[f] + groups[(f, k)] = g + continue + if k and not dim._defines & set(k): extras.setdefault(f, []).append((k, g)) continue @@ -462,9 +477,17 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs): ) groups[(f, k)] = g + if f in writeonly: + shared[f] = mapper[(f, k)] + # Alias deferred entries to a compatible f-buffer; create one if none fits for f, deferred in extras.items(): for k, g in deferred: + if f in shared: + mapper[(f, k)] = shared[f] + groups[(f, k)] = g + continue + # {candidate buffer -> Clusters already mapped to it} by_buf = defaultdict(list) for (ff, kk), b in mapper.items(): @@ -498,6 +521,8 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs): buf = _select_buffer( f, f.hdim, k, g, xds, async_degree, sregistry, callback ) + if f in writeonly: + shared[f] = buf mapper[(f, k)] = buf groups[(f, k)] = g @@ -1008,18 +1033,13 @@ def _overwritten(v, c, idxf, descriptors, order, d): 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. + both would be redundant and, for Functions storing each copy-back as a new + entry (e.g., compressed snapshots), would store the slot twice. """ 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 diff --git a/tests/test_buffering.py b/tests/test_buffering.py index 2a46203a85..29c864e848 100644 --- a/tests/test_buffering.py +++ b/tests/test_buffering.py @@ -1055,12 +1055,12 @@ def test_deferred_guard_aliasing(): assert np.allclose(f.data[:, 1, 1], 100. + np.arange(nt)) -def test_deferred_guard_no_fit(): +def test_deferred_guard_writeonly(): """ Like `test_deferred_guard_aliasing`, but the guarded write accesses - `f[time+1]`, which cannot be hosted by the main buffer (size 1, holding - `f[time]` only) without a slot collision; a dedicated buffer must be - created instead. + `f[time+1]`. Since `f` is only written to, the guarded write still shares + the main buffer: each slot is copied back right after being written, so + the two writes cannot collide. """ grid = Grid(shape=(5, 5)) nt = 5 @@ -1079,10 +1079,9 @@ def test_deferred_guard_no_fit(): op = Operator(eqs, opt='buffering') - # `f[time+1]` does not fit the main buffer's slots, so the guarded - # write gets a dedicated buffer + # A single buffer for the write-only `f` buffers = [i for i in FindSymbols().visit(op.body) if i.is_Array and i._mem_heap] - assert len(buffers) == 2 + assert len(buffers) == 1 op.apply(time_m=0, time_M=nt-2) # `f[t]` is overwritten by the main stream at iteration `t` for t < nt-1, @@ -1159,6 +1158,11 @@ def test_buffering_multi_cond(factor): # run operator with serialization op = Operator(eqs, opt='buffering') + + # Both writes to `f` share a single buffer + buffers = [i for i in FindSymbols().visit(op.body) if i.is_Array and i._mem_heap] + assert len(buffers) == 1 + op.apply(time_m=0, time_M=ntmod-2) # Now run backward as well with buffering