Cache mapped functions #531

majosm · 2024-07-26T04:44:24Z

Caches the results of mapping FunctionDefinitions. Caching is done globally to avoid re-traversing functions in different call stack frames.

Depends on ~~#503~~ (merged) and ~~#530~~ (merged).

majosm · 2024-08-02T22:25:58Z

@inducer Should be ready for a look. I think the mypy errors are unrelated, they're occurring on main too.

inducer · 2024-11-14T15:52:25Z

Uh-oh. The recent typing apocalypse has made a bit of a mess of this. I've started resolving conflicts, but I got bogged down in pytato.transform. Could you take a look? (Or, if you've lost patience with me, I can also try and finish.)

majosm · 2024-12-20T16:53:02Z

@inducer Should be ready for review again (mypy error looks unrelated?). In e57fe22 I added a few questions/comments about things I noticed while resolving conflicts. I'll remove them after you've had a chance to glance at them.

inducer · 2024-12-20T17:57:44Z

#571 should address the mypy issue.

inducer · 2025-01-08T19:12:37Z

pytato/transform/__init__.py

+        try:
+            method = self.map_function_definition  # type: ignore[attr-defined]
+        except AttributeError:
+            return cast("FunctionResultT", self.map_foreign(expr, *args, **kwargs))


Can this be an error instead?

inducer · 2025-01-08T19:13:54Z

pytato/transform/__init__.py

@@ -194,7 +193,7 @@ def handle_unsupported_array(self, expr: MappedT,
        raise UnsupportedArrayError(
                f"{type(self).__name__} cannot handle expressions of type {type(expr)}")

-    def map_foreign(self, expr: Any, *args: P.args, **kwargs: P.kwargs) -> ResultT:
+    def map_foreign(self, expr: Any, *args: P.args, **kwargs: P.kwargs) -> Any:


I think we can get rid of map_foreign.

Reprifier appears to override map_foreign to handle stringifying non-array fields. Seems like it needs to stay?

Let's move it to Reprifier then. I don't think it needs to be a feature of the generic machinery.

inducer · 2025-01-08T19:16:22Z

pytato/transform/__init__.py

+        else:
+            function_cache = {}
+
+        self._function_cache: dict[Hashable, FunctionResultT] = function_cache


Can this be an if expression?

inducer · 2025-01-08T19:25:31Z

pytato/transform/__init__.py

+            # Why multiple inheriting?
+            CachedMapper[ArrayOrNames, FunctionDefinition, P],
+            Mapper[ArrayOrNames, FunctionDefinition, P]


Suggested change

# Why multiple inheriting?

CachedMapper[ArrayOrNames, FunctionDefinition, P],

Mapper[ArrayOrNames, FunctionDefinition, P]

CachedMapper[ArrayOrNames, FunctionDefinition, P],

I don't think it should be.

Seems to work fine without it.

inducer · 2025-01-08T19:25:47Z

pytato/transform/__init__.py

    """
+    # FIXME: This inflates recursion depth


if TYPE_CHECKING: ... (same) else: @property def rec_ary(self): return self.rec

inducer · 2025-01-08T19:30:41Z

pytato/transform/__init__.py

    """
+    # FIXME: This inflates recursion depth


def is_array(expr: ArrayOrNames) -> Array: assert isinstance(expr, Array) return expr

(This is probably the saner way.)

Something like ca4f336?

(The other way appears to work too: majosm@34af8ef)

inducer · 2025-01-08T19:33:05Z

pytato/transform/__init__.py

+        # Don't need to pass function cache as argument here, because unlike
+        # CachedMapper we're not creating a new mapper for each call
+        self.function_cache: dict[FunctionDefinition, FunctionResultT] = {}


I think the behavior should be consistent across mapper types.

Added _function_cache to the __init__ args like the others. I didn't go the whole way and add clone_for_callee, etc., yet because I'm worried that's going to make a mess of my downstream changes.

inducer · 2025-01-08T19:33:39Z

pytato/transform/__init__.py

+        else:
+            visited_functions = set()
+
+        self._visited_functions: set[Any] = visited_functions


inducer · 2025-01-08T19:34:48Z

pytato/target/loopy/codegen.py

@@ -384,7 +384,7 @@ def update_t_unit(self, t_unit: lp.TranslationUnit) -> None:

 # {{{ codegen mapper

-class CodeGenMapper(Mapper[ImplementedResult, [CodeGenState]]):
+class CodeGenMapper(Mapper[ImplementedResult, None, [CodeGenState]]):


Suggested change

class CodeGenMapper(Mapper[ImplementedResult, None, [CodeGenState]]):

class CodeGenMapper(Mapper[ImplementedResult, Never, [CodeGenState]]):

Added here and in a few other places.

…t don't support functions

doesn't appear to be needed

the latter inflates recursion depth

majosm · 2025-01-09T22:18:41Z

Ready for another look I think @inducer.

inducer

Thanks! A few questions/comments below.

inducer · 2025-01-12T20:59:41Z

pytato/analysis/__init__.py

+    def __init__(
+            self,
+            count_duplicates: bool = False,
+            _visited_functions: set[Any] | None = None,


Can something more precise than Any be said here? (Also for other _visited_functions annotations.)

Looks like I did that to match _visited_arrays_or_names in CachedWalkMapper, which also uses Any. I guess I could change both to Hashable?

inducer · 2025-01-12T21:03:01Z

pytato/target/python/numpy_like.py

@@ -171,7 +172,7 @@ def _is_slice_trivial(slice_: NormalizedSlice,
 }


-class NumpyCodegenMapper(CachedMapper[str, []]):
+class NumpyCodegenMapper(CachedMapper[str, FunctionDefinition, []]):


Is this annotation right? Does the numpy target generate appropriate code for functions?

NumpyCodegenMapper doesn't support functions, so I guess this should be Never. Same for FancyDotWriter. Fixed these.

inducer · 2025-01-12T21:04:13Z

pytato/transform/calls.py

@@ -58,6 +58,7 @@ class PlaceholderSubstitutor(CopyMapper):
    """

    def __init__(self, substitutions: Mapping[str, Array]) -> None:
+        # Ignoring function cache, not needed


How come? Not using the function cache may lead to large amounts of redundant work?

(I'm OK with not acting on this right now, but maybe this comment should be a FIXME?)

PlaceholderSubstitutor should only be called on function bodies that don't have functions nested in them (since the placeholder names can overlap). I've disabled PlaceholderSubstitutor.map_named_call_result to make that more explicit.

pytato/transform/__init__.py

…r too

majosm · 2025-01-16T22:56:15Z

I changed CachedWalkMapper to use an alias of Hashable instead of Any, and I removed a few other uses of Any that I found. Should be ready for another look.

inducer

Thanks! This looks great. In it goes.

pytato/transform/__init__.py

majosm force-pushed the cache-mapped-function-results branch from 0ed120a to 0c73b20 Compare July 26, 2024 04:54

majosm mentioned this pull request Jul 26, 2024

More fixes for functions #503

Merged

majosm force-pushed the cache-mapped-function-results branch 6 times, most recently from dd6d54f to 2ff245d Compare August 2, 2024 21:33

majosm marked this pull request as ready for review August 2, 2024 22:24

majosm requested a review from inducer August 2, 2024 22:26

majosm force-pushed the cache-mapped-function-results branch from 2ff245d to f8c2c68 Compare August 12, 2024 20:45

inducer force-pushed the cache-mapped-function-results branch from f8c2c68 to 58c8dbf Compare August 25, 2024 04:54

inducer mentioned this pull request Aug 26, 2024

Parameter Study Transform #532

Draft

majosm force-pushed the cache-mapped-function-results branch 2 times, most recently from 4519ce2 to 913d3d1 Compare September 5, 2024 19:23

majosm mentioned this pull request Sep 24, 2024

Use a class for CachedMapper-derived mappers instead of a dict #549

Open

majosm force-pushed the cache-mapped-function-results branch 2 times, most recently from 312e8ff to 72ce101 Compare September 24, 2024 17:47

This was referenced Sep 24, 2024

Add duplication checks #550

Draft

Avoid duplicating arrays #515

Draft

majosm force-pushed the cache-mapped-function-results branch 4 times, most recently from b8cbbc5 to e57fe22 Compare December 20, 2024 02:36

majosm force-pushed the cache-mapped-function-results branch from e57fe22 to a4b609c Compare January 8, 2025 16:01

inducer reviewed Jan 8, 2025

View reviewed changes

majosm force-pushed the cache-mapped-function-results branch 4 times, most recently from 610d413 to 5d4b0e2 Compare January 9, 2025 21:30

majosm added 7 commits January 9, 2025 15:43

avoid traversing functions multiple times

eb75be0

condense function cache setup

786088c

pass function cache to CombineMapper.__init__ too

2108afd

use Never instead of None as the function result type for mappers tha…

1596db4

…t don't support functions

remove multiple inheritance for TransformMapperWithExtraArgs

2cac57f

doesn't appear to be needed

add _verify_is_array to avoid the need for rec_ary

40b47e9

the latter inflates recursion depth

remove map_foreign from Mapper

1890af8

majosm force-pushed the cache-mapped-function-results branch from 5d4b0e2 to 1890af8 Compare January 9, 2025 21:43

majosm requested a review from inducer January 9, 2025 22:18

inducer reviewed Jan 12, 2025

View reviewed changes

majosm added 6 commits January 12, 2025 18:40

use Never as FunctionResultT for NumpyCodegenMapper and FancyDotWrite…

4e9a4d0

…r too

make PlaceholderSubstitutor explicit about not supporting functions

d3671b4

tweak types in rec/rec_function_definition

a7e4797

use P.args/P.kwargs in CachedWalkMapper

0d18657

use a more specific type than Any for cache key in CachedWalkMapper

b704ede

remove some more Anys

d912360

majosm requested a review from inducer January 16, 2025 22:56

inducer approved these changes Jan 17, 2025

View reviewed changes

pytato/transform/__init__.py Show resolved Hide resolved

inducer merged commit 46431ee into inducer:main Jan 17, 2025
11 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Cache mapped functions #531

Cache mapped functions #531

majosm commented Jul 26, 2024 •

edited

Loading

majosm commented Aug 2, 2024

inducer commented Nov 14, 2024

majosm commented Dec 20, 2024

inducer commented Dec 20, 2024

inducer Jan 8, 2025

majosm Jan 9, 2025

inducer Jan 8, 2025

majosm Jan 9, 2025

inducer Jan 9, 2025

majosm Jan 9, 2025

inducer Jan 8, 2025

inducer Jan 8, 2025

majosm Jan 9, 2025

inducer Jan 8, 2025

inducer Jan 8, 2025

majosm Jan 9, 2025

inducer Jan 8, 2025

majosm Jan 9, 2025

inducer Jan 8, 2025

inducer Jan 8, 2025

majosm Jan 9, 2025

majosm commented Jan 9, 2025

inducer left a comment

inducer Jan 12, 2025

majosm Jan 13, 2025

inducer Jan 12, 2025

majosm Jan 13, 2025

inducer Jan 12, 2025

majosm Jan 13, 2025

majosm commented Jan 16, 2025

inducer left a comment

	class CodeGenMapper(Mapper[ImplementedResult, None, [CodeGenState]]):
	class CodeGenMapper(Mapper[ImplementedResult, Never, [CodeGenState]]):

Cache mapped functions #531

Cache mapped functions #531

Conversation

majosm commented Jul 26, 2024 • edited Loading

majosm commented Aug 2, 2024

inducer commented Nov 14, 2024

majosm commented Dec 20, 2024

inducer commented Dec 20, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

majosm commented Jan 9, 2025

inducer left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

majosm commented Jan 16, 2025

inducer left a comment

Choose a reason for hiding this comment

majosm commented Jul 26, 2024 •

edited

Loading