add PytatoKeyBuilder, persistent_dict test

pytato/analysis/__init__.py

@@ -28,6 +28,7 @@
 
 from typing import TYPE_CHECKING, Any, Mapping
 
+from loopy.tools import LoopyKeyBuilder
 from pymbolic.mapper.optimize import optimize_mapper
 from pytools import memoize_method
 
@@ -564,4 +565,47 @@ def get_num_call_sites(outputs: Array | DictOfNamedArrays) -> int:
 
 # }}}
 
+
+# {{{ PytatoKeyBuilder
+
+class PytatoKeyBuilder(LoopyKeyBuilder):
+    """A custom :class:`pytools.persistent_dict.KeyBuilder` subclass
+    for objects within :mod:`pytato`.
+    """
+
+    def update_for_ndarray(self, key_hash: Any, key: Any) -> None:
+        self.rec(key_hash, key.data.tobytes())
+
+    def update_for_TaggableCLArray(self, key_hash: Any, key: Any) -> None:
+        self.rec(key_hash, key.get())
+
+    def update_for_Array(self, key_hash: Any, key: Any) -> None:
+        # CL Array
+        self.rec(key_hash, key.get())
+
+    update_for_BitwiseAnd = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_BitwiseNot = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_BitwiseOr = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_BitwiseXor = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_Call = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_CallWithKwargs = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_Comparison = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_If = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_FloorDiv = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_LeftShift = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_LogicalAnd = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_LogicalNot = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_LogicalOr = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_Lookup = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_Power = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_Product = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_Quotient = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_Remainder = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_RightShift = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_Subscript = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_Sum = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    update_for_Variable = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+
+# }}}
+
 # vim: fdm=marker

test/test_pytato.py

@@ -1364,8 +1364,102 @@ def test_dot_visualizers():
     # }}}
 
 
+# {{{ Test PytatoKeyBuilder
+
+def run_test_with_new_python_invocation(f, *args, extra_env_vars=None) -> None:
+    import os
+    if extra_env_vars is None:
+        extra_env_vars = {}
+
+    from base64 import b64encode
+    from pickle import dumps
+    from subprocess import check_call
+
+    env_vars = {
+        "INVOCATION_INFO": b64encode(dumps((f, args))).decode(),
+    }
+    env_vars.update(extra_env_vars)
+
+    my_env = os.environ.copy()
+    my_env.update(env_vars)
+
+    check_call([sys.executable, __file__], env=my_env)
+
+
+def run_test_with_new_python_invocation_inner() -> None:
+    import os
+    from base64 import b64decode
+    from pickle import loads
+
+    f, args = loads(b64decode(os.environ["INVOCATION_INFO"].encode()))
+
+    f(*args)
+
+
+def test_persistent_hashing_and_persistent_dict() -> None:
+    import shutil
+    import tempfile
+
+    from pytools.persistent_dict import ReadOnlyEntryError, WriteOncePersistentDict
+
+    from pytato.analysis import PytatoKeyBuilder
+
+    try:
+        tmpdir = tempfile.mkdtemp()
+
+        pkb = PytatoKeyBuilder()
+
+        pd = WriteOncePersistentDict("test_persistent_dict",
+                                 key_builder=pkb,
+                                 container_dir=tmpdir)
+
+        for i in range(100):
+            rdagc = RandomDAGContext(np.random.default_rng(seed=i),
+                    axis_len=5, use_numpy=True)
+
+            dag = make_random_dag(rdagc)
+
+            # Make sure the PytatoKeyBuilder can handle 'dag'
+            pd[dag] = 42
+
+            # Make sure that the key stays the same within the same Python invocation
+            with pytest.raises(ReadOnlyEntryError):
+                pd[dag] = 42
+
+        # Make sure that the key stays the same across Python invocations
+        run_test_with_new_python_invocation(
+            _test_persistent_hashing_and_persistent_dict_stage2, tmpdir)
+    finally:
+        shutil.rmtree(tmpdir)
+
+
+def _test_persistent_hashing_and_persistent_dict_stage2(tmpdir) -> None:
+    from pytools.persistent_dict import ReadOnlyEntryError, WriteOncePersistentDict
+
+    from pytato.analysis import PytatoKeyBuilder
+    pkb = PytatoKeyBuilder()
+
+    pd = WriteOncePersistentDict("test_persistent_dict",
+                                 key_builder=pkb,
+                                 container_dir=tmpdir)
+
+    for i in range(100):
+        rdagc = RandomDAGContext(np.random.default_rng(seed=i),
+                axis_len=5, use_numpy=True)
+
+        dag = make_random_dag(rdagc)
+
+        with pytest.raises(ReadOnlyEntryError):
+            pd[dag] = 42
+
+# }}}
+
+
 if __name__ == "__main__":
-    if len(sys.argv) > 1:
+    import os
+    if "INVOCATION_INFO" in os.environ:
+        run_test_with_new_python_invocation_inner()
+    elif len(sys.argv) > 1:
         exec(sys.argv[1])
     else:
         from pytest import main