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)