Skip to content

Commit c6a977c

Browse files
authored
Merge pull request #3050 from devitocodes/fix-buffering-same-slot
compiler: Skip copy-backs overwritten by a later one to the same slot
2 parents 6bb6804 + 23502c0 commit c6a977c

4 files changed

Lines changed: 147 additions & 6 deletions

File tree

‎devito/passes/clusters/buffering.py‎

Lines changed: 68 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,9 @@
33
from itertools import chain, groupby
44

55
import numpy as np
6+
from sympy import And
67
from sympy import Eq as SympyEq
7-
from sympy import Ne, S, simplify
8+
from sympy import Ne, Not, Or, S, simplify, true
89

910
from devito.exceptions import CompilationError
1011
from devito.ir import (
@@ -13,7 +14,7 @@
1314
)
1415
from devito.logger import warning
1516
from devito.passes.clusters.utils import is_memcpy
16-
from devito.symbolics import IntDiv, retrieve_functions, uxreplace
17+
from devito.symbolics import CondEq, IntDiv, retrieve_functions, uxreplace
1718
from devito.tools import (
1819
Stamp, as_mapper, as_tuple, filter_ordered, flatten, frozendict, is_integer,
1920
timed_pass
@@ -184,6 +185,10 @@ def callback(self, clusters, prefix):
184185
# E.g., `usave[time+1, x+1, y+1] -> ub0[t1, x+1, y+1]`
185186
subs = make_buffer_subs(descriptors, mds, d)
186187

188+
# Program order of the Clusters, to spot copy-backs into a slot that a
189+
# later copy-back in the same iteration overwrites
190+
order = {c: i for i, c in enumerate(clusters)}
191+
187192
processed = []
188193
for c in clusters:
189194
# If a buffer is read but never written, then we need to add
@@ -259,6 +264,15 @@ def callback(self, clusters, prefix):
259264
else:
260265
guards = c.guards
261266

267+
# Skip the copy-back when a later one in the same iteration
268+
# writes the same slot, so that a single (possibly asynchronous)
269+
# copy-back per slot is performed, carrying the last value
270+
overwritten = _overwritten(v, c, idxf, descriptors, order, d)
271+
if overwritten is true:
272+
continue
273+
elif overwritten is not None:
274+
guards = guards.xandg(d, Not(overwritten))
275+
262276
properties = c.properties.sequentialize(d)
263277

264278
syncs = c.syncs
@@ -985,6 +999,58 @@ def _underlying_buffer(clusters):
985999
return None
9861000

9871001

1002+
def _overwritten(v, c, idxf, descriptors, order, d):
1003+
"""
1004+
The condition under which the copy-back of `v` following the Cluster `c`,
1005+
into the slot `idxf` of `v.f`, is overwritten by a later copy-back into the
1006+
same slot within the same iteration of `d`. None if it is never overwritten.
1007+
1008+
For example, given a subsampled save `Eq(usave, u)` and a save of the last
1009+
timestep `Eq(usave, u, implicit_dims=clast)`, both copy-backs target the same
1010+
slot whenever the last timestep falls onto the subsampled grid. Performing
1011+
both would be redundant and, with asynchronous copy-backs (e.g. to disk),
1012+
racy, as the two writes into the slot would be performed concurrently.
1013+
"""
1014+
conditions = []
1015+
for v1 in descriptors.flat_values():
1016+
if v1.f is not v.f or v1.is_readonly:
1017+
continue
1018+
# Copy-backs out of the same buffer are already serialized, since the
1019+
# buffer cannot be refilled before its previous copy-back is completed
1020+
if v1.b is v.b:
1021+
continue
1022+
1023+
for c1 in v1.lastwrite:
1024+
if order.get(c1, -1) <= order[c]:
1025+
continue
1026+
if c1.guards.get(d) != v1.guards.get(d):
1027+
continue
1028+
# The later copy-back must write the whole slot written by `v`
1029+
key = lambda i: i in v.bdims # noqa: B023
1030+
if c1.ispace.project(key) != c.ispace.project(key):
1031+
continue
1032+
# A guard over other Dimensions may restrict the later copy-back to
1033+
# a portion of the slot only
1034+
if any(not dd._defines & d._defines for dd in c1.guards):
1035+
continue
1036+
1037+
shift = v1.last_idx[c1] - idxf
1038+
if shift.is_Number:
1039+
if shift != 0:
1040+
continue
1041+
coincide = true
1042+
else:
1043+
coincide = CondEq(shift, 0)
1044+
1045+
# The later copy-back fires under all of its guards, which may be
1046+
# keyed on Dimensions derived from `d` (e.g., a layer Dimension)
1047+
conditions.append(And(*c1.guards.values(), coincide))
1048+
1049+
if not conditions:
1050+
return None
1051+
return Or(*conditions)
1052+
1053+
9881054
def _fits_buffer(f, clusters, b):
9891055
"""
9901056
True if the accesses to `f` in the given Clusters span no more slots than

‎devito/passes/clusters/cse.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
from devito.finite_differences.differentiable import IndexSum
1515
from devito.ir import Cluster, Scope, cluster_pass
1616
from devito.symbolics import (
17-
DefFunction, Reserved, estimate_cost, q_leaf, q_terminal, search
17+
DefFunction, IndexedPointer, Reserved, estimate_cost, q_leaf, q_terminal, search
1818
)
1919
from devito.symbolics.manipulation import _uxreplace
2020
from devito.tools import DAG, as_list, as_tuple, extract_dtype, frozendict
@@ -418,11 +418,17 @@ def _(expr):
418418

419419

420420
@_catch.register(Indexed)
421+
@_catch.register(IndexedPointer)
421422
@_catch.register(Symbol)
422423
@_catch.register(Reserved)
423424
def _(expr):
424425
"""
425426
Handler for objects preventing CSE to propagate through their arguments.
427+
428+
For IndexedPointer (e.g. `f_vec->sizes[(time / ct0f) + ...]`) this also
429+
keeps integer index expressions out of a floating-point cluster's CSE
430+
temporaries, which `cse_dtype` would otherwise declare as floating-point
431+
and thus produce non-integral array subscripts.
426432
"""
427433
return {}
428434

‎tests/test_buffering.py‎

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import numpy as np
22
import pytest
3-
from sympy import Or
3+
from sympy import Not, Or
44

55
from conftest import skipif
66
from devito import (
@@ -11,7 +11,8 @@
1111
from devito.arch.archinfo import AppleArm
1212
from devito.exceptions import CompilationError
1313
from devito.ir import (
14-
Cluster, FindSymbols, Interval, IterationSpace, lower_exprs, retrieve_iteration_tree
14+
Cluster, Conditional, FindNodes, FindSymbols, Interval, IterationSpace, lower_exprs,
15+
retrieve_iteration_tree
1516
)
1617
from devito.passes.clusters.buffering import BufferDimension, expand_halo_transfers
1718
from devito.types import Array
@@ -1090,6 +1091,47 @@ def test_deferred_guard_no_fit():
10901091
assert np.allclose(f.data[:, 1, 1], expected)
10911092

10921093

1094+
@pytest.mark.parametrize("factor", [1, 2, 3])
1095+
def test_buffering_same_slot_single_copyback(factor):
1096+
"""
1097+
A subsampled save and a CondEq save into the same Function land on the same
1098+
slot when the CondEq point is on the subsampled grid. Only one copy-back,
1099+
carrying the last value, must then be performed: the copy-back of the first
1100+
write is skipped whenever the second one overwrites its slot. Asynchronous
1101+
copy-backs (e.g., to disk) would otherwise write the slot concurrently.
1102+
"""
1103+
grid = Grid((4, 4))
1104+
1105+
nt = 5
1106+
ntmod = (nt - 1) * factor + 1
1107+
last = ntmod - 2
1108+
on_grid = last % factor == 0
1109+
nsave = nt - 1 if on_grid else nt
1110+
1111+
ct0 = ConditionalDimension(name="ct0", parent=grid.time_dim, factor=factor,
1112+
relation=Or)
1113+
ctend = ConditionalDimension(name="ctend", parent=grid.time_dim,
1114+
condition=CondEq(grid.time_dim, last),
1115+
relation='strict')
1116+
f = TimeFunction(grid=grid, name='f', time_order=0, space_order=0,
1117+
time_dim=ct0, save=nsave)
1118+
T = TimeFunction(grid=grid, name='T', time_order=0, space_order=0)
1119+
1120+
eqs = [Eq(T, grid.time_dim), Eq(f, T), Eq(f, T+1, implicit_dims=ctend)]
1121+
1122+
op = Operator(eqs, opt='buffering')
1123+
1124+
# The first copy-back is skipped when the second one targets the same slot
1125+
skips = [i for i in FindNodes(Conditional).visit(op) if i.condition.has(Not)]
1126+
assert len(skips) == 1
1127+
1128+
op.apply(time_m=0, time_M=last)
1129+
1130+
expected = factor * np.arange(nsave)
1131+
expected[-1] = last + 1
1132+
assert np.allclose(f.data[:, 1, 1], expected)
1133+
1134+
10931135
@pytest.mark.parametrize("factor", [1, 2, 3])
10941136
def test_buffering_multi_cond(factor):
10951137
grid = Grid((16, 16))

‎tests/test_cse.py‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
from devito.ir import Conditional, DummyEq, FindNodes, FindSymbols
1212
from devito.ir.support import generator
1313
from devito.passes.clusters.cse import CTemp, _cse
14-
from devito.symbolics import DefFunction, indexify
14+
from devito.symbolics import DefFunction, FieldFromPointer, IndexedPointer, indexify
1515
from devito.types import Array, Symbol, Temp
1616

1717

@@ -154,6 +154,33 @@ def test_deffunction_not_captured():
154154
assert processed == exprs
155155

156156

157+
def test_indexedpointer_index_not_captured():
158+
"""
159+
The index of an IndexedPointer (e.g. `f_vec->sizes[...]`) must not be
160+
CSE-captured: in a floating-point cluster `cse_dtype` would bind the
161+
integer index to a floating-point temporary, yielding a non-integral
162+
array subscript that does not compile (nvcc, hipcc).
163+
"""
164+
time = Symbol(name='time', dtype=np.int32)
165+
fac = Symbol(name='fac', dtype=np.int32)
166+
vec = Symbol(name='f_vec', dtype=np.int32)
167+
a = CTemp(name='a', dtype=np.float32)
168+
b = CTemp(name='b', dtype=np.float32)
169+
170+
index = time/fac + 2*time
171+
sizes = IndexedPointer(FieldFromPointer('sizes', vec), index)
172+
offs = IndexedPointer(FieldFromPointer('offs', vec), index)
173+
exprs = [DummyEq(a, sizes + 1.), DummyEq(b, offs + sizes)]
174+
175+
counter = generator()
176+
make = lambda e: CTemp(name=f'r{counter()}', dtype=np.float32)
177+
processed = _cse(exprs, make)
178+
179+
for e in processed:
180+
for i in e.find(IndexedPointer):
181+
assert not any(isinstance(s, CTemp) for s in i.index.free_symbols)
182+
183+
157184
def test_w_conditionals():
158185
grid = Grid(shape=(10, 10, 10))
159186
x, _, _ = grid.dimensions

0 commit comments

Comments
 (0)