avoid traversing functions multiple times

Author: majosm

pytato/analysis/__init__.py

@@ -404,14 +404,14 @@ class NodeCountMapper(CachedWalkMapper):
        The number of nodes.
     """
 
-    def __init__(self) -> None:
-        super().__init__()
+    def __init__(self, _visited_functions: set[Any] | None = None) -> None:
+        super().__init__(_visited_functions=_visited_functions)
         self.count = 0
 
     def get_cache_key(self, expr: ArrayOrNames) -> int:
         return id(expr)
 
-    def get_func_def_cache_key(self, expr: FunctionDefinition) -> int:
+    def get_function_definition_cache_key(self, expr: FunctionDefinition) -> int:
         return id(expr)
 
     def post_visit(self, expr: Any) -> None:
@@ -444,28 +444,25 @@ class CallSiteCountMapper(CachedWalkMapper):
        The number of nodes.
     """
 
-    def __init__(self) -> None:
-        super().__init__()
+    def __init__(self, _visited_functions: set[Any] | None = None) -> None:
+        super().__init__(_visited_functions=_visited_functions)
         self.count = 0
 
     def get_cache_key(self, expr: ArrayOrNames) -> int:
         return id(expr)
 
-    def get_func_def_cache_key(self, expr: FunctionDefinition) -> int:
+    def get_function_definition_cache_key(self, expr: FunctionDefinition) -> int:
         return id(expr)
 
     def map_function_definition(self, expr: FunctionDefinition) -> None:
-        cache_key = self.get_func_def_cache_key(expr)
-        if not self.visit(expr) or cache_key in self._visited_functions:
+        if not self.visit(expr):
             return
 
         new_mapper = self.clone_for_callee(expr)
         for subexpr in expr.returns.values():
             new_mapper(subexpr)
         self.count += new_mapper.count
 
-        self._visited_functions.add(cache_key)
-
         self.post_visit(expr)
 
     def post_visit(self, expr: Any) -> None:

pytato/codegen.py

@@ -24,7 +24,7 @@
 """
 
 import dataclasses
-from typing import Any, Mapping, Tuple
+from typing import Any, Hashable, Mapping, Tuple
 
 from immutabledict import immutabledict
 
@@ -118,10 +118,13 @@ class CodeGenPreprocessor(ToIndexLambdaMixin, CopyMapper):  # type: ignore[misc]
     ======================================  =====================================
     """
 
-    def __init__(self, target: Target,
-                 kernels_seen: dict[str, lp.LoopKernel] | None = None
-                 ) -> None:
-        super().__init__()
+    def __init__(
+            self,
+            target: Target,
+            kernels_seen: dict[str, lp.LoopKernel] | None = None,
+            _function_cache: dict[Hashable, FunctionDefinition] | None = None
+            ) -> None:
+        super().__init__(_function_cache=_function_cache)
         self.bound_arguments: dict[str, DataInterface] = {}
         self.var_name_gen: UniqueNameGenerator = UniqueNameGenerator()
         self.target = target
@@ -247,14 +250,14 @@ def normalize_outputs(
 
 @optimize_mapper(drop_args=True, drop_kwargs=True, inline_get_cache_key=True)
 class NamesValidityChecker(CachedWalkMapper):
-    def __init__(self) -> None:
+    def __init__(self, _visited_functions: set[Any] | None = None) -> None:
         self.name_to_input: dict[str, InputArgumentBase] = {}
-        super().__init__()
+        super().__init__(_visited_functions=_visited_functions)
 
     def get_cache_key(self, expr: ArrayOrNames) -> int:
         return id(expr)
 
-    def get_func_def_cache_key(self, expr: FunctionDefinition) -> int:
+    def get_function_definition_cache_key(self, expr: FunctionDefinition) -> int:
         return id(expr)
 
     def post_visit(self, expr: Any) -> None:

pytato/distributed/partition.py

@@ -292,8 +292,9 @@ def __init__(self,
                  recvd_ary_to_name: Mapping[Array, str],
                  sptpo_ary_to_name: Mapping[Array, str],
                  name_to_output: Mapping[str, Array],
+                 _function_cache: dict[Hashable, FunctionDefinition] | None = None,
                  ) -> None:
-        super().__init__()
+        super().__init__(_function_cache=_function_cache)
 
         self.recvd_ary_to_name = recvd_ary_to_name
         self.sptpo_ary_to_name = sptpo_ary_to_name
@@ -307,7 +308,7 @@ def clone_for_callee(
             self, function: FunctionDefinition) -> _DistributedInputReplacer:
         # Function definitions aren't allowed to contain receives,
         # stored arrays promoted to part outputs, or part outputs
-        return type(self)({}, {}, {})
+        return type(self)({}, {}, {}, _function_cache=self._function_cache)
 
     def map_placeholder(self, expr: Placeholder) -> Placeholder:
         self.user_input_names.add(expr.name)

pytato/distributed/verify.py

@@ -142,8 +142,8 @@ class MissingRecvError(DistributedPartitionVerificationError):
 
 @optimize_mapper(drop_args=True, drop_kwargs=True, inline_get_cache_key=True)
 class _SeenNodesWalkMapper(CachedWalkMapper):
-    def __init__(self) -> None:
-        super().__init__()
+    def __init__(self, _visited_functions: set[Any] | None = None) -> None:
+        super().__init__(_visited_functions=_visited_functions)
         self.seen_nodes: set[ArrayOrNames] = set()
 
     def get_cache_key(self, expr: ArrayOrNames) -> int:

pytato/equality.py

@@ -27,8 +27,6 @@
 
 from typing import TYPE_CHECKING, Any, Callable, Union
 
-from pytools import memoize_method
-
 from pytato.array import (
     AbstractResultWithNamedArrays,
     AdvancedIndexInContiguousAxes,
@@ -83,19 +81,23 @@ class EqualityComparer:
           more on this.
     """
     def __init__(self) -> None:
+        # Uses the same cache for both arrays and functions
         self._cache: dict[tuple[int, int], bool] = {}
 
-    def rec(self, expr1: ArrayOrNames, expr2: Any) -> bool:
+    def rec(self, expr1: ArrayOrNames | FunctionDefinition, expr2: Any) -> bool:
         cache_key = id(expr1), id(expr2)
         try:
             return self._cache[cache_key]
         except KeyError:
-
-            method: Callable[[Array | AbstractResultWithNamedArrays, Any],
-                             bool]
+            method: Callable[
+                [Array | AbstractResultWithNamedArrays | FunctionDefinition, Any],
+                bool]
 
             try:
-                method = getattr(self, expr1._mapper_method)
+                method = (
+                    getattr(self, expr1._mapper_method)
+                    if isinstance(expr1, (Array, AbstractResultWithNamedArrays))
+                    else self.map_function_definition)
             except AttributeError:
                 if isinstance(expr1, Array):
                     result = self.handle_unsupported_array(expr1, expr2)
@@ -293,7 +295,6 @@ def map_distributed_recv(self, expr1: DistributedRecv, expr2: Any) -> bool:
                 and expr1.tags == expr2.tags
                 )
 
-    @memoize_method
     def map_function_definition(self, expr1: FunctionDefinition, expr2: Any
                                 ) -> bool:
         return (expr1.__class__ is expr2.__class__
@@ -307,7 +308,7 @@ def map_function_definition(self, expr1: FunctionDefinition, expr2: Any
 
     def map_call(self, expr1: Call, expr2: Any) -> bool:
         return (expr1.__class__ is expr2.__class__
-                and self.map_function_definition(expr1.function, expr2.function)
+                and self.rec(expr1.function, expr2.function)
                 and frozenset(expr1.bindings) == frozenset(expr2.bindings)
                 and all(self.rec(bnd,
                                  expr2.bindings[name])

pytato/stringifier.py

@@ -31,8 +31,6 @@
 import numpy as np
 from immutabledict import immutabledict
 
-from pytools import memoize_method
-
 from pytato.array import (
     Array,
     Axis,
@@ -68,6 +66,7 @@ def __init__(self,
         self.truncation_depth = truncation_depth
         self.truncation_string = truncation_string
 
+        # Uses the same cache for both arrays and functions
         self._cache: dict[tuple[int, int], str] = {}
 
     def rec(self, expr: Any, depth: int) -> str:
@@ -79,6 +78,15 @@ def rec(self, expr: Any, depth: int) -> str:
             self._cache[cache_key] = result
             return result  # type: ignore[no-any-return]
 
+    def rec_function_definition(self, expr: FunctionDefinition, depth: int) -> str:
+        cache_key = (id(expr), depth)
+        try:
+            return self._cache[cache_key]
+        except KeyError:
+            result = super().rec_function_definition(expr, depth)
+            self._cache[cache_key] = result
+            return result  # type: ignore[no-any-return]
+
     def __call__(self, expr: Any, depth: int = 0) -> str:
         return self.rec(expr, depth)
 
@@ -168,7 +176,6 @@ def _get_field_val(field: str) -> str:
                         for field in attrs.fields(type(expr)))
                 + ")")
 
-    @memoize_method
     def map_function_definition(self, expr: FunctionDefinition, depth: int) -> str:
         if depth > self.truncation_depth:
             return self.truncation_string
@@ -191,7 +198,7 @@ def map_call(self, expr: Call, depth: int) -> str:
 
         def _get_field_val(field: str) -> str:
             if field == "function":
-                return self.map_function_definition(expr.function, depth+1)
+                return self.rec_function_definition(expr.function, depth+1)
             else:
                 return self.rec(getattr(expr, field), depth+1)
 

pytato/target/python/numpy_like.py

@@ -64,6 +64,7 @@
     SizeParam,
     Stack,
 )
+from pytato.function import FunctionDefinition
 from pytato.raising import BinaryOpType, C99CallOp
 from pytato.reductions import (
     AllReductionOperation,
@@ -193,7 +194,7 @@ def _is_slice_trivial(slice_: NormalizedSlice,
 }
 
 
-class NumpyCodegenMapper(CachedMapper[str]):
+class NumpyCodegenMapper(CachedMapper[str, FunctionDefinition]):
     """
     .. note::