Add collect_materialized_nodes#657
Merged
Merged
Conversation
2ecb14e to
5da2a0a
Compare
no idea why it was id(expr), doesn't seem to make sense...
…collect_materialized_nodes
a971e25 to
8a605b0
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
This PR introduces a unified mechanism for identifying “implicitly materialized” nodes in a Pytato DAG and wires it into both MPMS materialization and distributed partitioning, including improved handling of LoopyCall bindings/results.
Changes:
- Add
MaterializedNodeCollector/collect_materialized_nodesinpytato.analysis. - Update MPMS materialization to account for implicitly materialized nodes (e.g.
CSRMatmul,LoopyCallins/outs, distributed send/recv). - Replace the distributed partitioner’s bespoke materialized-node collector with the new analysis helper.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
pytato/transform/materialize.py |
MPMS materializer now consults collect_materialized_nodes and correctly traverses/handles LoopyCall bindings/results and other implicitly-materialized nodes. |
pytato/distributed/partition.py |
Removes the local materialized-node collector and uses collect_materialized_nodes to seed partition placement logic. |
pytato/analysis/__init__.py |
Adds the new collector + public API for determining which nodes are materialized. |
.basedpyright/baseline.json |
Updates typing baseline to reflect the new/modified typing surface. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
the idea is that one should be able to call collect_materialize_nodes on a subexpression, using include_outputs=False to indicate that the root is not a "true" output
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds a
collect_materialized_nodesfunction that identifies all materialized nodes, including those that are implicitly materialized based on their use by another node (e.g.,LoopyCallins/outs,DistributedSenddata). It also attempts to use this function to unify the identification of materialized nodes elsewhere in the code, namely MPMS materialization and distributed partitioning.Summary of changes:
MaterializedNodeCollectorandcollect_materialized_nodesto analysis._MaterializedArrayCollectorin distributed.MPMSMaterializernot traversingLoopyCall's bindings.MPMSMaterializerto be aware of implicitly-materialized arrays.