Reland "Arm backend: Add transpose propagation pass" (#20748)

Differential Revision: D110793024

Pull Request resolved: #20748

Author: rascani

backends/arm/_passes/__init__.py

@@ -151,6 +151,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,
@@ -133,11 +132,12 @@
     NormalizeTransformInputPlaceholdersPass,
     NormalizeWhileInitialArgsPass,
     PromoteBoolOperandsPass,
+    PropagateViewCopyPermuteDownPass,
+    PropagateViewCopyPermuteUpPass,
     QuantizeClampArgumentsPass,
     RemoveGetItemPass,
     RemoveGraphAssertsPass,
     RemoveNoopPass,
-    RemovePermutesAroundElementwiseTosaOps,
     ReplaceInfAndLimitValuesPass,
     ReplaceScalarWithTensorByProfilePass,
     RewriteAdaptiveAvgPool2dPass,
@@ -173,9 +173,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
@@ -606,6 +603,7 @@ def _tosa_pipeline(
                 RewriteAvgPool2dPass(),
                 ComputeConstantOpsAOTPass(exported_program),
                 FuseConstantArgsPass(exported_program),
+                CastInt64BuffersToInt32Pass(exported_program),
                 DecomposeSelectPass(),
                 ConvertSqueezesToViewPass(),
                 CastToInt32Pass(),
@@ -629,14 +627,14 @@ 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(),
                 FuseConsecutiveSlicesPass(),
                 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