Skip to content

Commit 1a9f86d

Browse files
authored
Transforms: Fix permute propagation through dependent pointwise branches (pytorch#22738)
Validate connected pointwise operations before moving a permute to their single output. Leave multi-output and layout-dependent groups unchanged. Fix a problem that popped up in pytorch#21065 ### Test plan Add regression coverage and re-enable the Arm mixed-precision pixel shuffle tests. Signed-off-by: Zingo Andersen <Zingo.Andersen@arm.com>
1 parent 1cd22ed commit 1a9f86d

4 files changed

Lines changed: 408 additions & 15 deletions

File tree

‎backends/arm/test/ops/test_pixel_shuffling.py‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -34,10 +34,6 @@
3434
"rand_4d_channels_last": "Known U55 partitioning limitation for large 4D pixel shuffle layouts.",
3535
}
3636

37-
mixed_precision_xfails = {
38-
"rand_4d_channels_last": "Permute propagation stops at f(x, g(x)) shapes such as the round decomposition.",
39-
}
40-
4137

4238
class PixelUnShuffle(nn.Module):
4339

@@ -114,9 +110,7 @@ def test_pixel_unshuffle_tosa_FP(test_data: input_t1):
114110
pipeline.run()
115111

116112

117-
@common.parametrize(
118-
"test_data", PixelUnShuffle.test_data_generators, xfails=mixed_precision_xfails
119-
)
113+
@common.parametrize("test_data", PixelUnShuffle.test_data_generators)
120114
def test_pixel_unshuffle_no_target_tosa_mixed_precision(test_data: input_t1):
121115
inputs, expected_transposes = test_data()
122116
pipeline = TosaPipelineINT[input_t1](
@@ -146,9 +140,7 @@ def test_pixel_shuffle_tosa_FP(test_data: input_t1):
146140
pipeline.run()
147141

148142

149-
@common.parametrize(
150-
"test_data", PixelShuffle.test_data_generators, xfails=mixed_precision_xfails
151-
)
143+
@common.parametrize("test_data", PixelShuffle.test_data_generators)
152144
def test_pixel_shuffle_no_target_tosa_mixed_precision(test_data: input_t1):
153145
inputs, expected_transposes = test_data()
154146
pipeline = TosaPipelineINT[input_t1](

‎backends/arm/test/passes/test_propagate_permutes_views_pass.py‎

Lines changed: 104 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
# This source code is licensed under the BSD-style license found in the
44
# LICENSE file in the root directory of this source tree.
55

6-
from collections.abc import Callable
6+
from collections.abc import Callable, Sequence
77
from typing import Tuple
88

99
import pytest
@@ -720,6 +720,109 @@ def test_down_pass_moves_matching_input_permutations_after_binary_op() -> None:
720720
assert targets.index(ADD) < targets.index(PERMUTE)
721721

722722

723+
def test_down_pass_moves_permute_through_single_exit_pointwise_region() -> None:
724+
"""Move a dimension reorder after related per-value operations.
725+
726+
Here ``P`` means permute, Both ``floor`` and ``sub`` process values independently:
727+
728+
Before: P(x) --+-----------------> sub -> output
729+
| ^
730+
+-> floor -----------+
731+
732+
After: x --+-----------------> sub -> P -> output
733+
| ^
734+
+-> floor -----------+
735+
736+
The first check confirms that the permute moved after both operations. The second
737+
check confirms that the rewritten graph still calculates the same values.
738+
739+
"""
740+
graph = torch.fx.Graph()
741+
x = graph.placeholder("x")
742+
x.meta["val"] = torch.empty((1, 2, 3, 4))
743+
permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1]))
744+
permute.meta["val"] = torch.empty((1, 3, 4, 2))
745+
floor = graph.call_function(FLOOR, args=(permute,))
746+
floor.meta["val"] = torch.empty((1, 3, 4, 2))
747+
sub = graph.call_function(SUB, args=(permute, floor))
748+
sub.meta["val"] = torch.empty((1, 3, 4, 2))
749+
graph.output(sub)
750+
751+
x_data = torch.randn(1, 2, 3, 4)
752+
expected = x_data.permute(0, 2, 3, 1)
753+
expected = expected - expected.floor()
754+
graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteDownPass)
755+
call_nodes = [
756+
node for node in graph_module.graph.nodes if node.op == "call_function"
757+
]
758+
targets = [node.target for node in call_nodes]
759+
760+
assert targets.index(FLOOR) < targets.index(SUB) < targets.index(PERMUTE)
761+
torch.testing.assert_close(graph_module(x_data), expected)
762+
763+
764+
class _BlockFloorToSubPass(PropagateViewCopyPermuteDownPass):
765+
def blocks_moving(
766+
self,
767+
moving_node: torch.fx.Node,
768+
frontier: torch.fx.Node,
769+
next_nodes: Sequence[torch.fx.Node],
770+
) -> bool:
771+
return frontier.target == FLOOR and any(
772+
next_node.target == SUB for next_node in next_nodes
773+
)
774+
775+
776+
def test_down_pass_checks_each_path_into_shared_pointwise_node() -> None:
777+
"""Keep a permute when a backend rejects one path into a shared node."""
778+
graph = torch.fx.Graph()
779+
x = graph.placeholder("x")
780+
x.meta["val"] = torch.empty((1, 2, 3, 4))
781+
permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1]))
782+
permute.meta["val"] = torch.empty((1, 3, 4, 2))
783+
floor = graph.call_function(FLOOR, args=(permute,))
784+
floor.meta["val"] = torch.empty((1, 3, 4, 2))
785+
sub = graph.call_function(SUB, args=(permute, floor))
786+
sub.meta["val"] = torch.empty((1, 3, 4, 2))
787+
graph.output(sub)
788+
789+
targets = _run_pass_on_graph(graph, _BlockFloorToSubPass)
790+
791+
assert targets.index(PERMUTE) < targets.index(FLOOR) < targets.index(SUB)
792+
793+
794+
def test_down_pass_keeps_permute_before_multiple_exit_pointwise_region() -> None:
795+
"""Keep a dimension reorder in place when two results are returned.
796+
797+
Here ``P`` means permute. Both ``floor`` and ``sub`` are returned from the graph:
798+
799+
P(x) --+-----------------> sub -> output 1
800+
| ^
801+
+-> floor -----------+
802+
|
803+
+---------------> output 0
804+
805+
Moving the permute would require adding a separate reorder before each output.
806+
This rewrite deliberately handles only one outgoing result, so the graph
807+
must remain unchanged.
808+
809+
"""
810+
graph = torch.fx.Graph()
811+
x = graph.placeholder("x")
812+
x.meta["val"] = torch.empty((1, 2, 3, 4))
813+
permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1]))
814+
permute.meta["val"] = torch.empty((1, 3, 4, 2))
815+
floor = graph.call_function(FLOOR, args=(permute,))
816+
floor.meta["val"] = torch.empty((1, 3, 4, 2))
817+
sub = graph.call_function(SUB, args=(permute, floor))
818+
sub.meta["val"] = torch.empty((1, 3, 4, 2))
819+
graph.output((floor, sub))
820+
821+
targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass)
822+
823+
assert targets.index(PERMUTE) < targets.index(FLOOR) < targets.index(SUB)
824+
825+
723826
def test_down_pass_keeps_lower_rank_singleton_input_before_binary_op() -> None:
724827
graph = torch.fx.Graph()
725828
x = graph.placeholder("x")

0 commit comments

Comments
 (0)