Merge back to back slices on the same dim

Summary: If we have back to back slices on the same dimension, we can remove the top slice and just perform the second one.

Differential Revision: D102278253

Author: Andrew Grebenisan

backends/cadence/aot/simplify_ops.py

@@ -13,10 +13,13 @@
 from typing import cast, Optional
 
 import torch
+
 from executorch.backends.cadence.aot.pass_utils import (
     CadencePassAttribute,
+    get_arg,
     register_cadence_pass,
     RemoveOrReplacePassInterface,
+    set_arg,
 )
 from executorch.backends.cadence.aot.utils import rebind
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -174,6 +177,61 @@ def maybe_remove_or_replace(self, node: Node) -> bool:
         return True
 
 
+@register_cadence_pass(CadencePassAttribute(opt_level=0))
+class BypassRedundantSliceChainPass(RemoveOrReplacePassInterface):
+    """Bypass chained slices on the same dim by slicing directly from the source.
+
+    When a slice_copy is followed by another slice_copy on the same dimension
+    with step=1, the child can read directly from the parent's input with
+    merged indices, eliminating the intermediate slice.
+    """
+
+    @property
+    def targets(self) -> list[EdgeOpOverload]:
+        return [exir_ops.edge.aten.slice_copy.Tensor]
+
+    def maybe_remove_or_replace(self, node: Node) -> bool:
+        parent_input = get_arg(node, "input", Node)
+        parent_dim = get_arg(node, "dim", int)
+        parent_start = get_arg(node, "start", Optional[int])
+        parent_end = get_arg(node, "end", Optional[int])
+        parent_step = get_arg(node, "step", int)
+
+        if parent_step != 1 or parent_start is None or parent_end is None:
+            return False
+
+        input_shape = parent_input.meta["val"].shape
+
+        modified = False
+        for child in list(node.users):
+            if child.target != exir_ops.edge.aten.slice_copy.Tensor:
+                continue
+
+            child_dim = get_arg(child, "dim", int)
+            if child_dim != parent_dim:
+                continue
+
+            child_start = get_arg(child, "start", Optional[int])
+            child_end = get_arg(child, "end", Optional[int])
+            child_step = get_arg(child, "step", int)
+
+            if child_step != 1 or child_start is None or child_end is None:
+                continue
+
+            new_start = parent_start + child_start
+            new_end = parent_start + child_end
+
+            if new_end > input_shape[parent_dim]:
+                continue
+
+            child.replace_input_with(node, parent_input)
+            set_arg(child, "start", new_start)
+            set_arg(child, "end", new_end)
+            modified = True
+
+        return modified
+
+
 @register_cadence_pass(CadencePassAttribute(opt_level=0))
 class BindOptionalArgsPass(ExportPass):
     """Bind all optional args and kwargs."""
@@ -217,5 +275,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 class CadenceSimplifyOpsInGraph:
     passes = [
         SimplifySliceOpPass,
+        BypassRedundantSliceChainPass,
         BindOptionalArgsPass,
     ]

backends/cadence/aot/tests/test_simplify_ops_passes.py

@@ -15,10 +15,11 @@
 from executorch.backends.cadence.aot.pass_utils import count_node
 from executorch.backends.cadence.aot.simplify_ops import (
     BindOptionalArgsPass,
+    BypassRedundantSliceChainPass,
     SimplifySliceOpPass,
 )
 from executorch.backends.cadence.aot.typing_stubs import expand
-from executorch.backends.test.graph_builder import single_op_builder
+from executorch.backends.test.graph_builder import GraphBuilder, single_op_builder
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 from torch.utils import _pytree as pytree
@@ -158,3 +159,137 @@ def test_simplify_slice_op_args(self) -> None:
         modified_slice_copy = list(gm.graph.nodes)[1]
         self.assertEqual(modified_slice_copy.args[1:], (1, None, 3, 1))
         self.assertEqual(modified_slice_copy.kwargs, {})
+
+
+class TestBypassRedundantSliceChainPass(unittest.TestCase):
+    def test_basic_chain_bypass(self) -> None:
+        """conv → slice(dim=3, 0:78) → slice(dim=3, 0:60) → direct slice(dim=3, 0:60)."""
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(2, 3, 4, 80))
+        parent = builder.call_operator(
+            op=exir_ops.edge.aten.slice_copy.Tensor,
+            args=(x, 3, 0, 78, 1),
+        )
+        child = builder.call_operator(
+            op=exir_ops.edge.aten.slice_copy.Tensor,
+            args=(parent, 3, 0, 60, 1),
+        )
+        builder.output([child])
+        original = builder.get_graph_module()
+
+        result = transform_and_check_numerics(
+            original,
+            (torch.randn(2, 3, 4, 80),),
+            BypassRedundantSliceChainPass(),
+            "BypassRedundantSliceChainPass",
+        )
+        self.assertTrue(result.modified)
+
+        # Parent slice should be dead-code-eliminated since child no longer uses it.
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.slice_copy.Tensor), 1
+        )
+
+    def test_chain_with_offset(self) -> None:
+        """slice(dim=1, 10:50) → slice(dim=1, 5:20) → direct slice(dim=1, 15:30)."""
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(4, 64))
+        parent = builder.call_operator(
+            op=exir_ops.edge.aten.slice_copy.Tensor,
+            args=(x, 1, 10, 50, 1),
+        )
+        child = builder.call_operator(
+            op=exir_ops.edge.aten.slice_copy.Tensor,
+            args=(parent, 1, 5, 20, 1),
+        )
+        builder.output([child])
+        original = builder.get_graph_module()
+
+        result = transform_and_check_numerics(
+            original,
+            (torch.randn(4, 64),),
+            BypassRedundantSliceChainPass(),
+            "BypassRedundantSliceChainPass",
+        )
+        self.assertTrue(result.modified)
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.slice_copy.Tensor), 1
+        )
+
+    def test_parent_kept_with_other_users(self) -> None:
+        """Parent slice has another user besides the child → parent stays."""
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(2, 3, 4, 80))
+        parent = builder.call_operator(
+            op=exir_ops.edge.aten.slice_copy.Tensor,
+            args=(x, 3, 0, 78, 1),
+        )
+        child = builder.call_operator(
+            op=exir_ops.edge.aten.slice_copy.Tensor,
+            args=(parent, 3, 0, 60, 1),
+        )
+        neg = builder.call_operator(op=exir_ops.edge.aten.neg.default, args=(parent,))
+        builder.output([child, neg])
+        original = builder.get_graph_module()
+
+        result = transform_and_check_numerics(
+            original,
+            (torch.randn(2, 3, 4, 80),),
+            BypassRedundantSliceChainPass(),
+            "BypassRedundantSliceChainPass",
+        )
+        self.assertTrue(result.modified)
+        # Parent kept (has neg user), child bypassed → 2 slices remain.
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.slice_copy.Tensor), 2
+        )
+
+    def test_different_dims_no_change(self) -> None:
+        """Chained slices on different dims → no change."""
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(8, 16, 32))
+        parent = builder.call_operator(
+            op=exir_ops.edge.aten.slice_copy.Tensor,
+            args=(x, 1, 0, 10, 1),
+        )
+        child = builder.call_operator(
+            op=exir_ops.edge.aten.slice_copy.Tensor,
+            args=(parent, 2, 0, 5, 1),
+        )
+        builder.output([child])
+        original = builder.get_graph_module()
+
+        result = cast(PassResult, BypassRedundantSliceChainPass()(original))
+        self.assertFalse(result.modified)
+
+    def test_step_not_one_no_change(self) -> None:
+        """Parent has step != 1 → no change."""
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(4, 64))
+        parent = builder.call_operator(
+            op=exir_ops.edge.aten.slice_copy.Tensor,
+            args=(x, 1, 0, 60, 2),
+        )
+        child = builder.call_operator(
+            op=exir_ops.edge.aten.slice_copy.Tensor,
+            args=(parent, 1, 0, 10, 1),
+        )
+        builder.output([child])
+        original = builder.get_graph_module()
+
+        result = cast(PassResult, BypassRedundantSliceChainPass()(original))
+        self.assertFalse(result.modified)
+
+    def test_no_chain_no_change(self) -> None:
+        """Single slice with no slice user → no change."""
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(4, 64))
+        sliced = builder.call_operator(
+            op=exir_ops.edge.aten.slice_copy.Tensor,
+            args=(x, 1, 0, 32, 1),
+        )
+        builder.output([sliced])
+        original = builder.get_graph_module()
+
+        result = cast(PassResult, BypassRedundantSliceChainPass()(original))
+        self.assertFalse(result.modified)