diff --git a/backends/arm/_passes/propagate_view_copy_permute_pass.py b/backends/arm/_passes/propagate_view_copy_permute_pass.py index a54c8f28312..657d3f9225b 100644 --- a/backends/arm/_passes/propagate_view_copy_permute_pass.py +++ b/backends/arm/_passes/propagate_view_copy_permute_pass.py @@ -14,7 +14,6 @@ import torch from executorch.backends.arm._passes.dim_maps import PermuteMap, ViewMap from executorch.backends.arm.tosa.mapping import TosaSpecialDtype -from executorch.backends.arm.tosa.specification import get_context_spec from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult @@ -97,10 +96,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: if result.modified: graph_module = self._retrace(graph_module) - # Do not run for Ethos-U85 since this exposes a numerical issue - # There is no target meta-data at this stage so use INT+cf as proxy - # To be removed after MLBEDSW-11805 - while not self._is_u85_like_tosa_int_cf(): + while True: iteration_modified = False for node in list(graph_module.graph.nodes): if node.target in self._TARGETS: @@ -129,21 +125,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: return PassResult(graph_module, modified) - def _is_u85_like_tosa_int_cf(self) -> bool: - if self.compile_spec is not None: - tosa_spec = self.compile_spec.tosa_spec - else: - try: - tosa_spec = get_context_spec() - except RuntimeError: - return False - - return ( - tosa_spec.support_integer() - and not tosa_spec.support_float() - and tosa_spec.support_extension("cf") - ) - def _retrace(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule: graph_module.graph.eliminate_dead_code() graph_module.graph.lint() diff --git a/backends/arm/test/passes/test_propagate_permutes_views_pass.py b/backends/arm/test/passes/test_propagate_permutes_views_pass.py index 79207eed9d4..9eae96aa7ad 100644 --- a/backends/arm/test/passes/test_propagate_permutes_views_pass.py +++ b/backends/arm/test/passes/test_propagate_permutes_views_pass.py @@ -373,7 +373,7 @@ def test_down_pass_moves_permute_after_transparent_chain() -> None: assert targets.index(RELU) < targets.index(NEG) < targets.index(PERMUTE) -def test_down_pass_skips_propagation_for_u85_like_tosa_int_cf() -> None: +def test_down_pass_propagates_for_u85_like_tosa_int_cf() -> None: graph = torch.fx.Graph() x = graph.placeholder("x") x.meta["val"] = torch.empty((1, 2, 3, 4)) @@ -388,7 +388,7 @@ def test_down_pass_skips_propagation_for_u85_like_tosa_int_cf() -> None: with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT+cf")): targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) - assert targets.index(PERMUTE) < targets.index(RELU) < targets.index(NEG) + assert targets.index(RELU) < targets.index(NEG) < targets.index(PERMUTE) def test_down_pass_still_canonicalizes_for_u85_like_tosa_int_cf() -> None: @@ -941,6 +941,50 @@ def test_down_pass_moves_permutation_after_reduction() -> None: assert transform.meta["val"].shape == torch.Size((1, 3, 4, 1)) +@pytest.mark.parametrize("mean_first", [False, True]) +def test_down_pass_keeps_permute_before_reduction_with_layout_dependent_user( + mean_first: bool, +) -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 512, 796)) + direct = graph.placeholder("direct") + direct.meta["val"] = torch.empty((1, 796, 512)) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 1])) + permute.meta["val"] = torch.empty((1, 796, 512)) + relu = graph.call_function(RELU, args=(permute,)) + relu.meta["val"] = torch.empty((1, 796, 512)) + mean = graph.call_function(MEAN, args=(relu, [1], True)) + mean.meta["val"] = torch.empty((1, 1, 512)) + sub_args = (mean, direct) if mean_first else (direct, mean) + sub = graph.call_function(SUB, args=sub_args) + sub.meta["val"] = torch.empty((1, 796, 512)) + graph.output(sub) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteDownPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + mean = next(node for node in call_nodes if node.target == MEAN) + sub = next(node for node in call_nodes if node.target == SUB) + mean_input_shape = mean.all_input_nodes[0].meta["val"].shape + mean_output_shape = mean.meta["val"].shape + sub_input_shapes = [ + input_node.meta["val"].shape for input_node in sub.all_input_nodes + ] + reduction_dims = [dim % len(mean_input_shape) for dim in mean.args[1]] + + assert all(mean_output_shape[dim] == 1 for dim in reduction_dims) + assert all( + output_dim == 1 if dim in reduction_dims else output_dim == input_dim + for dim, (input_dim, output_dim) in enumerate( + zip(mean_input_shape, mean_output_shape) + ) + ) + assert torch.broadcast_shapes(*sub_input_shapes) == sub.meta["val"].shape + assert torch.Size((1, 512, 1)) not in sub_input_shapes + + def test_down_pass_splits_permute_over_elementwise_fanout() -> None: graph = torch.fx.Graph() x = graph.placeholder("x") diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 6e916dfe50a..b37260ed764 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -162,9 +162,9 @@ def _is_permutation_sink_view(self, node: torch.fx.Node) -> bool: Flattening such a tensor -- e.g. the ``[1, C, 1, 1] -> [1, C]`` after a global pool -- is permutation-invariant: every layout of the input produces the identical output (the single non-unit run of elements is - contiguous regardless of which axis holds it). A permutation propagating - into it therefore simply dies, so the region can terminate here with no - compensating permute. + contiguous regardless of which axis holds it). The region may terminate + here without a compensating permute when downstream consumers do not use + the output shape for layout-dependent broadcasting. """ if node.target not in self._VIEW_OPS: return False @@ -176,6 +176,40 @@ def _is_permutation_sink_view(self, node: torch.fx.Node) -> bool: non_unit = [d for d in shape if not (isinstance(d, int) and d == 1)] return len(non_unit) <= 1 + def _sink_users_are_layout_invariant(self, sink: torch.fx.Node) -> bool: + """Return whether dropping layout at ``sink`` is safe for its consumers.""" + frontier = [(user, sink) for user in sink.users] + visited: set[torch.fx.Node] = set() + while frontier: + node, producer = frontier.pop() + if node in visited: + continue + visited.add(node) + + if node.op == "output": + continue + if node.target == exir_ops.edge.aten.permute_copy.default: + # This explicit transform re-establishes the downstream layout, + # so consumers beyond it do not depend on the sink's layout. + continue + if self._is_permutation_sink_view(node): + continue + + tensor_inputs = [ + input_node + for input_node in node.all_input_nodes + if input_node.meta.get("val") is not None + ] + if any( + input_node is not producer and input_node.meta["val"].numel() != 1 + for input_node in tensor_inputs + ): + return False + if not self.is_node_permutable(node): + return False + frontier.extend((user, node) for user in node.users) + return True + def _inserted_unit_dim(self, node: torch.fx.Node) -> int | None: """Position of the size-1 dim ``node`` inserts, else None. @@ -514,12 +548,10 @@ def visit( # noqa: C901 elif user.op == "output": return False elif self._is_permutation_sink_view(user): - # The permutation dies at this reshape (see - # _is_permutation_sink_view), so terminate the region here with - # no compensating permute and no further downstream traversal. - # Checked before the rank-change handling below: a sink always - # terminates cleanly, whereas crossing it would leave the region - # hunting for an end permute that layout-invariance made moot. + # The tensor's element order is invariant at this reshape, but + # its output shape can still carry broadcast-axis meaning. + if not self._sink_users_are_layout_invariant(user): + return False continue elif not self.visit( user, subgraph, processed_nodes, downstream_end, downstream_start diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index cf4583357d7..86018e16fcb 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -1556,6 +1556,110 @@ def test_permutation_sink_view_splitting_the_non_unit_dim(self) -> None: "permutation_sink_view_splitting_the_non_unit_dim", ) + def test_permutation_sink_view_preserves_broadcast_layout(self) -> None: + x_data = torch.randn(1, 4, 1, 1) + direct_data = torch.randn(1, 8, 4) + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + direct = builder.placeholder("direct", direct_data) + permute = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 3, 1]) + ) + mul = builder.call_operator( + op=exir_ops.edge.aten.mul.Tensor, args=(permute, permute) + ) + view = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(mul, [1, 1, 4]) + ) + sub = builder.call_operator( + op=exir_ops.edge.aten.sub.Tensor, args=(direct, view) + ) + builder.output([sub]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + self.assertFalse(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 1 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data, direct_data], + "permutation_sink_view_preserves_broadcast_layout", + ) + + def test_permutation_sink_view_preserves_cat_layout(self) -> None: + x_data = torch.randn(1, 4, 1, 1) + direct_data = torch.randn(1, 7, 4) + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + direct = builder.placeholder("direct", direct_data) + permute = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 3, 1]) + ) + mul = builder.call_operator( + op=exir_ops.edge.aten.mul.Tensor, args=(permute, permute) + ) + view = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(mul, [1, 1, 4]) + ) + cat = builder.call_operator( + op=exir_ops.edge.aten.cat.default, args=([direct, view], 1) + ) + builder.output([cat]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + self.assertFalse(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 1 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data, direct_data], + "permutation_sink_view_preserves_cat_layout", + ) + + def test_permutation_sink_view_preserves_keyword_broadcast_layout(self) -> None: + x_data = torch.randn(1, 4, 1, 1) + direct_data = torch.randn(1, 8, 4) + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + direct = builder.placeholder("direct", direct_data) + permute = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 3, 1]) + ) + mul = builder.call_operator( + op=exir_ops.edge.aten.mul.Tensor, args=(permute, permute) + ) + view = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(mul, [1, 1, 4]) + ) + sub = builder.call_operator( + op=exir_ops.edge.aten.sub.Tensor, + args=(view,), + kwargs={"other": direct}, + ) + builder.output([sub]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + self.assertFalse(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 1 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data, direct_data], + "permutation_sink_view_preserves_keyword_broadcast_layout", + ) + def test_upstream_squeeze_view_rank_mismatch_no_crash(self) -> None: """Regression test for IndexError when a squeeze view_copy is reached via upstream traversal with a permutation at the view's output rank.