From dc8bd28d706d2e55523bfc3098f5b644b93ccfc3 Mon Sep 17 00:00:00 2001 From: Adrian Lundell Date: Tue, 2 Jun 2026 10:26:48 +0200 Subject: [PATCH] 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 Change-Id: I31bc29540cebce9eb0ee3d5a3b2c016c19b15809 --- backends/arm/_passes/__init__.py | 4 + backends/arm/_passes/arm_pass_manager.py | 14 +- backends/arm/_passes/dim_maps.py | 269 +++- .../_passes/insert_data_layout_casts_pass.py | 7 + .../propagate_view_copy_permute_pass.py | 716 +++++++++ .../arm/test/misc/test_transpose_counts.py | 24 +- backends/arm/test/passes/test_dim_maps.py | 47 +- .../test_insert_data_layout_casts_pass.py | 29 + .../test_propagate_permutes_views_pass.py | 1418 +++++++++++++++++ 9 files changed, 2500 insertions(+), 28 deletions(-) create mode 100644 backends/arm/_passes/propagate_view_copy_permute_pass.py create mode 100644 backends/arm/test/passes/test_propagate_permutes_views_pass.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index b7855dbc9a6..5fa6ca61e65 100644 --- a/backends/arm/_passes/__init__.py +++ b/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 diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index eaa553507b6..0a46020804a 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/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(), ] ) diff --git a/backends/arm/_passes/dim_maps.py b/backends/arm/_passes/dim_maps.py index 07098035389..6fc9b8ac1f9 100644 --- a/backends/arm/_passes/dim_maps.py +++ b/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] diff --git a/backends/arm/_passes/insert_data_layout_casts_pass.py b/backends/arm/_passes/insert_data_layout_casts_pass.py index 07a2d186895..4e931396dab 100644 --- a/backends/arm/_passes/insert_data_layout_casts_pass.py +++ b/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 diff --git a/backends/arm/_passes/propagate_view_copy_permute_pass.py b/backends/arm/_passes/propagate_view_copy_permute_pass.py new file mode 100644 index 00000000000..421b425026e --- /dev/null +++ b/backends/arm/_passes/propagate_view_copy_permute_pass.py @@ -0,0 +1,716 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from abc import ABC, abstractmethod +from collections.abc import Iterable, Sequence +from typing import Any, cast, Set, Type + +import torch +from executorch.backends.arm._passes.arm_pass_utils import refresh_permute_view_meta +from executorch.backends.arm._passes.dim_maps import PermuteMap, ViewMap +from executorch.backends.arm.tosa.mapping import TosaSpecialDtype +from executorch.backends.arm.tosa.specification import get_context_spec +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + +from .arm_pass import ArmPass +from .canonicalize_view_copy_permute_pass import CanonicalizeViewCopyPermutePass +from .fuse_duplicate_users_pass import FuseDuplicateUsersPass +from .fuse_identical_input_transforms_pass import FuseIdenticalInputTransformsPass +from .remove_permutes_around_elementwise_tosa_ops import ( + RemovePermutesAroundElementwiseTosaOps, +) + +_Dim = int | torch.SymInt + + +class PropagateViewCopyPermutePass(ArmPass, ABC): + """Abstract implementation of a permute/view_copy propagation pass. + + To be used for upwards/downwards propagation by implementing the abstract + methods for the direction of propagation. + + """ + + _passes_required_after: Set[Type[ExportPass]] = set() + + _VIEW_TARGET = exir_ops.edge.aten.view_copy.default + _VIEW_DEFAULT_TARGET = exir_ops.edge.aten.view.default + _PERMUTE_TARGET = exir_ops.edge.aten.permute_copy.default + _TARGETS = {_VIEW_TARGET, _VIEW_DEFAULT_TARGET, _PERMUTE_TARGET} + _TRANSPARENT_TARGETS = { + exir_ops.edge.dim_order_ops._clone_dim_order.default, + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, + } + + _REDUCTION_TARGETS = { + exir_ops.edge.aten.mean.dim, + exir_ops.edge.aten.sum.dim_IntList, + } + _ARG_UPDATE_TARGETS = { + *_REDUCTION_TARGETS, + exir_ops.edge.aten.slice_copy.Tensor, + } + + def __init__( + self, + compile_spec: Any | None = None, + exported_program: ExportedProgram | None = None, + ) -> None: + super().__init__() + if isinstance(compile_spec, ExportedProgram) and exported_program is None: + exported_program = compile_spec + compile_spec = None + self.exported_program = exported_program + self.compile_spec = compile_spec + + @staticmethod + def _dim_arg(arg: Any) -> int | Sequence[int] | None: + if isinstance(arg, int): + return arg + if isinstance(arg, Sequence) and not isinstance(arg, (str, bytes)): + return cast(Sequence[int], arg) + return None + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + modified = False + + result = self.fuse_horizontal(graph_module) + graph_module = result.graph_module + modified |= result.modified + result = self.fuse_vertical(graph_module) + graph_module = result.graph_module + modified |= result.modified + if result.modified: + graph_module = self._retrace(graph_module) + + # Do not run for Ethos-U85 since this exposes a numerical issue + # There is no target meta-data at this stage so use INT+cf as proxy + # To be removed after MLBEDSW-11805 + while not self._is_u85_like_tosa_int_cf(): + iteration_modified = False + for node in list(graph_module.graph.nodes): + if node.target in self._TARGETS: + if len(node.users) == 0: + continue + iteration_modified |= self._propagate(node) + + if iteration_modified: + graph_module = self._retrace(graph_module) + result = self.fuse_horizontal(graph_module) + graph_module = result.graph_module + iteration_modified |= result.modified + result = self.fuse_vertical(graph_module) + graph_module = result.graph_module + iteration_modified |= result.modified + + modified |= iteration_modified + if not iteration_modified: + break + + if modified: + graph_module = self._retrace(graph_module) + graph_module.recompile() + + return PassResult(graph_module, modified) + + def _is_u85_like_tosa_int_cf(self) -> bool: + if self.compile_spec is not None: + tosa_spec = self.compile_spec.tosa_spec + else: + try: + tosa_spec = get_context_spec() + except RuntimeError: + return False + + return ( + tosa_spec.support_integer() + and not tosa_spec.support_float() + and tosa_spec.support_extension("cf") + ) + + def _retrace(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule: + graph_module.graph.eliminate_dead_code() + graph_module.graph.lint() + return super().call(graph_module).graph_module + + def _propagate(self, node: torch.fx.Node) -> bool: + """Propagate a single permute/view node.""" + + frontier = node + previous_frontier = None + moved = False + while True: + next_nodes = list(self._get_next_nodes(frontier)) + + if len(next_nodes) == 0: + assert node.op in ( + "placeholder", + "output", + ), f"{self.__class__.__name__} reached an endpoint node which is not a placeholder or output: {frontier}" + break + + if not self._can_cross_next_nodes(frontier, next_nodes): + break + + if len(next_nodes) > 1: + if self._maybe_split_downwards_slice_fanout(node, next_nodes): + return True + break + + next_node = next_nodes[0] + if self.is_elementwise(next_node) and self._is_unary_elementwise(next_node): + previous_frontier = frontier + frontier = next_node + moved = True + continue + + if self.is_swappable(next_node): + swapped_args = self._maybe_swap_args(node, next_node) + if swapped_args is None: + break + node.args = swapped_args[0] + next_node.args = swapped_args[1] + previous_frontier = frontier + frontier = next_node + moved = True + continue + + # Concats are a special case since they branch the graph. + # Perform the swap directly in this case and return. + # Otherwise break and move the node before the concat + if self._maybe_split_upwards_cat_fanout(node, next_node): + return True + + # Unhandled case, stop propagation + break + + if not moved: + return False + + assert previous_frontier is not None + self._move_node(node, frontier, previous_frontier) + refresh_permute_view_meta(node) + return True + + def fuse_vertical(self, graph_module: torch.fx.GraphModule) -> PassResult: + """Fuse consecutive permute/view nodes.""" + modified = False + + if self.exported_program is not None: + result = RemovePermutesAroundElementwiseTosaOps(self.exported_program).call( + graph_module + ) + graph_module = result.graph_module + modified |= result.modified + + result = CanonicalizeViewCopyPermutePass().call(graph_module) + graph_module = result.graph_module + modified |= result.modified + return PassResult(graph_module, modified) + + @abstractmethod + def fuse_horizontal(self, graph_module: torch.fx.GraphModule) -> PassResult: + """Fuse parallel permute/view nodes going into/ out a single node.""" + pass + + @abstractmethod + def _get_next_nodes(self, node: torch.fx.Node) -> Iterable[torch.fx.Node]: + """Return the next nodes in the direction of propagation.""" + pass + + @abstractmethod + def _get_prev_nodes(self, node: torch.fx.Node) -> Iterable[torch.fx.Node]: + """Return the previous nodes in the direction of propagation.""" + pass + + def _can_cross_next_nodes( + self, frontier: torch.fx.Node, next_nodes: Sequence[torch.fx.Node] + ) -> bool: + return True + + @abstractmethod + def _maybe_swap_permute_args( + self, node: torch.fx.Node, next_node: torch.fx.Node + ) -> Any | None: + pass + + @abstractmethod + def _maybe_swap_view_args( + self, node: torch.fx.Node, next_node: torch.fx.Node + ) -> Any | None: + pass + + def _maybe_split_upwards_cat_fanout( + self, node: torch.fx.Node, next_node: torch.fx.Node + ) -> bool: + """Swap cat([x1,x2]).permute(p) -> cat([x1.permute(p'), x2.permute(p')]) + if permutes before the concat are noops. + """ + return False + + def _maybe_split_downwards_slice_fanout( + self, node: torch.fx.Node, next_nodes: Sequence[torch.fx.Node] + ) -> bool: + """Swap x2 = x1.permute; y1 = x2.slice_copy[0]; y2 = x2.slice_copy[1] to + y1 = x1.permute.slice_copy[0]; y2 = x1.permute.slice_copy[1] Only if + permutes after slice are noops. + """ + return False + + def _maybe_swap_args( + self, node: torch.fx.Node, next_node: torch.fx.Node + ) -> Any | None: + """If the node can be swapped with its next_node, return the new args + for the next_node and new shape, otherwise return None. + """ + if node.target == self._PERMUTE_TARGET: + return self._maybe_swap_permute_args(node, next_node) + elif node.target in {self._VIEW_TARGET, self._VIEW_DEFAULT_TARGET}: + return self._maybe_swap_view_args(node, next_node) + else: + raise ValueError( + f"Unexpected node target {node.target} in {self.__class__.__name__}" + ) + + def _move_node( + self, + node: torch.fx.Node, + frontier: torch.fx.Node, + previous_frontier: torch.fx.Node, + ) -> None: + """Update the graph to move the node into its new position.""" + raise NotImplementedError() + + def is_elementwise(self, node: torch.fx.Node) -> bool: + if node.op != "call_function": + return False + + if node.target == exir_ops.backend.tosa.RESCALE.default: + return self._is_per_tensor_rescale(node) + + if node.target == exir_ops.backend.tosa.TABLE.default: + return True + + if node.target in self._TRANSPARENT_TARGETS: + return True + + op = getattr(node.target, "_op", None) + if op is not None and hasattr(op, "tags"): + return torch.Tag.pointwise in op.tags + return False + + def _is_per_tensor_rescale(self, node: torch.fx.Node) -> bool: + if len(node.args) < 3: + return False + input_nodes = node.all_input_nodes + if len(input_nodes) != 1: + return False + special_dtype_key = TosaSpecialDtype.meta_key() + if input_nodes[0].meta.get(special_dtype_key) != node.meta.get( + special_dtype_key + ): + return False + scales = node.args[2] + return not isinstance(scales, Sequence) or len(scales) == 1 + + def is_swappable(self, next_node: torch.fx.Node) -> bool: + if next_node.target not in self._ARG_UPDATE_TARGETS: + return False + if next_node.target in self._REDUCTION_TARGETS: + keep_dim = ( + next_node.args[2] + if len(next_node.args) > 2 + else next_node.kwargs.get("keepdim") + ) + if keep_dim is not True: + raise RuntimeError( + f"{self.__class__.__name__} expects keep_dim=True for reduction ops to simplify propagation logic, got {keep_dim} for node {next_node.name}." + ) + return True + + def _is_unary_elementwise(self, node: torch.fx.Node) -> bool: + if node.target == exir_ops.backend.tosa.TABLE.default: + return True + return len(node.all_input_nodes) == 1 + + @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) + ) + + +class PropagateViewCopyPermuteUpPass(PropagateViewCopyPermutePass): + """Implements PropagateViewCopyPermutePass for upwards propagation: + + - Next propagation nodes are the input of the current node + - Previous propagation nodes are the users of the current node + - Swaps are (op -> permute/view) to (permute/view -> op) + - Node is moved before the frontier next_node + - Horizontal fuses are performed on users + """ + + def fuse_horizontal(self, graph_module): + modified = False + result = FuseDuplicateUsersPass().call(graph_module) + graph_module = result.graph_module + modified |= result.modified + return PassResult(graph_module, modified) + + def _get_next_nodes(self, node: torch.fx.Node) -> Iterable[torch.fx.Node]: + return list(node.all_input_nodes) + + def _get_prev_nodes(self, node: torch.fx.Node) -> Iterable[torch.fx.Node]: + return list(node.users.keys()) + + def _can_cross_next_nodes( + self, frontier: torch.fx.Node, next_nodes: Sequence[torch.fx.Node] + ) -> bool: + if any( + user.target == exir_ops.backend.tosa.SCATTER.default + for user in frontier.users + ): + return False + return all( + all(prev_node is frontier for prev_node in self._get_prev_nodes(next_node)) + for next_node in next_nodes + ) + + def _maybe_swap_permute_args( + self, node: torch.fx.Node, next_node: torch.fx.Node + ) -> Any | None: + permute_map = PermuteMap(node) + args = self._dim_arg(next_node.args[1]) + if args is None: + return None + mapped_args = permute_map.map_dims(args) + new_args: int | list[int] = ( + mapped_args[0] if isinstance(args, int) else mapped_args + ) + return (node.args, (*next_node.args[:1], new_args, *next_node.args[2:])) + + def _maybe_swap_view_args( + self, node: torch.fx.Node, next_node: torch.fx.Node + ) -> Any | None: + view_map = ViewMap(node) + if not view_map.is_valid_map or len(next_node.all_input_nodes) != 1: + return None + + input_val = next_node.all_input_nodes[0].meta["val"] + input_shape = list(input_val.shape) + new_shape = view_map.remap_target_shape(input_shape) + + if next_node.target in self._REDUCTION_TARGETS: + return self._maybe_swap_reduction_view_args( + node, next_node, view_map, new_shape + ) + if next_node.target == exir_ops.edge.aten.slice_copy.Tensor: + return self._maybe_swap_slice_view_args( + node, next_node, view_map, input_shape, new_shape + ) + return None + + def _maybe_swap_reduction_view_args( + self, + node: torch.fx.Node, + next_node: torch.fx.Node, + view_map: ViewMap, + new_shape: list[_Dim] | None, + ) -> Any | None: + if new_shape is None: + return None + if len(next_node.args) <= 2 or next_node.args[2] is not True: + return None + reduction_dims = cast(int | Sequence[int], next_node.args[1]) + new_dims = view_map.map_dim(reduction_dims) + if new_dims is None or not self._is_contiguous_nonempty(new_dims): + return None + new_next_node_args = (*next_node.args[:1], new_dims, *next_node.args[2:]) + return ((*node.args[:1], new_shape), new_next_node_args) + + def _maybe_swap_slice_view_args( + self, + node: torch.fx.Node, + next_node: torch.fx.Node, + view_map: ViewMap, + input_shape: list[_Dim], + new_shape: list[_Dim] | None, + ) -> Any | None: + if new_shape is None: + return self._maybe_swap_unit_slice_view_args( + node, next_node, view_map, input_shape + ) + + slice_dim = cast(int, next_node.args[1]) + new_dim = self._map_slice_dim(view_map, slice_dim) + if new_dim is None: + return None + new_next_node_args = (*next_node.args[:1], new_dim, *next_node.args[2:]) + return ((*node.args[:1], new_shape), new_next_node_args) + + def _maybe_swap_unit_slice_view_args( + self, + node: torch.fx.Node, + next_node: torch.fx.Node, + view_map: ViewMap, + input_shape: list[_Dim], + ) -> Any | None: + if len(next_node.args) < 4: + return None + step = next_node.args[4] if len(next_node.args) > 4 else 1 + remapped_slice = view_map.remap_unit_slice( + input_shape, + cast(int, next_node.args[1]), + cast(_Dim, next_node.args[2]), + cast(_Dim, next_node.args[3]), + cast(_Dim, step), + ) + if remapped_slice is None: + return None + + new_shape, new_dim, new_start, new_end = remapped_slice + new_next_node_args = ( + *next_node.args[:1], + new_dim, + new_start, + new_end, + *next_node.args[4:], + ) + return ((*node.args[:1], new_shape), new_next_node_args) + + @staticmethod + def _map_slice_dim(view_map: ViewMap, slice_dim: int) -> int | None: + new_dims = view_map.map_source_dims_to_target_axes(slice_dim) + if new_dims is None or len(new_dims) != 1: + return None + + new_dim = new_dims[0] + normalized_slice_dim = slice_dim % view_map.source_rank + source_to_target_axes = view_map.source_to_target_axes() + target_source_axes = view_map.source_axes_for_target_axis( + new_dim, source_to_target_axes + ) + if any( + source_axis != normalized_slice_dim for source_axis in target_source_axes + ): + return None + return new_dim + + def _move_node( + self, + node: torch.fx.Node, + frontier: torch.fx.Node, + previous_frontier: torch.fx.Node, + ) -> None: + original_input = node.all_input_nodes[0] + if frontier.op == "placeholder": + # Nodes cannot be moved before placeholders + producer = frontier + frontier_user = previous_frontier + else: + producer = frontier.all_input_nodes[0] + frontier_user = frontier + + node.replace_input_with(original_input, producer) + frontier_user.replace_input_with(producer, node) + + for user in list(node.users): + if user is not frontier_user: + user.replace_input_with(node, original_input) + + frontier_user.prepend(node) + + def _maybe_split_upwards_cat_fanout( + self, node: torch.fx.Node, next_node: torch.fx.Node + ) -> bool: + """Swap cat([x1,x2]).permute(p) -> cat([x1.permute(p'), x2.permute(p')]) + if permutes before the concat are noops. + """ + if node.target != self._PERMUTE_TARGET: + return False + if next_node.target != exir_ops.edge.aten.cat.default: + return False + + cat_users = list(next_node.users) + if len(cat_users) == 0: + return False + if not all(n.target == self._PERMUTE_TARGET for n in cat_users): + return False + + permute_args = [self._dim_arg(n.args[1]) for n in cat_users] + if not isinstance(permute_args[0], Sequence) or not all( + p == permute_args[0] for p in permute_args + ): + return False + + cat_dim = ( + next_node.args[1] + if len(next_node.args) >= 2 + else next_node.kwargs.get("dim", 0) + ) + if not isinstance(cat_dim, int): + return False + new_cat_dim = PermuteMap(node).map_dims(cat_dim)[0] + + cat_inputs = list(next_node.all_input_nodes) + cat_input_shapes = [input_node.meta["val"].shape for input_node in cat_inputs] + + # Ensure all input permutes are noops + if not all( + CanonicalizeViewCopyPermutePass._is_singleton_permutation( + shape, permute_args[0] + ) + for shape in cat_input_shapes + ): + return False + + # Add permutes to all cat inputs, update cat arg, and remove old output permute + new_inputs = [] + for input_node in cat_inputs: + input_val = input_node.meta["val"] + output_shape = [input_val.shape[dim] for dim in permute_args[0]] + with next_node.graph.inserting_before(next_node): + permute = next_node.graph.call_function( + self._PERMUTE_TARGET, + args=(input_node, permute_args[0]), + ) + permute.meta = dict(input_node.meta) + permute.meta["val"] = input_val.new_empty(tuple(output_shape)) + new_inputs.append(permute) + + next_node.args = (new_inputs, new_cat_dim, *next_node.args[2:]) + next_node.meta = dict(node.meta) + for cat_user in cat_users: + cat_user.replace_all_uses_with(next_node) + for cat_user in cat_users: + if len(cat_user.users) == 0: + next_node.graph.erase_node(cat_user) + return True + + +class PropagateViewCopyPermuteDownPass(PropagateViewCopyPermutePass): + """Implements PropagateViewCopyPermutePass for downward propagation: + + - Next propagation nodes are the users of the current node + - Previous propagation nodes are the inputs of the current node + - Swaps are (permute/view -> op) to (op -> permute/view) + - Node is moved after the frontier next_node + - Horizontal fuses are performed on inputs + """ + + def fuse_horizontal(self, graph_module): + modified = False + result = FuseIdenticalInputTransformsPass().call(graph_module) + graph_module = result.graph_module + modified |= result.modified + return PassResult(graph_module, modified) + + def _get_next_nodes(self, node: torch.fx.Node) -> Iterable[torch.fx.Node]: + return list(node.users.keys()) + + def _get_prev_nodes(self, node: torch.fx.Node) -> Iterable[torch.fx.Node]: + return list(node.all_input_nodes) + + def _maybe_swap_permute_args( + self, node: torch.fx.Node, next_node: torch.fx.Node + ) -> Any | None: + permute_map = PermuteMap(node) + args = self._dim_arg(next_node.args[1]) + if args is None: + return None + mapped_args = permute_map.map_dims_inverse(args) + new_args: int | list[int] = ( + mapped_args[0] if isinstance(args, int) else mapped_args + ) + return (node.args, (*next_node.args[:1], new_args, *next_node.args[2:])) + + def _maybe_swap_view_args(self, node, next_node): + view_map = ViewMap(node) + if not view_map.is_valid_map: + return None + + if next_node.target in self._REDUCTION_TARGETS: + if len(next_node.args) <= 2 or next_node.args[2] is not True: + return None + new_dims = view_map.map_dim_inverse(next_node.args[1]) + if new_dims is None: + return None + elif next_node.target == exir_ops.edge.aten.slice_copy.Tensor: + new_dims = view_map.map_dim_inverse(next_node.args[1]) + if new_dims is None: + return None + if len(new_dims) != 1: + return None + new_dims = new_dims[0] + else: + return None + + output_val = next_node.meta["val"] + new_next_node_args = (*next_node.args[:1], new_dims, *next_node.args[2:]) + return ((*node.args[:1], list(output_val.shape)), new_next_node_args) + + def _maybe_split_downwards_slice_fanout( + self, node: torch.fx.Node, next_nodes: Sequence[torch.fx.Node] + ) -> bool: + """Duplicate a permute onto each slice branch. + + The duplicated permutes are left before the slices; later propagation + iterations handle swapping each one through its slice. + + """ + if node.target != self._PERMUTE_TARGET: + return False + if not all( + next_node.target == exir_ops.edge.aten.slice_copy.Tensor + and next_node.all_input_nodes == [node] + for next_node in next_nodes + ): + return False + + producer = node.all_input_nodes[0] + for next_node in next_nodes: + with next_node.graph.inserting_before(next_node): + branch_permute = next_node.graph.call_function( + self._PERMUTE_TARGET, + args=(producer, node.args[1]), + ) + branch_permute.meta = dict(node.meta) + next_node.replace_input_with(node, branch_permute) + + if len(node.users) == 0: + node.graph.erase_node(node) + return True + + def _move_node( + self, + node: torch.fx.Node, + frontier: torch.fx.Node, + previous_frontier: torch.fx.Node, + ) -> None: + original_user = next(iter(node.users)) + producer = node.all_input_nodes[0] + if frontier.op == "output": + # Nodes cannot be moved after output + frontier_input = previous_frontier + else: + frontier_input = frontier + frontier_users = list(frontier_input.users) + + original_user.replace_input_with(node, producer) + node.replace_input_with(producer, frontier_input) + + for user in frontier_users: + if user is not node: + user.replace_input_with(frontier_input, node) + + if frontier.op == "output": + frontier.prepend(node) + else: + frontier.append(node) diff --git a/backends/arm/test/misc/test_transpose_counts.py b/backends/arm/test/misc/test_transpose_counts.py index 086edc537ba..168dabe96b9 100644 --- a/backends/arm/test/misc/test_transpose_counts.py +++ b/backends/arm/test/misc/test_transpose_counts.py @@ -392,7 +392,7 @@ def forward(self, x: torch.Tensor): "grouped_conv": TransposeCountCase( GroupedConvModule(), (torch.randn(1, 4, 8, 8),), - 4, + 2, ), "transpose_conv": TransposeCountCase( TransposeConvModule(), @@ -413,7 +413,7 @@ def forward(self, x: torch.Tensor): "lstm": TransposeCountCase( LstmModule(), (torch.randn(2, 4, 8),), - 2, + 1, ), "groupnorm": TransposeCountCase( GroupNormModule(), @@ -428,7 +428,7 @@ def forward(self, x: torch.Tensor): "multihead_attention_rank3": TransposeCountCase( MultiheadAttentionModule(), (torch.randn(2, 4, 8),), - 7, + 6, ), "cumsum_rank3_dim0": TransposeCountCase( CumsumModule(), @@ -441,31 +441,31 @@ def forward(self, x: torch.Tensor): 0, ), "model_1_conv_maxpool_residual_linear": TransposeCountCase( - Model1ConvMaxPoolResidualLinear(), (torch.randn(2, 8, 64),), 5 + Model1ConvMaxPoolResidualLinear(), (torch.randn(2, 8, 64),), 1 ), "model_2_conv_mha_linear_layernorm": TransposeCountCase( - Model2ConvMhaLinearLayerNorm(), (torch.randn(2, 8, 32),), 8 + Model2ConvMhaLinearLayerNorm(), (torch.randn(2, 8, 32),), 7 ), "model_3_lstm_linear": TransposeCountCase( - Model3LstmLinear(), (torch.randn(2, 16, 8),), 2 + Model3LstmLinear(), (torch.randn(2, 16, 8),), 1 ), "model_4_conv_lstm_linear_layernorm": TransposeCountCase( - Model4ConvLstmLinearLayerNorm(), (torch.randn(2, 8, 32),), 3 + Model4ConvLstmLinearLayerNorm(), (torch.randn(2, 8, 32),), 2 ), "model_5_dwconv_gelu_layernorm_avgpool": TransposeCountCase( Model5DwConvGeluLayerNormAvgPool(), (torch.randn(1, 8, 16, 16),), 2 ), "model_6_gru_linear": TransposeCountCase( - Model6GruLinear(), (torch.randn(2, 16, 8),), 2 + Model6GruLinear(), (torch.randn(2, 16, 8),), 1 ), "model_7_dwconv_batchnorm_linear": TransposeCountCase( Model7DwConvBatchNormLinear(), (torch.randn(2, 8, 64),), 1 ), "model_8_conv_batchnorm_maxpool_residual": TransposeCountCase( - Model8ConvBatchNormMaxPoolResidual(), (torch.randn(1, 8, 16, 16),), 4 + Model8ConvBatchNormMaxPoolResidual(), (torch.randn(1, 8, 16, 16),), 2 ), "model_9_dilated_conv_batchnorm_avgpool_residual": TransposeCountCase( - Model9DilatedConvBatchNormAvgPoolResidual(), (torch.randn(1, 8, 16, 16),), 4 + Model9DilatedConvBatchNormAvgPoolResidual(), (torch.randn(1, 8, 16, 16),), 2 ), "model_10_dwconv_batchnorm_linear_cat": TransposeCountCase( Model10DwConvBatchNormLinearCat(), (torch.randn(2, 8, 64),), 1 @@ -495,7 +495,7 @@ def forward(self, x: torch.Tensor): "conv3d_rank5_channels_last": TransposeCountCase( Conv3dModule(), (torch.randn(1, 2, 6, 6, 6).to(memory_format=torch.channels_last_3d),), - 3, + 1, ), "linear_rank4_channels_last": TransposeCountCase( LinearModule(), @@ -538,7 +538,7 @@ def forward(self, x: torch.Tensor): "maxpool2d_dilation_channels_last": TransposeCountCase( MaxPool2dDilatedModule(), (torch.randn(1, 2, 8, 8).to(memory_format=torch.channels_last),), - 4, + 3, ), "groupnorm_channels_last": TransposeCountCase( GroupNormModule(), diff --git a/backends/arm/test/passes/test_dim_maps.py b/backends/arm/test/passes/test_dim_maps.py index e71c815c471..16a18720442 100644 --- a/backends/arm/test/passes/test_dim_maps.py +++ b/backends/arm/test/passes/test_dim_maps.py @@ -262,13 +262,13 @@ def test_dim_map_maps_split_and_merged_prime_factor_groups() -> None: view_map = ViewMap.from_shapes([1, 2, 3, 4], [1, 6, 2, 2]) assert view_map.is_valid_map - assert view_map.map_dim(0) is None + assert view_map.map_dim(0) == [0] assert view_map.map_dim(1) is None assert view_map.map_dim(2) is None assert view_map.map_dim(3) == [2, 3] assert view_map.map_dim([1, 2]) == [1] assert view_map.map_dim([3, 1]) is None - assert view_map.map_dim([3, 1, 2]) == [2, 3, 1] + assert view_map.map_dim([3, 1, 2]) == [1, 2, 3] assert view_map.map_dim_inverse(0) is None assert view_map.map_dim_inverse(1) == [1, 2] @@ -360,6 +360,49 @@ def test_dim_map_uses_strict_no_mapping_for_singletons() -> None: assert split_view_map.map_dim_inverse([0, 2]) == [0] +def test_dim_map_maps_reduced_singletons_only_when_unambiguous() -> None: + split_singleton_view_map = ViewMap.from_shapes([1, 4], [1, 1, 4]) + assert split_singleton_view_map.map_dim(0) == [0, 1] + + squeezed_singleton_view_map = ViewMap.from_shapes([1, 50, 10, 1], [1, 50, 10]) + assert squeezed_singleton_view_map.map_dim(-1) is None + assert squeezed_singleton_view_map.map_dim([0, -1]) == [0] + + +def test_dim_map_remaps_unit_slice_through_view() -> None: + view_map = ViewMap.from_shapes([5, 2, 1, 4, 6], [5, 2, 4, 6]) + + assert view_map.remap_unit_slice([5, 2, 3, 4, 6], 2, 0, 1) == ( + [5, 2, 12, 6], + 2, + 0, + 4, + ) + assert view_map.remap_unit_slice([5, 2, 3, 4, 6], 2, 1, 2) == ( + [5, 2, 12, 6], + 2, + 4, + 8, + ) + + +def test_dim_map_remaps_unit_slice_through_flattening_view() -> None: + view_map = ViewMap.from_shapes([5, 2, 1, 4, 6], [5, 2, 24]) + + assert view_map.remap_unit_slice([5, 2, 3, 4, 6], 2, 1, 2) == ( + [5, 2, 72], + 2, + 24, + 48, + ) + + +def test_dim_map_does_not_remap_unit_slice_into_previous_axis() -> None: + view_map = ViewMap.from_shapes([3, 3, 1], [3, 3]) + + assert view_map.remap_unit_slice([3, 3, 3], 2, 0, 1) is None + + def test_dim_map_preserves_symbolic_dimensions_as_prime_factors() -> None: shape_env = ShapeEnv() batch = _make_symint(shape_env, "batch", hint=4) diff --git a/backends/arm/test/passes/test_insert_data_layout_casts_pass.py b/backends/arm/test/passes/test_insert_data_layout_casts_pass.py index b4298977e5b..bdacf5d27db 100644 --- a/backends/arm/test/passes/test_insert_data_layout_casts_pass.py +++ b/backends/arm/test/passes/test_insert_data_layout_casts_pass.py @@ -39,6 +39,11 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return torch.cat([x, y], dim=1) +class SliceModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x[:, 1:3] + + def test_insert_data_layout_casts_no_target_view_fp_profile_inserts_casts() -> None: test_data = (torch.arange(4, dtype=torch.int32).reshape(1, 4),) @@ -109,3 +114,27 @@ def test_insert_data_layout_casts_no_target_cat_fp_profile_inserts_casts() -> No cast_dtypes = _collect_cast_dtypes(pipeline) assert cast_dtypes.count(torch.float32) == 2 assert cast_dtypes.count(torch.int32) == 1 + + +def test_insert_data_layout_casts_no_target_slice_bf16_profile_inserts_casts() -> None: + test_data = (torch.arange(4, dtype=torch.int32).reshape(1, 4),) + + pipeline = PassPipeline[tuple[torch.Tensor, ...]]( + SliceModule(), + test_data, + quantize=False, + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2, + }, + pass_list=[InsertDataLayoutCastsPass], + tosa_extensions=["bf16"], + ) + pipeline.run() + + cast_dtypes = _collect_cast_dtypes(pipeline) + assert cast_dtypes.count(torch.float32) == 1 + assert cast_dtypes.count(torch.int32) == 1 diff --git a/backends/arm/test/passes/test_propagate_permutes_views_pass.py b/backends/arm/test/passes/test_propagate_permutes_views_pass.py new file mode 100644 index 00000000000..0fba5fecc4e --- /dev/null +++ b/backends/arm/test/passes/test_propagate_permutes_views_pass.py @@ -0,0 +1,1418 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from collections.abc import Callable +from typing import Tuple + +import pytest +import torch +from executorch.backends.arm._passes import ( + PropagateViewCopyPermuteDownPass, + PropagateViewCopyPermuteUpPass, +) + +from executorch.backends.arm._passes.arm_pass import ArmPass +from executorch.backends.arm.test.tester.test_pipeline import PassPipeline +from executorch.backends.arm.tosa.mapping import TosaSpecialDtype +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops + +input_t = Tuple[torch.Tensor] + +PERMUTE = exir_ops.edge.aten.permute_copy.default +VIEW = exir_ops.edge.aten.view_copy.default +ADD = exir_ops.edge.aten.add.Tensor +RELU = exir_ops.edge.aten.relu.default +NEG = exir_ops.edge.aten.neg.default +MM = exir_ops.edge.aten.mm.default +RESCALE = exir_ops.backend.tosa.RESCALE.default +TABLE = exir_ops.backend.tosa.TABLE.default +SCATTER = exir_ops.backend.tosa.SCATTER.default +CAT = exir_ops.edge.aten.cat.default +SLICE = exir_ops.edge.aten.slice_copy.Tensor +SUM = exir_ops.edge.aten.sum.dim_IntList +MEAN = exir_ops.edge.aten.mean.dim + + +def _assert_call_targets( + predicate: Callable[[list[object]], None], +) -> Callable[[ExportedProgram], ExportedProgram]: + def check_order(exported_program: ExportedProgram) -> ExportedProgram: + targets = [ + node.target + for node in exported_program.graph_module.graph.nodes + if node.op == "call_function" + ] + predicate(targets) + return exported_program + + return check_order + + +class DownwardPermute(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.permute(0, 2, 3, 1).relu().neg() + + data = (torch.randn(1, 2, 3, 4),) + + +def test_propagate_permute_down_through_transparent_ops_tosa_FP() -> None: + def predicate(targets: list[object]) -> None: + assert targets.index(PERMUTE) < targets.index(RELU) < targets.index(NEG) + + pipeline = PassPipeline[input_t]( + DownwardPermute(), + DownwardPermute.data, + quantize=False, + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + }, + pass_list=[PropagateViewCopyPermuteUpPass], + pass_functions=[_assert_call_targets(predicate)], + ) + pipeline.run() + + +class DownwardBinaryPermute(torch.nn.Module): + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x.permute(0, 2, 3, 1) + y.permute(0, 2, 3, 1) + + data = (torch.randn(1, 2, 3, 4), torch.randn(1, 2, 3, 4)) + + +class DownwardView(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.view(2, 12).relu().neg() + + data = (torch.randn(2, 3, 4),) + + +def test_propagate_view_down_through_transparent_ops_tosa_FP() -> None: + def predicate(targets: list[object]) -> None: + assert targets.index(VIEW) < targets.index(RELU) < targets.index(NEG) + + pipeline = PassPipeline[input_t]( + DownwardView(), + DownwardView.data, + quantize=False, + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, + }, + pass_list=[PropagateViewCopyPermuteUpPass], + pass_functions=[_assert_call_targets(predicate)], + ) + pipeline.run() + + +class UpwardPermute(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.relu().neg().permute(0, 2, 3, 1) + + data = (torch.randn(1, 2, 3, 4),) + + +def test_propagate_permute_up_through_transparent_ops_tosa_FP() -> None: + def predicate(targets: list[object]) -> None: + assert targets.index(PERMUTE) < targets.index(RELU) < targets.index(NEG) + + pipeline = PassPipeline[input_t]( + UpwardPermute(), + UpwardPermute.data, + quantize=False, + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + }, + pass_list=[PropagateViewCopyPermuteUpPass], + pass_functions=[_assert_call_targets(predicate)], + ) + pipeline.run() + + +class UpwardBinaryPermute(torch.nn.Module): + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return (x + y).permute(0, 2, 3, 1) + + data = (torch.randn(1, 2, 3, 4), torch.randn(1, 2, 3, 4)) + + +def test_propagate_permute_up_swaps_with_binary_transparent_op_tosa_FP() -> None: + def predicate(targets: list[object]) -> None: + assert targets.count(PERMUTE) == 1 + assert targets.index(ADD) < targets.index(PERMUTE) + + pipeline = PassPipeline[Tuple[torch.Tensor, torch.Tensor]]( + UpwardBinaryPermute(), + UpwardBinaryPermute.data, + quantize=False, + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + }, + pass_list=[PropagateViewCopyPermuteUpPass], + pass_functions=[_assert_call_targets(predicate)], + ) + pipeline.run() + + +class UpwardView(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.relu().neg().view(2, 12) + + data = (torch.randn(2, 3, 4),) + + +def test_propagate_view_up_through_transparent_ops_tosa_FP() -> None: + def predicate(targets: list[object]) -> None: + assert targets.index(VIEW) < targets.index(RELU) < targets.index(NEG) + + pipeline = PassPipeline[input_t]( + UpwardView(), + UpwardView.data, + quantize=False, + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, + }, + pass_list=[PropagateViewCopyPermuteUpPass], + pass_functions=[_assert_call_targets(predicate)], + ) + pipeline.run() + + +class StopAtNonTransparent(torch.nn.Module): + def forward(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return x.permute(1, 0).mm(weight) + + data = (torch.randn(3, 2), torch.randn(3, 4)) + + +def test_propagate_stops_at_non_transparent_ops_tosa_FP() -> None: + def predicate(targets: list[object]) -> None: + assert targets.index(PERMUTE) < targets.index(MM) + + pipeline = PassPipeline[Tuple[torch.Tensor, torch.Tensor]]( + StopAtNonTransparent(), + StopAtNonTransparent.data, + quantize=False, + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + }, + pass_list=[PropagateViewCopyPermuteUpPass], + pass_functions=[_assert_call_targets(predicate)], + ) + pipeline.run() + + +class StopAtBranch(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = x.permute(0, 2, 3, 1) + return y.relu() + y.neg() + + data = (torch.randn(1, 2, 3, 4),) + + +def test_propagate_stops_at_branches_tosa_FP() -> None: + def predicate(targets: list[object]) -> None: + assert targets.index(PERMUTE) < targets.index(RELU) + assert targets.index(PERMUTE) < targets.index(NEG) + + pipeline = PassPipeline[input_t]( + StopAtBranch(), + StopAtBranch.data, + quantize=False, + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + }, + pass_list=[PropagateViewCopyPermuteUpPass], + pass_functions=[_assert_call_targets(predicate)], + ) + pipeline.run() + + +class StopAtSharedTransformInput(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = x.permute(0, 2, 3, 1) + return (y * y.sigmoid()).permute(0, 3, 1, 2) + + data = (torch.randn(1, 2, 3, 4),) + + +class StopAtParameter(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(1, 2, 3, 4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x + self.weight).permute(0, 2, 3, 1) + + data = (torch.randn(1, 2, 3, 4),) + + +def test_propagate_moves_before_parameter_tosa_FP() -> None: + def predicate(targets: list[object]) -> None: + assert targets.index(ADD) < targets.index(PERMUTE) + + pipeline = PassPipeline[input_t]( + StopAtParameter(), + StopAtParameter.data, + quantize=False, + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + }, + pass_list=[PropagateViewCopyPermuteUpPass], + pass_functions=[_assert_call_targets(predicate)], + ) + pipeline.run() + + +def _run_pass_on_graph_module( + graph: torch.fx.Graph, + pass_cls: type[ArmPass] = PropagateViewCopyPermuteUpPass, +) -> torch.fx.GraphModule: + graph.lint() + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + result = pass_cls().call(graph_module) + return result.graph_module + + +def _run_pass_on_graph( + graph: torch.fx.Graph, + pass_cls: type[ArmPass] = PropagateViewCopyPermuteUpPass, +) -> list[object]: + graph_module = _run_pass_on_graph_module(graph, pass_cls) + return [ + node.target for node in graph_module.graph.nodes if node.op == "call_function" + ] + + +def test_is_swappable_rejects_unnormalized_keep_dim_operator() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + sum_node = graph.call_function(SUM, args=(x, [1], False)) + + with pytest.raises( + RuntimeError, + match="expects keep_dim=True for reduction ops to simplify propagation logic, got", + ): + PropagateViewCopyPermuteUpPass().is_swappable(sum_node) + + +def test_down_pass_moves_permute_after_transparent_chain() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + relu = graph.call_function(RELU, args=(permute,)) + relu.meta["val"] = torch.empty((1, 3, 4, 2)) + neg = graph.call_function(NEG, args=(relu,)) + neg.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output(neg) + + targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + + assert targets.index(RELU) < targets.index(NEG) < targets.index(PERMUTE) + + +def test_down_pass_skips_propagation_for_u85_like_tosa_int_cf() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + relu = graph.call_function(RELU, args=(permute,)) + relu.meta["val"] = torch.empty((1, 3, 4, 2)) + neg = graph.call_function(NEG, args=(relu,)) + neg.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output(neg) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT+cf")): + targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + + assert targets.index(PERMUTE) < targets.index(RELU) < targets.index(NEG) + + +def test_down_pass_still_canonicalizes_for_u85_like_tosa_int_cf() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3)) + first_permute = graph.call_function(PERMUTE, args=(x, [0, 2, 1])) + first_permute.meta["val"] = torch.empty((1, 3, 2)) + second_permute = graph.call_function(PERMUTE, args=(first_permute, [0, 2, 1])) + second_permute.meta["val"] = torch.empty((1, 2, 3)) + graph.output(second_permute) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT+cf")): + targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + + assert PERMUTE not in targets + + +def test_down_pass_moves_view_after_transparent_chain() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((2, 3, 4)) + view = graph.call_function(VIEW, args=(x, [2, 12])) + view.meta["val"] = torch.empty((2, 12)) + relu = graph.call_function(RELU, args=(view,)) + relu.meta["val"] = torch.empty((2, 12)) + neg = graph.call_function(NEG, args=(relu,)) + neg.meta["val"] = torch.empty((2, 12)) + graph.output(neg) + + targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + + assert targets.index(RELU) < targets.index(NEG) < targets.index(VIEW) + + +def test_down_pass_moves_permute_to_graph_output() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + relu = graph.call_function(RELU, args=(permute,)) + relu.meta["val"] = torch.empty((1, 3, 4, 2)) + neg = graph.call_function(NEG, args=(relu,)) + neg.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output(neg) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteDownPass) + nodes = list(graph_module.graph.nodes) + output = next(node for node in nodes if node.op == "output") + moved_permute = next(node for node in nodes if node.target == PERMUTE) + moved_neg = next(node for node in nodes if node.target == NEG) + + assert output.args[0] is moved_permute + assert moved_permute.args[0] is moved_neg + assert nodes.index(moved_neg) < nodes.index(moved_permute) < nodes.index(output) + + +def test_down_pass_moves_permute_to_matching_output_branch() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + left = graph.call_function(RELU, args=(x,)) + left.meta["val"] = torch.empty((1, 2, 3, 4)) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + relu = graph.call_function(RELU, args=(permute,)) + relu.meta["val"] = torch.empty((1, 3, 4, 2)) + neg = graph.call_function(NEG, args=(relu,)) + neg.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output((left, neg)) + + graph_module = torch.fx.GraphModule({}, graph) + output = next(node for node in graph.nodes if node.op == "output") + PropagateViewCopyPermuteDownPass()._move_node(permute, output, neg) + graph.lint() + graph_module.recompile() + + assert output.args[0] == (left, permute) + assert relu.args[0] is x + assert permute.args[0] is neg + + +def test_up_pass_moves_permute_to_graph_input() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + relu = graph.call_function(RELU, args=(x,)) + relu.meta["val"] = torch.empty((1, 2, 3, 4)) + neg = graph.call_function(NEG, args=(relu,)) + neg.meta["val"] = torch.empty((1, 2, 3, 4)) + permute = graph.call_function(PERMUTE, args=(neg, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output(permute) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteUpPass) + nodes = list(graph_module.graph.nodes) + x = next(node for node in nodes if node.op == "placeholder") + moved_permute = next(node for node in nodes if node.target == PERMUTE) + moved_relu = next(node for node in nodes if node.target == RELU) + + assert moved_permute.args[0] is x + assert moved_relu.args[0] is moved_permute + assert nodes.index(x) < nodes.index(moved_permute) < nodes.index(moved_relu) + + +def test_up_pass_fuses_duplicate_permutes_at_placeholder() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 4, 3, 3)) + left_slice = graph.call_function(SLICE, args=(x, 1, 0, 2)) + left_slice.meta["val"] = torch.empty((1, 2, 3, 3)) + right_slice = graph.call_function(SLICE, args=(x, 1, 2, 4)) + right_slice.meta["val"] = torch.empty((1, 2, 3, 3)) + left_permute = graph.call_function(PERMUTE, args=(left_slice, [0, 2, 3, 1])) + left_permute.meta["val"] = torch.empty((1, 3, 3, 2)) + right_permute = graph.call_function(PERMUTE, args=(right_slice, [0, 2, 3, 1])) + right_permute.meta["val"] = torch.empty((1, 3, 3, 2)) + graph.output((left_permute, right_permute)) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteUpPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + permutes = [node for node in call_nodes if node.target == PERMUTE] + slices = [node for node in call_nodes if node.target == SLICE] + x = next(node for node in graph_module.graph.nodes if node.op == "placeholder") + + assert len(permutes) == 1 + assert len(slices) == 2 + assert permutes[0].args == (x, [0, 2, 3, 1]) + assert [slice_node.args for slice_node in slices] == [ + (permutes[0], 3, 0, 2), + (permutes[0], 3, 2, 4), + ] + + +def test_up_pass_refreshes_permute_meta_before_view_slice_swap() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((3, 2, 8, 16)) + slice_node = graph.call_function(SLICE, args=(x, 0, 0, 1)) + slice_node.meta["val"] = torch.empty((1, 2, 8, 16)) + permute = graph.call_function(PERMUTE, args=(slice_node, [0, 3, 1, 2])) + permute.meta["val"] = torch.empty((1, 16, 2, 8)) + view = graph.call_function(VIEW, args=(permute, [1, 32, 8])) + view.meta["val"] = torch.empty((1, 32, 8)) + graph.output(view) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteUpPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + permute = next(node for node in call_nodes if node.target == PERMUTE) + view = next(node for node in call_nodes if node.target == VIEW) + slice_node = next(node for node in call_nodes if node.target == SLICE) + graph_input = next( + node for node in graph_module.graph.nodes if node.op == "placeholder" + ) + + assert permute.args == (graph_input, [0, 3, 1, 2]) + assert permute.meta["val"].shape == torch.Size((3, 16, 2, 8)) + assert view.args == (permute, [3, 32, 8]) + assert slice_node.args == (view, 0, 0, 1) + + +def test_up_pass_keeps_scatter_input_view_after_slice() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((3, 16, 2, 8)) + indices = graph.placeholder("indices") + indices.meta["val"] = torch.empty((1, 4), dtype=torch.int32) + data = graph.placeholder("data") + data.meta["val"] = torch.empty((1, 4, 16)) + slice_node = graph.call_function(SLICE, args=(x, 0, 0, 1)) + slice_node.meta["val"] = torch.empty((1, 16, 2, 8)) + view = graph.call_function(VIEW, args=(slice_node, [1, 32, 8])) + view.meta["val"] = torch.empty((1, 32, 8)) + scatter = graph.call_function(SCATTER, args=(view, indices, data)) + scatter.meta["val"] = torch.empty((1, 32, 8)) + graph.output(scatter) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteUpPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + view = next(node for node in call_nodes if node.target == VIEW) + slice_node = next(node for node in call_nodes if node.target == SLICE) + scatter = next(node for node in call_nodes if node.target == SCATTER) + + assert view.args == (slice_node, [1, 32, 8]) + assert scatter.args[0] is view + + +def test_up_pass_hoists_matching_transform_chain_across_slice_fanout() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 4, 3, 3)) + left_slice = graph.call_function(SLICE, args=(x, 1, 0, 2)) + left_slice.meta["val"] = torch.empty((1, 2, 3, 3)) + right_slice = graph.call_function(SLICE, args=(x, 1, 2, 4)) + right_slice.meta["val"] = torch.empty((1, 2, 3, 3)) + left_view = graph.call_function(VIEW, args=(left_slice, [1, 2, 9])) + left_view.meta["val"] = torch.empty((1, 2, 9)) + right_view = graph.call_function(VIEW, args=(right_slice, [1, 2, 9])) + right_view.meta["val"] = torch.empty((1, 2, 9)) + left_permute = graph.call_function(PERMUTE, args=(left_view, [0, 2, 1])) + left_permute.meta["val"] = torch.empty((1, 9, 2)) + right_permute = graph.call_function(PERMUTE, args=(right_view, [0, 2, 1])) + right_permute.meta["val"] = torch.empty((1, 9, 2)) + graph.output((left_permute, right_permute)) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteUpPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + views = [node for node in call_nodes if node.target == VIEW] + permutes = [node for node in call_nodes if node.target == PERMUTE] + slices = [node for node in call_nodes if node.target == SLICE] + graph_input = next( + node for node in graph_module.graph.nodes if node.op == "placeholder" + ) + + assert len(views) == 1 + assert len(permutes) == 1 + assert len(slices) == 2 + assert views[0].args == (graph_input, [1, 4, 9]) + assert permutes[0].args == (views[0], [0, 2, 1]) + assert [slice_node.args for slice_node in slices] == [ + (permutes[0], 2, 0, 2), + (permutes[0], 2, 2, 4), + ] + + +def test_up_pass_hoists_unit_slice_views_with_different_args() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((5, 2, 3, 4, 6)) + left_slice = graph.call_function(SLICE, args=(x, 2, 0, 1)) + left_slice.meta["val"] = torch.empty((5, 2, 1, 4, 6)) + right_slice = graph.call_function(SLICE, args=(x, 2, 1, 2)) + right_slice.meta["val"] = torch.empty((5, 2, 1, 4, 6)) + left_view = graph.call_function(VIEW, args=(left_slice, [5, 2, 4, 6])) + left_view.meta["val"] = torch.empty((5, 2, 4, 6)) + right_view = graph.call_function(VIEW, args=(right_slice, [5, 2, 24])) + right_view.meta["val"] = torch.empty((5, 2, 24)) + graph.output((left_view, right_view)) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteUpPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + views = [node for node in call_nodes if node.target == VIEW] + slices = [node for node in call_nodes if node.target == SLICE] + graph_input = next( + node for node in graph_module.graph.nodes if node.op == "placeholder" + ) + + assert [view.args for view in views] == [ + (graph_input, [5, 2, 12, 6]), + (graph_input, [5, 2, 72]), + ] + assert [slice_node.args for slice_node in slices] == [ + (views[0], 2, 0, 4), + (views[1], 2, 24, 48), + ] + + +def test_up_pass_keeps_mismatched_transform_slice_fanout_split() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 4, 3, 3)) + left_slice = graph.call_function(SLICE, args=(x, 1, 0, 2)) + left_slice.meta["val"] = torch.empty((1, 2, 3, 3)) + right_slice = graph.call_function(SLICE, args=(x, 1, 2, 4)) + right_slice.meta["val"] = torch.empty((1, 2, 3, 3)) + left_view = graph.call_function(VIEW, args=(left_slice, [1, 2, 9])) + left_view.meta["val"] = torch.empty((1, 2, 9)) + right_permute = graph.call_function(PERMUTE, args=(right_slice, [0, 2, 3, 1])) + right_permute.meta["val"] = torch.empty((1, 3, 3, 2)) + graph.output((left_view, right_permute)) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteUpPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + slices = [node for node in call_nodes if node.target == SLICE] + + assert len(slices) == 2 + assert slices[0].args[0] is not slices[1].args[0] + + +def test_down_pass_moves_matching_input_permutations_after_binary_op() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((1, 2, 3, 4)) + x_permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + x_permute.meta["val"] = torch.empty((1, 3, 4, 2)) + y_permute = graph.call_function(PERMUTE, args=(y, [0, 2, 3, 1])) + y_permute.meta["val"] = torch.empty((1, 3, 4, 2)) + add = graph.call_function(ADD, args=(x_permute, y_permute)) + add.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output(add) + + targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + + assert targets.count(PERMUTE) == 1 + assert targets.index(ADD) < targets.index(PERMUTE) + + +def test_down_pass_keeps_sunk_view_before_rank_reducing_permute() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((2, 8, 1, 32)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((2, 8, 1, 32)) + x_view = graph.call_function(VIEW, args=(x, [2, 8, 32])) + x_view.meta["val"] = torch.empty((2, 8, 32)) + y_view = graph.call_function(VIEW, args=(y, [2, 8, 32])) + y_view.meta["val"] = torch.empty((2, 8, 32)) + add = graph.call_function(ADD, args=(x_view, y_view)) + add.meta["val"] = torch.empty((2, 8, 32)) + output_view = graph.call_function(VIEW, args=(add, [2, 8, 32])) + output_view.meta["val"] = torch.empty((2, 8, 32)) + permute = graph.call_function(PERMUTE, args=(output_view, [0, 2, 1])) + permute.meta["val"] = torch.empty((2, 32, 8)) + graph.output(permute) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteDownPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + add = next(node for node in call_nodes if node.target == ADD) + output_view = next(node for node in call_nodes if node.target == VIEW) + permute = next(node for node in call_nodes if node.target == PERMUTE) + + assert targets.count(VIEW) == 1 + assert targets.index(ADD) < targets.index(VIEW) < targets.index(PERMUTE) + assert add.meta["val"].shape == torch.Size((2, 8, 1, 32)) + assert output_view.args[0] is add + assert permute.args[0] is output_view + + +def test_down_pass_canonicalizes_horizontally_fused_singleton_permute() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 1, 1, 1)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((1, 1, 1, 1)) + x_permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + x_permute.meta["val"] = torch.empty((1, 1, 1, 1)) + y_permute = graph.call_function(PERMUTE, args=(y, [0, 2, 3, 1])) + y_permute.meta["val"] = torch.empty((1, 1, 1, 1)) + add = graph.call_function(ADD, args=(x_permute, y_permute)) + add.meta["val"] = torch.empty((1, 1, 1, 1)) + graph.output(add) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteDownPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + add = next(node for node in call_nodes if node.target == ADD) + + assert targets.count(PERMUTE) == 0 + assert targets.count(VIEW) == 0 + assert [input_node.name for input_node in add.all_input_nodes] == ["x", "y"] + assert next(iter(add.users)).op == "output" + + +def test_down_pass_moves_matching_input_permutations_after_cat() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((1, 2, 3, 4)) + x_permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + x_permute.meta["val"] = torch.empty((1, 3, 4, 2)) + y_permute = graph.call_function(PERMUTE, args=(y, [0, 2, 3, 1])) + y_permute.meta["val"] = torch.empty((1, 3, 4, 2)) + cat_node = graph.call_function(CAT, args=([x_permute, y_permute], 3)) + cat_node.meta["val"] = torch.empty((1, 3, 4, 4)) + graph.output(cat_node) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteDownPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + cat_node = next(node for node in call_nodes if node.target == CAT) + + assert targets.count(PERMUTE) == 1 + assert targets.index(CAT) < targets.index(PERMUTE) + assert cat_node.args[1] == 1 + + +def test_down_pass_swaps_concat_with_matching_input_permutations() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((1, 2, 3, 4)) + x_permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + x_permute.meta["val"] = torch.empty((1, 3, 4, 2)) + y_permute = graph.call_function(PERMUTE, args=(y, [0, 2, 3, 1])) + y_permute.meta["val"] = torch.empty((1, 3, 4, 2)) + cat_node = graph.call_function(CAT, args=([x_permute, y_permute], 3)) + cat_node.meta["val"] = torch.empty((1, 3, 4, 4)) + graph.output(cat_node) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteDownPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + cat_node = next(node for node in call_nodes if node.target == CAT) + permute = next(node for node in call_nodes if node.target == PERMUTE) + + assert targets.count(PERMUTE) == 1 + assert targets.index(CAT) < targets.index(PERMUTE) + assert [input_node.name for input_node in cat_node.args[0]] == ["x", "y"] + assert cat_node.args[1] == 1 + assert cat_node.meta["val"].shape == torch.Size((1, 4, 3, 4)) + assert permute.args == (cat_node, [0, 2, 3, 1]) + + +def test_up_pass_moves_noop_input_permutations_before_cat() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 1, 3, 4)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((1, 1, 3, 4)) + cat_node = graph.call_function(CAT, args=([x, y], 1)) + cat_node.meta["val"] = torch.empty((1, 2, 3, 4)) + permute = graph.call_function(PERMUTE, args=(cat_node, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output(permute) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteUpPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + cat_node = next(node for node in call_nodes if node.target == CAT) + + assert targets.count(PERMUTE) == 0 + assert targets.count(VIEW) == 2 + assert cat_node.args[1] == 3 + assert cat_node.meta["val"].shape == torch.Size((1, 3, 4, 2)) + assert all(input_node.target == VIEW for input_node in cat_node.args[0]) + + +def test_up_pass_swaps_concat_with_noop_output_permutation() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 1, 3, 4)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((1, 1, 3, 4)) + cat_node = graph.call_function(CAT, args=([x, y], 1)) + cat_node.meta["val"] = torch.empty((1, 2, 3, 4)) + permute = graph.call_function(PERMUTE, args=(cat_node, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output(permute) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteUpPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + cat_node = next(node for node in call_nodes if node.target == CAT) + + assert targets.count(PERMUTE) == 0 + assert targets.count(VIEW) == 2 + assert cat_node.args[1] == 3 + assert cat_node.meta["val"].shape == torch.Size((1, 3, 4, 2)) + assert all(input_node.target == VIEW for input_node in cat_node.args[0]) + + +def test_down_pass_keeps_shared_input_permutations_before_cat() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((1, 2, 3, 4)) + x_permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + x_permute.meta["val"] = torch.empty((1, 3, 4, 2)) + y_permute = graph.call_function(PERMUTE, args=(y, [0, 2, 3, 1])) + y_permute.meta["val"] = torch.empty((1, 3, 4, 2)) + relu = graph.call_function(RELU, args=(x_permute,)) + relu.meta["val"] = torch.empty((1, 3, 4, 2)) + cat_node = graph.call_function(CAT, args=([x_permute, y_permute], 3)) + cat_node.meta["val"] = torch.empty((1, 3, 4, 4)) + graph.output((cat_node, relu)) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteDownPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + cat_node = next(node for node in call_nodes if node.target == CAT) + + assert targets.count(PERMUTE) == 2 + assert [input_node.target for input_node in cat_node.args[0]] == [ + PERMUTE, + PERMUTE, + ] + assert cat_node.args[1] == 3 + + +def test_down_pass_moves_permutation_after_reduction() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + sum_node = graph.call_function(SUM, args=(permute, [3], True)) + sum_node.meta["val"] = torch.empty((1, 3, 4, 1)) + graph.output(sum_node) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteDownPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + sum_node = next(node for node in call_nodes if node.target == SUM) + transform = next(node for node in call_nodes if node.target in (PERMUTE, VIEW)) + + assert targets.index(SUM) < targets.index(transform.target) + assert sum_node.args[1] == [1] + assert transform.meta["val"].shape == torch.Size((1, 3, 4, 1)) + + +def test_down_pass_stops_when_fanout_does_not_converge() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + relu = graph.call_function(RELU, args=(permute,)) + relu.meta["val"] = torch.empty((1, 3, 4, 2)) + neg = graph.call_function(NEG, args=(permute,)) + neg.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output((relu, neg)) + + targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + + assert targets.count(PERMUTE) == 1 + assert targets.index(PERMUTE) < targets.index(RELU) + assert targets.index(PERMUTE) < targets.index(NEG) + + +def test_down_pass_splits_permute_over_slice_fanout() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 4, 3, 3)) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 3, 4)) + left_slice = graph.call_function(SLICE, args=(permute, 3, 0, 2)) + left_slice.meta["val"] = torch.empty((1, 3, 3, 2)) + right_slice = graph.call_function(SLICE, args=(permute, 3, 2, 4)) + right_slice.meta["val"] = torch.empty((1, 3, 3, 2)) + left_view = graph.call_function(VIEW, args=(left_slice, [1, 9, 2])) + left_view.meta["val"] = torch.empty((1, 9, 2)) + right_view = graph.call_function(VIEW, args=(right_slice, [1, 9, 2])) + right_view.meta["val"] = torch.empty((1, 9, 2)) + left_permute = graph.call_function(PERMUTE, args=(left_view, [0, 2, 1])) + left_permute.meta["val"] = torch.empty((1, 2, 9)) + right_permute = graph.call_function(PERMUTE, args=(right_view, [0, 2, 1])) + right_permute.meta["val"] = torch.empty((1, 2, 9)) + graph.output((left_permute, right_permute)) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteDownPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + slices = [node for node in call_nodes if node.target == SLICE] + permutes = [node for node in call_nodes if node.target == PERMUTE] + graph_input = next( + node for node in graph_module.graph.nodes if node.op == "placeholder" + ) + + assert [slice_node.args for slice_node in slices] == [ + (graph_input, 1, 0, 2), + (graph_input, 1, 2, 4), + ] + assert all(permute_node.args[0].target == SLICE for permute_node in permutes) + + +def test_down_pass_stops_when_fanout_branch_has_nontransparent_op() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((2, 3)) + weight = graph.placeholder("weight") + weight.meta["val"] = torch.empty((2, 2)) + permute = graph.call_function(PERMUTE, args=(x, [1, 0])) + permute.meta["val"] = torch.empty((3, 2)) + relu = graph.call_function(RELU, args=(permute,)) + relu.meta["val"] = torch.empty((3, 2)) + mm = graph.call_function(MM, args=(permute, weight)) + mm.meta["val"] = torch.empty((3, 2)) + add = graph.call_function(ADD, args=(relu, mm)) + add.meta["val"] = torch.empty((3, 2)) + graph.output(add) + + targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + + assert targets.count(PERMUTE) == 1 + assert targets.index(PERMUTE) < targets.index(RELU) + assert targets.index(PERMUTE) < targets.index(MM) + + +def test_down_pass_stops_when_convergence_has_untracked_input() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((1, 3, 4, 2)) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + relu = graph.call_function(RELU, args=(permute,)) + relu.meta["val"] = torch.empty((1, 3, 4, 2)) + neg = graph.call_function(NEG, args=(permute,)) + neg.meta["val"] = torch.empty((1, 3, 4, 2)) + cat_node = graph.call_function(CAT, args=([relu, neg, y], 3)) + cat_node.meta["val"] = torch.empty((1, 3, 4, 6)) + graph.output(cat_node) + + targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + + assert targets.count(PERMUTE) == 1 + assert targets.index(PERMUTE) < targets.index(RELU) + assert targets.index(PERMUTE) < targets.index(NEG) + + +def test_down_pass_stops_view_before_cat_converging_fanout() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + view = graph.call_function(VIEW, args=(x, [1, 3, 4, 2])) + view.meta["val"] = torch.empty((1, 3, 4, 2)) + relu = graph.call_function(RELU, args=(view,)) + relu.meta["val"] = torch.empty((1, 3, 4, 2)) + neg = graph.call_function(NEG, args=(view,)) + neg.meta["val"] = torch.empty((1, 3, 4, 2)) + cat_node = graph.call_function(CAT, args=([relu, neg], 3)) + cat_node.meta["val"] = torch.empty((1, 3, 4, 4)) + graph.output(cat_node) + + targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + + assert targets.count(VIEW) == 1 + assert targets.index(VIEW) < targets.index(RELU) + assert targets.index(VIEW) < targets.index(NEG) + + +def test_up_pass_fuses_equivalent_output_permutations_before_fan_out() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + relu = graph.call_function(RELU, args=(x,)) + relu.meta["val"] = torch.empty((1, 2, 3, 4)) + first_permute = graph.call_function(PERMUTE, args=(relu, [0, 2, 3, 1])) + first_permute.meta["val"] = torch.empty((1, 3, 4, 2)) + second_permute = graph.call_function(PERMUTE, args=(relu, [0, 2, 3, 1])) + second_permute.meta["val"] = torch.empty((1, 3, 4, 2)) + add = graph.call_function(ADD, args=(first_permute, second_permute)) + add.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output(add) + graph.lint() + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + result = PropagateViewCopyPermuteUpPass().call(graph_module) + targets = [ + node.target + for node in result.graph_module.graph.nodes + if node.op == "call_function" + ] + + assert targets.count(PERMUTE) == 1 + assert targets.index(PERMUTE) < targets.index(RELU) < targets.index(ADD) + + +def test_propagate_moves_before_dtype_changing_rescale() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) + rescale = graph.call_function(RESCALE, args=(x, torch.int8, [1.0], 0, 0)) + rescale.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int8) + permute = graph.call_function(PERMUTE, args=(rescale, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.int8) + graph.output(permute) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")): + targets = _run_pass_on_graph(graph) + + assert targets.index(PERMUTE) < targets.index(RESCALE) + + +def test_propagate_fuses_permute_view_around_table() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((2, 3, 4), dtype=torch.int8) + table = graph.placeholder("table") + table.meta["val"] = torch.empty((256,), dtype=torch.int8) + permute = graph.call_function(PERMUTE, args=(x, [1, 0, 2])) + permute.meta["val"] = torch.empty((3, 2, 4), dtype=torch.int8) + view = graph.call_function(VIEW, args=(permute, [3, 8])) + view.meta["val"] = torch.empty((3, 8), dtype=torch.int8) + table_node = graph.call_function(TABLE, args=(view, table)) + table_node.meta["val"] = torch.empty((3, 8), dtype=torch.int8) + output_view = graph.call_function(VIEW, args=(table_node, [3, 2, 4])) + output_view.meta["val"] = torch.empty((3, 2, 4), dtype=torch.int8) + output_permute = graph.call_function(PERMUTE, args=(output_view, [1, 0, 2])) + output_permute.meta["val"] = torch.empty((2, 3, 4), dtype=torch.int8) + graph.output(output_permute) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")): + graph_module = _run_pass_on_graph_module( + graph, PropagateViewCopyPermuteDownPass + ) + targets = [ + node.target for node in graph_module.graph.nodes if node.op == "call_function" + ] + + assert targets == [TABLE] + + +def test_propagate_stops_at_per_channel_rescale() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((2, 3, 4, 5), dtype=torch.int32) + rescale = graph.call_function( + RESCALE, args=(x, torch.int8, [1.0, 1.0, 1.0, 1.0, 1.0], 0, 0) + ) + rescale.meta["val"] = torch.empty((2, 3, 4, 5), dtype=torch.int8) + permute = graph.call_function(PERMUTE, args=(rescale, [0, 3, 1, 2])) + permute.meta["val"] = torch.empty((2, 5, 3, 4), dtype=torch.int8) + graph.output(permute) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")): + targets = _run_pass_on_graph(graph) + + assert targets.index(RESCALE) < targets.index(PERMUTE) + + +def test_propagate_stops_at_rescale_changing_special_dtype() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 1, 1, 15), dtype=torch.int32) + x.meta[TosaSpecialDtype.meta_key()] = TosaSpecialDtype.INT48 + rescale = graph.call_function(RESCALE, args=(x, torch.int32, [1.0], 0, 0)) + rescale.meta["val"] = torch.empty((1, 1, 1, 15), dtype=torch.int32) + view = graph.call_function(VIEW, args=(rescale, [15])) + view.meta["val"] = torch.empty((15,), dtype=torch.int32) + graph.output(view) + + targets = _run_pass_on_graph(graph) + + assert targets.index(RESCALE) < targets.index(VIEW) + + +def test_propagate_up_stops_at_shared_rescale_producer() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((10, 80, 16), dtype=torch.int8) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((10, 80, 16), dtype=torch.int32) + rescale = graph.call_function(RESCALE, args=(x, torch.int32, [1.0], 0, 0)) + rescale.meta["val"] = torch.empty((10, 80, 16), dtype=torch.int32) + permute = graph.call_function(PERMUTE, args=(rescale, [1, 0, 2])) + permute.meta["val"] = torch.empty((80, 10, 16), dtype=torch.int32) + add = graph.call_function(ADD, args=(rescale, y)) + add.meta["val"] = torch.empty((10, 80, 16), dtype=torch.int32) + graph.output((permute, add)) + + targets = _run_pass_on_graph(graph) + + assert targets.index(RESCALE) < targets.index(PERMUTE) + + +def test_propagate_moves_before_int48_special_dtype() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) + x.meta[TosaSpecialDtype.meta_key()] = TosaSpecialDtype.INT48 + relu = graph.call_function(RELU, args=(x,)) + relu.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) + relu.meta[TosaSpecialDtype.meta_key()] = TosaSpecialDtype.INT48 + permute = graph.call_function(PERMUTE, args=(relu, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.int32) + permute.meta[TosaSpecialDtype.meta_key()] = TosaSpecialDtype.INT48 + graph.output(permute) + + targets = _run_pass_on_graph(graph) + + assert targets.index(PERMUTE) < targets.index(RELU) + + +def test_propagate_moves_output_view_before_sum_with_split_dim_remap() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((6, 4)) + sum_node = graph.call_function(SUM, args=(x, [0], True)) + sum_node.meta["val"] = torch.empty((1, 4)) + view = graph.call_function(VIEW, args=(sum_node, [1, 1, 4])) + view.meta["val"] = torch.empty((1, 1, 4)) + graph.output(view) + + graph_module = _run_pass_on_graph_module(graph) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + sum_node = next(node for node in call_nodes if node.target == SUM) + view = next(node for node in call_nodes if node.target == VIEW) + + assert targets.index(VIEW) < targets.index(SUM) + assert sum_node.args[1] == [0, 1] + assert view.args[1] == [6, 1, 4] + + +def test_propagate_updates_view_map_between_arg_updates() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((6, 4)) + slice_node = graph.call_function(SLICE, args=(x, 0, 0, 4)) + slice_node.meta["val"] = torch.empty((4, 4)) + sum_node = graph.call_function(SUM, args=(slice_node, [0], True)) + sum_node.meta["val"] = torch.empty((1, 4)) + view = graph.call_function(VIEW, args=(sum_node, [1, 1, 4])) + view.meta["val"] = torch.empty((1, 1, 4)) + graph.output(view) + + graph_module = _run_pass_on_graph_module(graph) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + view = next(node for node in call_nodes if node.target == VIEW) + slice_node = next(node for node in call_nodes if node.target == SLICE) + sum_node = next(node for node in call_nodes if node.target == SUM) + + assert targets.index(VIEW) < targets.index(SLICE) < targets.index(SUM) + assert view.args[1] == [6, 1, 4] + assert slice_node.args[1] == 0 + assert sum_node.args[1] == [0, 1] + + +def test_propagate_moves_output_view_before_mean_with_split_dim_remap() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((6, 4)) + mean_node = graph.call_function(MEAN, args=(x, [0], True)) + mean_node.meta["val"] = torch.empty((1, 4)) + view = graph.call_function(VIEW, args=(mean_node, [1, 1, 4])) + view.meta["val"] = torch.empty((1, 1, 4)) + graph.output(view) + + graph_module = _run_pass_on_graph_module(graph) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + mean_node = next(node for node in call_nodes if node.target == MEAN) + view = next(node for node in call_nodes if node.target == VIEW) + + assert targets.index(VIEW) < targets.index(MEAN) + assert mean_node.args[1] == [0, 1] + assert view.args[1] == [6, 1, 4] + + +def test_propagate_keeps_reduction_squeeze_after_sum() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 50, 10, 20)) + sum_node = graph.call_function(SUM, args=(x, [-1], True)) + sum_node.meta["val"] = torch.empty((1, 50, 10, 1)) + view = graph.call_function(VIEW, args=(sum_node, [1, 50, 10])) + view.meta["val"] = torch.empty((1, 50, 10)) + graph.output(view) + + graph_module = _run_pass_on_graph_module(graph) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + sum_node = next(node for node in call_nodes if node.target == SUM) + view = next(node for node in call_nodes if node.target == VIEW) + + assert targets.index(SUM) < targets.index(VIEW) + assert sum_node.args[1] == [-1] + assert view.args[1] == [1, 50, 10] + + +def test_propagate_keeps_unit_slice_before_reordering_view() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 3, 1, 7)) + slice_node = graph.call_function(SLICE, args=(x, 3, 2, 3)) + slice_node.meta["val"] = torch.empty((1, 3, 1, 1)) + view = graph.call_function(VIEW, args=(slice_node, [1, 1, 1, 3])) + view.meta["val"] = torch.empty((1, 1, 1, 3)) + graph.output(view) + + graph_module = _run_pass_on_graph_module(graph) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + slice_node = next(node for node in call_nodes if node.target == SLICE) + view = next(node for node in call_nodes if node.target == VIEW) + + assert targets.index(SLICE) < targets.index(VIEW) + assert slice_node.args[1:4] == (3, 2, 3) + assert view.args == (slice_node, [1, 1, 1, 3]) + + +def test_propagate_stops_when_downward_inputs_are_not_equivalent_transforms() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((1, 3, 4, 2)) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + add = graph.call_function(ADD, args=(permute, y)) + add.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output(add) + + targets = _run_pass_on_graph(graph) + + assert targets.index(PERMUTE) < targets.index(ADD) + + +def test_propagate_stops_split_dim_view_at_slice() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((6, 4)) + slice_node = graph.call_function(SLICE, args=(x, 0, 0, 4)) + slice_node.meta["val"] = torch.empty((4, 4)) + view = graph.call_function(VIEW, args=(slice_node, [2, 2, 4])) + view.meta["val"] = torch.empty((2, 2, 4)) + graph.output(view) + + targets = _run_pass_on_graph(graph) + + assert targets.index(SLICE) < targets.index(VIEW) + + +def test_propagate_stops_merged_trailing_dim_view_at_slice() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((25, 5, 13, 7)) + slice_node = graph.call_function(SLICE, args=(x, 1, 0, 2)) + slice_node.meta["val"] = torch.empty((25, 2, 13, 7)) + view = graph.call_function(VIEW, args=(slice_node, [1, 25, 182])) + view.meta["val"] = torch.empty((1, 25, 182)) + graph.output(view) + + targets = _run_pass_on_graph(graph) + + assert targets.index(SLICE) < targets.index(VIEW) + + +def test_propagate_stops_split_dim_view_at_cat() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((3, 4)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((3, 4)) + cat_node = graph.call_function(CAT, args=([x, y], 0)) + cat_node.meta["val"] = torch.empty((6, 4)) + view = graph.call_function(VIEW, args=(cat_node, [2, 3, 4])) + view.meta["val"] = torch.empty((2, 3, 4)) + graph.output(view) + + targets = _run_pass_on_graph(graph) + + assert targets.index(CAT) < targets.index(VIEW) + + +def test_propagate_keeps_channel_unit_slice_before_reordering_view() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + slice_node = graph.call_function(SLICE, args=(x, 1, 0, 1)) + slice_node.meta["val"] = torch.empty((1, 1, 3, 4)) + view = graph.call_function(VIEW, args=(slice_node, [1, 3, 4, 1])) + view.meta["val"] = torch.empty((1, 3, 4, 1)) + graph.output(view) + + graph_module = _run_pass_on_graph_module(graph) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + targets = [node.target for node in call_nodes] + slice_node = next(node for node in call_nodes if node.target == SLICE) + view = next(node for node in call_nodes if node.target == VIEW) + + assert targets.index(SLICE) < targets.index(VIEW) + assert slice_node.args[1:4] == (1, 0, 1) + assert view.args == (slice_node, [1, 3, 4, 1]) + + +def test_propagate_up_stops_at_multiple_distinct_edge_nodes() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((1, 2, 3, 4)) + add = graph.call_function(ADD, args=(x, y)) + add.meta["val"] = torch.empty((1, 2, 3, 4)) + permute = graph.call_function(PERMUTE, args=(add, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output(permute) + + targets = _run_pass_on_graph(graph) + + assert targets.count(PERMUTE) == 1 + assert targets.index(ADD) < targets.index(PERMUTE) + + +def test_propagate_up_moves_to_top_node_before_distinct_edge_nodes() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + y = graph.placeholder("y") + y.meta["val"] = torch.empty((1, 2, 3, 4)) + add = graph.call_function(ADD, args=(x, y)) + add.meta["val"] = torch.empty((1, 2, 3, 4)) + relu = graph.call_function(RELU, args=(add,)) + relu.meta["val"] = torch.empty((1, 2, 3, 4)) + permute = graph.call_function(PERMUTE, args=(relu, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output(permute) + + targets = _run_pass_on_graph(graph) + + assert targets.count(PERMUTE) == 1 + assert targets.index(ADD) < targets.index(PERMUTE) < targets.index(RELU) + + +def test_propagate_stops_rank_changing_view_at_slice() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3)) + slice_node = graph.call_function(SLICE, args=(x, 0, 0, 1)) + slice_node.meta["val"] = torch.empty((1, 2, 3)) + view = graph.call_function(VIEW, args=(slice_node, [2, 3])) + view.meta["val"] = torch.empty((2, 3)) + graph.output(view) + + targets = _run_pass_on_graph(graph) + + assert targets.index(SLICE) < targets.index(VIEW)