Change node group partitioning to be with all nodes, to keep partition ids in top sort order #12871
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/12871
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 Cancelled Job, 1 Unrelated FailureAs of commit 24cfc84 with merge base 66e5591 ( CANCELLED JOB - The following job was cancelled. Please retry:
BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
This PR needs a
|
| return True | ||
|
|
||
| def _process_node_groups( | ||
| def _process_remaining_nodes( |
There was a problem hiding this comment.
let's change this to process_all_nodes or something like that
| if node in assignment or not self._is_node_supported(node): | ||
| continue | ||
|
|
||
| if node in processed_nodes: |
There was a problem hiding this comment.
what about this check? The purpose was that when traversing nodes in the graph module, we could get two nodes from the same group, and we could potentially add the group twice?
| # Update partition map | ||
| for node in partition.nodes: | ||
| for user in node.users: | ||
| target_id = assignment.get(user) |
There was a problem hiding this comment.
``
| target_id = assignment.get(user) | |
| target_id = assignment.get(user, None) |
Just for readability
leafs1
force-pushed
the
partitionerOptimize
branch
3 times, most recently
from
dd3542a to
988c4b2
Compare
…n ids in top sort order
leafs1
force-pushed
the
partitionerOptimize
branch
from
988c4b2 to
24cfc84
Compare
Summary
With the current partitioner implementation, we need to traverse the partitions in reverse top sort order for the merging algorithm, as we have some invariants when doing cycle detection and checking downstream dependencies. The current problem is that we first form the group partitions and give them incrementing ids. Since they aren't guaranteed to be in top sort order, when we assign the remaining partitions to the ungrouped nodes, their ids will start from where the grouped ids left off, breaking our invariant. This PR assigns ids in order of reverse top sort, where if we encounter a group node we create the partition for the whole group and if we just encounter ungrouped nodes we create the partition just for that node.
Test plan
All previous tests for functionality still pass