Fix stale dims and constants in downward permute propagation (#24739)

Summary:

As PropagateViewCopyPermuteDownPass walks a permute down, it rewrites each node it crosses for the permute's input layout: the dims of a slice or reduction, and a binary op's constant input, which it permutes in place. When the fork at the end of the walk could not be split, the pass fell back to a pointwise region move planned from the permute's original position. That region stops at the first node it cannot include, such as a slice, a reduction or a backend op that is not pointwise, and left the crossed nodes from there on applying their rewritten dims or constants to the permuted tensor.

For example, `permute(x, [0, 3, 1, 2]) -> clamp -> slice_copy(dim=1, 0, 4)` feeding both `permute(0, 2, 3, 1)` and `neg`, the pattern Arm's RewriteConvPass leaves around a channel split between convs, became `slice_copy(clamp(x), 2, 0, 4)`: for x of shape [1, 3, 5, 8] the outputs had shapes [1, 3, 4, 8] and [1, 8, 3, 4] instead of [1, 3, 5, 4] and [1, 4, 3, 5]. A keepdim sum in the same place made the pass raise while re-tracing, and a constant added after a crossed backend op that is not pointwise failed to broadcast or was broadcast along the wrong axis.

Plan the region move only when the walk has not advanced. Otherwise the permute moves past the frontier, where every rewrite is valid, and the next iteration retries the fork from there.

The test target now also depends on `//executorch/exir:lib`, for `to_edge`.

Differential Revision: D130233379

Pull Request resolved: #24739

Author: apullin

backends/transforms/propagate_view_copy_permute_pass.py

@@ -1198,6 +1198,13 @@ def _maybe_split_fork(
             self._apply_fork_split(node, frontier, producer, branch_splits)
             return True
 
+        if frontier is not node:
+            # The walk already rewrote the operations up to ``frontier`` for
+            # the permute's input, and a region planned from ``node`` may stop
+            # before some of them. Move ``node`` past ``frontier`` instead and
+            # retry the fork from there on the next iteration.
+            return False
+
         # If the branches use each other's results, try moving the dimension
         # reorder past the connected group of operations as one change.
         region_move = self._plan_pointwise_region_move(node)

backends/transforms/targets.bzl

@@ -405,6 +405,7 @@ def define_common_targets():
         deps = [
             "//caffe2:torch",
             ":propagate_view_copy_permute_pass",
+            "//executorch/exir:lib",
             "//executorch/exir/dialects:lib",
             "fbsource//third-party/pypi/pytest:pytest",
         ],

backends/transforms/test/test_propagate_view_copy_permute_pass.py

@@ -5,7 +5,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 
 import pytest
 import sympy  # type: ignore[import-untyped]
@@ -15,16 +15,19 @@
     PropagateViewCopyPermuteDownPass,
     PropagateViewCopyPermuteUpPass,
 )
+from executorch.exir import to_edge
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch._subclasses.fake_tensor import FakeTensorMode
 from torch.fx.experimental.symbolic_shapes import ShapeEnv
 
 PERMUTE = exir_ops.edge.aten.permute_copy.default
 ABS = exir_ops.edge.aten.abs.default
+EXP = exir_ops.edge.aten.exp.default
 NEG = exir_ops.edge.aten.neg.default
 SIGMOID = exir_ops.edge.aten.sigmoid.default
 ADD = exir_ops.edge.aten.add.Tensor
 SUM = exir_ops.edge.aten.sum.dim_IntList
+SLICE = exir_ops.edge.aten.slice_copy.Tensor
 WHERE = exir_ops.edge.aten.where.self
 
 
@@ -301,6 +304,96 @@ def test_downward_permute_crosses_reduction_over_all_dims() -> None:
     torch.testing.assert_close(result(values), expected)
 
 
+@pytest.mark.parametrize(
+    ("target", "args", "expected_dims"),
+    [(SLICE, (1, 0, 4), 3), (SUM, ([1], True), [3])],
+    ids=["slice", "keepdim_sum"],
+)
+def test_downward_permute_crosses_arg_update_op_before_unsplittable_fork(
+    target: Callable[..., object],
+    args: tuple[object, ...],
+    expected_dims: int | list[int],
+) -> None:
+    """Keep a crossed op's dims valid when the fork after it cannot be split.
+
+    Here ``P`` means permute:
+
+        P(x) -> abs -> op --+-> P -> output
+                            +-> neg -> output
+
+    Crossing ``op`` rewrites its dims for the unpermuted input, so the first
+    permute must not then move past ``abs`` alone.
+
+    """
+    values = torch.randn(1, 3, 5, 8)
+    graph = torch.fx.Graph()
+    x = graph.placeholder("x")
+    x.meta["val"] = values
+    permute = graph.call_function(PERMUTE, args=(x, [0, 3, 1, 2]))
+    permute.meta["val"] = values.permute(0, 3, 1, 2)
+    absolute = graph.call_function(ABS, args=(permute,))
+    absolute.meta["val"] = permute.meta["val"].abs()
+    crossed = graph.call_function(target, args=(absolute, *args))
+    crossed.meta["val"] = target(absolute.meta["val"], *args)
+    restored = graph.call_function(PERMUTE, args=(crossed, [0, 2, 3, 1]))
+    restored.meta["val"] = crossed.meta["val"].permute(0, 2, 3, 1)
+    neg = graph.call_function(NEG, args=(crossed,))
+    neg.meta["val"] = -crossed.meta["val"]
+    graph.output((restored, neg))
+    graph_module = torch.fx.GraphModule(torch.nn.Module(), graph)
+    expected = graph_module(values)
+
+    result = PropagateViewCopyPermuteDownPass().call(graph_module).graph_module
+    moved_op = next(node for node in result.graph.nodes if node.target == target)
+
+    assert moved_op.args[1] == expected_dims
+    torch.testing.assert_close(result(values), expected)
+
+
+class _BackendElementwiseExpPass(PropagateViewCopyPermuteDownPass):
+    def is_elementwise(self, node: torch.fx.Node) -> bool:
+        return node.target != EXP and super().is_elementwise(node)
+
+    def is_backend_elementwise(self, node: torch.fx.Node) -> bool:
+        return node.target == EXP
+
+
+class _BiasAfterExp(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.bias = torch.nn.Parameter(torch.randn(1, 8, 1, 1))
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        y = torch.exp(x.permute(0, 3, 1, 2).abs()) + self.bias
+        return y.permute(0, 2, 3, 1), -y
+
+
+def test_downward_permute_crosses_constant_before_unsplittable_fork() -> None:
+    """Keep a crossed constant valid when the fork after it cannot be split.
+
+    Here ``P`` means permute and ``exp`` stands in for a backend op that the
+    pass may cross but not include in a pointwise region:
+
+        P(x) -> abs -> exp -> add(bias) --+-> P -> output
+                                          +-> neg -> output
+
+    Crossing ``add`` permutes ``bias`` for the unpermuted input, so the first
+    permute must not then move past ``abs`` alone.
+
+    """
+    values = torch.randn(1, 3, 5, 8)
+    model = _BiasAfterExp()
+    expected = model(values)
+    exported_program = to_edge(torch.export.export(model, (values,))).exported_program()
+
+    result = _BackendElementwiseExpPass(exported_program=exported_program).call(
+        exported_program.graph_module
+    )
+    bias = exported_program.state_dict["bias"]
+
+    torch.testing.assert_close(result.graph_module(bias, values), expected)
+
+
 def test_upward_permute_crosses_rank_expanding_binary_after_reduction() -> None:
     values = torch.arange(6, dtype=torch.float32).reshape(2, 3)
     bias_value = torch.ones((1, 1, 1, 1), dtype=torch.float32)