From 91ffd27354e07fba07ec7738f065adf714f6892d Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Tue, 14 Jul 2026 14:54:32 -0700 Subject: [PATCH] Revert "Reland "Arm backend: Add transpose propagation pass" (#20748)" This reverts commit 18e2ff8d2664d8970f706bb6ada7676961797e3f. --- backends/arm/_passes/__init__.py | 4 - backends/arm/_passes/arm_pass_manager.py | 14 +- backends/arm/_passes/dim_maps.py | 269 +-- .../_passes/insert_data_layout_casts_pass.py | 7 - .../propagate_view_copy_permute_pass.py | 751 -------- .../arm/test/misc/test_transpose_counts.py | 24 +- backends/arm/test/passes/test_dim_maps.py | 47 +- .../test_insert_data_layout_casts_pass.py | 29 - .../test_propagate_permutes_views_pass.py | 1515 ----------------- 9 files changed, 28 insertions(+), 2632 deletions(-) delete mode 100644 backends/arm/_passes/propagate_view_copy_permute_pass.py delete mode 100644 backends/arm/test/passes/test_propagate_permutes_views_pass.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index b7d5efc6b50..4b853d95ee1 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -151,10 +151,6 @@ ) 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 bd1c5bb1aec..9d11212dec5 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -15,6 +15,7 @@ AccumulateIndexPutPass, BroadcastArgsPass, CanonicalizeGatherPass, + CanonicalizeViewCopyPermutePass, CastInt64BuffersToInt32Pass, CastToInt32Pass, ComputeConstantOpsAOTPass, @@ -132,12 +133,11 @@ NormalizeTransformInputPlaceholdersPass, NormalizeWhileInitialArgsPass, PromoteBoolOperandsPass, - PropagateViewCopyPermuteDownPass, - PropagateViewCopyPermuteUpPass, QuantizeClampArgumentsPass, RemoveGetItemPass, RemoveGraphAssertsPass, RemoveNoopPass, + RemovePermutesAroundElementwiseTosaOps, ReplaceInfAndLimitValuesPass, ReplaceScalarWithTensorByProfilePass, RewriteAdaptiveAvgPool2dPass, @@ -173,6 +173,9 @@ 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 @@ -603,7 +606,6 @@ def _tosa_pipeline( RewriteAvgPool2dPass(), ComputeConstantOpsAOTPass(exported_program), FuseConstantArgsPass(exported_program), - CastInt64BuffersToInt32Pass(exported_program), DecomposeSelectPass(), ConvertSqueezesToViewPass(), CastToInt32Pass(), @@ -627,14 +629,14 @@ def _tosa_pipeline( RewriteMatmulPass(), RewritePadPass(), FuseViewCopyTransformPass(), - PropagateViewCopyPermuteDownPass(self.compile_spec, exported_program), - PropagateViewCopyPermuteUpPass(self.compile_spec, exported_program), + RemovePermutesAroundElementwiseTosaOps(exported_program), + CanonicalizeViewCopyPermutePass(), + FuseCascadedTransposeOrPermuteOps(), RewriteHighRankSingletonPermutePass(), DecomposePermuteForU55Pass(), RewriteSlicePass(), FuseConsecutiveSlicesPass(), InsertConstShapesPass(), - InsertDataLayoutCastsPass(), ] ) diff --git a/backends/arm/_passes/dim_maps.py b/backends/arm/_passes/dim_maps.py index 6fc9b8ac1f9..07098035389 100644 --- a/backends/arm/_passes/dim_maps.py +++ b/backends/arm/_passes/dim_maps.py @@ -306,28 +306,15 @@ def map_dim( return None groups = self._valid_groups() - if not self._is_valid_reduction_or_singleton( - normalized_dims, groups.source_axis_to_groups - ): + if not self._is_valid_reduction(normalized_dims, groups.source_axis_to_groups): return None - 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] - ) + target_dims = self._map_dims( + normalized_dims, + groups.source_axis_to_groups, + groups.group_to_target_axes, ) - 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( + if not target_dims or not self._is_valid_reduction( target_dims, groups.target_axis_to_groups ): return None @@ -445,226 +432,6 @@ 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], @@ -786,30 +553,6 @@ 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 4e931396dab..07a2d186895 100644 --- a/backends/arm/_passes/insert_data_layout_casts_pass.py +++ b/backends/arm/_passes/insert_data_layout_casts_pass.py @@ -36,7 +36,6 @@ 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, @@ -45,12 +44,6 @@ 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 deleted file mode 100644 index e8f74edf9a8..00000000000 --- a/backends/arm/_passes/propagate_view_copy_permute_pass.py +++ /dev/null @@ -1,751 +0,0 @@ -# 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 frontier.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 - - def _would_strand_layout_op_on_wider_elements( - self, next_node: torch.fx.Node - ) -> bool: - """Whether crossing next_node leaves the layout op on wider elements. - - Crossing next_node upward moves the layout op onto next_node's input. When - next_node narrows the dtype (its input has wider elements than its output, - e.g. an int32 to int8 rescale), the layout op then has more bytes to move, so - block the crossing and keep it on the narrow side. The sole exception is a - placeholder input consumed only by next_node: moving the layout op onto such - a graph input folds it into the input's dim_order at no runtime cost. A - placeholder with other consumers cannot be relaid out for free, so the - crossing is still blocked. Valid for upward propagation only, where - next_node's single input is where the layout op would land. - - """ - input_nodes = next_node.all_input_nodes - if len(input_nodes) != 1: - return False - node_val = next_node.meta.get("val") - producer = input_nodes[0] - producer_val = producer.meta.get("val") - if not isinstance(node_val, torch.Tensor) or not isinstance( - producer_val, torch.Tensor - ): - return False - if producer_val.element_size() <= node_val.element_size(): - return False - return producer.op != "placeholder" or len(producer.users) != 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 - if any( - self._would_strand_layout_op_on_wider_elements(next_node) - for next_node in next_nodes - ): - 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 168dabe96b9..086edc537ba 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),), - 2, + 4, ), "transpose_conv": TransposeCountCase( TransposeConvModule(), @@ -413,7 +413,7 @@ def forward(self, x: torch.Tensor): "lstm": TransposeCountCase( LstmModule(), (torch.randn(2, 4, 8),), - 1, + 2, ), "groupnorm": TransposeCountCase( GroupNormModule(), @@ -428,7 +428,7 @@ def forward(self, x: torch.Tensor): "multihead_attention_rank3": TransposeCountCase( MultiheadAttentionModule(), (torch.randn(2, 4, 8),), - 6, + 7, ), "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),), 1 + Model1ConvMaxPoolResidualLinear(), (torch.randn(2, 8, 64),), 5 ), "model_2_conv_mha_linear_layernorm": TransposeCountCase( - Model2ConvMhaLinearLayerNorm(), (torch.randn(2, 8, 32),), 7 + Model2ConvMhaLinearLayerNorm(), (torch.randn(2, 8, 32),), 8 ), "model_3_lstm_linear": TransposeCountCase( - Model3LstmLinear(), (torch.randn(2, 16, 8),), 1 + Model3LstmLinear(), (torch.randn(2, 16, 8),), 2 ), "model_4_conv_lstm_linear_layernorm": TransposeCountCase( - Model4ConvLstmLinearLayerNorm(), (torch.randn(2, 8, 32),), 2 + Model4ConvLstmLinearLayerNorm(), (torch.randn(2, 8, 32),), 3 ), "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),), 1 + Model6GruLinear(), (torch.randn(2, 16, 8),), 2 ), "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),), 2 + Model8ConvBatchNormMaxPoolResidual(), (torch.randn(1, 8, 16, 16),), 4 ), "model_9_dilated_conv_batchnorm_avgpool_residual": TransposeCountCase( - Model9DilatedConvBatchNormAvgPoolResidual(), (torch.randn(1, 8, 16, 16),), 2 + Model9DilatedConvBatchNormAvgPoolResidual(), (torch.randn(1, 8, 16, 16),), 4 ), "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),), - 1, + 3, ), "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),), - 3, + 4, ), "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 16a18720442..e71c815c471 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) == [0] + assert view_map.map_dim(0) is None 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]) == [1, 2, 3] + assert view_map.map_dim([3, 1, 2]) == [2, 3, 1] assert view_map.map_dim_inverse(0) is None assert view_map.map_dim_inverse(1) == [1, 2] @@ -360,49 +360,6 @@ 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 bdacf5d27db..b4298977e5b 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,11 +39,6 @@ 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),) @@ -114,27 +109,3 @@ 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 deleted file mode 100644 index 7b140dea806..00000000000 --- a/backends/arm/test/passes/test_propagate_permutes_views_pass.py +++ /dev/null @@ -1,1515 +0,0 @@ -# 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 -CLONE = exir_ops.edge.dim_order_ops._clone_dim_order.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_up_stops_before_narrowing_rescale_fed_by_binary_op() -> None: - graph = torch.fx.Graph() - x = graph.placeholder("x") - x.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) - y = graph.placeholder("y") - y.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) - add = graph.call_function(ADD, args=(x, y)) - add.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) - rescale = graph.call_function(RESCALE, args=(add, 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(RESCALE) < targets.index(PERMUTE) - - -def test_propagate_up_crosses_same_width_rescale_fed_by_binary_op() -> None: - graph = torch.fx.Graph() - x = graph.placeholder("x") - x.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) - y = graph.placeholder("y") - y.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) - add = graph.call_function(ADD, args=(x, y)) - add.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) - rescale = graph.call_function(RESCALE, args=(add, torch.int32, [1.0], 0, 0)) - rescale.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) - permute = graph.call_function(PERMUTE, args=(rescale, [0, 2, 3, 1])) - permute.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.int32) - 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_up_stops_before_narrowing_rescale_from_shared_placeholder() -> 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) - other = graph.call_function(NEG, args=(x,)) - other.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) - 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, other)) - - 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_up_crosses_widening_rescale_fed_by_binary_op() -> None: - graph = torch.fx.Graph() - x = graph.placeholder("x") - x.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int8) - y = graph.placeholder("y") - y.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int8) - add = graph.call_function(ADD, args=(x, y)) - add.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int8) - rescale = graph.call_function(RESCALE, args=(add, torch.int32, [1.0], 0, 0)) - rescale.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) - permute = graph.call_function(PERMUTE, args=(rescale, [0, 2, 3, 1])) - permute.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.int32) - 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_up_stops_before_narrowing_rescale_behind_unary() -> None: - graph = torch.fx.Graph() - x = graph.placeholder("x") - x.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) - clone = graph.call_function(CLONE, args=(x,)) - clone.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) - rescale = graph.call_function(RESCALE, args=(clone, 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(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)