diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index bba54fbafc..2f98413bf1 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -7,15 +7,15 @@ from devito.ir.equations import ClusterizedEq from devito.ir.support import ( PARALLEL, PARALLEL_IF_PVT, BaseGuardBoundNext, DataSpace, Forward, Guards, Interval, - IntervalGroup, IterationSpace, PrefetchUpdate, Properties, Scope, WaitLock, WithLock, - detect_accesses, maximum, minimum, normalize_properties, normalize_syncs, null_ispace, - tailor_properties, update_properties + IntervalGroup, IterationSpace, PrefetchUpdate, Properties, Scope, TimedAccess, + WaitLock, WithLock, detect_accesses, maximum, minimum, normalize_properties, + normalize_syncs, null_ispace, tailor_properties, update_properties ) from devito.mpi.halo_scheme import HaloScheme, HaloTouch from devito.mpi.reduction_scheme import DistReduce from devito.symbolics import estimate_cost, uxreplace from devito.tools import ( - CacheInstances, as_tuple, cached_hash, filter_ordered, flatten, infer_dtype + CacheInstances, as_mapper, as_tuple, cached_hash, filter_ordered, flatten, infer_dtype ) from devito.types import ( CriticalRegion, Fence, Indexed, PhaseMarker, TensorMove, ThreadArrive, ThreadCommit, @@ -157,7 +157,8 @@ def local_sums(self): @cached_property def scope(self): - return Scope(self.exprs) + reads, writes = extract_accesses(self) + return Scope(self.exprs, reads=reads, writes=writes) @cached_property def functions(self): @@ -687,7 +688,8 @@ def local_sums(self): @cached_property def scope(self): - return Scope(exprs=self.exprs) + reads, writes = extract_accesses(self) + return Scope(exprs=self.exprs, reads=reads, writes=writes) @cached_property def functions(self): @@ -764,6 +766,23 @@ def meta(self): # *** Utils +def extract_accesses(cluster): + """ + The accesses performed by `cluster` on top of those of its expressions, + e.g. due to its synchronization operations. A wait, e.g. for a prefetch to + complete, is an index-free write of its target, so that no access to the + target may float above it upon topological sorting. + + Return the reads and writes as mappers from Functions to TimedAccesses, to + be fed to a Scope. + """ + accesses = [TimedAccess(s.target.indexed, 'W', 0, cluster.ispace) + for s in flatten(cluster.syncs.values()) if isinstance(s, WaitLock)] + + return (as_mapper([a for a in accesses if a.is_read], key=lambda i: i.function), + as_mapper([a for a in accesses if a.is_write], key=lambda i: i.function)) + + def reduce_properties(clusters): """ The normalized intersection (basically, a reduction) of the Properties in diff --git a/devito/ir/support/basic.py b/devito/ir/support/basic.py index 0a0d092a55..54091cfef8 100644 --- a/devito/ir/support/basic.py +++ b/devito/ir/support/basic.py @@ -1008,14 +1008,14 @@ def is_cross(source, sink): if accesses: writes[f] = accesses - return cls((), rules=is_cross, reads=reads.items(), writes=writes.items()) + return cls((), rules=is_cross, reads=reads, writes=writes) @classmethod def _preprocess_args(cls, exprs: Expr | Iterable[Expr], **kwargs) -> tuple[tuple, dict]: for i in ('reads', 'writes'): with suppress(KeyError): - kwargs[i] = tuple(kwargs[i]) + kwargs[i] = tuple((k, as_tuple(v)) for k, v in dict(kwargs[i]).items()) return (as_tuple(exprs),), kwargs @@ -1024,11 +1024,12 @@ def __init__(self, exprs: tuple[Expr], reads=None, writes=None) -> None: """ A Scope enables data dependence analysis on a totally ordered sequence - of expressions. + of expressions. `reads` and `writes`, mappers from objects to + TimedAccesses, are accesses performed on top of those in `exprs`. """ self.exprs = exprs - self._reads = dict(reads) if reads is not None else None - self._writes = dict(writes) if writes is not None else None + self._reads = dict(reads or ()) + self._writes = dict(writes or ()) # A set of rules to drive the collection of dependencies self.rules: tuple[Scope.Rule] = as_tuple(rules) # type: ignore[assignment] @@ -1067,10 +1068,8 @@ def writes(self): """ Create a mapper from functions to write accesses. """ - if self._writes is not None: - return self._writes - - return as_mapper(self.writes_gen(), key=lambda i: i.function) + return as_mapper(chain(self.writes_gen(), *self._writes.values()), + key=lambda i: i.function) @cached_property def writes_tensor(self): @@ -1209,10 +1208,8 @@ def reads(self): """ Create a mapper from functions to read accesses. """ - if self._reads is not None: - return self._reads - - return as_mapper(self.reads_gen(), key=lambda i: i.function) + return as_mapper(chain(self.reads_gen(), *self._reads.values()), + key=lambda i: i.function) @cached_property def read_only(self): diff --git a/tests/test_operator.py b/tests/test_operator.py index e6fc86463a..3f6e8e6d96 100644 --- a/tests/test_operator.py +++ b/tests/test_operator.py @@ -25,18 +25,21 @@ from devito.arch.archinfo import Device from devito.exceptions import InvalidOperator from devito.finite_differences.differentiable import diff2sympy -from devito.ir.equations import ClusterizedEq +from devito.ir.clusters import Cluster +from devito.ir.equations import ClusterizedEq, LoweredEq from devito.ir.equations.algorithms import lower_exprs from devito.ir.iet import ( Callable, Conditional, Expression, FindNodes, FindSymbols, IsPerfectIteration, Iteration, TimedList, retrieve_iteration_tree ) -from devito.ir.support import Any, Backward, Forward +from devito.ir.support import Any, Backward, Forward, WaitLock +from devito.passes.clusters import fuse from devito.passes.iet.languages.C import CDataManager from devito.symbolics import ListInitializer, indexify, retrieve_indexed from devito.tools import as_fp64_decimal, flatten, powerset, timed_region from devito.types import ( - Array, Barrier, ConditionalDimension, CustomDimension, Indirection, Scalar, Symbol + Array, Barrier, ConditionalDimension, CustomDimension, Indirection, Lock, Scalar, + Symbol ) from devito.warnings import DevitoWarning @@ -2255,6 +2258,33 @@ def test_topofuse_preserves_nonadjacent_fission(self): assert_structure(op, ['x,y', 'x', 'x,y']) + def test_wait_halts_topofuse(self): + """ + No reader of a Function is scheduled ahead of a wait on it, e.g. for its + prefetch to complete, even to fuse with a Cluster with the same guards. + """ + grid = Grid(shape=(4, 4)) + x, y = grid.dimensions + + a, b, f, g = [Function(name=i, grid=grid) for i in 'abfg'] + lock = Lock(name='lock', dimensions=CustomDimension(name='ld', symbolic_size=1)) + + exprs = [LoweredEq(i) for i in lower_exprs([Eq(a, 1), Eq(b, f), Eq(g, f)])] + c0, c1, c2 = [Cluster(i, i.ispace) for i in exprs] + + # `c2` has the same guards as `c0`, so it would be scheduled right after + # it, ahead of `c1` and its wait on `f` + guards = {y: Gt(x, 1)} + c0 = c0.rebuild(guards=guards) + c1 = c1.rebuild(syncs={y: (WaitLock(lock[0], f),)}) + c2 = c2.rebuild(guards=guards) + + with timed_region('x'): + clusters = fuse([c0, c1, c2], toposort=True) + + readers = [c for c in clusters if f in c.scope.reads] + assert readers[0].syncs == c1.syncs + def test_nonadjacent_constraints_preserve_directions(self): """ Constraints inferred from non-adjacent Clusters must not replace the