|
3 | 3 | from itertools import chain, groupby |
4 | 4 |
|
5 | 5 | import numpy as np |
| 6 | +from sympy import And |
6 | 7 | from sympy import Eq as SympyEq |
7 | | -from sympy import Ne, S, simplify |
| 8 | +from sympy import Ne, Not, Or, S, simplify, true |
8 | 9 |
|
9 | 10 | from devito.exceptions import CompilationError |
10 | 11 | from devito.ir import ( |
|
13 | 14 | ) |
14 | 15 | from devito.logger import warning |
15 | 16 | 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 |
17 | 18 | from devito.tools import ( |
18 | 19 | Stamp, as_mapper, as_tuple, filter_ordered, flatten, frozendict, is_integer, |
19 | 20 | timed_pass |
@@ -184,6 +185,10 @@ def callback(self, clusters, prefix): |
184 | 185 | # E.g., `usave[time+1, x+1, y+1] -> ub0[t1, x+1, y+1]` |
185 | 186 | subs = make_buffer_subs(descriptors, mds, d) |
186 | 187 |
|
| 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 | + |
187 | 192 | processed = [] |
188 | 193 | for c in clusters: |
189 | 194 | # If a buffer is read but never written, then we need to add |
@@ -259,6 +264,15 @@ def callback(self, clusters, prefix): |
259 | 264 | else: |
260 | 265 | guards = c.guards |
261 | 266 |
|
| 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 | + |
262 | 276 | properties = c.properties.sequentialize(d) |
263 | 277 |
|
264 | 278 | syncs = c.syncs |
@@ -985,6 +999,58 @@ def _underlying_buffer(clusters): |
985 | 999 | return None |
986 | 1000 |
|
987 | 1001 |
|
| 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 | + |
988 | 1054 | def _fits_buffer(f, clusters, b): |
989 | 1055 | """ |
990 | 1056 | True if the accesses to `f` in the given Clusters span no more slots than |
|
0 commit comments