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/abc.py

@@ -1,4 +1,5 @@
 import abc
+from collections.abc import Callable
 import weakref
 from concurrent.futures import Future
 from hashlib import sha1
@@ -289,8 +290,9 @@ def __repr__(self):
     __str__ = __repr__
 
 
-# Cached instance type for `WeakValueCache` (not covariant to avoid weird init logic)
-ValueType = TypeVar('ValueType')
+# Cached instance type for `WeakValueCache`
+ValueType = TypeVar('ValueType', covariant=True)
+Constructor = Callable[..., ValueType]
 
 
 class WeakValueCache(Generic[ValueType]):
@@ -304,34 +306,22 @@ class WeakValueCache(Generic[ValueType]):
     concurrent access while still allowing for threaded construction.
     """
 
-    def __init__(self, cls: type[ValueType]):
-        self._cls = cls
+    def __init__(self, constructor: Constructor[ValueType]):
+        self._constructor = constructor
         self._futures: dict[int, Future[ReferenceType[ValueType]]] = {}
         self._lock = RLock()
 
-    def _make_key(self, *args: Hashable, **kwargs: Hashable) -> int:
-        return hash((*args, frozenset(kwargs.items())))
+    def _make_key(self, cls: type[ValueType], *args: Hashable, **kwargs: Hashable) -> int:
+        return hash((cls, *args, frozenset(kwargs.items())))
 
-    def _create_instance(self, *args: Hashable, **kwargs: Hashable) -> ValueType:
-        if self._cls is object.__new__:
-            # If the constructor is object's __new__, we cannot pass any arguments
-            obj = self._cls()
-        else:
-            # Otherwise, forward all construction arguments
-            obj = self._cls(*args, **kwargs)
-
-        # Initialize the object so it's ready for consuming threads
-        obj.__init__(*args, **kwargs)
-
-        return obj
-
-    def get_or_create(self, *args: Hashable, **kwargs: Hashable) -> ValueType:
+    def get_or_create(self, cls: type[ValueType],
+                      *args: Hashable, **kwargs: Hashable) -> ValueType:
         """
         Gets an instance for the given construction arguments, creating it on this thread
         if it doesn't exist. If another thread is currently constructing it, blocks until
         the instance is available.
         """
-        key = self._make_key(*args, **kwargs)
+        key = self._make_key(cls, *args, **kwargs)
         future: Future[ReferenceType[ValueType]]
 
         while True:
@@ -360,7 +350,7 @@ def get_or_create(self, *args: Hashable, **kwargs: Hashable) -> ValueType:
         # If we got here, this is the thread that will create the new instance
         # Do this outside the lock to allow for concurrent construction
         try:
-            obj = self._create_instance(*args, **kwargs)
+            obj = self._constructor(cls, *args, **kwargs)
 
             # Listener for when the weak reference expires
             def on_obj_destroyed(k: int = key,

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