compiler: Cache Scopes + Dependences, remove QueueStateful

Author: enwask

devito/ir/clusters/algorithms.py

@@ -12,7 +12,8 @@
 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, QueueStateful, cluster_pass
+from devito.ir.clusters.visitors import Queue, cluster_pass
+from devito.ir.support import Scope
 from devito.mpi.halo_scheme import HaloScheme, HaloTouch
 from devito.mpi.reduction_scheme import DistReduce
 from devito.symbolics import (limits_mapper, retrieve_indexed, uxreplace,
@@ -77,7 +78,7 @@ def impose_total_ordering(clusters):
     return processed
 
 
-class Schedule(QueueStateful):
+class Schedule(Queue):
 
     """
     This special Queue produces a new sequence of "scheduled" Clusters, which
@@ -135,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 = self._fetch_scope(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

@@ -1,31 +1,39 @@
-from devito.ir.clusters.visitors import QueueStateful
+from devito.ir.clusters.cluster import Cluster
+from devito.ir.clusters.visitors import Queue
 from devito.ir.support import (AFFINE, PARALLEL, PARALLEL_INDEP, PARALLEL_IF_ATOMIC,
-                               SEQUENTIAL)
+                               SEQUENTIAL, Property, Scope)
+from devito.ir.support.space import IterationSpace
 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):
-    state = QueueStateful.State()
+    properties: Properties = {}
 
     # Collect properties
-    clusters = Parallelism(state).process(clusters)
-    clusters = Affiness(state).process(clusters)
+    clusters = Parallelism().process(clusters, properties=properties)
+    clusters = Affiness().process(clusters, properties=properties)
 
     # Reconstruct Clusters attaching the discovered properties
-    processed = [c.rebuild(properties=state.properties.get(c)) for c in clusters]
+    processed = [c.rebuild(properties=properties.get(c)) for c in clusters]
 
     return processed
 
 
-class Detector(QueueStateful):
+class Detector(Queue):
 
-    def process(self, elements):
-        return self._process_fatd(elements, 1)
+    def process(self, clusters: list[Cluster], properties: Properties) -> list[Cluster]:
+        return self._process_fatd(clusters, 1, properties=properties)
 
-    def callback(self, clusters, prefix):
+    def callback(self, clusters: list[Cluster], prefix: IterationSpace | None,
+                 properties: Properties) -> list[Cluster]:
         if not prefix:
             return clusters
 
@@ -41,11 +49,19 @@ def callback(self, clusters, prefix):
         # Update `self.state`
         if retval:
             for c in clusters:
-                properties = self.state.properties.setdefault(c, {})
-                properties.setdefault(d, set()).update(retval)
+                c_properties = properties.setdefault(c, {})
+                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):
 
@@ -72,27 +88,27 @@ class Parallelism(Detector):
             the 'write' is known to be an associative and commutative increment
     """
 
-    def _callback(self, clusters, d, prefix):
+    def _callback(self, clusters, dim, prefix):
         # Rule out if non-unitary increment Dimension (e.g., `t0=(time+1)%2`)
-        if any(c.sub_iterators[d] for c in clusters):
-            return SEQUENTIAL
+        if any(c.sub_iterators[dim] 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 = self._fetch_scope(clusters)
+        scope = Scope(flatten(c.exprs for c in clusters))
         for dep in scope.d_all_gen():
-            test00 = dep.is_indep(d) and not dep.is_storage_related(d)
+            test00 = dep.is_indep(dim) and not dep.is_storage_related(dim)
             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(d.root) == 0)
+                is_parallel_indep &= (dep.distance_mapper.get(dim.root) == 0)
                 continue
 
             if dep.function in scope.initialized:
@@ -103,14 +119,14 @@ def _callback(self, clusters, d, 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):
@@ -119,8 +135,11 @@ class Affiness(Detector):
     Detect the AFFINE Dimensions.
     """
 
-    def _callback(self, clusters, d, prefix):
-        scope = self._fetch_scope(clusters)
+    def _callback(self, clusters, dim, prefix):
+        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(d._defines) for a in accesses):
-            return AFFINE
+
+        if all(a.is_regular and a.affine_if_present(dim._defines) for a in accesses):
+            return {AFFINE}
+
+        return set()

devito/ir/clusters/visitors.py

@@ -1,12 +1,11 @@
-from collections import defaultdict
 from collections.abc import Iterable
 
 from itertools import groupby
 
-from devito.ir.support import IterationSpace, Scope, null_ispace
-from devito.tools import as_tuple, flatten, timed_pass
+from devito.ir.support import IterationSpace, null_ispace
+from devito.tools import flatten, timed_pass
 
-__all__ = ['Queue', 'QueueStateful', 'cluster_pass']
+__all__ = ['Queue', 'cluster_pass']
 
 
 class Queue:
@@ -113,48 +112,6 @@ 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,7 +1,9 @@
+from collections.abc import Iterable
 from itertools import chain, product
 from functools import cached_property
+from typing import Callable
 
-from sympy import S
+from sympy import S, Expr
 import sympy
 
 from devito.ir.support.space import Backward, null_ispace
@@ -12,7 +14,7 @@
                               uxreplace)
 from devito.tools import (Tag, as_mapper, as_tuple, is_integer, filter_sorted,
                           flatten, memoized_meth, memoized_generator, smart_gt,
-                          smart_lt)
+                          smart_lt, CacheInstances)
 from devito.types import (ComponentAccess, Dimension, DimensionTuple, Fence,
                           CriticalRegion, Function, Symbol, Temp, TempArray,
                           TBArray)
@@ -624,7 +626,7 @@ def is_imaginary(self):
         return S.ImaginaryUnit in self.distance
 
 
-class Dependence(Relation):
+class Dependence(Relation, CacheInstances):
 
     """
     A data dependence between two TimedAccess objects.
@@ -823,17 +825,26 @@ def project(self, function):
         return DependenceGroup(i for i in self if i.function is function)
 
 
-class Scope:
+class Scope(CacheInstances):
 
-    def __init__(self, exprs, rules=None):
+    # 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:
         """
         A Scope enables data dependence analysis on a totally ordered sequence
         of expressions.
         """
-        self.exprs = as_tuple(exprs)
+        self.exprs = exprs
 
         # A set of rules to drive the collection of dependencies
-        self.rules = as_tuple(rules)
+        self.rules: tuple[Scope.Rule] = as_tuple(rules)  # type: ignore[assignment]
         assert all(callable(i) for i in self.rules)
 
     @memoized_generator

devito/operator/operator.py

@@ -31,7 +31,8 @@
 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)
+                          split, timed_pass, timed_region, contains_val,
+                          CacheInstances)
 from devito.types import (Buffer, Evaluable, host_layer, device_layer,
                           disk_layer)
 from devito.types.dimension import Thickness
@@ -245,6 +246,9 @@ 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):

devito/tools/memoization.py

@@ -1,8 +1,9 @@
-from collections.abc import Hashable
-from functools import partial
+from collections.abc import Callable, Hashable
+from functools import lru_cache, partial
 from itertools import tee
+from typing import TypeVar
 
-__all__ = ['memoized_func', 'memoized_meth', 'memoized_generator']
+__all__ = ['memoized_func', 'memoized_meth', 'memoized_generator', 'CacheInstances']
 
 
 class memoized_func:
@@ -125,3 +126,63 @@ def __call__(self, *args, **kwargs):
         it = cache[key] if key in cache else self.func(*args, **kwargs)
         cache[key], result = tee(it)
         return result
+
+
+# Describes the type of a subclass of CacheInstances
+InstanceType = TypeVar('InstanceType', bound='CacheInstances', covariant=True)
+
+
+class CacheInstancesMeta(type):
+    """
+    Metaclass to wrap construction in an LRU cache.
+    """
+
+    _cached_types: set[type['CacheInstances']] = set()
+
+    def __init__(cls: type[InstanceType], *args) -> None:  # type: ignore
+        super().__init__(*args)
+
+        # Register the cached type
+        CacheInstancesMeta._cached_types.add(cls)
+
+    def __call__(cls: type[InstanceType],  # type: ignore
+                 *args, **kwargs) -> InstanceType:
+        if cls._instance_cache is None:
+            maxsize = cls._instance_cache_size
+            cls._instance_cache = lru_cache(maxsize=maxsize)(super().__call__)
+
+        args, kwargs = cls._preprocess_args(*args, **kwargs)
+        return cls._instance_cache(*args, **kwargs)
+
+    @classmethod
+    def clear_caches(cls: type['CacheInstancesMeta']) -> None:
+        """
+        Clear all caches for classes using this metaclass.
+        """
+        for cached_type in cls._cached_types:
+            if cached_type._instance_cache is not None:
+                cached_type._instance_cache.cache_clear()
+
+
+class CacheInstances(metaclass=CacheInstancesMeta):
+    """
+    Parent class that wraps construction in an LRU cache.
+    """
+
+    _instance_cache: Callable | None = None
+    _instance_cache_size: int = 128
+
+    @classmethod
+    def _preprocess_args(cls, *args, **kwargs):
+        """
+        Preprocess the arguments before caching. This can be overridden in subclasses
+        to customize argument handling (e.g. to convert to hashable types).
+        """
+        return args, kwargs
+
+    @staticmethod
+    def clear_caches() -> None:
+        """
+        Clears all IR instance caches.
+        """
+        CacheInstancesMeta.clear_caches()