From 95267042ffe29965b4d7dff3974e216e870f691d Mon Sep 17 00:00:00 2001 From: Adrian Lundell Date: Thu, 25 Jun 2026 15:17:17 +0200 Subject: [PATCH] Arm backend: Support mixed precision permute fusing Adds transpose counting to pixel-shuffle tests including a mixed precision case with decomposed Q/DQ nodes which the previous permute propagation did not handle. This is fixed by adding support for propagating views and permutes over across forks in the graph, as long as it can be proven that the branches will eventually converge. Also adds support for where to fuse_identiacal_input_transforms Signed-off-by: Adrian Lundell Change-Id: If66ff26d8020cdeab9ea88bc87a687a63a57dda0 --- backends/arm/_passes/__init__.py | 3 + backends/arm/_passes/arm_pass_manager.py | 2 + backends/arm/_passes/arm_pass_utils.py | 60 +-- .../canonicalize_view_copy_permute_pass.py | 4 +- backends/arm/_passes/decompose_var_pass.py | 2 +- backends/arm/_passes/dim_maps.py | 71 ++++ .../fuse_identical_input_transforms_pass.py | 128 ++++-- backends/arm/_passes/match_arg_ranks_pass.py | 17 + ...data_movement_ops_to_smaller_dtype_pass.py | 158 ++++++++ .../propagate_view_copy_permute_pass.py | 377 +++++++++++++----- .../arm/test/misc/test_transpose_counts.py | 2 +- backends/arm/test/ops/test_pixel_shuffling.py | 171 ++++++-- .../arm/test/passes/test_arm_pass_utils.py | 45 +++ backends/arm/test/passes/test_dim_maps.py | 52 ++- ...st_fuse_identical_input_transforms_pass.py | 51 ++- .../test/passes/test_match_arg_ranks_pass.py | 49 +++ ...data_movement_ops_to_smaller_dtype_pass.py | 113 ++++++ .../test_propagate_permutes_views_pass.py | 221 +++++++++- 18 files changed, 1273 insertions(+), 253 deletions(-) create mode 100644 backends/arm/_passes/move_data_movement_ops_to_smaller_dtype_pass.py create mode 100644 backends/arm/test/passes/test_arm_pass_utils.py create mode 100644 backends/arm/test/passes/test_match_arg_ranks_pass.py create mode 100644 backends/arm/test/passes/test_move_data_movement_ops_to_smaller_dtype_pass.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index 2d4ff8b2baf..613ebb1b421 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -146,6 +146,9 @@ from .match_arg_dtype_pass import MatchArgDtypePass # noqa from .match_arg_ranks_pass import MatchArgRanksPass # noqa from .mm_to_bmm_pass import ConvertMmToBmmPass # noqa +from .move_data_movement_ops_to_smaller_dtype_pass import ( # noqa + MoveDataMovementOpsToSmallerDtypePass, +) from .normalize_delegate_io_layout_pass import NormalizeDelegateIOLayoutPass # noqa from .normalize_index_put_bool_index_tensor_pass import ( # noqa NormalizeIndexPutBoolIndexTensorPass, diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index ebedc610864..49c53981c47 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -128,6 +128,7 @@ InsertTableOpsPass, MatchArgDtypePass, MatchArgRanksPass, + MoveDataMovementOpsToSmallerDtypePass, NormalizeDelegateIOLayoutPass, NormalizeIndexPutBoolIndexTensorPass, NormalizeIndexPutNoneIndicesPass, @@ -635,6 +636,7 @@ def _tosa_pipeline( PropagateViewCopyPermuteUpPass(self.compile_spec, exported_program), # Propagation can leave a binary op with mismatched operand ranks, # which TOSA rejects; re-match ranks before lowering. + MoveDataMovementOpsToSmallerDtypePass(), MatchArgRanksPass(exported_program), RewriteHighRankSingletonPermutePass(), DecomposePermuteForU55Pass(), diff --git a/backends/arm/_passes/arm_pass_utils.py b/backends/arm/_passes/arm_pass_utils.py index 85d4ed514ea..527fa42bf0e 100644 --- a/backends/arm/_passes/arm_pass_utils.py +++ b/backends/arm/_passes/arm_pass_utils.py @@ -9,15 +9,11 @@ import operator import traceback from inspect import isclass -from typing import cast, Optional, Sequence +from typing import Any, cast, Optional, Sequence import torch import torch.fx -from executorch.backends.arm._passes.dim_maps import ( - _normalize_dims, - normalize_view_shape, -) from executorch.backends.arm.common.debug import get_node_debug_info from executorch.backends.arm.common.type import ensure_type from executorch.backends.arm.tosa.mapping import TosaSpecialDtype @@ -36,8 +32,7 @@ from torch._ops import OpOverload from torch._subclasses.fake_tensor import FakeTensor from torch.export.graph_signature import InputKind - -_Dim = int | torch.SymInt +from torch.fx.node import map_arg def is_submodule_node(node: torch.fx.Node): @@ -252,42 +247,31 @@ def meta_without_qparams(meta: NodeMetadata) -> NodeMetadata: return NodeMetadata(plain_meta_dict) -def refresh_permute_view_meta(node: torch.fx.Node) -> None: - """Compute new meta-vals, specifically preserving SymInts for view/permute - nodes. +def refresh_node_meta(node: torch.fx.Node) -> None: + """Best-effort refresh of a call_function node's output value metadata. + + Node arguments are replaced by their metadata values and evaluated through + the operator's FakeTensor/meta implementation, which preserves symbolic + dimensions when the operator supports them. Missing input metadata or + unsupported meta evaluation leaves the existing output metadata unchanged. + """ - input_node = node.all_input_nodes[0] - input_val = input_node.meta.get("val") - if input_val is None or node.target not in { - exir_ops.edge.aten.view_copy.default, - exir_ops.edge.aten.permute_copy.default, - }: + if node.op != "call_function" or "val" not in node.meta: return - if not isinstance(input_val, torch.Tensor): - node.meta["val"] = node.target(input_val, *node.args[1:]) # type: ignore[operator] + try: + args = map_arg( + node.args, + lambda arg: arg.meta["val"] if isinstance(arg, torch.fx.Node) else arg, + ) + kwargs = map_arg( + node.kwargs, + lambda arg: arg.meta["val"] if isinstance(arg, torch.fx.Node) else arg, + ) + node.meta["val"] = cast(Any, node.target)(*args, **kwargs) + except Exception: return - # Compute new meta shapes to preserve SymInts. - match node.target: - case exir_ops.edge.aten.view_copy.default: - node.meta["val"] = input_val.new_empty( - tuple( - normalize_view_shape( - input_val.shape, cast(Sequence[_Dim], node.args[1]) - ) - ) - ) - case exir_ops.edge.aten.permute_copy.default: - dims = _normalize_dims( - cast(Sequence[int], node.args[1]), len(input_val.shape) - ) - node.meta["val"] = input_val.new_empty( - tuple(input_val.shape[dim] for dim in dims) - ) - case _: - node.meta["val"] = node.target(input_val, *node.args[1:]) # type: ignore[operator] - def insert_scalar( graph: torch.fx.Graph, diff --git a/backends/arm/_passes/canonicalize_view_copy_permute_pass.py b/backends/arm/_passes/canonicalize_view_copy_permute_pass.py index 1930ba0b5f9..985c1022ef9 100644 --- a/backends/arm/_passes/canonicalize_view_copy_permute_pass.py +++ b/backends/arm/_passes/canonicalize_view_copy_permute_pass.py @@ -10,7 +10,7 @@ import torch from executorch.backends.arm._passes.arm_pass import ArmPass -from executorch.backends.arm._passes.arm_pass_utils import refresh_permute_view_meta +from executorch.backends.arm._passes.arm_pass_utils import refresh_node_meta from executorch.backends.arm._passes.dim_maps import ( _dim_equals, _is_permutation, @@ -361,7 +361,7 @@ def _set_node_op( ) -> None: node.target = target node.args = (input_node, list(arg)) - refresh_permute_view_meta(node) + refresh_node_meta(node) def _permute_dims(self, node: Node) -> list[int]: assert node.target == self._PERMUTE_TARGET, "Expected permute node" diff --git a/backends/arm/_passes/decompose_var_pass.py b/backends/arm/_passes/decompose_var_pass.py index 90ea80b6b47..8c25f902e1d 100644 --- a/backends/arm/_passes/decompose_var_pass.py +++ b/backends/arm/_passes/decompose_var_pass.py @@ -74,7 +74,7 @@ def call_operator(self, op, args, kwargs, meta): shape = [1 for _ in input_shape] # Get dim from args based on argument type - dim = get_node_arg(args, key=list, default_value=list(range(len(shape)))) + dim = get_node_arg(args, key=list, default_value=list(range(len(input_shape)))) if op == torch.ops.aten.var.dim: keepdim = False diff --git a/backends/arm/_passes/dim_maps.py b/backends/arm/_passes/dim_maps.py index 6fc9b8ac1f9..0e3efbef45a 100644 --- a/backends/arm/_passes/dim_maps.py +++ b/backends/arm/_passes/dim_maps.py @@ -363,6 +363,53 @@ def map_dim_inverse( return None return source_dims + def map_reduction_after_view( + self, + source_shape: Sequence[_Dim], + source_dims: int | Sequence[int], + ) -> tuple[list[_Dim], list[int]] | None: + """Map ``reduce(view(x), dims)`` to ``view(reduce(x, mapped_dims))``. + + Returns the new view shape and reduction dims for: + + view(reduce(x, source_dims), self.target_shape) + == reduce(view(x, new_shape), target_dims) + + """ + target_shape = self.remap_target_shape(source_shape) + if target_shape is None: + return None + + target_dims = self.map_dim(source_dims) + if target_dims is None or not self._is_contiguous_nonempty(target_dims): + return None + return target_shape, target_dims + + def map_reduction_before_view( + self, + target_dims: int | Sequence[int], + ) -> tuple[list[int], list[_Dim]] | None: + """Map ``view(reduce(x, dims))`` to ``reduce(view(x), mapped_dims)``. + + Returns the reduction dims and output view shape for: + + reduce(view(x, self.target_shape), target_dims) + == view(reduce(x, source_dims), output_shape) + + """ + source_dims = self.map_dim_inverse(target_dims) + if source_dims is None or not self._is_contiguous_nonempty(source_dims): + return None + + try: + normalized_target_dims = _normalize_dims(target_dims, self.target_rank) + except AssertionError: + return None + + return source_dims, self._reduce_shape( + self.target_shape, normalized_target_dims + ) + def map_permutation( self, source_permutation: Sequence[int], @@ -446,6 +493,8 @@ def map_permutation_inverse( ) def remap_target_shape(self, source_shape: Sequence[_Dim]) -> list[_Dim] | None: + if not self.is_valid_map: + return None if len(source_shape) != self.source_rank: return None @@ -470,6 +519,8 @@ def remap_target_shape(self, source_shape: Sequence[_Dim]) -> list[_Dim] | None: if not same_numel(source_shape, target_shape): return None + if self._has_zero_dim(target_shape): + return None if not self._preserves_source_axis_order(source_shape, source_to_target_axes): return None return target_shape @@ -551,6 +602,8 @@ def remap_unit_slice( for target_axes in source_to_target_axes[:slice_dim] for target_axis in target_axes ] + if not prev_target_axes: + return None next_target_axes = [ target_axis for target_axes in source_to_target_axes[slice_dim + 1 :] @@ -810,6 +863,24 @@ def _is_valid_reduction_or_singleton( group_to_axes[group].issubset(normalized_dims) for group in selected_groups ) + @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) + ) + + @staticmethod + def _reduce_shape(shape: Sequence[_Dim], dims: Sequence[int]) -> list[_Dim]: + reduced_shape = list(shape) + for dim in dims: + reduced_shape[dim] = 1 + return reduced_shape + + @staticmethod + def _has_zero_dim(shape: Sequence[_Dim]) -> bool: + return any(_dim_equals(dim, 0) for dim in shape) + @classmethod def _build_groups( cls, source_shape: Sequence[_Dim], target_shape: Sequence[_Dim] diff --git a/backends/arm/_passes/fuse_identical_input_transforms_pass.py b/backends/arm/_passes/fuse_identical_input_transforms_pass.py index e46624301d8..f614048704f 100644 --- a/backends/arm/_passes/fuse_identical_input_transforms_pass.py +++ b/backends/arm/_passes/fuse_identical_input_transforms_pass.py @@ -10,8 +10,13 @@ import torch from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass -from executorch.backends.arm._passes.arm_pass_utils import refresh_permute_view_meta -from executorch.backends.arm._passes.dim_maps import PermuteMap, same_numel, ViewMap +from executorch.backends.arm._passes.arm_pass_utils import refresh_node_meta +from executorch.backends.arm._passes.dim_maps import ( + _dim_equals, + PermuteMap, + same_numel, + ViewMap, +) from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult from torch.export.exported_program import ExportedProgram @@ -142,8 +147,12 @@ class FuseIdenticalInputTransformsPass(ArmOpTargetedPass): exir_ops.edge.aten.bitwise_xor.Tensor, exir_ops.edge.aten.remainder.Tensor, } + _NARY_ELEMENTWISE_OPS = { + exir_ops.edge.aten.where.self, + } + _ELEMENTWISE_OPS = _BINARY_ELEMENTWISE_OPS | _NARY_ELEMENTWISE_OPS - target_ops = _BINARY_ELEMENTWISE_OPS | _CONCAT_OPS + target_ops = _ELEMENTWISE_OPS | _CONCAT_OPS def __init__(self, exported_program: ExportedProgram | None = None) -> None: super().__init__() @@ -180,31 +189,43 @@ def _sink_identical_input_transforms(self, node: Node) -> bool: if node.target not in self.target_ops: return False - input_transforms = self._input_transforms(node) - if input_transforms is None: + input_nodes = list(node.all_input_nodes) + if len(input_nodes) < 2: return False node_val = node.meta.get("val", None) if node_val is None: return False - transform = input_transforms[0] - updated_args = self._updated_node_args( - node, transform, node_val, input_transforms - ) + transforms = [n for n in input_nodes if n.target in self._TARGETS] + if not transforms: + return False + transform = transforms[0] + if not self._inputs_share_transform_or_are_layout_invariant( + node, transform, input_nodes + ): + return False + + updated_args = self._updated_node_args(node, transform, node_val, input_nodes) if updated_args is None: return False node_args, node_kwargs, transform_args, node_output_shape = updated_args # Remove input transforms - producers = [n.all_input_nodes[0] for n in input_transforms] + producers = [ + n.all_input_nodes[0] if n.target in self._TARGETS else n + for n in input_nodes + ] node.args = node_args node.kwargs = node_kwargs - for input_transform, producer in zip(input_transforms, producers): + for input_transform, producer in zip(input_nodes, producers): node.replace_input_with(input_transform, producer) - for input_transform in dict.fromkeys(input_transforms): - if len(input_transform.users) == 0: + for input_transform in dict.fromkeys(input_nodes): + if ( + input_transform.target in self._TARGETS + and len(input_transform.users) == 0 + ): node.graph.erase_node(input_transform) node.meta = copy.copy(node.meta) @@ -218,7 +239,7 @@ def _sink_identical_input_transforms(self, node: Node) -> bool: kwargs=dict(transform.kwargs), ) new_node.meta = self._new_transform_meta(node, transform) - refresh_permute_view_meta(new_node) + refresh_node_meta(new_node) for user in list(node.users): if user is not new_node: @@ -235,27 +256,52 @@ def _new_transform_meta(self, node: Node, transform: Node) -> dict[str, Any]: return meta def _updated_node_args( - self, node: Node, transform: Node, node_val: Any, input_transforms: list[Node] + self, node: Node, transform: Node, node_val: Any, input_nodes: list[Node] ) -> ( tuple[tuple[Any, ...], dict[str, Any], tuple[Any, ...], tuple[Any, ...]] | None ): - if not self._transforms_are_identical(input_transforms): - return None - if not self._transforms_only_used_by_node(node, input_transforms): - return None - if node.target in self._BINARY_ELEMENTWISE_OPS: - return self._update_node_args_binary( - node, transform, node_val, input_transforms - ) + return self._update_node_args_binary(node, transform, node_val, input_nodes) if node.target in self._CONCAT_OPS: - return self._update_node_args_concat( - node, transform, node_val, input_transforms - ) + return self._update_node_args_concat(node, transform, node_val, input_nodes) + + if node.target in self._NARY_ELEMENTWISE_OPS: + return self._update_node_args_binary(node, transform, node_val, input_nodes) return None + def _inputs_share_transform_or_are_layout_invariant( + self, node: Node, transform: Node, input_nodes: list[Node] + ) -> bool: + transforms = [n for n in input_nodes if n.target in self._TARGETS] + if not self._transforms_are_identical(transforms): + return False + if not self._transforms_only_used_by_node(node, transforms): + return False + if len(transforms) == len(input_nodes): + return True + if node.target not in self._ELEMENTWISE_OPS: + return False + + transform_val = transform.meta.get("val") + if not isinstance(transform_val, torch.Tensor): + return False + rank = len(transform_val.shape) + return all( + input_node in transforms or self.is_layout_invariant(input_node, rank) + for input_node in input_nodes + ) + + @staticmethod + def is_layout_invariant(node: Node, rank: int) -> bool: + value = node.meta.get("val") + return ( + isinstance(value, torch.Tensor) + and len(value.shape) == rank + and all(_dim_equals(dim, 1) for dim in value.shape) + ) + def _transforms_are_identical(self, input_transforms: list[Node]) -> bool: target = input_transforms[0].target if target not in self._TARGETS: @@ -281,7 +327,15 @@ def _transforms_only_used_by_node( def _update_node_args_binary(self, node, transform, node_val, input_transforms): producer_shapes = [ - tuple(input_node.all_input_nodes[0].meta["val"].shape) + tuple( + ( + input_node.all_input_nodes[0] + if input_node.target in self._TARGETS + else input_node + ) + .meta["val"] + .shape + ) for input_node in input_transforms ] @@ -293,6 +347,9 @@ def _update_node_args_binary(self, node, transform, node_val, input_transforms): transform_args = (node, *transform.args[1:]) if transform.target == self._VIEW_TARGET: transform_args = (node, list(node_val.shape)) + # Reshaping before an elementwise op can change which dimensions + # broadcast. Sinking is safe only when the broadcast in the source + # layout already has exactly one of the producer shapes. if node_output_shape not in producer_shapes: return None @@ -382,20 +439,3 @@ def _mapped_concat_dim(self, transform: Node, concat_dim: int) -> int | None: if mapped_dims is None or len(mapped_dims) != 1: return None return mapped_dims[0] - - def _input_transforms(self, node: Node) -> list[Node] | None: - if node.target in self._BINARY_ELEMENTWISE_OPS: - input_transforms = list(node.args[:2]) - elif node.target in self._CONCAT_OPS: - if len(node.args) == 0 or not isinstance(node.args[0], Sequence): - return None - input_transforms = list(node.args[0]) - else: - return None - - if len(input_transforms) < 2 or not all( - isinstance(n, Node) for n in input_transforms - ): - return None - - return cast(list[Node], input_transforms) diff --git a/backends/arm/_passes/match_arg_ranks_pass.py b/backends/arm/_passes/match_arg_ranks_pass.py index 5e4463d5ee6..e2fe1135d1a 100644 --- a/backends/arm/_passes/match_arg_ranks_pass.py +++ b/backends/arm/_passes/match_arg_ranks_pass.py @@ -66,6 +66,23 @@ def __init__(self, exported_program: ExportedProgram, *args, **kwargs) -> None: exir_ops.edge.aten.bitwise_or.Tensor, exir_ops.edge.aten.maximum.default, exir_ops.edge.aten.minimum.default, + exir_ops.backend.tosa.ADD.default, + exir_ops.backend.tosa.ARITHMETIC_RIGHT_SHIFT.default, + exir_ops.backend.tosa.BITWISE_AND.default, + exir_ops.backend.tosa.BITWISE_OR.default, + exir_ops.backend.tosa.BITWISE_XOR.default, + exir_ops.backend.tosa.EQUAL.default, + exir_ops.backend.tosa.GREATER.default, + exir_ops.backend.tosa.GREATER_EQUAL.default, + exir_ops.backend.tosa.LOGICAL_AND.default, + exir_ops.backend.tosa.LOGICAL_LEFT_SHIFT.default, + exir_ops.backend.tosa.LOGICAL_OR.default, + exir_ops.backend.tosa.LOGICAL_XOR.default, + exir_ops.backend.tosa.MAXIMUM.default, + exir_ops.backend.tosa.MINIMUM.default, + exir_ops.backend.tosa.MUL.default, + exir_ops.backend.tosa.POW.default, + exir_ops.backend.tosa.SUB.default, ] def _match_op_rank(self, graph_module, node, arg, max_rank): diff --git a/backends/arm/_passes/move_data_movement_ops_to_smaller_dtype_pass.py b/backends/arm/_passes/move_data_movement_ops_to_smaller_dtype_pass.py new file mode 100644 index 00000000000..7ae4ad27455 --- /dev/null +++ b/backends/arm/_passes/move_data_movement_ops_to_smaller_dtype_pass.py @@ -0,0 +1,158 @@ +# 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 Sequence +from typing import Any, cast, Set, Type + +import torch +from executorch.backends.arm._passes.arm_pass_utils import refresh_node_meta +from executorch.backends.arm.tosa.mapping import TosaSpecialDtype +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + +from .arm_pass import ArmPass + + +class MoveDataMovementOpsToSmallerDtypePass(ArmPass): + """Move layout operations next to rescales onto the smaller element type.""" + + _RESCALE = exir_ops.backend.tosa.RESCALE.default + _passes_required_after: Set[Type[ExportPass]] = set() + _DATA_MOVEMENT_OPS = { + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.dim_order_ops._clone_dim_order.default, + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, + } + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + modified = False + iteration_modified = True + while iteration_modified: + iteration_modified = False + for node in list(graph_module.graph.nodes): + if node.target not in self._DATA_MOVEMENT_OPS: + continue + input_nodes = node.all_input_nodes + if len(input_nodes) != 1: + continue + + producer = input_nodes[0] + if producer.target == self._RESCALE and self._uses_wider_dtype( + node, producer, producer.all_input_nodes[0] + ): + self._move_before_rescale(node, producer) + iteration_modified = True + break + + narrowing_rescale = next( + ( + user + for user in node.users + if user.target == self._RESCALE + and self._uses_wider_dtype(node, user, user) + ), + None, + ) + if narrowing_rescale is not None: + if len(node.users) == 1: + self._move_after_rescale(node, narrowing_rescale) + else: + self._split_after_rescale(node, narrowing_rescale) + iteration_modified = True + break + modified |= iteration_modified + + if modified: + graph_module.graph.lint() + graph_module = super().call(graph_module).graph_module + graph_module.recompile() + return PassResult(graph_module, modified) + + def _uses_wider_dtype( + self, + data_movement: torch.fx.Node, + rescale: torch.fx.Node, + smaller_side: torch.fx.Node, + ) -> bool: + if not self._is_per_tensor_rescale(rescale): + return False + movement_val = data_movement.meta.get("val") + smaller_val = smaller_side.meta.get("val") + return ( + isinstance(movement_val, torch.Tensor) + and isinstance(smaller_val, torch.Tensor) + and movement_val.element_size() > smaller_val.element_size() + ) + + def _is_per_tensor_rescale(self, rescale: torch.fx.Node) -> bool: + if len(rescale.args) < 3 or len(rescale.all_input_nodes) != 1: + return False + special_dtype_key = TosaSpecialDtype.meta_key() + if rescale.all_input_nodes[0].meta.get(special_dtype_key) != rescale.meta.get( + special_dtype_key + ): + return False + scales = rescale.args[2] + return not isinstance(scales, Sequence) or len(scales) == 1 + + def _move_before_rescale( + self, data_movement: torch.fx.Node, rescale: torch.fx.Node + ) -> None: + producer = rescale.all_input_nodes[0] + movement_users = list(data_movement.users) + data_movement.replace_input_with(rescale, producer) + rescale.replace_input_with(producer, data_movement) + for user in movement_users: + user.replace_input_with(data_movement, rescale) + data_movement.append(rescale) + + old_movement_val = cast(torch.Tensor, data_movement.meta["val"]) + producer_val = cast(torch.Tensor, producer.meta["val"]) + data_movement.meta = dict(data_movement.meta) + data_movement.meta["val"] = old_movement_val.new_empty( + old_movement_val.shape, dtype=producer_val.dtype + ) + rescale.meta = dict(rescale.meta) + rescale.meta["val"] = old_movement_val + + def _move_after_rescale( + self, data_movement: torch.fx.Node, rescale: torch.fx.Node + ) -> None: + producer = data_movement.all_input_nodes[0] + rescale_users = list(rescale.users) + rescale.replace_input_with(data_movement, producer) + data_movement.replace_input_with(producer, rescale) + for user in rescale_users: + user.replace_input_with(rescale, data_movement) + rescale.append(data_movement) + + old_rescale_val = cast(torch.Tensor, rescale.meta["val"]) + producer_val = cast(torch.Tensor, producer.meta["val"]) + rescale.meta = dict(rescale.meta) + rescale.meta["val"] = producer_val.new_empty( + producer_val.shape, dtype=old_rescale_val.dtype + ) + refresh_node_meta(data_movement) + + def _split_after_rescale( + self, data_movement: torch.fx.Node, rescale: torch.fx.Node + ) -> None: + producer = data_movement.all_input_nodes[0] + rescale_users = list(rescale.users) + rescale.replace_input_with(data_movement, producer) + + with rescale.graph.inserting_after(rescale): + branch_movement = rescale.graph.call_function( + cast(Any, data_movement.target), + args=(rescale, *data_movement.args[1:]), + kwargs=dict(data_movement.kwargs), + ) + branch_movement.meta = dict(data_movement.meta) + refresh_node_meta(branch_movement) + for user in rescale_users: + user.replace_input_with(rescale, branch_movement) + + if not data_movement.users: + data_movement.graph.erase_node(data_movement) diff --git a/backends/arm/_passes/propagate_view_copy_permute_pass.py b/backends/arm/_passes/propagate_view_copy_permute_pass.py index e8f74edf9a8..982cd2fc143 100644 --- a/backends/arm/_passes/propagate_view_copy_permute_pass.py +++ b/backends/arm/_passes/propagate_view_copy_permute_pass.py @@ -5,12 +5,14 @@ # pyre-unsafe +import copy from abc import ABC, abstractmethod from collections.abc import Iterable, Sequence +from dataclasses import dataclass 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.arm_pass_utils import refresh_node_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 @@ -29,6 +31,13 @@ _Dim = int | torch.SymInt +@dataclass(frozen=True) +class _ForkBranchSplit: + next_node: torch.fx.Node + source_shape: tuple[_Dim, ...] + arg_update: tuple[Any, Any] | None + + class PropagateViewCopyPermutePass(ArmPass, ABC): """Abstract implementation of a permute/view_copy propagation pass. @@ -159,12 +168,14 @@ def _propagate(self, node: torch.fx.Node) -> bool: break if len(next_nodes) > 1: - if self._maybe_split_downwards_slice_fanout(node, next_nodes): + if self._maybe_split_fork( + node, frontier, previous_frontier, next_nodes + ): return True break next_node = next_nodes[0] - if self.is_elementwise(next_node) and self._is_unary_elementwise(next_node): + if self._can_move_through_elementwise(node, frontier, next_node): previous_frontier = frontier frontier = next_node moved = True @@ -195,7 +206,8 @@ def _propagate(self, node: torch.fx.Node) -> bool: assert previous_frontier is not None self._move_node(node, frontier, previous_frontier) - refresh_permute_view_meta(node) + refresh_node_meta(frontier) + refresh_node_meta(node) return True def fuse_vertical(self, graph_module: torch.fx.GraphModule) -> PassResult: @@ -254,12 +266,15 @@ def _maybe_split_upwards_cat_fanout( """ return False - def _maybe_split_downwards_slice_fanout( - self, node: torch.fx.Node, next_nodes: Sequence[torch.fx.Node] + def _maybe_split_fork( + self, + node: torch.fx.Node, + frontier: torch.fx.Node, + previous_frontier: torch.fx.Node | None, + 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. + """Optionally split a transform over a fork in the propagation + direction. """ return False @@ -334,46 +349,48 @@ def is_swappable(self, next_node: torch.fx.Node) -> bool: ) 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 + def _can_move_through_elementwise( + self, + moving_node: torch.fx.Node, + frontier: torch.fx.Node, + 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. + """Return whether ``moving_node`` can cross an elementwise operation. + + Only simple single-input, single-output operations are handled here. + Multi-input elementwise operations are handled by horizontal fusion. """ - 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 + if ( + not self.is_elementwise(next_node) + or ( + len(next_node.all_input_nodes) != 1 + and next_node.target != exir_ops.backend.tosa.TABLE.default + ) + or not isinstance(next_node.meta.get("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) + if moving_node.target == self._VIEW_TARGET: + view_map = ViewMap(moving_node) + if not view_map.is_valid_map: + return False + if ( + next_node.target in self._TRANSPARENT_TARGETS + or next_node.target == exir_ops.backend.tosa.RESCALE.default + ) and view_map.source_rank != view_map.target_rank: + return False + + if moving_node.target != self._PERMUTE_TARGET: + return True + + dims = self._dim_arg(moving_node.args[1]) + source_rank = len(dims) if isinstance(dims, Sequence) else None + next_val = next_node.meta.get("val") + return ( + source_rank is not None + and isinstance(next_val, torch.Tensor) + and len(next_val.shape) == source_rank ) @@ -408,16 +425,37 @@ def _can_cross_next_nodes( 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 _can_move_through_elementwise( + self, + moving_node: torch.fx.Node, + frontier: torch.fx.Node, + next_node: torch.fx.Node, + ) -> bool: + if super()._can_move_through_elementwise(moving_node, frontier, next_node): + return True + if moving_node.target != self._PERMUTE_TARGET or not self.is_elementwise( + next_node + ): + return False + + frontier_val = frontier.meta.get("val") + if not isinstance(frontier_val, torch.Tensor): + return False + rank = len(frontier_val.shape) + layout_dependent_inputs = [ + input_node + for input_node in next_node.all_input_nodes + if not FuseIdenticalInputTransformsPass.is_layout_invariant( + input_node, rank + ) + ] + return len(layout_dependent_inputs) == 1 + def _maybe_swap_permute_args( self, node: torch.fx.Node, next_node: torch.fx.Node ) -> Any | None: @@ -443,9 +481,7 @@ def _maybe_swap_view_args( 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 - ) + return self._maybe_swap_reduction_view_args(node, next_node, view_map) 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 @@ -457,16 +493,15 @@ def _maybe_swap_reduction_view_args( 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): + input_val = next_node.all_input_nodes[0].meta["val"] + swap = view_map.map_reduction_after_view(input_val.shape, reduction_dims) + if swap is None: return None + new_shape, new_dims = swap new_next_node_args = (*next_node.args[:1], new_dims, *next_node.args[2:]) return ((*node.args[:1], new_shape), new_next_node_args) @@ -477,47 +512,46 @@ def _maybe_swap_slice_view_args( 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( + unit_slice_swap = 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: + if unit_slice_swap is not None: + new_shape, new_dim, new_start, new_end = unit_slice_swap + if not self._valid_slice_interval(new_shape, new_dim, new_start, new_end): + return None + 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) + + if new_shape 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:], - ) + slice_dim = cast(int, next_node.args[1]) + mapped_dim = self._map_slice_dim(view_map, slice_dim) + if mapped_dim is None: + return None + if not self._valid_slice_interval( + new_shape, + mapped_dim, + cast(_Dim, next_node.args[2]), + cast(_Dim, next_node.args[3]), + ): + return None + new_next_node_args = (*next_node.args[:1], mapped_dim, *next_node.args[2:]) return ((*node.args[:1], new_shape), new_next_node_args) @staticmethod @@ -538,6 +572,16 @@ def _map_slice_dim(view_map: ViewMap, slice_dim: int) -> int | None: return None return new_dim + @staticmethod + def _valid_slice_interval( + shape: Sequence[_Dim], dim: int, start: _Dim, end: _Dim + ) -> bool: + try: + dim = dim % len(shape) + return 0 <= start < end <= shape[dim] + except (RuntimeError, TypeError): + return False + def _move_node( self, node: torch.fx.Node, @@ -674,9 +718,10 @@ def _maybe_swap_view_args(self, node, next_node): 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: + swap = view_map.map_reduction_before_view(next_node.args[1]) + if swap is None: return None + new_dims, output_shape = swap 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: @@ -684,44 +729,156 @@ def _maybe_swap_view_args(self, node, next_node): if len(new_dims) != 1: return None new_dims = new_dims[0] + output_shape = list(next_node.meta["val"].shape) 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) + return ((*node.args[:1], output_shape), new_next_node_args) - def _maybe_split_downwards_slice_fanout( - self, node: torch.fx.Node, next_nodes: Sequence[torch.fx.Node] + def _maybe_split_fork( + self, + node: torch.fx.Node, + frontier: torch.fx.Node, + previous_frontier: torch.fx.Node | None, + next_nodes: Sequence[torch.fx.Node], ) -> bool: - """Duplicate a permute onto each slice branch. + if frontier is not node and previous_frontier is None: + return False + plan = self._plan_fork_split(node, frontier, next_nodes) + if plan is None: + return False + producer, branch_splits = plan + self._apply_fork_split(node, frontier, producer, branch_splits) + return True + + def _plan_fork_split( + self, + node: torch.fx.Node, + frontier: torch.fx.Node, + next_nodes: Sequence[torch.fx.Node], + ) -> tuple[torch.fx.Node, tuple[_ForkBranchSplit, ...]] | None: + """Validate every fork branch and return an all-or-nothing rewrite plan. + + Fork propagation is only implemented for permutes. The permutation must be + valid, every branch must support the same propagation step, and every branch + output must retain the permutation rank. Planning all branches before editing + is essential: discovering one unsupported branch after mutating earlier ones + would leave a partially rewritten graph. - The duplicated permutes are left before the slices; later propagation - iterations handle swapping each one through its slice. + Each planned branch records both its pre-permute shape and any argument update + needed when crossing reductions or slices. """ - if node.target != self._PERMUTE_TARGET: + if node.target != self._PERMUTE_TARGET or len(node.all_input_nodes) != 1: + return None + producer = node.all_input_nodes[0] + if not isinstance(producer.meta.get("val"), torch.Tensor): + return None + permute_dims = self._dim_arg(node.args[1]) + if not isinstance(permute_dims, Sequence): + return None + rank = len(permute_dims) + normalized_dims = [dim if dim >= 0 else dim + rank for dim in permute_dims] + if sorted(normalized_dims) != list(range(rank)): + return None + + branch_splits = [] + for next_node in next_nodes: + if self._can_move_through_elementwise( + node, frontier, next_node + ) or self._can_split_through_elementwise(node, frontier, next_node): + arg_update = None + elif self.is_swappable(next_node): + arg_update = self._maybe_swap_args(node, next_node) + if arg_update is None: + return None + else: + return None + + next_val = next_node.meta.get("val") + if not isinstance(next_val, torch.Tensor) or len(next_val.shape) != rank: + return None + source_shape: list[_Dim] = [1] * rank + for output_axis, source_axis in enumerate(normalized_dims): + source_shape[source_axis] = next_val.shape[output_axis] + branch_splits.append( + _ForkBranchSplit(next_node, tuple(source_shape), arg_update) + ) + return producer, tuple(branch_splits) + + def _can_split_through_elementwise( + self, + moving_node: torch.fx.Node, + frontier: torch.fx.Node, + next_node: torch.fx.Node, + ) -> bool: + if not self.is_elementwise(next_node): 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 - ): + frontier_val = frontier.meta.get("val") + if not isinstance(frontier_val, torch.Tensor): return False + rank = len(frontier_val.shape) + return all( + input_node is frontier + or FuseIdenticalInputTransformsPass.is_layout_invariant(input_node, rank) + for input_node in next_node.all_input_nodes + ) - 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]), + def _apply_fork_split( + self, + node: torch.fx.Node, + frontier: torch.fx.Node, + producer: torch.fx.Node, + branch_splits: Sequence[_ForkBranchSplit], + ) -> None: + """Apply a validated fork plan and place one transform after each + branch. + + If propagation already crossed operators before reaching the fork, the + original path is first detached from the permute. Each branch is then + rewired to the pre-permute layout, its metadata is updated to that + layout, and a copy of the original permute is inserted after it. + Argument-changing operations use the transform arguments captured during + planning. + + """ + if frontier is not node: + original_user = next(iter(node.users)) + original_user.replace_input_with(node, producer) + + for branch_split in branch_splits: + next_node = branch_split.next_node + old_next_meta = copy.copy(next_node.meta) + if frontier is node: + next_node.replace_input_with(node, producer) + if branch_split.arg_update is not None: + branch_input = ( + producer if frontier is node else branch_split.arg_update[1][0] ) - branch_permute.meta = dict(node.meta) - next_node.replace_input_with(node, branch_permute) + next_node.args = (branch_input, *branch_split.arg_update[1][1:]) + next_node.meta = copy.copy(next_node.meta) + next_node.meta["val"] = old_next_meta["val"].new_empty( + branch_split.source_shape + ) + transform_suffix = ( + node.args[1:] + if branch_split.arg_update is None + else branch_split.arg_update[0][1:] + ) + with next_node.graph.inserting_after(next_node): + branch_transform = next_node.graph.call_function( + cast(Any, node.target), + args=(next_node, *transform_suffix), + kwargs=dict(node.kwargs), + ) + branch_transform.meta = old_next_meta + for user in list(next_node.users): + if user is not branch_transform: + user.replace_input_with(next_node, branch_transform) - if len(node.users) == 0: + if not node.users: node.graph.erase_node(node) - return True def _move_node( self, diff --git a/backends/arm/test/misc/test_transpose_counts.py b/backends/arm/test/misc/test_transpose_counts.py index 168dabe96b9..bd73ddfe0cb 100644 --- a/backends/arm/test/misc/test_transpose_counts.py +++ b/backends/arm/test/misc/test_transpose_counts.py @@ -399,7 +399,7 @@ def forward(self, x: torch.Tensor): (torch.randn(1, 2, 8, 8),), 2, ), - "views": TransposeCountCase(ViewsModule(), (torch.rand(1, 2, 2, 4),), 4), + "views": TransposeCountCase(ViewsModule(), (torch.rand(1, 2, 2, 4),), 2), "transposes": TransposeCountCase( TransposesModule(), (torch.randn(1, 2, 3, 4),), diff --git a/backends/arm/test/ops/test_pixel_shuffling.py b/backends/arm/test/ops/test_pixel_shuffling.py index 8f89b7284c4..4980e24bab3 100644 --- a/backends/arm/test/ops/test_pixel_shuffling.py +++ b/backends/arm/test/ops/test_pixel_shuffling.py @@ -3,15 +3,11 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. - from typing import Tuple import torch - from executorch.backends.arm.constants import MAX_RANK - from executorch.backends.arm.test import common - from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, @@ -33,14 +29,27 @@ max_rank_input_supported = MAX_RANK - 2 +u55_pixel_xfails = { + "rand_4d_contiguous": "Known U55 partitioning limitation for large 4D pixel shuffle layouts.", + "rand_4d_channels_last": "Known U55 partitioning limitation for large 4D pixel shuffle layouts.", +} + class PixelUnShuffle(nn.Module): upscale_factor = 2 test_data_generators = { - "rand_4d": lambda: (torch.randn(1, 12, 64, 64),), - "test_4d": lambda: (torch.tensor([[[[10.0, 20.0], [30.0, 40.0]]]]),), - "test_3d": lambda: (torch.tensor([[[10.0, 20.0], [30.0, 40.0]]]),), + "rand_4d": lambda: ((torch.randn(1, 12, 64, 64),), None), + "rand_4d_contiguous": lambda: ((torch.randn(1, 12, 270, 480),), 1), + "rand_4d_channels_last": lambda: ( + (torch.randn(1, 12, 270, 480).to(memory_format=torch.channels_last),), + 1, + ), + "test_4d": lambda: ( + (torch.tensor([[[[10.0, 20.0], [30.0, 40.0]]]]),), + None, + ), + "test_3d": lambda: ((torch.tensor([[[10.0, 20.0], [30.0, 40.0]]]),), None), } def __init__(self, *args, **kwargs): @@ -57,16 +66,26 @@ def forward(self, inputs: torch.Tensor) -> torch.Tensor: class PixelShuffle(nn.Module): - upscale_factor = 2 test_data_generators = { - "rand_4d": lambda: (torch.randn(1, 12, 64, 64),), - "test_4d": lambda: (torch.tensor([[[[10.0]], [[20.0]], [[30.0]], [[40.0]]]]),), - "test_3d": lambda: (torch.tensor([[[10.0]], [[20.0]], [[30.0]], [[40.0]]]),), + "rand_4d": lambda: ((torch.randn(1, 12, 64, 64),), 1), + "test_4d": lambda: ( + (torch.tensor([[[[10.0]], [[20.0]], [[30.0]], [[40.0]]]]),), + 0, + ), + "test_3d": lambda: ( + (torch.tensor([[[10.0]], [[20.0]], [[30.0]], [[40.0]]]),), + 0, + ), + "rand_4d_contiguous": lambda: ((torch.randn(1, 12, 270, 480),), 1), + "rand_4d_channels_last": lambda: ( + (torch.randn(1, 12, 270, 480).to(memory_format=torch.channels_last),), + 1, + ), } - def __init__(self, *args, **kwargs): + def __init__(self, upscale_factor: int = 2, *args, **kwargs): super().__init__(*args, **kwargs) - + self.upscale_factor = upscale_factor self.depth_to_space = nn.PixelShuffle(self.upscale_factor) def forward(self, inputs: torch.Tensor) -> torch.Tensor: @@ -79,58 +98,91 @@ def forward(self, inputs: torch.Tensor) -> torch.Tensor: @common.parametrize("test_data", PixelUnShuffle.test_data_generators) def test_pixel_unshuffle_tosa_FP(test_data: input_t1): + inputs, expected_transposes = test_data() pipeline = TosaPipelineFP[input_t1]( PixelUnShuffle(), - test_data(), + inputs, aten_op_pixel_unshuffle, exir_op_pixel_unshuffle, ) + if expected_transposes is not None: + pipeline.count_tosa_ops({"TRANSPOSE": expected_transposes}) pipeline.run() @common.parametrize("test_data", PixelUnShuffle.test_data_generators) -def test_pixel_unshuffle_tosa_INT(test_data: input_t1): +def test_pixel_unshuffle_no_target_tosa_mixed_precision(test_data: input_t1): + inputs, expected_transposes = test_data() pipeline = TosaPipelineINT[input_t1]( PixelUnShuffle(), - test_data(), + inputs, aten_op_pixel_unshuffle, exir_op_pixel_unshuffle, + tosa_extensions=["FP"], + qtol=1, ) + if expected_transposes is not None: + pipeline.count_tosa_ops({"TRANSPOSE": expected_transposes}) pipeline.run() @common.parametrize("test_data", PixelShuffle.test_data_generators) def test_pixel_shuffle_tosa_FP(test_data: input_t1): + inputs, expected_transposes = test_data() pipeline = TosaPipelineFP[input_t1]( PixelShuffle(), - test_data(), + inputs, aten_op_pixel_shuffle, exir_op_pixel_shuffle, ) + if expected_transposes is not None: + pipeline.count_tosa_ops({"TRANSPOSE": expected_transposes}) pipeline.run() @common.parametrize("test_data", PixelShuffle.test_data_generators) -def test_pixel_shuffle_tosa_INT(test_data: input_t1): +def test_pixel_shuffle_no_target_tosa_mixed_precision(test_data: input_t1): + inputs, expected_transposes = test_data() pipeline = TosaPipelineINT[input_t1]( PixelShuffle(), - test_data(), + inputs, aten_op_pixel_shuffle, exir_op_pixel_shuffle, + tosa_extensions=["FP"], + qtol=1, ) + if expected_transposes is not None: + pipeline.count_tosa_ops({"TRANSPOSE": expected_transposes}) pipeline.run() @common.parametrize("test_data", PixelUnShuffle.test_data_generators) @common.SkipIfNoModelConverter def test_pixel_unshuffle_vgf_no_quant(test_data: input_t1): + inputs, expected_transposes = test_data() + pipeline = VgfPipeline[input_t1]( + PixelUnShuffle(), + inputs, + aten_op_pixel_unshuffle, + exir_op_pixel_unshuffle, + run_on_vulkan_runtime=True, + quantize=False, + ) + pipeline.run() + + +@common.parametrize("test_data", PixelUnShuffle.test_data_generators) +@common.SkipIfNoModelConverter +def test_pixel_unshuffle_vgf_FP(test_data: input_t1): + inputs, _ = test_data() pipeline = VgfPipeline[input_t1]( PixelUnShuffle(), - test_data(), + inputs, aten_op_pixel_unshuffle, exir_op_pixel_unshuffle, run_on_vulkan_runtime=True, quantize=False, + tosa_spec="TOSA-1.0+FP", ) pipeline.run() @@ -138,9 +190,10 @@ def test_pixel_unshuffle_vgf_no_quant(test_data: input_t1): @common.parametrize("test_data", PixelUnShuffle.test_data_generators) @common.SkipIfNoModelConverter def test_pixel_unshuffle_vgf_quant(test_data: input_t1): + inputs, expected_transposes = test_data() pipeline = VgfPipeline[input_t1]( PixelUnShuffle(), - test_data(), + inputs, aten_op_pixel_unshuffle, exir_op_pixel_unshuffle, run_on_vulkan_runtime=True, @@ -152,13 +205,30 @@ def test_pixel_unshuffle_vgf_quant(test_data: input_t1): @common.parametrize("test_data", PixelShuffle.test_data_generators) @common.SkipIfNoModelConverter def test_pixel_shuffle_vgf_no_quant(test_data: input_t1): + inputs, expected_transposes = test_data() + pipeline = VgfPipeline[input_t1]( + PixelShuffle(), + inputs, + aten_op_pixel_shuffle, + exir_op_pixel_shuffle, + run_on_vulkan_runtime=True, + quantize=False, + ) + pipeline.run() + + +@common.parametrize("test_data", PixelShuffle.test_data_generators) +@common.SkipIfNoModelConverter +def test_pixel_shuffle_vgf_FP(test_data: input_t1): + inputs, _ = test_data() pipeline = VgfPipeline[input_t1]( PixelShuffle(), - test_data(), + inputs, aten_op_pixel_shuffle, exir_op_pixel_shuffle, run_on_vulkan_runtime=True, quantize=False, + tosa_spec="TOSA-1.0+FP", ) pipeline.run() @@ -166,9 +236,10 @@ def test_pixel_shuffle_vgf_no_quant(test_data: input_t1): @common.parametrize("test_data", PixelShuffle.test_data_generators) @common.SkipIfNoModelConverter def test_pixel_shuffle_vgf_quant(test_data: input_t1): + inputs, expected_transposes = test_data() pipeline = VgfPipeline[input_t1]( PixelShuffle(), - test_data(), + inputs, aten_op_pixel_shuffle, exir_op_pixel_shuffle, run_on_vulkan_runtime=True, @@ -178,11 +249,48 @@ def test_pixel_shuffle_vgf_quant(test_data: input_t1): @common.parametrize("test_data", PixelUnShuffle.test_data_generators) +@common.SkipIfNoModelConverter +def test_pixel_unshuffle_vgf_INT(test_data: input_t1): + inputs, _ = test_data() + pipeline = VgfPipeline[input_t1]( + PixelUnShuffle(), + inputs, + aten_op_pixel_unshuffle, + exir_op_pixel_unshuffle, + run_on_vulkan_runtime=True, + quantize=True, + tosa_spec="TOSA-1.0+INT", + ) + pipeline.run() + + +@common.parametrize("test_data", PixelShuffle.test_data_generators) +@common.SkipIfNoModelConverter +def test_pixel_shuffle_vgf_INT(test_data: input_t1): + inputs, _ = test_data() + pipeline = VgfPipeline[input_t1]( + PixelShuffle(), + inputs, + aten_op_pixel_shuffle, + exir_op_pixel_shuffle, + run_on_vulkan_runtime=True, + quantize=True, + tosa_spec="TOSA-1.0+INT", + ) + pipeline.run() + + +@common.parametrize( + "test_data", + PixelUnShuffle.test_data_generators, + xfails=u55_pixel_xfails, +) @common.XfailIfNoCorstone300 def test_pixel_unshuffle_u55_INT(test_data: input_t1): + inputs, _ = test_data() pipeline = EthosU55PipelineINT[input_t1]( PixelUnShuffle(), - test_data(), + inputs, aten_op_pixel_unshuffle, exir_op_pixel_unshuffle, run_on_fvp=True, @@ -196,9 +304,10 @@ def test_pixel_unshuffle_u55_INT(test_data: input_t1): ) @common.XfailIfNoCorstone320 def test_pixel_unshuffle_u85_INT(test_data: input_t1): + inputs, _ = test_data() pipeline = EthosU85PipelineINT[input_t1]( PixelUnShuffle(), - test_data(), + inputs, aten_op_pixel_unshuffle, exir_op_pixel_unshuffle, run_on_fvp=True, @@ -206,12 +315,17 @@ def test_pixel_unshuffle_u85_INT(test_data: input_t1): pipeline.run() -@common.parametrize("test_data", PixelShuffle.test_data_generators) +@common.parametrize( + "test_data", + PixelShuffle.test_data_generators, + xfails=u55_pixel_xfails, +) @common.XfailIfNoCorstone300 def test_pixel_shuffle_u55_INT(test_data: input_t1): + inputs, _ = test_data() pipeline = EthosU55PipelineINT[input_t1]( PixelShuffle(), - test_data(), + inputs, aten_op_pixel_shuffle, exir_op_pixel_shuffle, run_on_fvp=True, @@ -225,9 +339,10 @@ def test_pixel_shuffle_u55_INT(test_data: input_t1): ) @common.XfailIfNoCorstone320 def test_pixel_shuffle_u85_INT(test_data: input_t1): + inputs, _ = test_data() pipeline = EthosU85PipelineINT[input_t1]( PixelShuffle(), - test_data(), + inputs, aten_op_pixel_shuffle, exir_op_pixel_shuffle, run_on_fvp=True, diff --git a/backends/arm/test/passes/test_arm_pass_utils.py b/backends/arm/test/passes/test_arm_pass_utils.py new file mode 100644 index 00000000000..0ab44329952 --- /dev/null +++ b/backends/arm/test/passes/test_arm_pass_utils.py @@ -0,0 +1,45 @@ +# 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. + +import sympy # type: ignore[import-untyped] +import torch +from executorch.backends.arm._passes.arm_pass_utils import refresh_node_meta +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.symbolic_shapes import ShapeEnv + + +def test_refresh_node_meta_preserves_symbolic_shape() -> None: + shape_env = ShapeEnv() + batch = shape_env.create_symintnode(sympy.Symbol("batch"), hint=2) + assert isinstance(batch, torch.SymInt) + shape_env.constrain_symbol_range(batch.node.expr, compiler_min=1, compiler_max=8) + with FakeTensorMode(shape_env=shape_env, allow_non_fake_inputs=True): + x_val = torch.empty((batch, 3)) + y_val = torch.empty((1, 3)) + + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = x_val + y = graph.placeholder("y") + y.meta["val"] = y_val + add = graph.call_function(torch.ops.aten.add.Tensor, args=(x, y)) + add.meta["val"] = torch.empty((2, 3)) + + refresh_node_meta(add) + + assert isinstance(add.meta["val"].shape[0], torch.SymInt) + assert sympy.simplify(add.meta["val"].shape[0].node.expr - batch.node.expr) == 0 + + +def test_refresh_node_meta_leaves_value_when_input_meta_is_missing() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + relu = graph.call_function(torch.ops.aten.relu.default, args=(x,)) + original_val = torch.empty((2, 3)) + relu.meta["val"] = original_val + + refresh_node_meta(relu) + + assert relu.meta["val"] is original_val diff --git a/backends/arm/test/passes/test_dim_maps.py b/backends/arm/test/passes/test_dim_maps.py index 16a18720442..a6c9052a24b 100644 --- a/backends/arm/test/passes/test_dim_maps.py +++ b/backends/arm/test/passes/test_dim_maps.py @@ -170,10 +170,11 @@ def _propose_reduction_view_swap( if not view_map.is_valid_map: return None - target_dims = view_map.map_dim(source_dims) - if target_dims is None: + proposal = view_map.map_reduction_after_view(input_shape, source_dims) + if proposal is None: return None - return list(view_shape), target_dims + view_shape, target_dims = proposal + return cast(list[_DimT], view_shape), target_dims def _propose_view_reduction_swap( @@ -185,10 +186,11 @@ def _propose_view_reduction_swap( if not view_map.is_valid_map: return None - source_dims = view_map.map_dim_inverse(target_dims) - if source_dims is None: + proposal = view_map.map_reduction_before_view(target_dims) + if proposal is None: return None - return source_dims, _reduce_shape(view_shape, target_dims) + source_dims, output_shape = proposal + return source_dims, cast(list[_DimT], output_shape) def _bruteforce_permute_view_swaps( @@ -303,6 +305,42 @@ def test_dim_map_matches_view_map_docstring_example_reduction_dims() -> None: assert view_map.map_dim_inverse(2) is None +def test_dim_map_allows_singleton_rank_change_reduction_swap() -> None: + x = _tensor([6, 4]) + proposal = _propose_reduction_view_swap([6, 4], [0], [6, 1, 4]) + candidates = _bruteforce_reduction_view_swaps(x, [0], [6, 1, 4]) + + assert proposal == ([6, 1, 4], [0]) + assert proposal in candidates + + +def test_dim_map_allows_inverse_singleton_rank_change_reduction_swap() -> None: + x = _tensor([6, 4]) + proposal = _propose_view_reduction_swap([6, 4], [6, 1, 4], [0]) + candidates = _bruteforce_view_reduction_swaps(x, [6, 1, 4], [0]) + + assert proposal == ([0], [1, 1, 4]) + assert proposal in candidates + + +def test_dim_map_allows_output_singleton_rank_change_reduction_swap() -> None: + x = _tensor([6, 4]) + view_map = ViewMap.from_shapes([1, 4], [1, 1, 4]) + proposal = view_map.map_reduction_after_view([6, 4], [0]) + + assert proposal == ([6, 1, 4], [0, 1]) + view_shape, target_dims = proposal + original = x.sum(dim=0, keepdim=True).reshape(1, 1, 4) + candidate = x.reshape(view_shape).sum(dim=tuple(target_dims), keepdim=True) + assert _same(original, candidate) + + +def test_dim_map_rejects_invalid_zero_dim_remap() -> None: + view_map = ViewMap.from_shapes([5, 2, 60], [2, 60]) + + assert view_map.remap_target_shape([1, 2, 60]) is None + + def test_dim_map_matches_view_map_docstring_example_permutation_dims() -> None: view_map = ViewMap.from_shapes([4, 3, 10], [2, 2, 1, 5, 3, 2]) @@ -505,7 +543,7 @@ def test_dim_map_reduction_view_swap_preserves_symbolic_view_shape_dims() -> Non assert proposal is not None view_shape, target_dims = proposal assert isinstance(view_shape[0], torch.SymInt) - assert view_shape[0] is batch + assert sympy.simplify(view_shape[0].node.expr - batch.node.expr) == 0 assert view_shape[1:] == [2, 3] assert target_dims == [1, 2] diff --git a/backends/arm/test/passes/test_fuse_identical_input_transforms_pass.py b/backends/arm/test/passes/test_fuse_identical_input_transforms_pass.py index 0b7a0b575cd..4a5112975ee 100644 --- a/backends/arm/test/passes/test_fuse_identical_input_transforms_pass.py +++ b/backends/arm/test/passes/test_fuse_identical_input_transforms_pass.py @@ -32,6 +32,7 @@ _ADD = exir_ops.edge.aten.add.Tensor _CAT = exir_ops.edge.aten.cat.default +_WHERE = exir_ops.edge.aten.where.self _VIEW = exir_ops.edge.aten.view_copy.default _PERMUTE = exir_ops.edge.aten.permute_copy.default @@ -221,6 +222,32 @@ def test_fuse_identical_input_transforms_sunk_permute_replaces_stale_tag() -> No _validate_numerics(original, result.graph_module, (x_data, y_data)) +def test_fuse_identical_input_transforms_sinks_where_input_permutes() -> None: + condition_data = torch.randn(2, 3, 4) > 0 + x_data = torch.randn(2, 3, 4) + y_data = torch.randn(2, 3, 4) + + builder = GraphBuilder() + condition = builder.placeholder("condition", condition_data) + x = builder.placeholder("x", x_data) + y = builder.placeholder("y", y_data) + condition_permute = builder.call_operator(op=_PERMUTE, args=(condition, [0, 2, 1])) + x_permute = builder.call_operator(op=_PERMUTE, args=(x, [0, 2, 1])) + y_permute = builder.call_operator(op=_PERMUTE, args=(y, [0, 2, 1])) + where = builder.call_operator( + op=_WHERE, args=(condition_permute, x_permute, y_permute) + ) + builder.output([where]) + graph_module = builder.get_graph_module() + + original, result = _apply_pass(graph_module) + + assert result.modified + assert _count_node(result.graph_module, _PERMUTE) == 1 + assert _compute_nodes(result.graph_module) == [_WHERE, _PERMUTE] + _validate_numerics(original, result.graph_module, (condition_data, x_data, y_data)) + + def test_fuse_identical_input_transforms_sinks_cat_input_views() -> None: x_data = torch.randn(2, 3, 4) y_data = torch.randn(2, 3, 4) @@ -790,16 +817,26 @@ def test_fuse_identical_input_transforms_sinks_views_through_rank_broadcasts( _validate_numerics(original, result.graph_module, (x_data, y_data)) -def test_fuse_identical_input_transforms_rejects_views_through_expanding_broadcast() -> ( - None -): - x_data = torch.randn(1, 6) - y_data = torch.randn(6, 1) +@pytest.mark.parametrize( + "x_shape, y_shape, view_shape", + [ + ((1, 6), (6, 1), (2, 3)), + ((1, 1, 6), (1, 6, 1), (3, 2)), + ((6,), (6, 1), (3, 2)), + ], +) +def test_fuse_identical_input_transforms_rejects_views_through_expanding_broadcast( + x_shape: tuple[int, ...], + y_shape: tuple[int, ...], + view_shape: tuple[int, ...], +) -> None: + x_data = torch.randn(x_shape) + y_data = torch.randn(y_shape) graph_module = _binary_transform_graph( x_data, y_data, - [(_VIEW, [2, 3])], - [(_VIEW, [2, 3])], + [(_VIEW, view_shape)], + [(_VIEW, view_shape)], ) original, result = _apply_pass(graph_module) diff --git a/backends/arm/test/passes/test_match_arg_ranks_pass.py b/backends/arm/test/passes/test_match_arg_ranks_pass.py new file mode 100644 index 00000000000..b2bd453b692 --- /dev/null +++ b/backends/arm/test/passes/test_match_arg_ranks_pass.py @@ -0,0 +1,49 @@ +# 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 typing import cast + +import torch +from executorch.backends.arm._passes.match_arg_ranks_pass import MatchArgRanksPass +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from torch._subclasses.fake_tensor import FakeTensorMode + + +def test_match_arg_ranks_handles_tosa_binary_ops() -> None: + fake_mode = FakeTensorMode() + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = fake_mode.from_tensor(torch.empty((1, 270, 480, 3, 2, 2))) + y = graph.placeholder("y") + y.meta["val"] = fake_mode.from_tensor(torch.empty((1, 1, 1, 1))) + add = graph.call_function(exir_ops.backend.tosa.ADD.default, args=(x, y)) + add.meta["val"] = fake_mode.from_tensor(torch.empty((1, 270, 480, 3, 2, 2))) + graph.output(add) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP")): + result = MatchArgRanksPass(cast(ExportedProgram, None)).call(graph_module) + call_nodes = [ + node for node in result.graph_module.graph.nodes if node.op == "call_function" + ] + view = next( + node + for node in call_nodes + if node.target == exir_ops.edge.aten.view_copy.default + ) + add = next( + node for node in call_nodes if node.target == exir_ops.backend.tosa.ADD.default + ) + + assert result.modified + assert view.all_input_nodes[0].target == "y" + assert view.args[1] == [1, 1, 1, 1, 1, 1] + assert add.all_input_nodes[0].target == "x" + assert add.all_input_nodes[1] is view diff --git a/backends/arm/test/passes/test_move_data_movement_ops_to_smaller_dtype_pass.py b/backends/arm/test/passes/test_move_data_movement_ops_to_smaller_dtype_pass.py new file mode 100644 index 00000000000..7b570659f8a --- /dev/null +++ b/backends/arm/test/passes/test_move_data_movement_ops_to_smaller_dtype_pass.py @@ -0,0 +1,113 @@ +# 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. + +import torch +from executorch.backends.arm._passes import ( + MoveDataMovementOpsToSmallerDtypePass, + PropagateViewCopyPermuteDownPass, + PropagateViewCopyPermuteUpPass, +) +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from executorch.exir.dialects._ops import ops as exir_ops + + +PERMUTE = exir_ops.edge.aten.permute_copy.default +RESCALE = exir_ops.backend.tosa.RESCALE.default + + +def _apply_pass(graph: torch.fx.Graph) -> torch.fx.GraphModule: + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")): + return MoveDataMovementOpsToSmallerDtypePass().call(graph_module).graph_module + + +def _call_nodes(graph_module: torch.fx.GraphModule) -> list[torch.fx.Node]: + return [node for node in graph_module.graph.nodes if node.op == "call_function"] + + +def test_moves_permute_after_narrowing_rescale() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.int32) + rescale = graph.call_function(RESCALE, args=(permute, torch.int8, [1.0], 0, 0)) + rescale.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.int8) + graph.output(rescale) + + graph_module = _apply_pass(graph) + rescale, permute = _call_nodes(graph_module) + + assert rescale.target == RESCALE + assert permute.target == PERMUTE + assert rescale.meta["val"].shape == torch.Size((1, 2, 3, 4)) + assert permute.meta["val"].dtype == torch.int8 + + +def test_moves_permute_before_widening_rescale() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int8) + rescale = graph.call_function(RESCALE, args=(x, 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) + + graph_module = _apply_pass(graph) + permute, rescale = _call_nodes(graph_module) + + assert permute.target == PERMUTE + assert rescale.target == RESCALE + assert permute.meta["val"].dtype == torch.int8 + assert rescale.meta["val"].shape == torch.Size((1, 3, 4, 2)) + + +def test_keeps_per_channel_rescale_order() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.int32) + rescale = graph.call_function(RESCALE, args=(permute, torch.int8, [1.0, 2.0], 0, 0)) + rescale.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.int8) + graph.output(rescale) + + graph_module = _apply_pass(graph) + + assert [node.target for node in _call_nodes(graph_module)] == [PERMUTE, RESCALE] + + +def test_splits_merged_permute_at_distinct_narrowing_rescales() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.int32) + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.int32) + rescale_1 = graph.call_function(RESCALE, args=(permute, torch.int8, [1.0], 0, 0)) + rescale_1.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.int8) + rescale_2 = graph.call_function(RESCALE, args=(permute, torch.int8, [2.0], 0, 0)) + rescale_2.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.int8) + graph.output((rescale_1, rescale_2)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")): + graph_module = ( + PropagateViewCopyPermuteDownPass().call(graph_module).graph_module + ) + graph_module = PropagateViewCopyPermuteUpPass().call(graph_module).graph_module + assert sum(node.target == PERMUTE for node in _call_nodes(graph_module)) == 1 + graph_module = ( + MoveDataMovementOpsToSmallerDtypePass().call(graph_module).graph_module + ) + + call_nodes = _call_nodes(graph_module) + permutes = [node for node in call_nodes if node.target == PERMUTE] + assert len(permutes) == 2 + assert all(node.meta["val"].dtype == torch.int8 for node in permutes) + assert all(node.all_input_nodes[0].target == RESCALE for node in permutes) diff --git a/backends/arm/test/passes/test_propagate_permutes_views_pass.py b/backends/arm/test/passes/test_propagate_permutes_views_pass.py index 7b140dea806..79207eed9d4 100644 --- a/backends/arm/test/passes/test_propagate_permutes_views_pass.py +++ b/backends/arm/test/passes/test_propagate_permutes_views_pass.py @@ -9,6 +9,7 @@ import pytest import torch from executorch.backends.arm._passes import ( + MoveDataMovementOpsToSmallerDtypePass, PropagateViewCopyPermuteDownPass, PropagateViewCopyPermuteUpPass, ) @@ -28,6 +29,12 @@ PERMUTE = exir_ops.edge.aten.permute_copy.default VIEW = exir_ops.edge.aten.view_copy.default ADD = exir_ops.edge.aten.add.Tensor +SUB = exir_ops.edge.aten.sub.Tensor +MUL = exir_ops.edge.aten.mul.Tensor +GE = exir_ops.edge.aten.ge.Tensor +FLOOR = exir_ops.edge.aten.floor.default +CEIL = exir_ops.edge.aten.ceil.default +WHERE = exir_ops.edge.aten.where.self RELU = exir_ops.edge.aten.relu.default NEG = exir_ops.edge.aten.neg.default MM = exir_ops.edge.aten.mm.default @@ -117,6 +124,29 @@ def predicate(targets: list[object]) -> None: pipeline.run() +def test_propagate_scalar_view_down_through_unary_elementwise() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty(()) + view = graph.call_function(VIEW, args=(x, [1])) + view.meta["val"] = torch.empty((1,)) + relu = graph.call_function(RELU, args=(view,)) + relu.meta["val"] = torch.empty((1,)) + graph.output(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] + relu = next(node for node in call_nodes if node.target == RELU) + view = next(node for node in call_nodes if node.target == VIEW) + + assert targets.index(RELU) < targets.index(VIEW) + assert relu.meta["val"].shape == torch.Size(()) + assert view.meta["val"].shape == torch.Size((1,)) + + class UpwardPermute(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: return x.relu().neg().permute(0, 2, 3, 1) @@ -671,6 +701,23 @@ def test_down_pass_moves_matching_input_permutations_after_binary_op() -> None: assert targets.index(ADD) < targets.index(PERMUTE) +def test_down_pass_keeps_lower_rank_singleton_input_before_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, 1)) + 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, PropagateViewCopyPermuteDownPass) + + assert targets.index(PERMUTE) < targets.index(ADD) + + def test_down_pass_keeps_sunk_view_before_rank_reducing_permute() -> None: graph = torch.fx.Graph() x = graph.placeholder("x") @@ -894,7 +941,7 @@ def test_down_pass_moves_permutation_after_reduction() -> None: assert transform.meta["val"].shape == torch.Size((1, 3, 4, 1)) -def test_down_pass_stops_when_fanout_does_not_converge() -> None: +def test_down_pass_splits_permute_over_elementwise_fanout() -> None: graph = torch.fx.Graph() x = graph.placeholder("x") x.meta["val"] = torch.empty((1, 2, 3, 4)) @@ -906,11 +953,18 @@ def test_down_pass_stops_when_fanout_does_not_converge() -> None: neg.meta["val"] = torch.empty((1, 3, 4, 2)) graph.output((relu, neg)) - targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + 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] + branch_transforms = [node for node in call_nodes if node.target == PERMUTE] - assert targets.count(PERMUTE) == 1 - assert targets.index(PERMUTE) < targets.index(RELU) - assert targets.index(PERMUTE) < targets.index(NEG) + assert targets.count(PERMUTE) == 2 + assert {transform.args[0].target for transform in branch_transforms} == { + RELU, + NEG, + } def test_down_pass_splits_permute_over_slice_fanout() -> None: @@ -989,11 +1043,18 @@ def test_down_pass_stops_when_convergence_has_untracked_input() -> None: cat_node.meta["val"] = torch.empty((1, 3, 4, 6)) graph.output(cat_node) - targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + 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] + branch_transforms = [node for node in call_nodes if node.target == PERMUTE] - assert targets.count(PERMUTE) == 1 - assert targets.index(PERMUTE) < targets.index(RELU) - assert targets.index(PERMUTE) < targets.index(NEG) + assert targets.count(PERMUTE) == 2 + assert {transform.args[0].target for transform in branch_transforms} == { + RELU, + NEG, + } def test_down_pass_stops_view_before_cat_converging_fanout() -> None: @@ -1060,6 +1121,112 @@ def test_propagate_moves_before_dtype_changing_rescale() -> None: assert targets.index(PERMUTE) < targets.index(RESCALE) +def test_propagate_moves_through_elementwise_fork() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.float32) + scale = graph.placeholder("scale") + scale.meta["val"] = torch.empty((1, 1, 1, 1), dtype=torch.float32) + bias = graph.placeholder("bias") + bias.meta["val"] = torch.empty((1, 1, 1, 1), dtype=torch.float32) + + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.float32) + scaled = graph.call_function(MUL, args=(permute, scale)) + scaled.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.float32) + condition = graph.call_function(GE, args=(scaled, bias)) + condition.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.bool) + add = graph.call_function(ADD, args=(scaled, bias)) + add.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.float32) + sub = graph.call_function(SUB, args=(scaled, bias)) + sub.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.float32) + floor = graph.call_function(FLOOR, args=(add,)) + floor.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.float32) + ceil = graph.call_function(CEIL, args=(sub,)) + ceil.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.float32) + where = graph.call_function(WHERE, args=(condition, floor, ceil)) + where.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.float32) + graph.output(where) + + targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + + assert targets.count(PERMUTE) == 1 + assert targets.index(WHERE) < targets.index(PERMUTE) + + +def test_propagate_moves_through_reduction_fork() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4), dtype=torch.float32) + bias = graph.placeholder("bias") + bias.meta["val"] = torch.empty((1, 1, 1, 1), dtype=torch.float32) + + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.empty((1, 3, 4, 2), dtype=torch.float32) + condition_sum = graph.call_function(SUM, args=(permute, [1], True)) + condition_sum.meta["val"] = torch.empty((1, 1, 4, 2), dtype=torch.float32) + condition = graph.call_function(GE, args=(condition_sum, bias)) + condition.meta["val"] = torch.empty((1, 1, 4, 2), dtype=torch.bool) + data_sum = graph.call_function(SUM, args=(permute, [1], True)) + data_sum.meta["val"] = torch.empty((1, 1, 4, 2), dtype=torch.float32) + data_mean = graph.call_function(MEAN, args=(permute, [1], True)) + data_mean.meta["val"] = torch.empty((1, 1, 4, 2), dtype=torch.float32) + floor = graph.call_function(FLOOR, args=(data_sum,)) + floor.meta["val"] = torch.empty((1, 1, 4, 2), dtype=torch.float32) + ceil = graph.call_function(CEIL, args=(data_mean,)) + ceil.meta["val"] = torch.empty((1, 1, 4, 2), dtype=torch.float32) + where = graph.call_function(WHERE, args=(condition, floor, ceil)) + where.meta["val"] = torch.empty((1, 1, 4, 2), dtype=torch.float32) + graph.output(where) + + 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] + reductions = [node for node in call_nodes if node.target in {SUM, MEAN}] + + assert targets.count(PERMUTE) == 1 + assert targets.index(WHERE) < targets.index(PERMUTE) + assert all(reduction.args[1] == [2] for reduction in reductions) + + +def test_propagate_view_moves_through_elementwise_fork() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 12), dtype=torch.float32) + scale = graph.placeholder("scale") + scale.meta["val"] = torch.empty((1, 1, 1), dtype=torch.float32) + bias = graph.placeholder("bias") + bias.meta["val"] = torch.empty((1, 1, 1), dtype=torch.float32) + + view = graph.call_function(VIEW, args=(x, [1, 3, 8])) + view.meta["val"] = torch.empty((1, 3, 8), dtype=torch.float32) + scaled = graph.call_function(MUL, args=(view, scale)) + scaled.meta["val"] = torch.empty((1, 3, 8), dtype=torch.float32) + condition = graph.call_function(GE, args=(scaled, bias)) + condition.meta["val"] = torch.empty((1, 3, 8), dtype=torch.bool) + add = graph.call_function(ADD, args=(scaled, bias)) + add.meta["val"] = torch.empty((1, 3, 8), dtype=torch.float32) + sub = graph.call_function(SUB, args=(scaled, bias)) + sub.meta["val"] = torch.empty((1, 3, 8), dtype=torch.float32) + floor = graph.call_function(FLOOR, args=(add,)) + floor.meta["val"] = torch.empty((1, 3, 8), dtype=torch.float32) + ceil = graph.call_function(CEIL, args=(sub,)) + ceil.meta["val"] = torch.empty((1, 3, 8), dtype=torch.float32) + where = graph.call_function(WHERE, args=(condition, floor, ceil)) + where.meta["val"] = torch.empty((1, 3, 8), dtype=torch.float32) + graph.output(where) + + targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) + + assert targets.count(VIEW) == 1 + assert targets.index(MUL) < targets.index(VIEW) + assert targets.index(VIEW) < targets.index(GE) + assert targets.index(VIEW) < targets.index(ADD) + assert targets.index(VIEW) < targets.index(SUB) + + def test_propagate_fuses_permute_view_around_table() -> None: graph = torch.fx.Graph() x = graph.placeholder("x") @@ -1142,7 +1309,9 @@ def test_propagate_up_stops_at_shared_rescale_producer() -> None: assert targets.index(RESCALE) < targets.index(PERMUTE) -def test_propagate_up_stops_before_narrowing_rescale_fed_by_binary_op() -> None: +def test_smaller_dtype_pass_restores_permute_after_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) @@ -1157,7 +1326,13 @@ def test_propagate_up_stops_before_narrowing_rescale_fed_by_binary_op() -> None: graph.output(permute) with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")): - targets = _run_pass_on_graph(graph) + graph_module = _run_pass_on_graph_module(graph) + result = MoveDataMovementOpsToSmallerDtypePass().call(graph_module) + targets = [ + node.target + for node in result.graph_module.graph.nodes + if node.op == "call_function" + ] assert targets.index(RESCALE) < targets.index(PERMUTE) @@ -1182,7 +1357,9 @@ def test_propagate_up_crosses_same_width_rescale_fed_by_binary_op() -> None: assert targets.index(PERMUTE) < targets.index(RESCALE) -def test_propagate_up_stops_before_narrowing_rescale_from_shared_placeholder() -> None: +def test_smaller_dtype_pass_restores_permute_after_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) @@ -1195,7 +1372,13 @@ def test_propagate_up_stops_before_narrowing_rescale_from_shared_placeholder() - graph.output((permute, other)) with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")): - targets = _run_pass_on_graph(graph) + graph_module = _run_pass_on_graph_module(graph) + result = MoveDataMovementOpsToSmallerDtypePass().call(graph_module) + targets = [ + node.target + for node in result.graph_module.graph.nodes + if node.op == "call_function" + ] assert targets.index(RESCALE) < targets.index(PERMUTE) @@ -1220,7 +1403,9 @@ def test_propagate_up_crosses_widening_rescale_fed_by_binary_op() -> None: assert targets.index(PERMUTE) < targets.index(RESCALE) -def test_propagate_up_stops_before_narrowing_rescale_behind_unary() -> None: +def test_smaller_dtype_pass_restores_permute_after_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) @@ -1233,7 +1418,13 @@ def test_propagate_up_stops_before_narrowing_rescale_behind_unary() -> None: graph.output(permute) with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")): - targets = _run_pass_on_graph(graph) + graph_module = _run_pass_on_graph_module(graph) + result = MoveDataMovementOpsToSmallerDtypePass().call(graph_module) + targets = [ + node.target + for node in result.graph_module.graph.nodes + if node.op == "call_function" + ] assert targets.index(RESCALE) < targets.index(PERMUTE)