Arm backend: Add transpose propagation pass

Adds abstract pass PropagateViewCopyPermutePass and
two implementations PropagateViewCopyPermuteUpPass
and PropagateViewCopyPermutePass

Adds a helper function to ViewMap to update the view-shape
after a view-op swap + additional validity checks.

Adds new InsertDataLayoutCastsPass under the propagation passes
to not move permutes/views outside of their intended dtype cast

The move exposed an issue where tosa ops
require the same dtype on both inputs which is not guaranteed at
that stage in the pipeline. This is fixed by adding
CastInt64BuffersToInt32Pass before TOSA dialect rewrite.

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

Author: AdrianLundell

backends/arm/_passes/__init__.py

@@ -148,6 +148,10 @@
 )
 from .normalize_while_initial_args_pass import NormalizeWhileInitialArgsPass  # noqa
 from .promote_bool_operands_pass import PromoteBoolOperandsPass  # noqa
+from .propagate_view_copy_permute_pass import (  # noqa
+    PropagateViewCopyPermuteDownPass,
+    PropagateViewCopyPermuteUpPass,
+)
 from .remove_getitem_pass import RemoveGetItemPass  # noqa
 from .remove_graph_asserts_pass import RemoveGraphAssertsPass  # noqa
 from .remove_noop_pass import RemoveNoopPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -15,7 +15,6 @@
     AccumulateIndexPutPass,
     BroadcastArgsPass,
     CanonicalizeGatherPass,
-    CanonicalizeViewCopyPermutePass,
     CastInt64BuffersToInt32Pass,
     CastToInt32Pass,
     ComputeConstantOpsAOTPass,
@@ -130,11 +129,12 @@
     NormalizeIndexPutNoneIndicesPass,
     NormalizeWhileInitialArgsPass,
     PromoteBoolOperandsPass,
+    PropagateViewCopyPermuteDownPass,
+    PropagateViewCopyPermuteUpPass,
     QuantizeClampArgumentsPass,
     RemoveGetItemPass,
     RemoveGraphAssertsPass,
     RemoveNoopPass,
-    RemovePermutesAroundElementwiseTosaOps,
     ReplaceInfAndLimitValuesPass,
     ReplaceScalarWithTensorByProfilePass,
     RewriteAdaptiveAvgPool2dPass,
@@ -167,9 +167,6 @@
     TosaLoweringContext,
     TosaSpecification,
 )
-from executorch.backends.transforms.fuse_cascaded_transpose_or_permute_ops import (
-    FuseCascadedTransposeOrPermuteOps,
-)
 
 from executorch.exir import ExportedProgram
 from executorch.exir._program_utils import _get_updated_graph_signature
@@ -597,6 +594,7 @@ def _tosa_pipeline(
                 RewriteAvgPool2dPass(),
                 ComputeConstantOpsAOTPass(exported_program),
                 FuseConstantArgsPass(exported_program),
+                CastInt64BuffersToInt32Pass(exported_program),
                 DecomposeSelectPass(),
                 ConvertSqueezesToViewPass(),
                 CastToInt32Pass(),
@@ -620,13 +618,13 @@ def _tosa_pipeline(
                 RewriteMatmulPass(),
                 RewritePadPass(),
                 FuseViewCopyTransformPass(),
-                RemovePermutesAroundElementwiseTosaOps(exported_program),
-                CanonicalizeViewCopyPermutePass(),
-                FuseCascadedTransposeOrPermuteOps(),
+                PropagateViewCopyPermuteDownPass(self.compile_spec, exported_program),
+                PropagateViewCopyPermuteUpPass(self.compile_spec, exported_program),
                 RewriteHighRankSingletonPermutePass(),
                 DecomposePermuteForU55Pass(),
                 RewriteSlicePass(),
                 InsertConstShapesPass(),
+                InsertDataLayoutCastsPass(),
             ]
         )
 

backends/arm/_passes/dim_maps.py

@@ -306,15 +306,28 @@ def map_dim(
             return None
 
         groups = self._valid_groups()
-        if not self._is_valid_reduction(normalized_dims, groups.source_axis_to_groups):
+        if not self._is_valid_reduction_or_singleton(
+            normalized_dims, groups.source_axis_to_groups
+        ):
             return None
 
-        target_dims = self._map_dims(
-            normalized_dims,
-            groups.source_axis_to_groups,
-            groups.group_to_target_axes,
+        source_to_target_axes = self.source_to_target_axes()
+        target_dims = sorted(
+            _dedupe(
+                target_axis
+                for source_dim in normalized_dims
+                for target_axis in source_to_target_axes[source_dim]
+            )
         )
-        if not target_dims or not self._is_valid_reduction(
+        if not target_dims or any(
+            source_axis not in normalized_dims
+            for target_axis in target_dims
+            for source_axis in self.source_axes_for_target_axis(
+                target_axis, source_to_target_axes
+            )
+        ):
+            return None
+        if not self._is_valid_reduction_or_singleton(
             target_dims, groups.target_axis_to_groups
         ):
             return None
@@ -432,6 +445,226 @@ def map_permutation_inverse(
             else None
         )
 
+    def remap_target_shape(self, source_shape: Sequence[_Dim]) -> list[_Dim] | None:
+        if len(source_shape) != self.source_rank:
+            return None
+
+        source_to_target_axes = self.source_to_target_axes()
+        target_to_source_axes = [
+            self.source_axes_for_target_axis(target_axis, source_to_target_axes)
+            for target_axis in range(self.target_rank)
+        ]
+        target_shape: list[_Dim] = [1] * self.target_rank
+
+        for source_axis, target_axes in enumerate(source_to_target_axes):
+            updates = self._target_axis_updates_for_source_axis(
+                source_shape,
+                source_axis,
+                target_axes,
+                target_to_source_axes,
+            )
+            if updates is None:
+                return None
+            for target_axis, target_dim in updates:
+                target_shape[target_axis] = target_dim
+
+        if not same_numel(source_shape, target_shape):
+            return None
+        if not self._preserves_source_axis_order(source_shape, source_to_target_axes):
+            return None
+        return target_shape
+
+    def _target_axis_updates_for_source_axis(
+        self,
+        source_shape: Sequence[_Dim],
+        source_axis: int,
+        target_axes: Sequence[int],
+        target_to_source_axes: Sequence[Sequence[int]],
+    ) -> list[tuple[int, _Dim]] | None:
+        if not target_axes:
+            return []
+
+        if len(target_axes) == 1:
+            target_axis = target_axes[0]
+            source_axes = target_to_source_axes[target_axis]
+            if source_axis != source_axes[0]:
+                return []
+            target_dim = numel(source_shape[source_axis] for source_axis in source_axes)
+            return [(target_axis, target_dim)]
+
+        if any(
+            len(target_to_source_axes[target_axis]) > 1 for target_axis in target_axes
+        ):
+            return []
+
+        target_dims = [self.target_shape[target_axis] for target_axis in target_axes]
+        if _dim_equals(source_shape[source_axis], self.source_shape[source_axis]):
+            return list(zip(target_axes, target_dims))
+        if _dim_equals(numel(target_dims), 1):
+            return [(target_axes[0], source_shape[source_axis])]
+        if _dim_equals(numel(target_dims), self.source_shape[source_axis]):
+            return list(zip(target_axes, target_dims))
+        return None
+
+    def remap_unit_slice(
+        self,
+        producer_shape: Sequence[_Dim],
+        slice_dim: int,
+        start: _Dim,
+        end: _Dim,
+        step: _Dim = 1,
+    ) -> tuple[list[_Dim], int, _Dim, _Dim] | None:
+        """Move a view before a unit slice.
+
+        Returns the new view shape and slice interval for:
+
+            view(slice(x, dim, start, end), self.target_shape)
+            == slice(view(x, new_shape), new_dim, new_start, new_end)
+
+        This handles the case where a unit slice produces a singleton source
+        axis that the view removes, so normal source-to-target dim mapping has
+        no target axis for the slice dim.
+
+        """
+        if (
+            len(producer_shape) != self.source_rank
+            or not isinstance(slice_dim, int)
+            or not isinstance(start, (int, torch.SymInt))
+            or not isinstance(end, (int, torch.SymInt))
+            or not isinstance(step, (int, torch.SymInt))
+        ):
+            return None
+        if not _dim_equals(step, 1) or not _dim_equals(end - start, 1):
+            return None
+
+        try:
+            slice_dim = _normalize_dim(slice_dim, self.source_rank)
+        except AssertionError:
+            return None
+
+        source_to_target_axes = self.source_to_target_axes()
+        if source_to_target_axes[slice_dim]:
+            return None
+
+        prev_target_axes = [
+            target_axis
+            for target_axes in source_to_target_axes[:slice_dim]
+            for target_axis in target_axes
+        ]
+        next_target_axes = [
+            target_axis
+            for target_axes in source_to_target_axes[slice_dim + 1 :]
+            for target_axis in target_axes
+        ]
+        fold_axes = [
+            target_axes[0]
+            for target_axes in source_to_target_axes[slice_dim + 1 :]
+            if target_axes
+        ]
+        fold_axes = [
+            target_axis
+            for target_axis in fold_axes
+            if all(
+                prev_target_axis <= target_axis for prev_target_axis in prev_target_axes
+            )
+            and all(
+                target_axis <= next_target_axis for next_target_axis in next_target_axes
+            )
+        ]
+        if not fold_axes:
+            return None
+
+        fold_axis = fold_axes[0]
+        target_shape = list(self.target_shape)
+        chunk = target_shape[fold_axis]
+        target_shape[fold_axis] = chunk * producer_shape[slice_dim]
+        return target_shape, fold_axis, start * chunk, end * chunk
+
+    def source_to_target_axes(self) -> list[list[int]]:
+        groups = self._valid_groups()
+        source_to_target_axes = [
+            self._map_dims(
+                [source_axis],
+                groups.source_axis_to_groups,
+                groups.group_to_target_axes,
+            )
+            for source_axis in range(self.source_rank)
+        ]
+
+        self._add_singleton_axes(source_to_target_axes)
+        return source_to_target_axes
+
+    def map_source_dims_to_target_axes(
+        self, source_dims: int | Sequence[int]
+    ) -> list[int] | None:
+        try:
+            normalized_dims = _normalize_dims(source_dims, self.source_rank)
+        except AssertionError:
+            return None
+        source_to_target_axes = self.source_to_target_axes()
+        return _dedupe(
+            target_axis
+            for source_dim in normalized_dims
+            for target_axis in source_to_target_axes[source_dim]
+        )
+
+    @staticmethod
+    def source_axes_for_target_axis(
+        target_axis: int, source_to_target_axes: Sequence[Sequence[int]]
+    ) -> list[int]:
+        return [
+            source_axis
+            for source_axis, target_axes in enumerate(source_to_target_axes)
+            if target_axis in target_axes
+        ]
+
+    def _add_singleton_axes(self, source_to_target_axes: list[list[int]]) -> None:
+        mapped_source_axes = {
+            source_axis
+            for source_axis, target_axes in enumerate(source_to_target_axes)
+            if target_axes
+        }
+        mapped_target_axes = {
+            target_axis
+            for target_axes in source_to_target_axes
+            for target_axis in target_axes
+        }
+        source_singletons = [
+            axis
+            for axis, dim in enumerate(self.source_shape)
+            if axis not in mapped_source_axes and _dim_equals(dim, 1)
+        ]
+        target_singletons = [
+            axis
+            for axis, dim in enumerate(self.target_shape)
+            if axis not in mapped_target_axes and _dim_equals(dim, 1)
+        ]
+
+        if len(source_singletons) == len(target_singletons):
+            pairs = zip(source_singletons, target_singletons)
+        elif len(source_singletons) == 1:
+            pairs = zip(source_singletons * len(target_singletons), target_singletons)
+        elif len(target_singletons) == 1:
+            pairs = zip(source_singletons, target_singletons * len(source_singletons))
+        else:
+            pairs = zip(source_singletons, target_singletons)
+
+        for source_axis, target_axis in pairs:
+            source_to_target_axes[source_axis].append(target_axis)
+
+    @staticmethod
+    def _preserves_source_axis_order(
+        source_shape: Sequence[_Dim],
+        source_to_target_axes: Sequence[Sequence[int]],
+    ) -> bool:
+        target_axes = [
+            target_axis
+            for source_axis, axes in enumerate(source_to_target_axes)
+            if not _dim_equals(source_shape[source_axis], 1)
+            for target_axis in axes
+        ]
+        return target_axes == sorted(target_axes)
+
     @staticmethod
     def _map_dims(
         source_dims: Iterable[int],
@@ -553,6 +786,30 @@ def _is_valid_reduction(
             group_to_axes[group].issubset(normalized_dims) for group in selected_groups
         )
 
+    @staticmethod
+    def _is_valid_reduction_or_singleton(
+        normalized_dims: Iterable[int],
+        axis_to_groups: Sequence[Sequence[int]],
+    ) -> bool:
+        """Return whether dims cover complete groups, allowing singleton
+        axes.
+        """
+        normalized_dims = set(normalized_dims)
+        if not normalized_dims:
+            return False
+
+        group_to_axes: dict[int, set[int]] = defaultdict(set)
+        selected_groups: set[int] = set()
+        for axis, groups in enumerate(axis_to_groups):
+            for group in groups:
+                group_to_axes[group].add(axis)
+                if axis in normalized_dims:
+                    selected_groups.add(group)
+
+        return all(
+            group_to_axes[group].issubset(normalized_dims) for group in selected_groups
+        )
+
     @classmethod
     def _build_groups(
         cls, source_shape: Sequence[_Dim], target_shape: Sequence[_Dim]

backends/arm/_passes/insert_data_layout_casts_pass.py

@@ -36,6 +36,7 @@ class InsertDataLayoutCastsPass(ArmOpTargetedPass):
     _concat_ops = {
         exir_ops.edge.aten.cat.default,
         exir_ops.edge.aten.concatenate.default,
+        exir_ops.backend.tosa.CONCAT.default,
     }
     _single_input_ops = {
         exir_ops.edge.aten.constant_pad_nd.default,
@@ -44,6 +45,12 @@ class InsertDataLayoutCastsPass(ArmOpTargetedPass):
         exir_ops.edge.aten.permute_copy.default,
         exir_ops.edge.aten.slice_copy.Tensor,
         exir_ops.edge.aten.flip.default,
+        exir_ops.backend.tosa.PAD.default,
+        exir_ops.backend.tosa.RESHAPE.default,
+        exir_ops.backend.tosa.TILE.default,
+        exir_ops.backend.tosa.TRANSPOSE.default,
+        exir_ops.backend.tosa.SLICE.default,
+        exir_ops.backend.tosa.REVERSE.default,
     }
     target_ops = _concat_ops | _single_input_ops