Normalize permute dims in PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView

Summary:
When the pass moves a permute below a view that inserts or removes a size-1 dim, it shifts each permute dim by comparing it with the inserted or removed axis. That comparison assumes non-negative dims. A negative dim that points below the axis keeps its value, so it lands on a different axis once the rank changes and the rewritten permute is invalid: `permute(-3, 2, 1).unsqueeze(2)` becomes `[1, 3, 2, 1]`, and `permute(-4, 2, 1, 3).squeeze(1)` keeps a `-4` at rank 3.

Normalize the dims to their non-negative form before rewriting.

Differential Revision: D128123955

Pull Request resolved: #24248

Author: pssrawat

backends/transforms/postpone_permute_below_squeeze_view.py

@@ -12,6 +12,7 @@
 import torch
 import torch.fx
 from executorch.backends.transforms.channels_last_layout import PERMUTE_COPY_TARGETS
+from executorch.backends.transforms.dim_maps import _normalize_dims
 from executorch.backends.transforms.permute_pass_utils import (
     get_shape,
     RemoveOrReplacePassInterface,
@@ -78,7 +79,9 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool:
 
         permute_node_shape = [*cast(list, get_shape(node.graph.owning_module, node))]
 
-        permute_dims = cast(list, node.args[1])
+        permute_dims = _normalize_dims(
+            cast(list, node.args[1]), len(permute_node_shape)
+        )
         view_node = users[0]
 
         if view_node.meta.get("val") is None:

backends/transforms/targets.bzl

@@ -750,6 +750,7 @@ def define_common_targets():
         ],
         deps = [
             ":channels_last_layout",
+            ":dim_maps",
             "//caffe2:torch",
             "//executorch/exir:pass_base",
             "//executorch/exir/dialects:lib",

backends/transforms/test/test_permute_optimization_passes.py

@@ -591,6 +591,84 @@ def test_negative_not_squeeze_like(self) -> None:
         self.assertEqual(targets[0], exir_ops.edge.aten.view_copy.default)
         self.assertEqual(targets[1], exir_ops.edge.aten.permute_copy.default)
 
+    def test_negative_permute_dims_below_unsqueeze_axis(self) -> None:
+        builder = GraphBuilder()
+        x_data = torch.randn(2, 3, 4)
+        x = builder.placeholder("x", x_data)
+        p = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(x, [-3, 2, 1])
+        )
+        v = builder.call_operator(
+            op=exir_ops.edge.aten.view_copy.default, args=(p, [2, 4, 1, 3])
+        )
+        builder.output([v])
+        original = builder.get_graph_module()
+        gm_before = copy.deepcopy(original)
+
+        result = cast(
+            PassResult,
+            PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView().call(original),
+        )
+
+        self.assertTrue(result.modified)
+        gm = result.graph_module
+        self.assertEqual(
+            get_compute_nodes(gm),
+            [
+                exir_ops.edge.aten.view_copy.default,
+                exir_ops.edge.aten.permute_copy.default,
+            ],
+        )
+        (permute,) = gm.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.permute_copy.default
+        )
+        self.assertEqual(permute.args[1], [0, 3, 2, 1])
+        validate_numerics(
+            gm_before,
+            gm,
+            [x_data],
+            "PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView",
+        )
+
+    def test_negative_permute_dims_below_squeeze_axis(self) -> None:
+        builder = GraphBuilder()
+        x_data = torch.randn(2, 3, 1, 4)
+        x = builder.placeholder("x", x_data)
+        p = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(x, [-4, 2, 1, 3])
+        )
+        v = builder.call_operator(
+            op=exir_ops.edge.aten.view_copy.default, args=(p, [2, 3, 4])
+        )
+        builder.output([v])
+        original = builder.get_graph_module()
+        gm_before = copy.deepcopy(original)
+
+        result = cast(
+            PassResult,
+            PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView().call(original),
+        )
+
+        self.assertTrue(result.modified)
+        gm = result.graph_module
+        self.assertEqual(
+            get_compute_nodes(gm),
+            [
+                exir_ops.edge.aten.view_copy.default,
+                exir_ops.edge.aten.permute_copy.default,
+            ],
+        )
+        (permute,) = gm.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.permute_copy.default
+        )
+        self.assertEqual(permute.args[1], [0, 1, 2])
+        validate_numerics(
+            gm_before,
+            gm,
+            [x_data],
+            "PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView",
+        )
+
 
 class FuseTransposeOrPermuteOpPairsTest(unittest.TestCase):
     def test_per_tensor_qdq_is_bypassed(self) -> None: