Move split/concat chain folding to OSS (pytorch#22628)

Summary:

Move `MergeSplitConcatChainPass` into `executorch.backends.transforms` package and implement it as a standard `PassBase`. The pass recognizes complete, ordered outputs from ATen `chunk`, `split`, and `split_with_sizes`, as well as edge `split_with_sizes_copy`, followed by `cat`.

Replace the chain with `view_copy` only when the input is contiguous, element counts match, and every axis crossed between the split and concat dimensions is singleton. This preserves flattened storage order for both pre-edge and post-edge consumers.

Reviewed By: DrJessop

Differential Revision: D119242414

Pull Request resolved: pytorch#22628

Author: ethansfng

backends/transforms/merge_split_concat_chain.py

@@ -0,0 +1,153 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import operator
+from collections.abc import Callable, Sequence, Set
+
+import torch
+from executorch.backends.transforms.permute_pass_utils import get_arg
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import PassResult
+from torch.fx import GraphModule, Node
+from torch.fx.passes.infra.pass_base import PassBase
+
+
+_RAW_SPLIT_TARGETS: frozenset[Callable[..., object]] = frozenset(
+    {
+        torch.ops.aten.chunk.default,
+        torch.ops.aten.split.Tensor,
+        torch.ops.aten.split_with_sizes.default,
+    }
+)
+_EDGE_SPLIT_TARGETS: frozenset[Callable[..., object]] = frozenset(
+    {exir_ops.edge.aten.split_with_sizes_copy.default}
+)
+
+
+def _normalize_dim(dim: int, rank: int) -> int:
+    normalized = dim + rank if dim < 0 else dim
+    assert 0 <= normalized < rank
+    return normalized
+
+
+def _ordered_split_inputs(
+    cat_node: Node,
+    split_targets: Set[Callable[..., object]],
+) -> tuple[Node, list[Node]] | None:
+    cat_inputs = get_arg(cat_node, "tensors")
+    if not isinstance(cat_inputs, Sequence) or not cat_inputs:
+        return None
+
+    getitem_nodes: list[Node] = []
+    for inp in cat_inputs:
+        if not isinstance(inp, Node) or inp.target != operator.getitem:
+            return None
+        getitem_nodes.append(inp)
+
+    split_node = getitem_nodes[0].args[0]
+    if not isinstance(split_node, Node) or split_node.target not in split_targets:
+        return None
+
+    split_outputs = split_node.meta["val"]
+    if not isinstance(split_outputs, (tuple, list)) or len(getitem_nodes) != len(
+        split_outputs
+    ):
+        return None
+    for index, getitem_node in enumerate(getitem_nodes):
+        if getitem_node.args[0] != split_node or getitem_node.args[1] != index:
+            return None
+    return split_node, getitem_nodes
+
+
+def _is_view_equivalent(
+    cat_node: Node,
+    split_node: Node,
+    getitem_nodes: Sequence[Node],
+) -> bool:
+    split_input = get_arg(split_node, "input", Node)
+    input_val = split_input.meta["val"]
+    output_val = cat_node.meta["val"]
+    if not isinstance(input_val, torch.Tensor) or not isinstance(
+        output_val, torch.Tensor
+    ):
+        return False
+    if not input_val.is_contiguous() or input_val.numel() != output_val.numel():
+        return False
+
+    rank = input_val.ndim
+    split_dim = _normalize_dim(get_arg(split_node, "dim", int), rank)
+    cat_dim = _normalize_dim(get_arg(cat_node, "dim", int), rank)
+
+    first_moved_dim, last_moved_dim = sorted((split_dim, cat_dim))
+    # Moving split parts across axes preserves flattened storage order only
+    # when every crossed axis is singleton.
+    for getitem_node in getitem_nodes:
+        value = getitem_node.meta["val"]
+        if not isinstance(value, torch.Tensor) or any(
+            size != 1 for size in value.shape[first_moved_dim:last_moved_dim]
+        ):
+            return False
+    return True
+
+
+class MergeSplitConcatChainPass(PassBase):
+    """Replace view-equivalent split/getitem/cat chains with a view.
+
+    For example, splitting a contiguous ``[2, 1, 6, 4, 4]`` tensor into
+    three parts along dimension 2 and concatenating them along dimension 1
+    is equivalent to a view with shape ``[2, 3, 2, 4, 4]``.
+    """
+
+    def _replace_cats(
+        self,
+        graph_module: GraphModule,
+        cat_target: Callable[..., object],
+        split_targets: Set[Callable[..., object]],
+        view_target: Callable[..., object],
+    ) -> bool:
+        modified = False
+        for cat_node in graph_module.graph.find_nodes(
+            op="call_function", target=cat_target
+        ):
+            match = _ordered_split_inputs(cat_node, split_targets)
+            if match is None:
+                continue
+            split_node, getitem_nodes = match
+            if not _is_view_equivalent(cat_node, split_node, getitem_nodes):
+                continue
+
+            output_val = cat_node.meta["val"]
+            assert isinstance(output_val, torch.Tensor)
+            split_input = get_arg(split_node, "input", Node)
+            with graph_module.graph.inserting_before(cat_node):
+                replacement_view = graph_module.graph.call_function(
+                    view_target,
+                    (split_input, list(output_val.shape)),
+                )
+                replacement_view.meta = cat_node.meta.copy()
+            cat_node.replace_all_uses_with(replacement_view)
+            modified = True
+        return modified
+
+    def call(self, graph_module: GraphModule) -> PassResult:
+        modified = self._replace_cats(
+            graph_module,
+            torch.ops.aten.cat.default,
+            _RAW_SPLIT_TARGETS,
+            torch.ops.aten.view_copy.default,
+        )
+        modified |= self._replace_cats(
+            graph_module,
+            exir_ops.edge.aten.cat.default,
+            _EDGE_SPLIT_TARGETS,
+            exir_ops.edge.aten.view_copy.default,
+        )
+        if modified:
+            graph_module.graph.eliminate_dead_code()
+            graph_module.recompile()
+        return PassResult(graph_module, modified)

backends/transforms/targets.bzl

@@ -521,6 +521,30 @@ def define_common_targets():
         ],
     )
 
+    runtime.python_library(
+        name = "merge_split_concat_chain",
+        srcs = ["merge_split_concat_chain.py"],
+        visibility = ["PUBLIC"],
+        deps = [
+            ":permute_pass_utils",
+            "//caffe2:torch",
+            "//executorch/exir:pass_base",
+            "//executorch/exir/dialects:lib",
+        ],
+    )
+
+    runtime.python_test(
+        name = "test_merge_split_concat_chain",
+        srcs = ["test/test_merge_split_concat_chain.py"],
+        deps = [
+            ":merge_split_concat_chain",
+            "//caffe2:torch",
+            "//executorch/backends/test:graph_builder",
+            "//executorch/exir:pass_base",
+            "//executorch/exir/dialects:lib",
+        ],
+    )
+
     runtime.python_library(
         name = "fuse_cascaded_transpose_or_permute_ops",
         srcs = ["fuse_cascaded_transpose_or_permute_ops.py"],

backends/transforms/test/test_merge_split_concat_chain.py

@@ -0,0 +1,197 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import copy
+import operator
+import unittest
+from collections.abc import Callable, Sequence
+from typing import cast
+
+import torch
+from executorch.backends.test.graph_builder import GraphBuilder
+from executorch.backends.transforms.merge_split_concat_chain import (
+    MergeSplitConcatChainPass,
+)
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import PassResult
+from torch.fx import GraphModule
+
+
+def _build_split_cat_graph(
+    input_value: torch.Tensor,
+    split_target: Callable[..., object],
+    split_spec: int | list[int],
+    split_dim: int,
+    output_count: int,
+    cat_target: Callable[..., object],
+    cat_dim: int,
+    order: Sequence[int] | None = None,
+) -> GraphModule:
+    builder = GraphBuilder()
+    input_node = builder.placeholder("input", input_value)
+    split = builder.call_operator(
+        split_target,
+        (input_node, split_spec, split_dim),
+    )
+    output_order = range(output_count) if order is None else order
+    split_outputs = [
+        builder.call_operator(operator.getitem, (split, index))
+        for index in output_order
+    ]
+    cat = builder.call_operator(cat_target, (split_outputs, cat_dim))
+    builder.output([cat])
+    return builder.get_graph_module()
+
+
+_VIEW_EQUIVALENT_CASES = (
+    (
+        "split",
+        torch.ops.aten.split.Tensor,
+        2,
+        (2, 1, 6, 4, 4),
+        2,
+        3,
+        torch.ops.aten.cat.default,
+        1,
+        torch.ops.aten.view_copy.default,
+    ),
+    (
+        "chunk",
+        torch.ops.aten.chunk.default,
+        3,
+        (2, 1, 5, 4, 4),
+        2,
+        3,
+        torch.ops.aten.cat.default,
+        2,
+        torch.ops.aten.view_copy.default,
+    ),
+    (
+        "split_with_sizes",
+        torch.ops.aten.split_with_sizes.default,
+        [1, 1, 1, 1],
+        (1, 4, 3, 2),
+        1,
+        4,
+        torch.ops.aten.cat.default,
+        0,
+        torch.ops.aten.view_copy.default,
+    ),
+    (
+        "edge_split_with_sizes",
+        exir_ops.edge.aten.split_with_sizes_copy.default,
+        [1, 1, 1, 1],
+        (4, 1, 3, 2),
+        0,
+        4,
+        exir_ops.edge.aten.cat.default,
+        1,
+        exir_ops.edge.aten.view_copy.default,
+    ),
+)
+
+_UNSAFE_CASES = (
+    ("reordered_outputs", [1, 0, 2, 3], 0, False),
+    ("non_singleton_crossing", [0, 1, 2, 3], 3, False),
+    ("non_contiguous_input", [0, 1, 2, 3], 0, True),
+)
+
+_SPLIT_CAT_DIALECTS = (
+    (
+        "aten",
+        torch.ops.aten.split_with_sizes.default,
+        torch.ops.aten.cat.default,
+    ),
+    (
+        "edge",
+        exir_ops.edge.aten.split_with_sizes_copy.default,
+        exir_ops.edge.aten.cat.default,
+    ),
+)
+
+
+class MergeSplitConcatChainPassTest(unittest.TestCase):
+    def test_merges_view_equivalent_chains(self) -> None:
+        for (
+            name,
+            split_target,
+            split_spec,
+            input_shape,
+            split_dim,
+            output_count,
+            cat_target,
+            cat_dim,
+            view_target,
+        ) in _VIEW_EQUIVALENT_CASES:
+            with self.subTest(name=name):
+                input_value = torch.randn(input_shape)
+                graph_module = _build_split_cat_graph(
+                    input_value,
+                    split_target,
+                    split_spec,
+                    split_dim,
+                    output_count,
+                    cat_target,
+                    cat_dim,
+                )
+                reference = copy.deepcopy(graph_module)
+
+                result = cast(PassResult, MergeSplitConcatChainPass()(graph_module))
+
+                self.assertTrue(result.modified)
+                torch.testing.assert_close(
+                    reference(input_value), result.graph_module(input_value)
+                )
+                for target, expected_count in (
+                    (split_target, 0),
+                    (cat_target, 0),
+                    (view_target, 1),
+                ):
+                    self.assertEqual(
+                        expected_count,
+                        len(
+                            result.graph_module.graph.find_nodes(
+                                op="call_function", target=target
+                            )
+                        ),
+                    )
+
+    def test_does_not_merge_unsafe_chains(self) -> None:
+        for dialect, split_target, cat_target in _SPLIT_CAT_DIALECTS:
+            for name, order, cat_dim, non_contiguous in _UNSAFE_CASES:
+                with self.subTest(dialect=dialect, name=name):
+                    input_value = torch.randn(1, 4, 3, 2)
+                    if non_contiguous:
+                        input_value = input_value.transpose(2, 3)
+                    graph_module = _build_split_cat_graph(
+                        input_value,
+                        split_target,
+                        [1, 1, 1, 1],
+                        1,
+                        4,
+                        cat_target,
+                        cat_dim,
+                        order,
+                    )
+                    reference = copy.deepcopy(graph_module)
+
+                    result = cast(PassResult, MergeSplitConcatChainPass()(graph_module))
+
+                    self.assertFalse(result.modified)
+                    torch.testing.assert_close(
+                        reference(input_value), result.graph_module(input_value)
+                    )
+                    self.assertEqual(
+                        1,
+                        len(
+                            result.graph_module.graph.find_nodes(
+                                op="call_function",
+                                target=cat_target,
+                            )
+                        ),
+                    )