compiler: Rework Scope + Dependence caching

Author: enwask

devito/ir/clusters/algorithms.py

@@ -136,7 +136,7 @@ def callback(self, clusters, prefix, backlog=None, known_break=None):
         # `clusters` are supposed to share it
         candidates = prefix[-1].dim._defines
 
-        scope = Scope.fetch(flatten(c.exprs for c in clusters))
+        scope = Scope(flatten(c.exprs for c in clusters))
 
         # Handle the nastiest case -- ambiguity due to the presence of both a
         # flow- and an anti-dependence.

devito/ir/clusters/analysis.py

@@ -99,7 +99,7 @@ def _callback(self, clusters, dim, prefix):
         is_parallel_indep = True
         is_parallel_atomic = False
 
-        scope = Scope.fetch(flatten(c.exprs for c in clusters))
+        scope = Scope(flatten(c.exprs for c in clusters))
         for dep in scope.d_all_gen():
             test00 = dep.is_indep(dim) and not dep.is_storage_related(dim)
             test01 = all(dep.is_reduce_atmost(i) for i in prev)
@@ -136,7 +136,7 @@ class Affiness(Detector):
     """
 
     def _callback(self, clusters, dim, prefix):
-        scope = Scope.fetch(flatten(c.exprs for c in clusters))
+        scope = Scope(flatten(c.exprs for c in 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):

devito/ir/clusters/cluster.py

@@ -183,7 +183,7 @@ def dist_dimensions(self):
 
     @cached_property
     def scope(self):
-        return Scope.fetch(self.exprs)
+        return Scope(self.exprs)
 
     @cached_property
     def functions(self):
@@ -473,7 +473,7 @@ def exprs(self):
 
     @cached_property
     def scope(self):
-        return Scope.fetch(exprs=self.exprs)
+        return Scope(exprs=self.exprs)
 
     @cached_property
     def ispace(self):

devito/ir/support/basic.py

@@ -1,11 +1,12 @@
 from collections.abc import Iterable
 from itertools import chain, product
-from functools import cached_property, lru_cache
+from functools import cached_property
 from typing import Callable
 
 from sympy import S, Expr
 import sympy
 
+from devito.ir.support.caching import CacheInstances
 from devito.ir.support.space import Backward, null_ispace
 from devito.ir.support.utils import AccessMode, extrema
 from devito.ir.support.vector import LabeledVector, Vector
@@ -626,21 +627,12 @@ def is_imaginary(self):
         return S.ImaginaryUnit in self.distance
 
 
-class Dependence(Relation):
+class Dependence(Relation, CacheInstances):
 
     """
     A data dependence between two TimedAccess objects.
     """
 
-    @classmethod
-    @lru_cache(maxsize=128)
-    def fetch(cls: type['Dependence'],
-              source: TimedAccess, sink: TimedAccess) -> 'Dependence':
-        """
-        Obtain a (potentially cached) Dependence for analysis.
-        """
-        return cls(source, sink)
-
     def __repr__(self):
         return "%s -> %s" % (self.source, self.sink)
 
@@ -834,14 +826,18 @@ def project(self, function):
         return DependenceGroup(i for i in self if i.function is function)
 
 
-class Scope:
+class Scope(CacheInstances):
 
     # Describes a rule for dependencies
     Rule = Callable[[TimedAccess, TimedAccess], bool]
 
+    @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:
+                 rules: Rule | tuple[Rule] | None = None) -> None:
         """
         A Scope enables data dependence analysis on a totally ordered sequence
         of expressions.
@@ -852,24 +848,6 @@ def __init__(self, exprs: tuple[Expr],
         self.rules: tuple[Scope.Rule] = as_tuple(rules)  # type: ignore[assignment]
         assert all(callable(i) for i in self.rules)
 
-    @classmethod
-    @lru_cache(maxsize=128)
-    def _fetch(cls: type['Scope'], exprs: tuple[Expr],
-               rules: Rule | tuple[Rule] | None = None) -> 'Scope':
-        """
-        Obtains a (potentially cached) Scope from a sequence of expressions.
-        Helper function called with hashable arguments.
-        """
-        return cls(exprs, rules=rules)
-
-    @classmethod
-    def fetch(cls: type['Scope'], exprs: Expr | Iterable[Expr],
-              rules: Rule | tuple[Rule] | None = None) -> 'Scope':
-        """
-        Obtains a (potentially cached) Scope from a sequence of expressions.
-        """
-        return cls._fetch(as_tuple(exprs), rules=rules)
-
     @memoized_generator
     def writes_gen(self):
         """
@@ -1111,7 +1089,7 @@ def d_flow_gen(self):
                     if any(not rule(w, r) for rule in self.rules):
                         continue
 
-                    dependence = Dependence.fetch(w, r)
+                    dependence = Dependence(w, r)
 
                     if dependence.is_imaginary:
                         continue
@@ -1141,7 +1119,7 @@ def d_anti_gen(self):
                     if any(not rule(r, w) for rule in self.rules):
                         continue
 
-                    dependence = Dependence.fetch(r, w)
+                    dependence = Dependence(r, w)
 
                     if dependence.is_imaginary:
                         continue
@@ -1171,7 +1149,7 @@ def d_output_gen(self):
                     if any(not rule(w2, w1) for rule in self.rules):
                         continue
 
-                    dependence = Dependence.fetch(w2, w1)
+                    dependence = Dependence(w2, w1)
 
                     if dependence.is_imaginary:
                         continue

devito/mpi/halo_scheme.py

@@ -519,7 +519,7 @@ def classify(exprs, ispace):
     # performed before (reads) or after (writes) the OWNED region is computed
     loc_indices_from_reads = configuration['mpi'] not in ('dual',)
 
-    scope = Scope.fetch(exprs)
+    scope = Scope(exprs)
 
     mapper = {}
     for f, r in scope.reads.items():

devito/operator/operator.py

@@ -13,7 +13,7 @@
 from devito.data import default_allocator
 from devito.exceptions import (CompilationError, ExecutionError, InvalidArgument,
                                InvalidOperator)
-from devito.ir.support.basic import Dependence, Scope
+from devito.ir.support.caching import CacheInstances
 from devito.logger import (debug, info, perf, warning, is_log_enabled_for,
                            switch_log_level)
 from devito.ir.equations import LoweredEq, lower_exprs, concretize_subdims
@@ -247,8 +247,7 @@ def _build(cls, expressions, **kwargs):
         op._profiler = profiler
 
         # Clear Scope + Dependence caches
-        Scope._fetch.cache_clear()
-        Dependence.fetch.cache_clear()
+        CacheInstances.clear_caches()
 
         return op
 

devito/passes/clusters/blocking.py

@@ -279,7 +279,7 @@ def callback(self, clusters, prefix):
         if len(clusters) > 1:
             # Heuristic: same as above if it induces dynamic bounds
             exprs = flatten(c.exprs for c in as_tuple(clusters))
-            scope = Scope.fetch(exprs)
+            scope = Scope(exprs)
             if any(i.is_lex_non_stmt for i in scope.d_all_gen()):
                 return clusters
         else:

devito/passes/clusters/cse.py

@@ -124,10 +124,10 @@ def _cse(maybe_exprs, make, min_cost=1, mode='basic'):
         maybe_exprs = as_list(maybe_exprs)
         if all(e.is_Equality for e in maybe_exprs):
             exprs = maybe_exprs
-            scope = Scope.fetch(maybe_exprs)
+            scope = Scope(maybe_exprs)
         else:
             exprs = [Eq(make(e), e) for e in maybe_exprs]
-            scope = Scope.fetch([])
+            scope = Scope([])
 
     # Some sub-expressions aren't really "common" -- that's the case of Dimension-
     # independent data dependences. For example:

devito/passes/clusters/misc.py

@@ -355,7 +355,7 @@ def is_cross(source, sink):
             for n1, cg1 in enumerate(cgroups[n+1:], start=n+1):
 
                 # A Scope to compute all cross-ClusterGroup anti-dependences
-                scope = Scope.fetch(exprs=cg0.exprs + cg1.exprs, rules=is_cross)
+                scope = Scope(exprs=cg0.exprs + cg1.exprs, rules=is_cross)
 
                 # Anti-dependences along `prefix` break the execution flow
                 # (intuitively, "the loop nests are to be kept separated")
@@ -444,7 +444,7 @@ def callback(self, clusters, prefix):
             return clusters
 
         # Analyze and abort if fissioning would break a dependence
-        scope = Scope.fetch(flatten(c.exprs for c in clusters))
+        scope = Scope(flatten(c.exprs for c in clusters))
         if any(d._defines & dep.cause or dep.is_reduce(d) for dep in scope.d_all_gen()):
             return clusters
 

devito/passes/iet/mpi.py

@@ -40,7 +40,7 @@ def _drop_reduction_halospots(iet):
 
     # If all HaloSpot reads pertain to reductions, then the HaloSpot is useless
     for hs, expressions in MapNodes(HaloSpot, Expression).visit(iet).items():
-        scope = Scope.fetch(i.expr for i in expressions)
+        scope = Scope(i.expr for i in expressions)
         for k, v in hs.fmapper.items():
             f = v.bundle or k
             if f not in scope.reads:
@@ -82,7 +82,7 @@ def _hoist_redundant_from_conditionals(iet):
 
     mapper = HaloSpotMapper()
     for it, halo_spots in iter_mapper.items():
-        scope = Scope.fetch(e.expr for e in FindNodes(Expression).visit(it))
+        scope = Scope(e.expr for e in FindNodes(Expression).visit(it))
 
         for hs0 in halo_spots:
             conditions = cond_mapper[hs0]
@@ -282,7 +282,7 @@ def _mark_overlappable(iet):
         if not expressions:
             continue
 
-        scope = Scope.fetch(i.expr for i in expressions)
+        scope = Scope(i.expr for i in expressions)
 
         # Comp/comm overlaps is legal only if the OWNED regions can grow
         # arbitrarly, which means all of the dependences must be carried
@@ -462,7 +462,7 @@ def _derive_scope(it, hs0, hs1):
     and ends at the HaloSpot `hs1`.
     """
     expressions = FindWithin(Expression, hs0, stop=hs1).visit(it)
-    return Scope.fetch(e.expr for e in expressions)
+    return Scope(e.expr for e in expressions)
 
 
 def _check_control_flow(hs0, hs1, cond_mapper):