compiler: Schedule the readers of a waited-on Function after the wait

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

devito/ir/support/basic.py

@@ -999,14 +999,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
 
@@ -1015,11 +1015,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]
@@ -1058,10 +1059,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):
@@ -1200,10 +1199,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):

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