Avoid duplicating arrays #515

majosm · 2024-07-12T04:09:06Z

Depends on ~~#503~~ (merged), ~~#512~~ (merged), ~~#531~~ (merged), ~~#549~~ (merged), ~~#585~~ (merged), ~~#583~~ (merged), and ~~#550~~ (merged).

~~@inducer note to self: Reviewed up to (but not including) "tweak Inliner/PlaceholderSubstitutor implementations".~~

Draft because:

⚠️ Squash "corresponding" commits before merging.

Closes #590.

majosm · 2025-02-28T13:35:16Z

Ready for a look after #585, #583, and #550. I've left the duplication checks turned off by default for now, because some things in arraycontext, grudge, etc. still duplicate.

pytato/transform/__init__.py

pytato/transform/calls.py

pytato/array.py

pytato/codegen.py

inducer

LGTM!

inducer · 2025-03-31T21:18:46Z

pytato/transform/calls.py

+        # Must disable collision/duplication checking because we're combining
+        # expressions that were previously in two different call stack frames
+        # (and were thus cached separately)


Add to doc: Inliner may produce DAGs with duplicates.

inducer · 2025-03-31T21:39:16Z

pytato/transform/remove_broadcasts_einsum.py

@@ -97,6 +169,6 @@ def rewrite_einsums_with_no_broadcasts(expr: MappedT) -> MappedT:
        alter its value.
    """
    mapper = EinsumWithNoBroadcastsRewriter()
-    return cast("MappedT", mapper(expr))
+    return cast("MappedT", mapper(expr, None))


return cast("MappedT", mapper(expr, ())

and save the None checking above?

inducer · 2025-03-31T21:45:43Z

pytato/transform/remove_broadcasts_einsum.py

-                      expr.redn_axis_to_redn_descr,
-                      tags=expr.tags,
-                      axes=expr.axes,)
+        if (


Maybe tackle #590? (OK if it doesn't pan out.)

Took a stab at it. Seemed to work out OK.

inducer · 2025-04-16T19:01:36Z

pytato/transform/__init__.py

+        return all(b_i is a_i for a_i, b_i in zip(a, b, strict=True))
+
+
+class IdentityMappingHelperMixin:


It struck me that it would likely be better to bake these into the node type. (I.e. Placeholder._if_modified(self, other).) We can probably generate them automatically, too. As it stands, dispatch is name-based, even though it's only "single", i.e. based on the type of the node.

How does 69c4e3c look?

I'm not sure how to make it work for DictOfNamedArrays, though. See here:

pytato/pytato/transform/__init__.py

Lines 881 to 883 in 69c4e3c

# FIXME: Doesn't work because __init__ argument is "data" but member

# is "_data"

# return expr.replace_if_different(_data=new_data)

Special-casing DictOfNamedArrays is OK, for the reason you mention.

I'm OK with the approach, though I think I'd prefer generated code for this, since it's bound to be a fairly hot code path. There's already some code generation in @array_dataclass, it should be straightforward to add there.

I attempted to implement this in a687544, but in doing so I've seemingly triggered a bunch of unrelated(?) mypy errors. I fixed some in d27da88, but some of the errors that remain don't make a lot of sense. I've also triggered some new pylint errors with my fixes. Any idea what's going on here?

Edit: I guess this is similar to the issues encountered in #598?

Edit: I guess this is similar to the issues encountered in #598?

I don't think the errors you're seeing here are all that related. Some just outright don't make any sense, IMO. For example:

pytato/transform/__init__.py:1115:25: error: Invalid index type "str" for "Call"; expected type "Call" [index]

I can't make heads or tails of that.

The Pylint issues are similarly weird, maybe it's time to drop Pylint, too?

I think the best proposal I've got is to switch to pyright (the "based" version, to avoid drowning in noise). I've pushed changes that do that. At least most of what pyright has to say seems to make some sense.

Unfortunately, the "baseline" feature of basedpyright seems a bit imperfect, too. It seems to complain about errors that are already contained in the baseline. I don't see this locally, but it does seem to happen in CI. 🤔

add note about Inliner duplication deduplicate in inline_calls

now inherits from CachedMapper add comments explaining result deduplication to MPMSMaterializerCache

simplify EinsumWithNoBroadcastsRewriter by using () instead of None

…lication avoidance logic move duplication-preventing boilerplate into array types move replace_if_different code into array_dataclass

majosm force-pushed the avoid-array-duplication branch from e813eb2 to 1d49631 Compare July 15, 2024 20:20

majosm mentioned this pull request Jul 25, 2024

More fixes for functions #503

Merged

majosm force-pushed the avoid-array-duplication branch 2 times, most recently from a5dbc5a to c55dee2 Compare August 27, 2024 21:33

majosm force-pushed the avoid-array-duplication branch from c55dee2 to f6a3014 Compare September 24, 2024 20:36

majosm force-pushed the avoid-array-duplication branch 3 times, most recently from dd6cd37 to 8ac0b8d Compare February 7, 2025 04:13

majosm mentioned this pull request Feb 21, 2025

Avoid double caching in mappers that derive from CachedMapper #585

Merged

majosm force-pushed the avoid-array-duplication branch 3 times, most recently from ae9c939 to 2ca7a65 Compare February 28, 2025 02:42

majosm marked this pull request as ready for review February 28, 2025 13:33

majosm requested a review from inducer February 28, 2025 13:35

majosm force-pushed the avoid-array-duplication branch 4 times, most recently from 8bb1a1e to 8789290 Compare March 18, 2025 14:34

inducer reviewed Mar 18, 2025

View reviewed changes

majosm force-pushed the avoid-array-duplication branch from 8789290 to e00662b Compare March 31, 2025 20:24

inducer marked this pull request as draft March 31, 2025 21:41

inducer approved these changes Mar 31, 2025

View reviewed changes

majosm force-pushed the avoid-array-duplication branch from e00662b to d21de9d Compare April 15, 2025 04:36

majosm requested a review from inducer April 15, 2025 15:07

inducer reviewed Apr 16, 2025

View reviewed changes

majosm force-pushed the avoid-array-duplication branch 4 times, most recently from 4af1241 to d27da88 Compare May 7, 2025 20:32

limit PlaceholderSubstitutor to one call stack frame

5fef7d9

inducer mentioned this pull request May 14, 2025

Get rid of sorted in bindings rewrites #601

Open

inducer force-pushed the avoid-array-duplication branch from d98c27d to a36694e Compare May 14, 2025 22:10

majosm and others added 24 commits May 14, 2025 17:14

tweak Inliner/PlaceholderSubstitutor implementations

0d6c8c0

add note about Inliner duplication deduplicate in inline_calls

use context manager to avoid leaking traceback tag setting in test

39aa718

refactor FFTRealizationMapper to avoid resetting cache in __init__

eb4bd4b

add allow_duplicate_nodes option to RandomDAGContext in tests

bdf6395

fix some more tests

94acf22

don't check for collisions in ArrayToDotNodeInfoMapper

78d8ba0

avoid duplication in MPMSMaterializer

6ea7d8b

now inherits from CachedMapper add comments explaining result deduplication to MPMSMaterializerCache

avoid duplicates in EinsumWithNoBroadcastsRewriter

2c76a7c

simplify EinsumWithNoBroadcastsRewriter by using () instead of None

forbid DependencyMapper from being called on functions

b826922

deduplicate in advection example

b3bb849

add FIXME

8ee1d3f

Add replace_if_different to reduce the amount of boilerplate from dup…

997fbf0

…lication avoidance logic move duplication-preventing boilerplate into array types move replace_if_different code into array_dataclass

add doc for PartId to fix build error

4d49af9

Drop license specifier in pyproject

2331d91

Add a cast/FIXME in _is_non_negative

fc4f1a7

Add some @OverRide's to placate pyright

3e2c521

Fix some import issues flagged by pyright

65011de

Add missing superclass constructor call

98b5129

Add missing field_specifiers to @dataclass_transform

c05af45

Fix imports of Never, Self from typing->typing_extensions

c7a2fde

Better type einsum

0c13260

CI: Switch from basedmypy to basedpyright

27a7e3b

Fire pylint

f3ce4b5

Work around bpr errors only shown in CI

20bd9f0

inducer force-pushed the avoid-array-duplication branch from a36694e to 20bd9f0 Compare May 15, 2025 14:56

inducer enabled auto-merge (rebase) May 15, 2025 14:59

inducer merged commit e3b9065 into inducer:main May 15, 2025
10 checks passed

		return all(b_i is a_i for a_i, b_i in zip(a, b, strict=True))


		class IdentityMappingHelperMixin:

	# FIXME: Doesn't work because __init__ argument is "data" but member
	# is "_data"
	# return expr.replace_if_different(_data=new_data)

Avoid duplicating arrays #515

Avoid duplicating arrays #515

Uh oh!

Conversation

majosm commented Jul 12, 2024 • edited by inducer Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

majosm commented Feb 28, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

inducer left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

majosm May 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

inducer May 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

majosm commented Jul 12, 2024 •

edited by inducer

Loading

majosm May 7, 2025 •

edited

Loading

inducer May 8, 2025 •

edited

Loading