misc: memoized_instances -> weak_instance_cache

Author: enwask

devito/ir/support/basic.py

@@ -1,6 +1,6 @@
 from itertools import chain, product
 from functools import cached_property
-from typing import Any
+from typing import Any, TypeVar
 
 from sympy import S
 import sympy
@@ -13,8 +13,7 @@
                               uxreplace)
 from devito.tools import (Tag, as_mapper, as_tuple, is_integer, filter_sorted,
                           flatten, memoized_meth, memoized_generator, smart_gt,
-                          smart_lt)
-from devito.tools.memoization import _memoized_instances
+                          smart_lt, weak_instance_cache)
 from devito.types import (ComponentAccess, Dimension, DimensionTuple, Fence,
                           CriticalRegion, Function, Symbol, Temp, TempArray,
                           TBArray)
@@ -511,13 +510,26 @@ def touched_halo(self, findex):
         return (touch_halo_left, touch_halo_right)
 
 
-@_memoized_instances
+# Type variable for subclass of `Relation` returned from `maybe_cached`
+RelationType = TypeVar('RelationType', bound='Relation')
+
+
 class Relation:
 
     """
     A relation between two TimedAccess objects.
     """
 
+    @classmethod
+    @weak_instance_cache
+    def maybe_cached(cls: type[RelationType],
+                     source: TimedAccess, sink: TimedAccess) -> RelationType:
+        """
+        Constructs a new Relation or retrieves one from the cache if it
+        already/still exists.
+        """
+        return cls(source, sink)
+
     def __init__(self, source, sink):
         assert isinstance(source, TimedAccess) and isinstance(sink, TimedAccess)
         assert source.function == sink.function
@@ -628,7 +640,6 @@ def is_imaginary(self):
         return S.ImaginaryUnit in self.distance
 
 
-@_memoized_instances
 class Dependence(Relation):
 
     """
@@ -1082,7 +1093,7 @@ def d_flow_gen(self):
                     if any(not rule(w, r) for rule in self.rules):
                         continue
 
-                    dependence = Dependence(w, r)
+                    dependence = Dependence.maybe_cached(w, r)
 
                     if dependence.is_imaginary:
                         continue
@@ -1112,7 +1123,7 @@ def d_anti_gen(self):
                     if any(not rule(r, w) for rule in self.rules):
                         continue
 
-                    dependence = Dependence(r, w)
+                    dependence = Dependence.maybe_cached(r, w)
 
                     if dependence.is_imaginary:
                         continue
@@ -1142,7 +1153,7 @@ def d_output_gen(self):
                     if any(not rule(w2, w1) for rule in self.rules):
                         continue
 
-                    dependence = Dependence(w2, w1)
+                    dependence = Dependence.maybe_cached(w2, w1)
 
                     if dependence.is_imaginary:
                         continue
@@ -1204,7 +1215,7 @@ def r_gen(self):
                 if a0 is a1:
                     continue
 
-                r = Relation(a0, a1)
+                r = Relation.maybe_cached(a0, a1)
                 if r.is_imaginary:
                     continue
 

devito/tools/memoization.py

@@ -6,7 +6,7 @@
 from devito.tools import WeakValueCache
 
 
-__all__ = ['memoized_func', 'memoized_meth', 'memoized_generator']
+__all__ = ['memoized_func', 'memoized_meth', 'memoized_generator', 'weak_instance_cache']
 
 
 class memoized_func:
@@ -131,79 +131,28 @@ def __call__(self, *args, **kwargs):
         return result
 
 
-# Describes the type of an object cached by `memoized_constructor`
+# Describes the type of an object cached by `weak_instance_cache`
 InstanceType = TypeVar('InstanceType')
-Constructor = Callable[..., InstanceType]
+Constructor = TypeVar('Constructor', bound=Callable[..., InstanceType])
 
 
-def _memoized_instances(cls: type[InstanceType]) -> type[InstanceType]:
+def weak_instance_cache(fun: Constructor[InstanceType]) -> Constructor[InstanceType]:
     """
-    Decorator for a class that caches instances based on the hash values of
-    constructing arguments. The constructed values are stored weakly and
-    evicted from the cache when no longer referenced.
+    Decorator for a class method that caches instances based on the hash
+    values of constructing arguments. The constructed values are stored
+    weakly and evicted from the cache when no longer referenced.
 
-    We need to override both __new__ and __init__ to ensure initialization
-    only happens once for a cached instance.
+    See `Dependence` for a usage example.
     """
 
-    # Check if we already decorated a parent class
-    already_applied = getattr(cls, '_memoized_instances__exists', False)
-    cls._memoized_instances__exists = True
-
-    new = cls.__new__
-    init = cls.__init__
-    cache: WeakValueCache[InstanceType] = WeakValueCache(cls)
-
-    @wraps(new)
-    def _new(_cls: type[InstanceType], *args: Hashable,
-             _memoized_instances__use_cache: bool = True,
-             **kwargs: Hashable) -> InstanceType:
-        # The decorator must be reapplied to a child class, so we make sure cls matches
-        if _memoized_instances__use_cache and _cls is not cls:
-            raise TypeError(f"_memoized_instances must be applied to {_cls.__name__}, "
-                            f"not (just) {cls.__name__}")
-
-        # The cache called the constructor; avoid infinite recursion
-        if not _memoized_instances__use_cache:
-            # If the class doesn't define __new__, we can't pass any args
-            if new is object.__new__:
-                obj = new(_cls)
-
-            # Otherwise forward all arguments
-            else:
-                # If we applied the decorator to a parent class, forward the caching flag
-                if already_applied:
-                    kwargs['_memoized_instances__use_cache'] = False
-                obj = new(_cls, *args, **kwargs)
-
-            # Set our initialization flag and return
-            obj._memoized_instances__initialized = False
-            return obj
-
-        return cache.get_or_create(*args, _memoized_instances__use_cache=False, **kwargs)
-
-    @wraps(init)
-    def _init(self: InstanceType, *args: Hashable, **kwargs: Hashable) -> None:
-        # Skip reinitialization if this object was obtained from the cache
-        try:
-            if self._memoized_instances__initialized:
-                return
-        except AttributeError:
-            # If the attribute doesn't exist, this is a new instance
-            self._memoized_instances__initialized = False
-
-        # Don't forward our extra argument to the original __init__
-        kwargs.pop('_memoized_instances__use_cache', None)
-        init(self, *args, **kwargs)
-        self._memoized_instances__initialized = True
+    # Initialize a weak value cache bound to this constructor
+    cache = WeakValueCache(fun)
 
-    def _copy(self: InstanceType) -> InstanceType:
-        # Copy should just return the cached instance; bypass the cache machinery
-        return self
+    @wraps(fun)
+    def _wrapper(cls: type[InstanceType], *args: Hashable, **kwargs: Hashable) \
+            -> InstanceType:
+        # Retrieve from the cache or create a new instance
+        return cache.get_or_create(cls, *args, **kwargs)
 
-    # Update the class's methods
-    cls.__new__ = _new
-    cls.__init__ = _init
-    cls.__copy__ = _copy
+    return _wrapper
 
-    return cls