Merge pull request #2675 from devitocodes/revert-2669-cached-scopes

Revert "compiler: Cache `Scope` + `Dependence` instances"

Author: mloubout

devito/ir/clusters/algorithms.py

@@ -12,8 +12,7 @@
 from devito.ir.equations import OpMin, OpMax, identity_mapper
 from devito.ir.clusters.analysis import analyze
 from devito.ir.clusters.cluster import Cluster, ClusterGroup
-from devito.ir.clusters.visitors import Queue, cluster_pass
-from devito.ir.support import Scope
+from devito.ir.clusters.visitors import Queue, QueueStateful, cluster_pass
 from devito.mpi.halo_scheme import HaloScheme, HaloTouch
 from devito.mpi.reduction_scheme import DistReduce
 from devito.symbolics import (limits_mapper, retrieve_indexed, uxreplace,
@@ -78,7 +77,7 @@ def impose_total_ordering(clusters):
     return processed
 
 
-class Schedule(Queue):
+class Schedule(QueueStateful):
 
     """
     This special Queue produces a new sequence of "scheduled" Clusters, which
@@ -136,7 +135,7 @@ def callback(self, clusters, prefix, backlog=None, known_break=None):
         # `clusters` are supposed to share it
         candidates = prefix[-1].dim._defines
 
-        scope = Scope(flatten(c.exprs for c in clusters))
+        scope = self._fetch_scope(clusters)
 
         # Handle the nastiest case -- ambiguity due to the presence of both a
         # flow- and an anti-dependence.

devito/ir/clusters/analysis.py

@@ -1,39 +1,31 @@
-from devito.ir.clusters.cluster import Cluster
-from devito.ir.clusters.visitors import Queue
+from devito.ir.clusters.visitors import QueueStateful
 from devito.ir.support import (AFFINE, PARALLEL, PARALLEL_INDEP, PARALLEL_IF_ATOMIC,
-                               SEQUENTIAL, Property, Scope)
-from devito.ir.support.space import IterationSpace
+                               SEQUENTIAL)
 from devito.tools import as_tuple, flatten, timed_pass
-from devito.types.dimension import Dimension
 
 __all__ = ['analyze']
 
 
-# Describes properties fetched by a `Detector`
-Properties = dict[Cluster, dict[Dimension, set[Property]]]
-
-
 @timed_pass()
 def analyze(clusters):
-    properties: Properties = {}
+    state = QueueStateful.State()
 
     # Collect properties
-    clusters = Parallelism().process(clusters, properties=properties)
-    clusters = Affiness().process(clusters, properties=properties)
+    clusters = Parallelism(state).process(clusters)
+    clusters = Affiness(state).process(clusters)
 
     # Reconstruct Clusters attaching the discovered properties
-    processed = [c.rebuild(properties=properties.get(c)) for c in clusters]
+    processed = [c.rebuild(properties=state.properties.get(c)) for c in clusters]
 
     return processed
 
 
-class Detector(Queue):
+class Detector(QueueStateful):
 
-    def process(self, clusters: list[Cluster], properties: Properties) -> list[Cluster]:
-        return self._process_fatd(clusters, 1, properties=properties)
+    def process(self, elements):
+        return self._process_fatd(elements, 1)
 
-    def callback(self, clusters: list[Cluster], prefix: IterationSpace | None,
-                 properties: Properties) -> list[Cluster]:
+    def callback(self, clusters, prefix):
         if not prefix:
             return clusters
 
@@ -49,19 +41,11 @@ def callback(self, clusters: list[Cluster], prefix: IterationSpace | None,
         # Update `self.state`
         if retval:
             for c in clusters:
-                c_properties = properties.setdefault(c, {})
-                c_properties.setdefault(d, set()).update(retval)
+                properties = self.state.properties.setdefault(c, {})
+                properties.setdefault(d, set()).update(retval)
 
         return clusters
 
-    def _callback(self, clusters: list[Cluster], dim: Dimension,
-                  prefix: IterationSpace | None) -> set[Property]:
-        """
-        Callback to be implemented by subclasses. It should return a set of
-        properties for the given dimension.
-        """
-        raise NotImplementedError()
-
 
 class Parallelism(Detector):
 
@@ -88,27 +72,27 @@ class Parallelism(Detector):
             the 'write' is known to be an associative and commutative increment
     """
 
-    def _callback(self, clusters, dim, prefix):
+    def _callback(self, clusters, d, prefix):
         # Rule out if non-unitary increment Dimension (e.g., `t0=(time+1)%2`)
-        if any(c.sub_iterators[dim] for c in clusters):
-            return {SEQUENTIAL}
+        if any(c.sub_iterators[d] for c in clusters):
+            return SEQUENTIAL
 
         # All Dimensions up to and including `i-1`
         prev = flatten(i.dim._defines for i in prefix[:-1])
 
         is_parallel_indep = True
         is_parallel_atomic = False
 
-        scope = Scope(flatten(c.exprs for c in clusters))
+        scope = self._fetch_scope(clusters)
         for dep in scope.d_all_gen():
-            test00 = dep.is_indep(dim) and not dep.is_storage_related(dim)
+            test00 = dep.is_indep(d) and not dep.is_storage_related(d)
             test01 = all(dep.is_reduce_atmost(i) for i in prev)
             if test00 and test01:
                 continue
 
             test1 = len(prev) > 0 and any(dep.is_carried(i) for i in prev)
             if test1:
-                is_parallel_indep &= (dep.distance_mapper.get(dim.root) == 0)
+                is_parallel_indep &= (dep.distance_mapper.get(d.root) == 0)
                 continue
 
             if dep.function in scope.initialized:
@@ -119,14 +103,14 @@ def _callback(self, clusters, dim, prefix):
                 is_parallel_atomic = True
                 continue
 
-            return {SEQUENTIAL}
+            return SEQUENTIAL
 
         if is_parallel_atomic:
-            return {PARALLEL_IF_ATOMIC}
+            return PARALLEL_IF_ATOMIC
         elif is_parallel_indep:
             return {PARALLEL, PARALLEL_INDEP}
         else:
-            return {PARALLEL}
+            return PARALLEL
 
 
 class Affiness(Detector):
@@ -135,11 +119,8 @@ class Affiness(Detector):
     Detect the AFFINE Dimensions.
     """
 
-    def _callback(self, clusters, dim, prefix):
-        scope = Scope(flatten(c.exprs for c in clusters))
+    def _callback(self, clusters, d, prefix):
+        scope = self._fetch_scope(clusters)
         accesses = [a for a in scope.accesses if not a.is_scalar]
-
-        if all(a.is_regular and a.affine_if_present(dim._defines) for a in accesses):
-            return {AFFINE}
-
-        return set()
+        if all(a.is_regular and a.affine_if_present(d._defines) for a in accesses):
+            return AFFINE

devito/ir/clusters/visitors.py

@@ -1,11 +1,12 @@
+from collections import defaultdict
 from collections.abc import Iterable
 
 from itertools import groupby
 
-from devito.ir.support import IterationSpace, null_ispace
-from devito.tools import flatten, timed_pass
+from devito.ir.support import IterationSpace, Scope, null_ispace
+from devito.tools import as_tuple, flatten, timed_pass
 
-__all__ = ['Queue', 'cluster_pass']
+__all__ = ['Queue', 'QueueStateful', 'cluster_pass']
 
 
 class Queue:
@@ -112,6 +113,48 @@ def _process_fatd(self, clusters, level, prefix=None, **kwargs):
         return processed
 
 
+class QueueStateful(Queue):
+
+    """
+    A Queue carrying along some state. This is useful when one wants to avoid
+    expensive re-computations of information.
+    """
+
+    class State:
+
+        def __init__(self):
+            self.properties = {}
+            self.scopes = {}
+
+    def __init__(self, state=None):
+        super().__init__()
+        self.state = state or QueueStateful.State()
+
+    def _fetch_scope(self, clusters):
+        exprs = flatten(c.exprs for c in as_tuple(clusters))
+        key = tuple(exprs)
+        if key not in self.state.scopes:
+            self.state.scopes[key] = Scope(exprs)
+        return self.state.scopes[key]
+
+    def _fetch_properties(self, clusters, prefix):
+        # If the situation is:
+        #
+        # t
+        #   x0
+        #     <some clusters>
+        #   x1
+        #     <some other clusters>
+        #
+        # then retain only the "common" properties, that is those along `t`
+        properties = defaultdict(set)
+        for c in clusters:
+            v = self.state.properties.get(c, {})
+            for i in prefix:
+                properties[i.dim].update(v.get(i.dim, set()))
+        return properties
+
+
 class Prefix(IterationSpace):
 
     def __init__(self, ispace, guards, properties, syncs):

devito/ir/support/basic.py

@@ -1,9 +1,7 @@
-from collections.abc import Iterable
 from itertools import chain, product
 from functools import cached_property
-from typing import Callable
 
-from sympy import S, Expr
+from sympy import S
 import sympy
 
 from devito.ir.support.space import Backward, null_ispace
@@ -14,7 +12,7 @@
                               uxreplace)
 from devito.tools import (Tag, as_mapper, as_tuple, is_integer, filter_sorted,
                           flatten, memoized_meth, memoized_generator, smart_gt,
-                          smart_lt, CacheInstances)
+                          smart_lt)
 from devito.types import (ComponentAccess, Dimension, DimensionTuple, Fence,
                           CriticalRegion, Function, Symbol, Temp, TempArray,
                           TBArray)
@@ -248,7 +246,7 @@ def __eq__(self, other):
                 self.ispace == other.ispace)
 
     def __hash__(self):
-        return hash((self.access, self.mode, self.timestamp, self.ispace))
+        return super().__hash__()
 
     @property
     def function(self):
@@ -626,7 +624,7 @@ def is_imaginary(self):
         return S.ImaginaryUnit in self.distance
 
 
-class Dependence(Relation, CacheInstances):
+class Dependence(Relation):
 
     """
     A data dependence between two TimedAccess objects.
@@ -825,26 +823,17 @@ def project(self, function):
         return DependenceGroup(i for i in self if i.function is function)
 
 
-class Scope(CacheInstances):
-
-    # Describes a rule for dependencies
-    Rule = Callable[[TimedAccess, TimedAccess], bool]
+class Scope:
 
-    @classmethod
-    def _preprocess_args(cls, exprs: Expr | Iterable[Expr],
-                         **kwargs) -> tuple[tuple, dict]:
-        return (as_tuple(exprs),), kwargs
-
-    def __init__(self, exprs: tuple[Expr],
-                 rules: Rule | tuple[Rule] | None = None) -> None:
+    def __init__(self, exprs, rules=None):
         """
         A Scope enables data dependence analysis on a totally ordered sequence
         of expressions.
         """
-        self.exprs = exprs
+        self.exprs = as_tuple(exprs)
 
         # A set of rules to drive the collection of dependencies
-        self.rules: tuple[Scope.Rule] = as_tuple(rules)  # type: ignore[assignment]
+        self.rules = as_tuple(rules)
         assert all(callable(i) for i in self.rules)
 
     @memoized_generator
@@ -1183,10 +1172,12 @@ def d_from_access_gen(self, accesses):
         Generate all flow, anti, and output dependences involving any of
         the given TimedAccess objects.
         """
-        accesses = set(as_tuple(accesses))
+        accesses = as_tuple(accesses)
         for d in self.d_all_gen():
-            if accesses & {d.source, d.sink}:
-                yield d
+            for i in accesses:
+                if d.source == i or d.sink == i:
+                    yield d
+                    break
 
     @memoized_meth
     def d_from_access(self, accesses):

devito/operator/operator.py

@@ -31,8 +31,7 @@
 from devito.symbolics import estimate_cost, subs_op_args
 from devito.tools import (DAG, OrderedSet, Signer, ReducerMap, as_mapper, as_tuple,
                           flatten, filter_sorted, frozendict, is_integer,
-                          split, timed_pass, timed_region, contains_val,
-                          CacheInstances)
+                          split, timed_pass, timed_region, contains_val)
 from devito.types import (Buffer, Evaluable, host_layer, device_layer,
                           disk_layer)
 from devito.types.dimension import Thickness
@@ -246,9 +245,6 @@ def _build(cls, expressions, **kwargs):
         op._dtype, op._dspace = irs.clusters.meta
         op._profiler = profiler
 
-        # Clear build-scoped instance caches
-        CacheInstances.clear_caches()
-
         return op
 
     def __init__(self, *args, **kwargs):