Reland "Arm backend: Add transpose propagation pass" (#20748) - #20748
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20748
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 4aec2e4 with merge base a588b26 ( 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. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
|
@rascani has imported this pull request. If you are a Meta employee, you can view this in D110793024. |
meta-codesync
Bot
changed the title
a0680ed to
8ac44af
Compare
|
@rascani has exported this pull request. If you are a Meta employee, you can view the originating Diff in D110793024. |
8ac44af to
6557ec2
Compare
6557ec2 to
65c1f05
Compare
65c1f05 to
4aec2e4
Compare
| def _would_strand_layout_op_on_wider_elements( | ||
| self, next_node: torch.fx.Node | ||
| ) -> bool: | ||
| """Whether crossing next_node leaves the layout op on wider elements. | ||
|
|
||
| Crossing next_node upward moves the layout op onto next_node's input. When | ||
| next_node narrows the dtype (its input has wider elements than its output, | ||
| e.g. an int32 to int8 rescale), the layout op then has more bytes to move, so | ||
| block the crossing and keep it on the narrow side. The sole exception is a | ||
| placeholder input consumed only by next_node: moving the layout op onto such | ||
| a graph input folds it into the input's dim_order at no runtime cost. A | ||
| placeholder with other consumers cannot be relaid out for free, so the | ||
| crossing is still blocked. Valid for upward propagation only, where | ||
| next_node's single input is where the layout op would land. | ||
|
|
||
| """ | ||
| input_nodes = next_node.all_input_nodes | ||
| if len(input_nodes) != 1: | ||
| return False | ||
| node_val = next_node.meta.get("val") | ||
| producer = input_nodes[0] | ||
| producer_val = producer.meta.get("val") | ||
| if not isinstance(node_val, torch.Tensor) or not isinstance( | ||
| producer_val, torch.Tensor | ||
| ): | ||
| return False | ||
| if producer_val.element_size() <= node_val.element_size(): | ||
| return False | ||
| return producer.op != "placeholder" or len(producer.users) != 1 |
There was a problem hiding this comment.
@AdrianLundell - This is the main significant change from the original (plus its usage at line 412). We had some cases where this pass would insert a transpose between an Add -> Rescale, which wound up blocking a fusion and the resulting graph hit a regor pass bug. I'll fix that bug separately, but figured it would still be good to block the propagation here to minimize the data movement.
There was a problem hiding this comment.
Makes sense, thanks!
| next_nodes = list(self._get_next_nodes(frontier)) | ||
|
|
||
| if len(next_nodes) == 0: | ||
| assert frontier.op in ( |
There was a problem hiding this comment.
Fixed a nascent bug here. It was previously assert node.op, which doesn't match the comment or the intention. I don't think it caused any problems, but wanted to call it out.
digantdesai
left a comment
digantdesai
left a comment
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
Differential Revision: D112005256 Pull Request resolved: #20947
3 participants
Summary:
Reverts #20742 and relands #20625.
Differential Revision: D110793024
Pulled By: rascani