misc: Move IR caches to tool/memoization

Author: enwask

devito/ir/support/__init__.py

@@ -1,7 +1,6 @@
 from .vector import *  # noqa
 from .utils import *  # noqa
 from .basic import *  # noqa
-from .caching import *  # noqa
 from .space import *  # noqa
 from .guards import *  # noqa
 from .syncs import *  # noqa

devito/ir/support/basic.py

@@ -6,7 +6,6 @@
 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
@@ -15,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)

devito/ir/support/caching.py

@@ -13,7 +13,6 @@
 from devito.data import default_allocator
 from devito.exceptions import (CompilationError, ExecutionError, InvalidArgument,
                                InvalidOperator)
-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
@@ -32,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
@@ -246,7 +246,7 @@ def _build(cls, expressions, **kwargs):
         op._dtype, op._dspace = irs.clusters.meta
         op._profiler = profiler
 
-        # Clear Scope + Dependence caches
+        # Clear build-scoped instance caches
         CacheInstances.clear_caches()
 
         return op

devito/operator/operator.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()

devito/tools/memoization.py

@@ -12,7 +12,6 @@
 from devito.ir.support.basic import (IterationInstance, TimedAccess, Scope,
                                      Vector, AFFINE, REGULAR, IRREGULAR, mocksym0,
                                      mocksym1)
-from devito.ir.support.caching import CacheInstances
 from devito.ir.support.space import (NullInterval, Interval, Forward, Backward,
                                      IntervalGroup, IterationSpace)
 from devito.ir.support.guards import GuardOverflow
@@ -1102,65 +1101,3 @@ def test_guard_overflow(self):
         guard = GuardOverflow(freespace, size)
 
         assert ccode(guard) == 'freespace >= f_vec->size[0]*f_vec->size[1]'
-
-
-class TestCacheInstances:
-
-    def test_caching(self):
-        """
-        Tests basic functionality of cached instances.
-        """
-        class Object(CacheInstances):
-            def __init__(self, value: int):
-                self.value = value
-
-        obj1 = Object(1)
-        obj2 = Object(1)
-        obj3 = Object(2)
-
-        assert obj1 is obj2
-        assert obj1 is not obj3
-
-    def test_cache_size(self):
-        """
-        Tests specifying the size of the instance cache.
-        """
-        class Object(CacheInstances):
-            _instance_cache_size = 2
-
-            def __init__(self, value: int):
-                self.value = value
-
-        obj1 = Object(1)
-        obj2 = Object(2)
-        obj3 = Object(3)
-        obj4 = Object(1)
-        obj5 = Object(3)
-
-        # obj1 should have been evicted before obj4 was created
-        assert obj1 is not obj4
-        assert obj1 is not obj2
-        assert obj3 is obj5
-
-        hits, _, _, cursize = Object._instance_cache.cache_info()
-        assert hits == 1  # obj5 hit the cache
-        assert cursize == 2
-
-    def test_cleared_after_build(self):
-        """
-        Tests that instance caches are cleared after building an Operator.
-        """
-        class Object(CacheInstances):
-            def __init__(self, value: int):
-                self.value = value
-
-        obj1 = Object(1)
-        cache_size = Object._instance_cache.cache_info()[-1]
-        assert cache_size == 1
-
-        x = Symbol('x')
-        Operator(Eq(x, obj1.value))
-
-        # Cache should be cleared after Operator construction
-        cache_size = Object._instance_cache.cache_info()[-1]
-        assert cache_size == 0

tests/test_ir.py

@@ -4,8 +4,10 @@
 
 import time
 
+from devito import Operator, Eq
 from devito.tools import (UnboundedMultiTuple, ctypes_to_cstr, toposort,
-                          filter_ordered, transitive_closure, UnboundTuple)
+                          filter_ordered, transitive_closure, UnboundTuple,
+                          CacheInstances)
 from devito.types.basic import Symbol
 
 
@@ -145,3 +147,65 @@ def test_unbound_tuple():
     assert ub.next() == 2
     ub.iter()
     assert ub.next() == 1
+
+
+class TestCacheInstances:
+
+    def test_caching(self):
+        """
+        Tests basic functionality of cached instances.
+        """
+        class Object(CacheInstances):
+            def __init__(self, value: int):
+                self.value = value
+
+        obj1 = Object(1)
+        obj2 = Object(1)
+        obj3 = Object(2)
+
+        assert obj1 is obj2
+        assert obj1 is not obj3
+
+    def test_cache_size(self):
+        """
+        Tests specifying the size of the instance cache.
+        """
+        class Object(CacheInstances):
+            _instance_cache_size = 2
+
+            def __init__(self, value: int):
+                self.value = value
+
+        obj1 = Object(1)
+        obj2 = Object(2)
+        obj3 = Object(3)
+        obj4 = Object(1)
+        obj5 = Object(3)
+
+        # obj1 should have been evicted before obj4 was created
+        assert obj1 is not obj4
+        assert obj1 is not obj2
+        assert obj3 is obj5
+
+        hits, _, _, cursize = Object._instance_cache.cache_info()
+        assert hits == 1  # obj5 hit the cache
+        assert cursize == 2
+
+    def test_cleared_after_build(self):
+        """
+        Tests that instance caches are cleared after building an Operator.
+        """
+        class Object(CacheInstances):
+            def __init__(self, value: int):
+                self.value = value
+
+        obj1 = Object(1)
+        cache_size = Object._instance_cache.cache_info()[-1]
+        assert cache_size == 1
+
+        x = Symbol('x')
+        Operator(Eq(x, obj1.value))
+
+        # Cache should be cleared after Operator construction
+        cache_size = Object._instance_cache.cache_info()[-1]
+        assert cache_size == 0