Arm backend: Improve permute/view fusing

backends/arm/_passes/__init__.py

@@ -146,6 +146,9 @@
 from .match_arg_dtype_pass import MatchArgDtypePass  # noqa
 from .match_arg_ranks_pass import MatchArgRanksPass  # noqa
 from .mm_to_bmm_pass import ConvertMmToBmmPass  # noqa
+from .move_data_movement_ops_to_smaller_dtype_pass import (  # noqa
+    MoveDataMovementOpsToSmallerDtypePass,
+)
 from .normalize_delegate_io_layout_pass import NormalizeDelegateIOLayoutPass  # noqa
 from .normalize_index_put_bool_index_tensor_pass import (  # noqa
     NormalizeIndexPutBoolIndexTensorPass,

backends/arm/_passes/arm_pass_manager.py

@@ -128,6 +128,7 @@
     InsertTableOpsPass,
     MatchArgDtypePass,
     MatchArgRanksPass,
+    MoveDataMovementOpsToSmallerDtypePass,
     NormalizeDelegateIOLayoutPass,
     NormalizeIndexPutBoolIndexTensorPass,
     NormalizeIndexPutNoneIndicesPass,
@@ -635,6 +636,7 @@ def _tosa_pipeline(
                 PropagateViewCopyPermuteUpPass(self.compile_spec, exported_program),
                 # Propagation can leave a binary op with mismatched operand ranks,
                 # which TOSA rejects; re-match ranks before lowering.
+                MoveDataMovementOpsToSmallerDtypePass(),
                 MatchArgRanksPass(exported_program),
                 RewriteHighRankSingletonPermutePass(),
                 DecomposePermuteForU55Pass(),

backends/arm/_passes/arm_pass_utils.py

@@ -9,15 +9,11 @@
 import operator
 import traceback
 from inspect import isclass
-from typing import cast, Optional, Sequence
+from typing import Any, cast, Optional, Sequence
 
 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
@@ -36,8 +32,7 @@
 from torch._ops import OpOverload
 from torch._subclasses.fake_tensor import FakeTensor
 from torch.export.graph_signature import InputKind
-
-_Dim = int | torch.SymInt
+from torch.fx.node import map_arg
 
 
 def is_submodule_node(node: torch.fx.Node):
@@ -252,42 +247,31 @@ 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.
+def refresh_node_meta(node: torch.fx.Node) -> None:
+    """Best-effort refresh of a call_function node's output value metadata.
+
+    Node arguments are replaced by their metadata values and evaluated through
+    the operator's FakeTensor/meta implementation, which preserves symbolic
+    dimensions when the operator supports them. Missing input metadata or
+    unsupported meta evaluation leaves the existing output metadata unchanged.
+
     """
-    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,
-    }:
+    if node.op != "call_function" or "val" not in node.meta:
         return
 
-    if not isinstance(input_val, torch.Tensor):
-        node.meta["val"] = node.target(input_val, *node.args[1:])  # type: ignore[operator]
+    try:
+        args = map_arg(
+            node.args,
+            lambda arg: arg.meta["val"] if isinstance(arg, torch.fx.Node) else arg,
+        )
+        kwargs = map_arg(
+            node.kwargs,
+            lambda arg: arg.meta["val"] if isinstance(arg, torch.fx.Node) else arg,
+        )
+        node.meta["val"] = cast(Any, node.target)(*args, **kwargs)
+    except Exception:
         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,

backends/arm/_passes/canonicalize_view_copy_permute_pass.py

@@ -10,7 +10,7 @@
 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.arm_pass_utils import refresh_node_meta
 from executorch.backends.arm._passes.dim_maps import (
     _dim_equals,
     _is_permutation,
@@ -361,7 +361,7 @@ def _set_node_op(
     ) -> None:
         node.target = target
         node.args = (input_node, list(arg))
-        refresh_permute_view_meta(node)
+        refresh_node_meta(node)
 
     def _permute_dims(self, node: Node) -> list[int]:
         assert node.target == self._PERMUTE_TARGET, "Expected permute node"

backends/arm/_passes/decompose_var_pass.py

@@ -74,7 +74,7 @@ def call_operator(self, op, args, kwargs, meta):
             shape = [1 for _ in input_shape]
 
         # Get dim from args based on argument type
-        dim = get_node_arg(args, key=list, default_value=list(range(len(shape))))
+        dim = get_node_arg(args, key=list, default_value=list(range(len(input_shape))))
 
         if op == torch.ops.aten.var.dim:
             keepdim = False

backends/arm/_passes/dim_maps.py

@@ -363,6 +363,53 @@ def map_dim_inverse(
             return None
         return source_dims
 
+    def map_reduction_after_view(
+        self,
+        source_shape: Sequence[_Dim],
+        source_dims: int | Sequence[int],
+    ) -> tuple[list[_Dim], list[int]] | None:
+        """Map ``reduce(view(x), dims)`` to ``view(reduce(x, mapped_dims))``.
+
+        Returns the new view shape and reduction dims for:
+
+            view(reduce(x, source_dims), self.target_shape)
+            == reduce(view(x, new_shape), target_dims)
+
+        """
+        target_shape = self.remap_target_shape(source_shape)
+        if target_shape is None:
+            return None
+
+        target_dims = self.map_dim(source_dims)
+        if target_dims is None or not self._is_contiguous_nonempty(target_dims):
+            return None
+        return target_shape, target_dims
+
+    def map_reduction_before_view(
+        self,
+        target_dims: int | Sequence[int],
+    ) -> tuple[list[int], list[_Dim]] | None:
+        """Map ``view(reduce(x, dims))`` to ``reduce(view(x), mapped_dims)``.
+
+        Returns the reduction dims and output view shape for:
+
+            reduce(view(x, self.target_shape), target_dims)
+            == view(reduce(x, source_dims), output_shape)
+
+        """
+        source_dims = self.map_dim_inverse(target_dims)
+        if source_dims is None or not self._is_contiguous_nonempty(source_dims):
+            return None
+
+        try:
+            normalized_target_dims = _normalize_dims(target_dims, self.target_rank)
+        except AssertionError:
+            return None
+
+        return source_dims, self._reduce_shape(
+            self.target_shape, normalized_target_dims
+        )
+
     def map_permutation(
         self,
         source_permutation: Sequence[int],
@@ -446,6 +493,8 @@ def map_permutation_inverse(
         )
 
     def remap_target_shape(self, source_shape: Sequence[_Dim]) -> list[_Dim] | None:
+        if not self.is_valid_map:
+            return None
         if len(source_shape) != self.source_rank:
             return None
 
@@ -470,6 +519,8 @@ def remap_target_shape(self, source_shape: Sequence[_Dim]) -> list[_Dim] | None:
 
         if not same_numel(source_shape, target_shape):
             return None
+        if self._has_zero_dim(target_shape):
+            return None
         if not self._preserves_source_axis_order(source_shape, source_to_target_axes):
             return None
         return target_shape
@@ -551,6 +602,8 @@ def remap_unit_slice(
             for target_axes in source_to_target_axes[:slice_dim]
             for target_axis in target_axes
         ]
+        if not prev_target_axes:
+            return None
         next_target_axes = [
             target_axis
             for target_axes in source_to_target_axes[slice_dim + 1 :]
@@ -810,6 +863,24 @@ def _is_valid_reduction_or_singleton(
             group_to_axes[group].issubset(normalized_dims) for group in selected_groups
         )
 
+    @staticmethod
+    def _is_contiguous_nonempty(dims: Sequence[int]) -> bool:
+        sorted_dims = sorted(set(dims))
+        return bool(sorted_dims) and sorted_dims == list(
+            range(sorted_dims[0], sorted_dims[-1] + 1)
+        )
+
+    @staticmethod
+    def _reduce_shape(shape: Sequence[_Dim], dims: Sequence[int]) -> list[_Dim]:
+        reduced_shape = list(shape)
+        for dim in dims:
+            reduced_shape[dim] = 1
+        return reduced_shape
+
+    @staticmethod
+    def _has_zero_dim(shape: Sequence[_Dim]) -> bool:
+        return any(_dim_equals(dim, 0) for dim in shape)
+
     @classmethod
     def _build_groups(
         cls, source_shape: Sequence[_Dim], target_shape: Sequence[_Dim]