Make layout propagation fixed-point and broadcast-safe (#21956)

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()

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")

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

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.