Arm backend: Add FuseIdenticalInputTransformsPass (#20514)

Sinks identical data movement ops on all inputs
of a binary op, e.g.
   add(permute(x), permute(y))
becomes:
   permute(add(x, y))

Adds support for -1 indexing to view-map using the normalize_view_shape
helper, and collapsing of known SymInts using _simplify_dim.



Signed-off-by: Adrian Lundell <adrian.lundell@arm.com>

Author: AdrianLundell

backends/arm/_passes/__init__.py

@@ -119,6 +119,9 @@
 )
 from .fuse_duplicate_users_pass import FuseDuplicateUsersPass  # noqa
 from .fuse_equal_placeholders_pass import FuseEqualPlaceholdersPass  # noqa
+from .fuse_identical_input_transforms_pass import (  # noqa
+    FuseIdenticalInputTransformsPass,
+)
 from .fuse_quantized_activation_pass import FuseQuantizedActivationPass  # noqa
 from .fuse_view_copy_transform_pass import FuseViewCopyTransformPass  # noqa
 from .insert_const_shapes import InsertConstShapesPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -112,6 +112,7 @@
     FuseConstantArgsPass,
     FuseDuplicateUsersPass,
     FuseEqualPlaceholdersPass,
+    FuseIdenticalInputTransformsPass,
     FuseQuantizedActivationPass,
     FuseViewCopyTransformPass,
     InsertConstShapesPass,
@@ -500,6 +501,7 @@ def _tosa_pipeline(
                 # TODO: DecomposeLinearPass should run after InsertRescaleInt32Pass or
                 # before FoldAndAnnotateQParamsPass but is unable to at the moment.
                 # Ticket: MLETORCH-1539
+                FuseIdenticalInputTransformsPass(),
                 DecomposeLinearPass(),
                 InsertRescaleInt32Pass(),
                 FuseConsecutiveRescalesPass(),

backends/arm/_passes/arm_pass_utils.py

@@ -13,14 +13,18 @@
 
 import torch
 import torch.fx
+
+from executorch.backends.arm._passes.dim_maps import (
+    _normalize_dims,
+    normalize_view_shape,
+)
 from executorch.backends.arm.common.debug import get_node_debug_info
 from executorch.backends.arm.common.type import ensure_type
 from executorch.backends.arm.tosa.mapping import TosaSpecialDtype
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.pass_base import NodeMetadata
-
 from torch._export.utils import (
     get_buffer,
     get_lifted_tensor_constant,
@@ -33,6 +37,8 @@
 from torch._subclasses.fake_tensor import FakeTensor
 from torch.export.graph_signature import InputKind
 
+_Dim = int | torch.SymInt
+
 
 def is_submodule_node(node: torch.fx.Node):
     if node.op not in ("get_attr", "placeholder"):
@@ -243,6 +249,43 @@ def meta_without_qparams(meta: NodeMetadata) -> NodeMetadata:
     return NodeMetadata(plain_meta_dict)
 
 
+def refresh_permute_view_meta(node: torch.fx.Node) -> None:
+    """Compute new meta-vals, specifically preserving SymInts for view/permute
+    nodes.
+    """
+    input_node = node.all_input_nodes[0]
+    input_val = input_node.meta.get("val")
+    if input_val is None or node.target not in {
+        exir_ops.edge.aten.view_copy.default,
+        exir_ops.edge.aten.permute_copy.default,
+    }:
+        return
+
+    if not isinstance(input_val, torch.Tensor):
+        node.meta["val"] = node.target(input_val, *node.args[1:])  # type: ignore[operator]
+        return
+
+    # Compute new meta shapes to preserve SymInts.
+    match node.target:
+        case exir_ops.edge.aten.view_copy.default:
+            node.meta["val"] = input_val.new_empty(
+                tuple(
+                    normalize_view_shape(
+                        input_val.shape, cast(Sequence[_Dim], node.args[1])
+                    )
+                )
+            )
+        case exir_ops.edge.aten.permute_copy.default:
+            dims = _normalize_dims(
+                cast(Sequence[int], node.args[1]), len(input_val.shape)
+            )
+            node.meta["val"] = input_val.new_empty(
+                tuple(input_val.shape[dim] for dim in dims)
+            )
+        case _:
+            node.meta["val"] = node.target(input_val, *node.args[1:])  # type: ignore[operator]
+
+
 def insert_scalar(
     graph: torch.fx.Graph,
     value: int | float,

backends/arm/_passes/canonicalize_view_copy_permute_pass.py

@@ -8,7 +8,9 @@
 from typing import cast, Sequence, Set, Type
 
 import torch
+
 from executorch.backends.arm._passes.arm_pass import ArmPass
+from executorch.backends.arm._passes.arm_pass_utils import refresh_permute_view_meta
 from executorch.backends.arm._passes.dim_maps import (
     _dim_equals,
     _is_permutation,
@@ -18,7 +20,6 @@
 )
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
-
 from torch.fx import GraphModule, Node
 from torch.fx.node import Target
 from torch.fx.passes.infra.pass_base import PassResult
@@ -360,39 +361,12 @@ def _set_node_op(
     ) -> None:
         node.target = target
         node.args = (input_node, list(arg))
-        self._refresh_meta(node)
+        refresh_permute_view_meta(node)
 
     def _permute_dims(self, node: Node) -> list[int]:
         assert node.target == self._PERMUTE_TARGET, "Expected permute node"
         return list(cast(Sequence[int], node.args[1]))
 
-    @classmethod
-    def _refresh_meta(cls, node: Node) -> None:
-        input_node = node.args[0]
-        assert isinstance(input_node, Node)
-        input_val = input_node.meta.get("val")
-        if input_val is None or node.target not in cls._TARGETS:
-            return
-
-        # Compute new meta shapes to preserve SymInts.
-        if isinstance(input_val, torch.Tensor):
-            if node.target == cls._VIEW_TARGET:
-                node.meta["val"] = input_val.new_empty(
-                    tuple(cast(Sequence[_Dim], node.args[1]))
-                )
-                return
-
-            if node.target == cls._PERMUTE_TARGET:
-                dims = _normalize_dims(
-                    cast(Sequence[int], node.args[1]), len(input_val.shape)
-                )
-                node.meta["val"] = input_val.new_empty(
-                    tuple(input_val.shape[dim] for dim in dims)
-                )
-                return
-
-        node.meta["val"] = node.target(input_val, *node.args[1:])  # type: ignore[operator]
-
     @staticmethod
     def _shape(node: Node) -> list[_Dim]:
         return cast(list[_Dim], list(node.meta["val"].shape))

backends/arm/_passes/dim_maps.py

@@ -91,7 +91,21 @@ def _dim_expr(dim: _Dim) -> sympy.Basic:
     return sympy.Integer(dim) if isinstance(dim, int) else dim.node.expr
 
 
+def _simplify_dim(dim: _Dim) -> _Dim:
+    if isinstance(dim, int):
+        return dim
+
+    maybe_int = dim.node.maybe_as_int()
+    return maybe_int if maybe_int is not None else dim
+
+
+def _simplify_shape(shape: Iterable[_Dim]) -> list[_Dim]:
+    return [_simplify_dim(dim) for dim in shape]
+
+
 def _dim_equals(lhs: _Dim, rhs: _Dim) -> bool:
+    lhs = _simplify_dim(lhs)
+    rhs = _simplify_dim(rhs)
     if isinstance(lhs, int) and isinstance(rhs, int):
         return lhs == rhs
     return sympy.simplify(_dim_expr(lhs) - _dim_expr(rhs)) == 0
@@ -113,6 +127,7 @@ def _factor_int(dim: int) -> list[_FactorKey] | None:
 
 
 def _factor_dim(dim: _Dim) -> list[_FactorKey] | None:
+    dim = _simplify_dim(dim)
     if _dim_equals(dim, 1):
         return []
     if isinstance(dim, int):
@@ -143,14 +158,39 @@ def _dedupe(items: Iterable[int]) -> list[int]:
 def numel(shape: Iterable[_Dim]) -> _Dim:
     numel: _Dim = 1
     for dim in shape:
-        numel *= dim
+        numel = _simplify_dim(numel * _simplify_dim(dim))
     return numel
 
 
 def same_numel(first_shape: Iterable[_Dim], second_shape: Iterable[_Dim]) -> bool:
     return _dim_equals(numel(first_shape), numel(second_shape))
 
 
+def normalize_view_shape(
+    source_shape: Sequence[_Dim], target_shape: Sequence[_Dim]
+) -> list[_Dim]:
+    """Normalize a view shape with <=1 unknown dim, indicated by -1."""
+    source_shape = _simplify_shape(source_shape)
+    normalized_shape = _simplify_shape(target_shape)
+    inferred_dims = [
+        index for index, dim in enumerate(normalized_shape) if _dim_equals(dim, -1)
+    ]
+    if not inferred_dims:
+        return normalized_shape
+
+    assert len(inferred_dims) == 1, f"Invalid view shape {target_shape}"
+    inferred_dim = inferred_dims[0]
+    known_shape = [
+        dim for index, dim in enumerate(normalized_shape) if index != inferred_dim
+    ]
+    source_numel = numel(source_shape)
+    known_numel = numel(known_shape)
+    normalized_shape[inferred_dim] = _simplify_dim(
+        source_numel if _dim_equals(known_numel, 1) else source_numel // known_numel
+    )
+    return normalized_shape
+
+
 class _UnionFind:
     def __init__(self, size: int) -> None:
         self.parents = list(range(size))
@@ -210,8 +250,10 @@ def __init__(self, view_node: Node) -> None:
         input_val = input_node.meta["val"]
         assert isinstance(input_val, torch.Tensor)
 
-        self.source_shape = cast(list[_Dim], list(input_val.shape))
-        self.target_shape = list(cast(Sequence[_Dim], view_node.args[1]))
+        self.source_shape = _simplify_shape(cast(Sequence[_Dim], input_val.shape))
+        self.target_shape = normalize_view_shape(
+            self.source_shape, cast(Sequence[_Dim], view_node.args[1])
+        )
         self._groups = self._build_groups(self.source_shape, self.target_shape)
 
     @classmethod
@@ -220,8 +262,10 @@ def from_shapes(
     ) -> ViewMap:
         """Build a view map directly from source and target shapes."""
         view_map = cls.__new__(cls)
-        view_map.source_shape = list(source_shape)
-        view_map.target_shape = list(target_shape)
+        view_map.source_shape = _simplify_shape(source_shape)
+        view_map.target_shape = normalize_view_shape(
+            view_map.source_shape, target_shape
+        )
         view_map._groups = cls._build_groups(
             view_map.source_shape, view_map.target_shape
         )