diff --git a/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py b/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py index b241038f7a9..e84343f8cac 100644 --- a/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py +++ b/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py @@ -28,7 +28,9 @@ def __init__(self, exported_program: ExportedProgram) -> None: def _is_constant(self, node: torch.fx.Node) -> bool: # Override fragile string match check with exported program check - return super()._is_constant(node) or is_param_node(self.exported_program, node) + exported_program = self.exported_program + assert exported_program is not None + return super()._is_constant(node) or is_param_node(exported_program, node) def permute_subgraph(self, subgraph) -> bool: # TABLE lookup inputs are already tied to the table layout. diff --git a/backends/transforms/channels_last_ops.py b/backends/transforms/channels_last_ops.py index 0a90f0d1889..d7bfd76fb50 100644 --- a/backends/transforms/channels_last_ops.py +++ b/backends/transforms/channels_last_ops.py @@ -175,6 +175,7 @@ def _permute_copy(input, dims): lib.impl("max_pool2d", _max_pool2d, "CompositeExplicitAutograd") register_fake("channels_last::max_pool2d", _max_pool2d, lib=lib) + lib.define( "grid_sampler_2d(Tensor input, Tensor grid, int interpolation_mode, " "int padding_mode, bool align_corners) -> Tensor" diff --git a/backends/transforms/decompose_channels_last_pass.py b/backends/transforms/decompose_channels_last_pass.py index 05ec247c30a..2a28600eed8 100644 --- a/backends/transforms/decompose_channels_last_pass.py +++ b/backends/transforms/decompose_channels_last_pass.py @@ -27,6 +27,10 @@ exir_ops.edge.channels_last.grid_sampler_2d.default: exir_ops.edge.aten.grid_sampler_2d.default, } +_DIRECT_DECOMPOSITIONS = { + exir_ops.edge.channels_last.permute_copy.default: exir_ops.edge.aten.permute_copy.default, +} + class DecomposeChannelsLastPass(ExportPass): """Decompose channels_last dialect ops into permute + aten op + permute. @@ -39,6 +43,10 @@ class DecomposeChannelsLastPass(ExportPass): """ def call_operator(self, op, args, kwargs, meta): + direct_op = _DIRECT_DECOMPOSITIONS.get(op) + if direct_op is not None: + return super().call_operator(direct_op, args, kwargs, meta) + aten_op = _DECOMPOSITIONS.get(op) if aten_op is not None: nchw_in = super().call_operator( @@ -90,8 +98,4 @@ def call_operator(self, op, args, kwargs, meta): meta, ) return values, indices - if op == exir_ops.edge.channels_last.permute_copy.default: - return super().call_operator( - exir_ops.edge.aten.permute_copy.default, args, kwargs, meta - ) return super().call_operator(op, args, kwargs, meta) diff --git a/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py b/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py index 83b33533995..62183da2ebd 100644 --- a/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py +++ b/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py @@ -6,6 +6,7 @@ # pyre-unsafe +from collections import deque from typing import Any, Callable, cast import torch @@ -41,6 +42,42 @@ class FuseTransposeOrPermuteOpPairsPass(FuseOpPairsAcrossBranchesPass): exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, } + def __init__( + self, + can_propagate: Callable[[torch.fx.Node], bool] | None = None, + ) -> None: + super().__init__() + self.can_propagate = can_propagate + + def get_fuse_candidates( + self, + producer: torch.fx.Node, + consumer_op_packets: set[EdgeOpOverloadPacket], + bypass_ops: set[EdgeOpOverload], + ) -> list[torch.fx.Node]: + if self.can_propagate is None: + return super().get_fuse_candidates( + producer, consumer_op_packets, bypass_ops + ) + + users = deque(producer.users) + visited: set[torch.fx.Node] = set() + removal_candidates = [] + while users: + user = users.popleft() + if user in visited: + continue + visited.add(user) + if user.target in bypass_ops: + if not self.can_propagate(user): + return [] + users.extend(user.users) + elif self.can_fuse_for_chain(producer, user, consumer_op_packets): + removal_candidates.append(user) + else: + return [] + return removal_candidates + def can_fuse_for_chain( self, producer: torch.fx.Node, diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 34068e97ecd..a6d55deb4f1 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -7,16 +7,21 @@ # pyre-unsafe +from collections.abc import Callable from dataclasses import dataclass, field from typing import cast import torch import torch.fx from executorch.backends.transforms.channels_last_layout import ( + ATEN_PERMUTE_COPY, + is_layout_copy, is_permute_copy, + LAYOUT_PERMUTE_COPY, PERMUTE_COPY_TARGETS, ) from executorch.backends.transforms.permute_pass_utils import get_arg, set_arg +from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult @@ -29,6 +34,9 @@ class RemovePermutesAroundElementwiseOps(ExportPass): based on the permute's parameter such as mean, cat, and slice. The repeat_interleave idiom (unsqueeze -> expand_copy -> merging view_copy) is recognised as a single rank-preserving unit; see _interleave_triple. + + ``extra_permutable_ops`` must be layout-equivariant without argument remapping. + Layout-boundary propagation applies only to layout-owned dialect copies. """ @dataclass() @@ -45,6 +53,14 @@ class Subgraph: constant_edges_in: set[tuple[torch.fx.Node, torch.fx.Node]] = field( default_factory=set ) + # Open boundaries used only for structural layout-copy propagation. + input_boundaries: set[tuple[torch.fx.Node, torch.fx.Node, tuple[int, ...]]] = ( + field(default_factory=set) + ) + output_boundaries: set[tuple[torch.fx.Node, torch.fx.Node, tuple[int, ...]]] = ( + field(default_factory=set) + ) + layout_region: bool = False # Per-node expected end permutation (may differ from end_permute # when the subgraph contains rank-changing views). node_end_permute: dict[torch.fx.Node, list[int]] = field(default_factory=dict) @@ -56,8 +72,18 @@ class Subgraph: torch.fx.Node, tuple[int, int, torch.fx.Node, torch.fx.Node] ] = field(default_factory=dict) - def __init__(self, extra_permutable_ops: set | None = None) -> None: + def __init__( + self, + extra_permutable_ops: set | None = None, + *, + exported_program: ExportedProgram | None = None, + allow_layout_boundary_propagation: bool = False, + can_propagate: Callable[[torch.fx.Node], bool] | None = None, + ) -> None: super().__init__() + self.exported_program = exported_program + self.allow_layout_boundary_propagation = allow_layout_boundary_propagation + self.can_propagate = can_propagate self._permutable_ops = { exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.mul.Tensor, @@ -329,15 +355,39 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 self._interleave_cache.clear() subgraphs_found: list[RemovePermutesAroundElementwiseOps.Subgraph] = [] processed_nodes: set[torch.fx.Node] = set() - for node in graph_module.graph.nodes: - if not is_permute_copy(node): - continue + permute_nodes = [ + node for node in graph_module.graph.nodes if is_permute_copy(node) + ] + for node in permute_nodes: start_permute = self.get_permutation(node) if start_permute is None: continue + layout_region = self.allow_layout_boundary_propagation and is_layout_copy( + node + ) # Expected end permutation for the subgraph. end_permute = [start_permute.index(i) for i in range(len(start_permute))] + if layout_region: + users = list(node.users) + if users and all( + self.is_node_permutable(user) + or self._interleave_triple(user) is not None + for user in users + ): + subgraph = self.Subgraph( + start_permute, + end_permute, + layout_region=True, + ) + if all( + self.visit(user, subgraph, processed_nodes) for user in users + ): + subgraphs_found.append(subgraph) + processed_nodes.update(subgraph.nodes) + # Layout boundary movement is atomic across every direct user. + continue + # Try direct users first (same-rank matching) for user in node.users: if ( @@ -345,7 +395,11 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 and self._interleave_triple(user) is None ): continue - subgraph = self.Subgraph(start_permute, end_permute) + subgraph = self.Subgraph( + start_permute, + end_permute, + layout_region=layout_region, + ) if self.visit(user, subgraph, processed_nodes): subgraphs_found.append(subgraph) for n in subgraph.nodes: @@ -375,7 +429,11 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 and self._interleave_triple(view_user) is None ): continue - subgraph = self.Subgraph(adapted_start, adapted_end) + subgraph = self.Subgraph( + adapted_start, + adapted_end, + layout_region=layout_region, + ) # Include the view in the subgraph subgraph.nodes.add(view_node) subgraph.node_end_permute[view_node] = adapted_end @@ -395,6 +453,31 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 for n in subgraph.nodes: processed_nodes.add(n) + if self.allow_layout_boundary_propagation: + for node in permute_nodes: + end_permute = self.get_permutation(node) + if end_permute is None: + continue + if not is_layout_copy(node): + continue + producer = node.args[0] if node.args else None + if not isinstance(producer, torch.fx.Node): + continue + if ( + not self.is_node_permutable(producer) + and self._interleave_triple(producer) is None + ): + continue + start_permute = [end_permute.index(i) for i in range(len(end_permute))] + subgraph = self.Subgraph( + start_permute, + end_permute, + layout_region=True, + ) + if self.visit(producer, subgraph, processed_nodes): + subgraphs_found.append(subgraph) + processed_nodes.update(subgraph.nodes) + modified = False for subgraph in subgraphs_found: if self.permute_subgraph(subgraph): @@ -516,7 +599,14 @@ def visit( # noqa: C901 continue return False elif user.op == "output": - return False + if ( + not self.allow_layout_boundary_propagation + or not subgraph.layout_region + ): + return False + subgraph.output_boundaries.add( + (users_source, user, tuple(downstream_start)) + ) elif self._is_permutation_sink_view(user): # The permutation dies at this reshape (see # _is_permutation_sink_view), so terminate the region here with @@ -525,6 +615,15 @@ def visit( # noqa: C901 # terminates cleanly, whereas crossing it would leave the region # hunting for an end permute that layout-invariance made moot. continue + elif ( + self.allow_layout_boundary_propagation + and subgraph.layout_region + and self.can_propagate is not None + and not self.can_propagate(user) + ): + subgraph.output_boundaries.add( + (users_source, user, tuple(downstream_start)) + ) elif not self.visit( user, subgraph, processed_nodes, downstream_end, downstream_start ): @@ -553,6 +652,22 @@ def visit( # noqa: C901 if const_rank < permute_rank and inp.meta.get("val") is None: return False subgraph.constant_edges_in.add((inp, node)) + elif self._is_user_input(inp): + if ( + not self.allow_layout_boundary_propagation + or not subgraph.layout_region + or self._get_node_rank(inp) != len(current_end_permute) + ): + return False + subgraph.input_boundaries.add((inp, node, tuple(current_end_permute))) + elif ( + self.allow_layout_boundary_propagation + and subgraph.layout_region + and self.can_propagate is not None + and not self.can_propagate(inp) + and self._get_node_rank(inp) == len(current_end_permute) + ): + subgraph.input_boundaries.add((inp, node, tuple(current_end_permute))) elif not self.visit( inp, subgraph, @@ -602,6 +717,13 @@ def _is_constant(self, node: torch.fx.Node) -> bool: return True return False + def _is_user_input(self, node: torch.fx.Node) -> bool: + if node.op != "placeholder": + return False + if self.exported_program is not None: + return node.name in self.exported_program.graph_signature.user_inputs + return not self._is_constant(node) + def _get_node_rank(self, node: torch.fx.Node) -> int | None: """Return the tensor rank of a node's output, or None if unknown.""" val = node.meta.get("val") @@ -618,6 +740,8 @@ def _is_pointwise(target) -> bool: return False def is_node_permutable(self, node: torch.fx.Node) -> bool: + if self.can_propagate is not None and not self.can_propagate(node): + return False if node.target in self._PAD_OPS and not self._is_constant_pad(node): return False if node.target in self._permutable_ops: @@ -654,6 +778,11 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 # Ensure that the subgraph's edges have not been modified by an earlier rewrite before applying changes. if not self._subgraph_edges_are_current(subgraph): return False + if subgraph.layout_region and ( + not self._boundary_permutations_are_layout_copies(subgraph) + or not self._boundary_rewrite_is_cost_safe(subgraph) + ): + return False # Nodes belonging to a repeat_interleave triple are rewritten as a unit # below, so they must skip the per-node dim handling and the view rank @@ -736,9 +865,14 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 if const_rank is not None and const_rank == permute_rank: new_node = graph.create_node( "call_function", - exir_ops.edge.aten.permute_copy.default, + ( + LAYOUT_PERMUTE_COPY + if subgraph.layout_region + else ATEN_PERMUTE_COPY + ), args=(const_node, node_end_perm), ) + new_node.meta = {} elif ( const_rank is not None and const_rank < permute_rank @@ -762,9 +896,12 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 assert out.target in PERMUTE_COPY_TARGETS out.replace_all_uses_with(inp) + self._insert_input_boundary_permutations(subgraph) + self._insert_output_boundary_permutations(subgraph) + return True - def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: + def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: # noqa: C901 """Return false if an earlier rewrite invalidated this candidate.""" for inp, out in subgraph.edges_in: if inp.target not in PERMUTE_COPY_TARGETS or inp not in out.all_input_nodes: @@ -778,6 +915,32 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: if const_node not in user_node.all_input_nodes: return False + for input_node, user_node, _ in subgraph.input_boundaries: + if input_node not in user_node.all_input_nodes: + return False + future_occurrences = self._node_argument_count(user_node, input_node) + future_occurrences += sum( + self._node_argument_count(user_node, permute) + for permute, user in subgraph.edges_in + if user is user_node + and len(permute.args) >= 1 + and permute.args[0] is input_node + ) + if future_occurrences != 1: + return False + + for producer, output_node, _ in subgraph.output_boundaries: + if producer not in output_node.all_input_nodes: + return False + future_occurrences = self._node_argument_count(output_node, producer) + future_occurrences += sum( + self._node_argument_count(output_node, permute) + for source, permute in subgraph.edges_out + if source is producer + ) + if future_occurrences != 1: + return False + for head, (_, _, expand_node, view_node) in subgraph.interleaves.items(): if ( len(head.users) != 1 @@ -789,6 +952,127 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: return True + @staticmethod + def _node_argument_count(node: torch.fx.Node, target: torch.fx.Node) -> int: + count = 0 + + def visit(argument): + nonlocal count + if argument is target: + count += 1 + return argument + + torch.fx.map_arg((node.args, node.kwargs), visit) + return count + + def _boundary_permutations_are_layout_copies(self, subgraph: Subgraph) -> bool: + return all(is_layout_copy(permute) for permute, _ in subgraph.edges_in) and all( + is_layout_copy(permute) for _, permute in subgraph.edges_out + ) + + def _boundary_rewrite_is_cost_safe(self, subgraph: Subgraph) -> bool: + for const_node, user_node in subgraph.constant_edges_in: + node_end_perm = subgraph.node_end_permute.get( + user_node, subgraph.end_permute + ) + if self._constant_transform_requires_data_copy(const_node, node_end_perm): + return False + + removed_copies = { + permute + for permute, _ in subgraph.edges_in + if all(user in subgraph.nodes for user in permute.users) + } | {permute for _, permute in subgraph.edges_out} + new_copy_sources = { + (input_node, permutation) + for input_node, _, permutation in subgraph.input_boundaries + } | { + (producer, permutation) + for producer, _, permutation in subgraph.output_boundaries + } + + if not new_copy_sources: + return True + + removed_bytes = [self._static_tensor_bytes(node) for node in removed_copies] + new_bytes = [self._static_tensor_bytes(node) for node, _ in new_copy_sources] + if any(size is None for size in removed_bytes + new_bytes): + return False + return sum(cast(int, size) for size in new_bytes) <= sum( + cast(int, size) for size in removed_bytes + ) + + def _constant_transform_requires_data_copy( + self, node: torch.fx.Node, permutation: list[int] + ) -> bool: + val = node.meta.get("val") + if not isinstance(val, torch.Tensor): + return True + shape = list(val.shape) + if len(shape) > len(permutation) or not all( + isinstance(dim, int) for dim in shape + ): + return True + rank_difference = len(permutation) - len(shape) + padded_shape = [1] * rank_difference + shape + output_shape = [padded_shape[dim] for dim in permutation] + if any(size != 1 for size in output_shape[:rank_difference]): + return True + old_order = [dim for dim, size in enumerate(padded_shape) if size != 1] + new_order = [dim for dim, size in zip(permutation, output_shape) if size != 1] + return old_order != new_order + + @staticmethod + def _static_tensor_bytes(node: torch.fx.Node) -> int | None: + val = node.meta.get("val") + if not isinstance(val, torch.Tensor) or not all( + isinstance(dim, int) for dim in val.shape + ): + return None + return val.numel() * val.element_size() + + def _insert_input_boundary_permutations(self, subgraph: Subgraph) -> None: + if not subgraph.input_boundaries: + return + assert subgraph.layout_region + groups: dict[tuple[torch.fx.Node, tuple[int, ...]], list[torch.fx.Node]] = {} + for input_node, user_node, permutation in subgraph.input_boundaries: + groups.setdefault((input_node, permutation), []).append(user_node) + + graph = next(iter(subgraph.input_boundaries))[0].graph + node_order = {node: index for index, node in enumerate(graph.nodes)} + for (input_node, permutation), users in groups.items(): + first_user = min(users, key=node_order.__getitem__) + with input_node.graph.inserting_before(first_user): + new_permute = input_node.graph.call_function( + LAYOUT_PERMUTE_COPY, + args=(input_node, list(permutation)), + ) + new_permute.meta = dict(input_node.meta) + for user in users: + user.replace_input_with(input_node, new_permute) + + def _insert_output_boundary_permutations(self, subgraph: Subgraph) -> None: + if not subgraph.output_boundaries: + return + assert subgraph.layout_region + groups: dict[tuple[torch.fx.Node, tuple[int, ...]], list[torch.fx.Node]] = {} + for producer, output_node, permutation in subgraph.output_boundaries: + groups.setdefault((producer, permutation), []).append(output_node) + + graph = next(iter(subgraph.output_boundaries))[0].graph + node_order = {node: index for index, node in enumerate(graph.nodes)} + for (producer, permutation), outputs in groups.items(): + first_output = min(outputs, key=node_order.__getitem__) + with producer.graph.inserting_before(first_output): + new_permute = producer.graph.call_function( + LAYOUT_PERMUTE_COPY, + args=(producer, list(permutation)), + ) + new_permute.meta = dict(producer.meta) + for output in outputs: + output.replace_input_with(producer, new_permute) + def update_interleave( self, head: torch.fx.Node, diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index c774852cf63..fc6aec148d5 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -101,6 +101,79 @@ def get_compute_nodes( class FuseCascadedTransposeOrPermuteOpsTest(unittest.TestCase): + def test_structural_permute_composition_preserves_provenance(self) -> None: + for second_target, expected_target in ( + ( + exir_ops.edge.channels_last.permute_copy.default, + exir_ops.edge.channels_last.permute_copy.default, + ), + ( + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.permute_copy.default, + ), + ): + with self.subTest(second_target=second_target): + builder = GraphBuilder() + x_data = torch.randn(1, 2, 3, 4) + x = builder.placeholder("x", x_data) + first = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + second = builder.call_operator( + op=second_target, + args=(first, [0, 1, 3, 2]), + ) + builder.output([second]) + graph_module = builder.get_graph_module() + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + FuseCascadedTransposeOrPermuteOps()(graph_module), + ) + + self.assertTrue(result.modified) + self.assertEqual(count_node(result.graph_module, expected_target), 1) + validate_numerics( + before, + result.graph_module, + [x_data], + "FuseCascadedTransposeOrPermuteOps", + ) + + def test_ordinary_mixed_inverse_permutes_are_fused(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 2, 3, 4)) + first = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + second = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(first, [0, 3, 1, 2]), + ) + builder.output([second]) + graph_module = builder.get_graph_module() + + result = cast( + PassResult, + FuseCascadedTransposeOrPermuteOps()(graph_module), + ) + + self.assertTrue(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 0, + ) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), + 0, + ) + def test_permute_transpose_fusion(self) -> None: builder = GraphBuilder() x = builder.placeholder("x", torch.randn(3, 1, 3, 1, 4)) @@ -557,6 +630,7 @@ def test_negative_not_squeeze_like(self) -> None: class FuseTransposeOrPermuteOpPairsTest(unittest.TestCase): + def test_per_tensor_qdq_is_bypassed(self) -> None: for op, x_data in ( ( @@ -742,10 +816,579 @@ def test_per_channel_branch_blocks_shared_permute_fusion(self) -> None: # ────────────────────────────────────────────────────────────────────── -# Tests for ReplaceNopTransposeOrPermuteWithViewPass +# Tests for structural layout boundary propagation # ────────────────────────────────────────────────────────────────────── +class StructuralLayoutBoundaryPropagationTest(unittest.TestCase): + @staticmethod + def _layout_add_graph( + bias_name: str, bias_data: torch.Tensor + ) -> tuple[torch.fx.GraphModule, torch.Tensor]: + builder = GraphBuilder() + x_data = torch.randn(1, 8, 8, 4) + x = builder.placeholder("x", x_data) + bias = builder.placeholder(bias_name, bias_data) + to_nchw = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 3, 1, 2]), + ) + add = builder.call_operator( + op=exir_ops.edge.aten.add.Tensor, + args=(to_nchw, bias), + ) + to_nhwc = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(add, [0, 2, 3, 1]), + ) + builder.output([to_nhwc]) + return builder.get_graph_module(), x_data + + @staticmethod + def _layout_pad_graph( + shape: tuple[int, ...], + to_inner: list[int], + to_outer: list[int], + pad: list[int], + ) -> tuple[torch.fx.GraphModule, torch.Tensor]: + builder = GraphBuilder() + x_data = torch.randn(*shape) + x = builder.placeholder("x", x_data) + inner = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, to_inner), + ) + padded = builder.call_operator( + op=exir_ops.edge.aten.constant_pad_nd.default, + args=(inner, pad, 0.0), + ) + outer = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(padded, to_outer), + ) + builder.output([outer]) + return builder.get_graph_module(), x_data + + def test_layout_pad_argument_is_remapped(self) -> None: + for shape, to_inner, to_outer, pad in ( + ((1, 8, 8, 3), [0, 3, 1, 2], [0, 2, 3, 1], [0, 0, 0, 0, 0, 1]), + ((2, 8, 3), [0, 2, 1], [0, 2, 1], [0, 0, 0, 1]), + ): + with self.subTest(shape=shape): + graph_module, x_data = self._layout_pad_graph( + shape, to_inner, to_outer, pad + ) + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps( + allow_layout_boundary_propagation=True, + )(graph_module), + ) + + self.assertTrue(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.aten.constant_pad_nd.default, + ), + 1, + ) + validate_numerics( + before, + result.graph_module, + [x_data], + "RemovePermutesAroundElementwiseOps", + ) + + def test_existing_layout_pad_is_remapped(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 3, 8, 8)) + pad = builder.call_operator( + op=exir_ops.edge.aten.constant_pad_nd.default, + args=(x, [0, 0, 0, 0, 0, 1], 0.0), + ) + builder.output([pad]) + + RemovePermutesAroundElementwiseOps().update_pad(pad.node, [0, 3, 1, 2]) + + self.assertEqual(pad.node.args[1], [0, 1]) + + def test_pair_fusion_recognizes_structural_permutes(self) -> None: + builder = GraphBuilder() + x_data = torch.randn(1, 2, 3, 4) + x = builder.placeholder("x", x_data) + to_nhwc = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + quantize = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(to_nhwc, 0.25, 0, -128, 127, torch.int8), + ) + to_nchw = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(quantize, [0, 3, 1, 2]), + ) + builder.output([to_nchw]) + graph_module = builder.get_graph_module() + before = copy.deepcopy(graph_module) + + result = cast(PassResult, FuseTransposeOrPermuteOpPairsPass()(graph_module)) + + self.assertTrue(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 0, + ) + validate_numerics( + before, + result.graph_module, + [x_data], + "FuseTransposeOrPermuteOpPairsPass", + ) + + def test_pair_fusion_respects_backend_propagation_barrier(self) -> None: + builder = GraphBuilder() + x_data = torch.randn(1, 2, 3, 4) + x = builder.placeholder("x", x_data) + to_nhwc = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + quantize = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(to_nhwc, 0.25, 0, -128, 127, torch.int8), + ) + to_nchw = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(quantize, [0, 3, 1, 2]), + ) + builder.output([to_nchw]) + graph_module = builder.get_graph_module() + + result = cast( + PassResult, + FuseTransposeOrPermuteOpPairsPass( + can_propagate=lambda node: node.target + != exir_ops.edge.quantized_decomposed.quantize_per_tensor.default + )(graph_module), + ) + + self.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 2, + ) + + def test_pair_fusion_does_not_bypass_structural_per_channel_qdq(self) -> None: + for op, x_data in ( + ( + exir_ops.edge.quantized_decomposed.quantize_per_channel.default, + torch.randn(1, 2, 3, 4), + ), + ( + exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, + torch.randint(-128, 127, (1, 2, 3, 4), dtype=torch.int8), + ), + ): + with self.subTest(op=op): + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + scales = builder.placeholder("scales", torch.tensor([0.25, 0.5])) + zero_points = builder.placeholder( + "zero_points", torch.tensor([0, 0], dtype=torch.int64) + ) + to_nhwc = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + qdq = builder.call_operator( + op=op, + args=(to_nhwc, scales, zero_points, 3, -128, 127, torch.int8), + ) + to_nchw = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(qdq, [0, 3, 1, 2]), + ) + builder.output([to_nchw]) + graph_module = builder.get_graph_module() + + result = cast( + PassResult, FuseTransposeOrPermuteOpPairsPass()(graph_module) + ) + + self.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 2, + ) + + def test_layout_copy_moves_to_static_output_boundary(self) -> None: + builder = GraphBuilder() + x_data = torch.randn(1, 2, 3, 4) + x = builder.placeholder("x", x_data) + permute = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + output = builder.call_operator( + op=exir_ops.edge.aten.hardtanh.default, + args=(permute,), + ) + builder.output([output]) + graph_module = builder.get_graph_module() + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( + graph_module + ), + ) + + self.assertTrue(result.modified) + surviving_permute = result.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.channels_last.permute_copy.default, + )[0] + self.assertEqual( + surviving_permute.args[0].target, + exir_ops.edge.aten.hardtanh.default, + ) + validate_numerics( + before, + result.graph_module, + [x_data], + "RemovePermutesAroundElementwiseOps", + ) + + def test_layout_region_terminates_at_backend_barrier(self) -> None: + builder = GraphBuilder() + x_data = torch.randn(1, 2, 3, 4) + x = builder.placeholder("x", x_data) + to_nhwc = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + activation = builder.call_operator( + op=exir_ops.edge.aten.hardtanh.default, + args=(to_nhwc,), + ) + to_nchw = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(activation, [0, 3, 1, 2]), + ) + barrier = builder.call_operator( + op=exir_ops.edge.aten.relu.default, + args=(activation,), + ) + builder.output([to_nchw, barrier]) + graph_module = builder.get_graph_module() + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps( + allow_layout_boundary_propagation=True, + can_propagate=lambda node: node.target + != exir_ops.edge.aten.relu.default, + )(graph_module), + ) + + self.assertTrue(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 1, + ) + validate_numerics( + before, + result.graph_module, + [x_data], + "RemovePermutesAroundElementwiseOps", + ) + + def test_layout_copy_does_not_fork_to_more_boundaries(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 2, 3, 4)) + permute = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + first = builder.call_operator( + op=exir_ops.edge.aten.hardtanh.default, + args=(permute,), + ) + second = builder.call_operator( + op=exir_ops.edge.aten.mul.Tensor, + args=(permute, permute), + ) + builder.output([first, second]) + graph_module = builder.get_graph_module() + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( + graph_module + ), + ) + + self.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 1, + ) + + def test_shared_incoming_layout_copy_is_not_credited_as_removed(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 2, 3, 4)) + permute = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + main = builder.call_operator( + op=exir_ops.edge.aten.hardtanh.default, + args=(permute,), + ) + auxiliary = builder.call_operator( + op=exir_ops.edge.aten._softmax.default, + args=(permute, -1, False), + ) + builder.output([main, auxiliary]) + graph_module = builder.get_graph_module() + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( + graph_module + ), + ) + + self.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 1, + ) + + def test_input_boundary_does_not_alias_permuted_and_unpermuted_input(self) -> None: + builder = GraphBuilder() + x_data = torch.randn(1, 2, 2, 2) + x = builder.placeholder("x", x_data) + permute = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + add = builder.call_operator( + op=exir_ops.edge.aten.add.Tensor, + args=(permute, x), + ) + builder.output([add]) + graph_module = builder.get_graph_module() + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( + graph_module + ), + ) + + self.assertFalse(result.modified) + validate_numerics( + before, + result.graph_module, + [x_data], + "RemovePermutesAroundElementwiseOps", + ) + + def test_output_boundary_does_not_alias_two_output_edges(self) -> None: + builder = GraphBuilder() + x_data = torch.randn(1, 2, 3, 4) + x = builder.placeholder("x", x_data) + to_nhwc = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + activation = builder.call_operator( + op=exir_ops.edge.aten.hardtanh.default, + args=(to_nhwc,), + ) + to_nchw = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(activation, [0, 3, 1, 2]), + ) + builder.output([to_nchw, activation]) + graph_module = builder.get_graph_module() + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( + graph_module + ), + ) + + self.assertFalse(result.modified) + validate_numerics( + before, + result.graph_module, + [x_data], + "RemovePermutesAroundElementwiseOps", + ) + + def test_layout_copy_does_not_cross_unknown_cost_boundary(self) -> None: + class DynamicDequantize(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + permuted = torch.ops.channels_last.permute_copy(x, [0, 2, 3, 1]) + return torch.ops.quantized_decomposed.dequantize_per_tensor.default( + permuted, 0.1, 0, -128, 127, torch.int8 + ) + + inputs = (torch.randint(-128, 127, (1, 4, 8, 10), dtype=torch.int8),) + exported = torch.export.export( + DynamicDequantize(), + inputs, + dynamic_shapes={"x": {2: torch.export.Dim("height", min=2, max=16)}}, + ) + edge = to_edge( + exported, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=True, + ), + ) + graph_module = edge.exported_program().graph_module + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps( + exported_program=edge.exported_program(), + allow_layout_boundary_propagation=True, + )(graph_module), + ) + + self.assertFalse(result.modified) + permute = result.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.channels_last.permute_copy.default, + )[0] + dequantize = result.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + )[0] + self.assertIs(dequantize.args[0], permute) + self.assertEqual(permute.meta["val"].dtype, torch.int8) + + def test_layout_copy_rejects_rank_mismatched_runtime_input(self) -> None: + bias_data = torch.randn(4, 1, 1) + graph_module, x_data = self._layout_add_graph("bias", bias_data) + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( + graph_module + ), + ) + + self.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 2, + ) + validate_numerics( + before, + result.graph_module, + [x_data, bias_data], + "RemovePermutesAroundElementwiseOps", + ) + + def test_layout_copy_rejects_spatial_constant_reordering(self) -> None: + bias_data = torch.randn(4, 8, 8) + graph_module, x_data = self._layout_add_graph("b_bias", bias_data) + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( + graph_module + ), + ) + + self.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 2, + ) + validate_numerics( + before, + result.graph_module, + [x_data, bias_data], + "RemovePermutesAroundElementwiseOps", + ) + + def test_layout_copy_reshapes_channel_constant_without_copy(self) -> None: + bias_data = torch.randn(4, 1, 1) + graph_module, x_data = self._layout_add_graph("b_bias", bias_data) + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( + graph_module + ), + ) + + self.assertTrue(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 0, + ) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.view_copy.default), + 1, + ) + validate_numerics( + before, + result.graph_module, + [x_data, bias_data], + "RemovePermutesAroundElementwiseOps", + ) + + +# ───────────────────────────────────── +# Tests for ReplaceNopTransposeOrPermuteWithViewPass +# ───────────────────────────────────── + + class ReplaceNopTransposeOrPermuteWithViewTest(unittest.TestCase): def test_replace_nop_transpose_with_view_float(self) -> None: x = torch.randn(2, 1, 3, 1)