tests: Add memoized method tests

Author: enwask

tests/test_tools.py

@@ -1,3 +1,5 @@
+from concurrent.futures import ThreadPoolExecutor
+from threading import RLock
 import numpy as np
 import pytest
 from sympy.abc import a, b, c, d, e
@@ -7,7 +9,7 @@
 from devito import Operator, Eq
 from devito.tools import (UnboundedMultiTuple, ctypes_to_cstr, toposort,
                           filter_ordered, transitive_closure, UnboundTuple,
-                          CacheInstances)
+                          CacheInstances, memoized_meth, memoized_generator)
 from devito.types.basic import Symbol
 
 
@@ -209,3 +211,156 @@ def __init__(self, value: int):
         # Cache should be cleared after Operator construction
         cache_size = Object._instance_cache.cache_info()[-1]
         assert cache_size == 0
+
+
+class TestMemoizedMethods:
+
+    def test_memoized_meth(self):
+        """
+        Tests basic functionality of memoized_meth
+        """
+        class Object:
+            def __init__(self):
+                self.misses = 0
+
+            @memoized_meth
+            def compute(self, x):
+                self.misses += 1
+                return x * 2
+
+        obj = Object()
+        obj.compute(2)
+        obj.compute(4)
+        assert obj.compute(2) == 4
+        assert obj.compute(4) == 8
+        assert obj.misses == 2  # Only two unique calls
+
+    def test_unhashable_args(self):
+        """
+        Tests that memoized_meth raises an error for unhashable arguments.
+        """
+        class Object:
+            def __init__(self):
+                self.misses = 0
+
+            @memoized_meth
+            def compute(self, x: list[int]):
+                self.misses += 1
+                return sum(x)
+
+        obj = Object()
+        with pytest.raises(TypeError):
+            obj.compute([1, 2, 3])
+
+    @pytest.mark.parametrize('num_threads', [5, 11, 17])
+    def test_memoized_meth_concurrency(self, num_threads: int):
+        """
+        Tests concurrent calls to a memoized method
+        """
+        # Each thread should have its own cache; the calls should not block
+        class Object:
+            def __init__(self):
+                self.misses = 0
+                self.lock = RLock()
+
+            @memoized_meth
+            def compute(self, x):
+                # print ID of the running thread
+                with self.lock:
+                    self.misses += 1
+
+                # Simulate some computation
+                time.sleep(0.2)
+                return x * 2
+
+        obj = Object()
+        def worker(x: int) -> int:
+            a = obj.compute(x)
+            b = obj.compute(x)
+            assert a == b
+            return a
+
+        with ThreadPoolExecutor(max_workers=num_threads) as executor:
+            stime = time.perf_counter()
+            futures = [executor.submit(worker, i % 4) for i in range(num_threads)]
+            results = [f.result() for f in futures]
+            etime = time.perf_counter()
+
+        assert len(set(results)) == 4  # Should have gotten four unique results
+        assert obj.misses == num_threads  # Each thread should have missed once
+
+        # Ensure that the total time is approximately 0.2 seconds (one miss per thread)
+        expected = 0.2
+        assert abs(etime - stime - expected) < 0.1 * expected
+
+    def test_memoized_generator(self):
+        """
+        Tests basic functionality of memoized_generator
+        """
+        class Object:
+            def __init__(self):
+                self.misses = 0
+
+            @memoized_generator
+            def compute(self, x):
+                self.misses += 1
+                yield x * 2
+                yield x * 3
+
+        obj = Object()
+        list(obj.compute(2))
+        assert tuple(obj.compute(2)) == (4, 6)
+        assert obj.misses == 1  # Only one unique call
+
+    @pytest.mark.parametrize('num_threads', [5, 11, 17])
+    def test_memoized_generator_concurrency(self, num_threads: int):
+        """
+        Tests concurrent calls to a memoized generator
+        """
+        class Object:
+            def __init__(self):
+                self.misses = 0
+                self.lock = RLock()
+
+            @memoized_generator
+            def compute(self, x):
+                with self.lock:
+                    self.misses += 1
+
+                time.sleep(0.25)
+                yield x * 2
+
+                time.sleep(0.25)
+                yield x * 3
+
+        # With memoized_generator, the initial construction should block but iteration
+        # should be concurrent and reuse the same iterator.
+
+        obj = Object()
+        def worker(x: int) -> list[int]:
+            return list(obj.compute(x))
+
+        # If one thread consumes the generator, subsequent iteration shouldn't block
+        # First we iterate concurrently; all but one thread should block to wait for
+        # the producing thread, so all will take ~0.5 seconds
+        with ThreadPoolExecutor(max_workers=num_threads) as executor:
+            stime = time.perf_counter()
+            futures = [executor.submit(worker, i % 4) for i in range(num_threads)]
+            results = [f.result() for f in futures]
+            etime = time.perf_counter()
+
+        expected = 0.5
+        assert abs(etime - stime - expected) < 0.1 * expected
+        assert set(tuple(r) for r in results) == {(0, 0), (2, 3), (4, 6), (6, 9)}
+        assert obj.misses == 4  # One miss per unique call
+
+        # Now iterating the same calls should use buffered generators from the cache
+        with ThreadPoolExecutor(max_workers=num_threads) as executor:
+            stime = time.perf_counter()
+            futures = [executor.submit(worker, i % 4) for i in range(num_threads)]
+            results = [f.result() for f in futures]
+            etime = time.perf_counter()
+
+        assert etime - stime < 0.1  # Should take epsilon time
+        assert set(tuple(r) for r in results) == {(0, 0), (2, 3), (4, 6), (6, 9)}
+        assert obj.misses == 4  # No new misses; all calls reused cached generators