From f0c8c8fa95da45c6dad95bda00272b6f56f90cbc Mon Sep 17 00:00:00 2001 From: Matthias Cremon Date: Thu, 20 Aug 2026 12:21:39 -0700 Subject: [PATCH] Make layout propagation fixed-point and broadcast-safe (#21956) Summary: Run view/permute propagation to a fixed point for every supported TOSA specification instead of conditionally disabling iterative propagation for one specification family. A layout transform may be element-order invariant at a reshape while the reshape output shape remains semantically significant to downstream broadcasting. Make permutation-sink termination conditional on downstream layout invariance, preserving the transform when another consumer depends on its logical axis placement. This allows all targets to use the same propagation algorithm while retaining conservative correctness checks around elementwise inputs, reshape sinks, and keep-dimension reductions. Reviewed By: rascani Differential Revision: D116683130 --- .../propagate_view_copy_permute_pass.py | 21 +--- .../test_propagate_permutes_views_pass.py | 48 +++++++- .../remove_permutes_around_elementwise_ops.py | 50 +++++++-- .../test/test_permute_optimization_passes.py | 104 ++++++++++++++++++ 4 files changed, 192 insertions(+), 31 deletions(-) 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.