From 8ecb850ad9d78f681833633a92d22f209d24fd50 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 19 Aug 2026 21:50:34 -0700 Subject: [PATCH 1/2] Update [ghstack-poisoned] --- ...replace_ops_with_channels_last_variants.py | 7 ++++++- ...replace_ops_with_channels_last_variants.py | 20 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/backends/transforms/replace_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py index 5540fe55f22..5e8a4612163 100644 --- a/backends/transforms/replace_ops_with_channels_last_variants.py +++ b/backends/transforms/replace_ops_with_channels_last_variants.py @@ -133,6 +133,11 @@ class ReplaceOpsWithChannelsLastVariants(ExportPass): By default, all currently implemented channels_last dialect ops are replaced. Pass a custom op_map to restrict or extend the set of replacements. + + Metadata from each replaced operator is preserved so provenance and backend + annotations survive the rewrite. ExportPass recomputes shape metadata after + retracing. Callers must reject or remap semantic metadata tied to dimensions + changed by ``input_indices`` or ``output_indices``, such as per-channel axes. """ def __init__( @@ -229,7 +234,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: args=tuple(args), kwargs=node.kwargs, ) - nhwc_node.meta = {} + nhwc_node.meta = dict(node.meta) users = list(node.users) if all( diff --git a/backends/transforms/test/test_replace_ops_with_channels_last_variants.py b/backends/transforms/test/test_replace_ops_with_channels_last_variants.py index e131463320e..3c8fdd8309f 100644 --- a/backends/transforms/test/test_replace_ops_with_channels_last_variants.py +++ b/backends/transforms/test/test_replace_ops_with_channels_last_variants.py @@ -132,6 +132,26 @@ def forward(self, x): class TestReplaceOpsWithChannelsLastVariants: + def test_preserves_metadata_and_recomputes_shape(self): + ep = _export_to_edge(Conv2dModule(), (torch.randn(1, 4, 8, 8),)) + conv = _find_nodes(ep.graph_module, exir_ops.edge.aten.convolution.default)[0] + metadata = { + "debug_handle": 1234, + "from_node": [("source", "convolution")], + "input_qparams": {0: "input"}, + "output_qparams": {0: "output"}, + } + conv.meta.update(metadata) + + result = ReplaceOpsWithChannelsLastVariants(ep)(ep.graph_module) + replaced = _find_nodes( + result.graph_module, exir_ops.edge.channels_last.convolution.default + )[0] + + for key, value in metadata.items(): + assert replaced.meta[key] == value + assert tuple(replaced.meta["val"].shape) == (1, 8, 8, 4) + def test_conv2d(self): ep = _export_to_edge(Conv2dModule(bias=True), (torch.randn(1, 4, 8, 8),)) assert _count(ep.graph_module, exir_ops.edge.aten.convolution.default) == 1 From cdf9043cf019e19660763cf52e770741097574c9 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 19 Aug 2026 21:50:40 -0700 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- ...ve_permutes_around_elementwise_tosa_ops.py | 4 +- backends/transforms/channels_last_layout.py | 50 ++ backends/transforms/channels_last_ops.py | 11 + .../decompose_channels_last_pass.py | 13 +- .../fuse_cascaded_transpose_or_permute_ops.py | 15 +- ...fuse_transpose_or_permute_op_pairs_pass.py | 50 +- backends/transforms/permute_pass_utils.py | 3 +- .../postpone_permute_below_squeeze_view.py | 5 +- .../remove_permutes_around_elementwise_ops.py | 357 +++++++- ...lace_nop_transpose_or_permute_with_view.py | 5 +- ...replace_ops_with_channels_last_variants.py | 5 +- backends/transforms/targets.bzl | 23 + .../transforms/test/test_channels_last_ops.py | 10 + .../test/test_decompose_channels_last_pass.py | 12 + .../test/test_permute_optimization_passes.py | 759 +++++++++++++++++- 15 files changed, 1280 insertions(+), 42 deletions(-) create mode 100644 backends/transforms/channels_last_layout.py 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_layout.py b/backends/transforms/channels_last_layout.py new file mode 100644 index 00000000000..65dd2fd00ee --- /dev/null +++ b/backends/transforms/channels_last_layout.py @@ -0,0 +1,50 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import executorch.backends.transforms.channels_last_ops # noqa: F401 + +import torch + +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx.node import Target + +ATEN_PERMUTE_COPY = exir_ops.edge.aten.permute_copy.default +LAYOUT_PERMUTE_COPY = exir_ops.edge.channels_last.permute_copy.default +PERMUTE_COPY_TARGETS: frozenset[Target] = frozenset( + (ATEN_PERMUTE_COPY, LAYOUT_PERMUTE_COPY) +) + + +def is_permute_copy(node: torch.fx.Node) -> bool: + return node.op == "call_function" and node.target in PERMUTE_COPY_TARGETS + + +def is_layout_copy(node: torch.fx.Node) -> bool: + return node.op == "call_function" and node.target == LAYOUT_PERMUTE_COPY + + +def is_channels_last_input_normalization_pair( + first: torch.fx.Node, second: torch.fx.Node +) -> bool: + if first.target != ATEN_PERMUTE_COPY or second.target != LAYOUT_PERMUTE_COPY: + return False + input_node = first.args[0] if first.args else None + val = input_node.meta.get("val") if isinstance(input_node, torch.fx.Node) else None + return ( + input_node is not None + and input_node.op == "placeholder" + and isinstance(val, torch.Tensor) + and val.dim() == 4 + and tuple(val.dim_order()) == (0, 2, 3, 1) + and list(first.args[1]) == [0, 2, 3, 1] + and list(second.args[1]) == [0, 3, 1, 2] + ) + + +def composed_permute_target(first: torch.fx.Node, second: torch.fx.Node) -> Target: + if is_layout_copy(first) and is_layout_copy(second): + return LAYOUT_PERMUTE_COPY + return ATEN_PERMUTE_COPY diff --git a/backends/transforms/channels_last_ops.py b/backends/transforms/channels_last_ops.py index 0a90f0d1889..9a8d94f859f 100644 --- a/backends/transforms/channels_last_ops.py +++ b/backends/transforms/channels_last_ops.py @@ -11,6 +11,9 @@ opposed to the implicit dim-order handling used elsewhere. They let layout-handling passes (see RFC #19299) make channels-last regions explicit in the graph. +The ``constant_pad_nd`` padding list follows the contiguous NHWC tensor axes, +last axis first, after the layout transform has remapped it from NCHW. + Efficiency is a non-goal: kernels are implemented as ``permute -> aten op -> permute``. Importing this module registers the dialect. """ @@ -112,6 +115,10 @@ def _max_pool2d( return out.permute(0, 2, 3, 1).contiguous() +def _constant_pad_nd(input, pad, value=0): + return torch.ops.aten.constant_pad_nd(input, pad, value) + + def _grid_sampler_2d(input, grid, interpolation_mode, padding_mode, align_corners): nchw = input.permute(0, 3, 1, 2) out = torch.ops.aten.grid_sampler_2d( @@ -175,6 +182,10 @@ def _permute_copy(input, dims): lib.impl("max_pool2d", _max_pool2d, "CompositeExplicitAutograd") register_fake("channels_last::max_pool2d", _max_pool2d, lib=lib) +lib.define("constant_pad_nd(Tensor input, SymInt[] pad, Scalar value=0) -> Tensor") +lib.impl("constant_pad_nd", _constant_pad_nd, "CompositeExplicitAutograd") +register_fake("channels_last::constant_pad_nd", _constant_pad_nd, 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..477fa9fb552 100644 --- a/backends/transforms/decompose_channels_last_pass.py +++ b/backends/transforms/decompose_channels_last_pass.py @@ -27,6 +27,11 @@ exir_ops.edge.channels_last.grid_sampler_2d.default: exir_ops.edge.aten.grid_sampler_2d.default, } +_DIRECT_DECOMPOSITIONS = { + exir_ops.edge.channels_last.constant_pad_nd.default: exir_ops.edge.aten.constant_pad_nd.default, + 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 +44,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 +99,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_cascaded_transpose_or_permute_ops.py b/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py index f350120e7eb..ab75a6e0b5f 100644 --- a/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py +++ b/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py @@ -6,6 +6,11 @@ # pyre-unsafe +from executorch.backends.transforms.channels_last_layout import ( + composed_permute_target, + is_channels_last_input_normalization_pair, + PERMUTE_COPY_TARGETS, +) from executorch.backends.transforms.permute_pass_utils import ( get_arg, get_permuted_dims, @@ -26,7 +31,7 @@ class FuseCascadedTransposeOrPermuteOps(RemoveOrReplacePassInterface): transpose_or_permute_target = { exir_ops.edge.aten.transpose_copy.int, - exir_ops.edge.aten.permute_copy.default, + *PERMUTE_COPY_TARGETS, } _VIEW_OPS = { @@ -67,14 +72,16 @@ def _fuse_direct(self, node: Node, parent_node: Node) -> bool: dims = get_permuted_dims(node, dims) if dims == sorted(dims): + if is_channels_last_input_normalization_pair(parent_node, node): + return False node.replace_all_uses_with(input_of_parent) else: with node.graph.inserting_before(node): new_permute = node.graph.call_function( - exir_ops.edge.aten.permute_copy.default, + composed_permute_target(parent_node, node), args=(input_of_parent, dims), ) - new_permute.meta = node.meta + new_permute.meta = dict(node.meta) node.replace_all_uses_with(new_permute) return True @@ -141,7 +148,7 @@ def _fuse_across_view(self, node: Node, view_node: Node) -> bool: # noqa: C901 node_dims = list(range(len(dims))) node_dims = get_transposed_dims(node, node_dims) dims = [dims[d] for d in node_dims] - elif node.target == exir_ops.edge.aten.permute_copy.default: + elif node.target in PERMUTE_COPY_TARGETS: perm = get_arg(node, "dims") dims = [dims[d] for d in perm] else: 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 7270736f5bd..24c56ed2552 100644 --- a/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py +++ b/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py @@ -6,10 +6,16 @@ # pyre-unsafe +from collections import deque from typing import Any, Callable, cast import torch import torch.fx +from executorch.backends.transforms.channels_last_layout import ( + ATEN_PERMUTE_COPY, + is_channels_last_input_normalization_pair, + LAYOUT_PERMUTE_COPY, +) from executorch.backends.transforms.permute_pass_utils import ( FuseOpPairsAcrossBranchesPass, get_permuted_dims, @@ -37,6 +43,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, @@ -45,6 +87,8 @@ def can_fuse_for_chain( ) -> bool: if not super().can_fuse_for_chain(producer, consumer, consumer_op_packets): return False + if is_channels_last_input_normalization_pair(producer, consumer): + return False # checking that permut2(permut1(identity)) == identity, modulo unitary dimensions producer_input = cast(torch.fx.Node, producer.args[0]) @@ -55,7 +99,8 @@ def can_fuse_for_chain( # this mapping helps to handle both transpose and permutations f: dict[Any, Callable] = { exir_ops.edge.aten.transpose_copy.int: get_transposed_dims, - exir_ops.edge.aten.permute_copy.default: get_permuted_dims, + ATEN_PERMUTE_COPY: get_permuted_dims, + LAYOUT_PERMUTE_COPY: get_permuted_dims, } in_dims = f[producer.target](producer, ident_dims) out_dims = f[consumer.target](consumer, in_dims) @@ -80,6 +125,7 @@ def get_fused_node( (consumer.args[0], output_shape), {}, ) + view.meta = dict(consumer.meta) return view def call(self, graph_module: torch.fx.GraphModule) -> PassResult: @@ -89,10 +135,12 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: producer_op_packets={ exir_ops.edge.aten.transpose_copy, exir_ops.edge.aten.permute_copy, + exir_ops.edge.channels_last.permute_copy, }, consumer_op_packets={ exir_ops.edge.aten.transpose_copy, exir_ops.edge.aten.permute_copy, + exir_ops.edge.channels_last.permute_copy, }, bypass_ops=self.bypass_ops, ) diff --git a/backends/transforms/permute_pass_utils.py b/backends/transforms/permute_pass_utils.py index fca8946165e..97588beb423 100644 --- a/backends/transforms/permute_pass_utils.py +++ b/backends/transforms/permute_pass_utils.py @@ -18,6 +18,7 @@ import torch import torch.fx +from executorch.backends.transforms.channels_last_layout import PERMUTE_COPY_TARGETS from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.dialects.edge._ops import EdgeOpOverload, EdgeOpOverloadPacket from executorch.exir.pass_base import ExportPass, PassResult @@ -76,7 +77,7 @@ def get_transposed_dims( def get_permuted_dims(node: torch.fx.Node, dims: List[int]) -> List[int]: """Applies the permutation as given by node onto the dimensions given in input.""" - assert node.target == exir_ops.edge.aten.permute_copy.default + assert node.target in PERMUTE_COPY_TARGETS # pyre-fixme[6]: This combined typecheck isn't supported yet. permute_dims: List[int] = list(node.args[1]) assert all(isinstance(x, int) for x in permute_dims) diff --git a/backends/transforms/postpone_permute_below_squeeze_view.py b/backends/transforms/postpone_permute_below_squeeze_view.py index e0e9a3ec198..226924cbea7 100644 --- a/backends/transforms/postpone_permute_below_squeeze_view.py +++ b/backends/transforms/postpone_permute_below_squeeze_view.py @@ -11,6 +11,7 @@ import torch import torch.fx +from executorch.backends.transforms.channels_last_layout import PERMUTE_COPY_TARGETS from executorch.backends.transforms.permute_pass_utils import ( get_shape, RemoveOrReplacePassInterface, @@ -36,7 +37,7 @@ class PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView(RemoveOrReplacePassInterf @property def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.permute_copy.default] + return list(PERMUTE_COPY_TARGETS) # If list1 and list2 are same (same values and in same order) except # list1 has one more element with value of 1. Return index of the extra 1. @@ -182,7 +183,7 @@ def _insert_nodes( permute_target, args=(new_view_node, new_permute_dims), ) - new_permute_node.meta = view_node.meta + new_permute_node.meta = dict(view_node.meta) view_node.replace_all_uses_with(new_permute_node) # view_node is user of permute_node, so must erase view_node first diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 6e916dfe50a..01d387a4cc7 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -7,14 +7,24 @@ # 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 +from torch.fx.node import Target class RemovePermutesAroundElementwiseOps(ExportPass): @@ -25,6 +35,11 @@ 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. + ``layout_pad_target`` opts into retargeting rank-4 constant pads after their + layout-dependent pad argument has been remapped. """ @dataclass() @@ -41,6 +56,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) @@ -52,8 +75,20 @@ 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, + layout_pad_target: Target | None = None, + 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.layout_pad_target = layout_pad_target + self.can_propagate = can_propagate self._permutable_ops = { exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.mul.Tensor, @@ -114,6 +149,7 @@ def _view_shapes(self, node: torch.fx.Node) -> tuple[list[int], list[int]] | Non return in_shape, out_shape _PAD_OPS = ( + exir_ops.edge.channels_last.constant_pad_nd.default, exir_ops.edge.aten.constant_pad_nd.default, exir_ops.edge.aten.pad.default, ) @@ -325,15 +361,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.find_nodes( - op="call_function", target=exir_ops.edge.aten.permute_copy.default - ): + 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 ( @@ -341,7 +401,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: @@ -371,7 +435,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 @@ -391,6 +459,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): @@ -483,7 +576,7 @@ def visit( # noqa: C901 # Traverse downstream: for user in users_source.users: - if user.target == exir_ops.edge.aten.permute_copy.default: + if user.target in PERMUTE_COPY_TARGETS: user_perm = self.get_permutation(user) if user_perm == downstream_end: subgraph.edges_out.add((users_source, user)) @@ -512,7 +605,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 @@ -521,6 +621,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 ): @@ -528,7 +637,7 @@ def visit( # noqa: C901 # Traverse upstream: for inp in node.all_input_nodes: - if inp.target == exir_ops.edge.aten.permute_copy.default: + if inp.target in PERMUTE_COPY_TARGETS: if self.get_permutation(inp) != current_start_permute: return False subgraph.edges_in.add((inp, node)) @@ -549,6 +658,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, @@ -598,6 +723,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") @@ -614,6 +746,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: @@ -650,6 +784,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 @@ -698,8 +837,13 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 self.update_mean_dim(node, node_start_perm) elif node.target == exir_ops.edge.aten.slice_copy.Tensor: self.update_slice_copy(node, node_start_perm) + elif node.target in ( + exir_ops.edge.aten._softmax.default, + exir_ops.edge.aten.softmax.int, + ): + self.update_dim(node, node_start_perm) elif node.target in self._PAD_OPS: - self.update_pad(node, node_start_perm) + self.update_pad(node, node_start_perm, subgraph.layout_region) elif node.target in self._VIEW_OPS: self.update_view_copy(node, node_start_perm) @@ -712,7 +856,7 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 # Skip incoming permutes. for inp, out in subgraph.edges_in: - assert inp.target == exir_ops.edge.aten.permute_copy.default + assert inp.target in PERMUTE_COPY_TARGETS if len(inp.args) >= 1: out.replace_input_with(inp, cast(torch.fx.Node, inp.args[0])) else: @@ -732,9 +876,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 @@ -755,31 +904,54 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 # Skip outgoing permutes. for inp, out in subgraph.edges_out: - assert out.target == exir_ops.edge.aten.permute_copy.default + 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 != exir_ops.edge.aten.permute_copy.default - or inp not in out.all_input_nodes - ): + if inp.target not in PERMUTE_COPY_TARGETS or inp not in out.all_input_nodes: return False for inp, out in subgraph.edges_out: - if ( - out.target != exir_ops.edge.aten.permute_copy.default - or out not in inp.users - ): + if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users: return False for const_node, user_node in subgraph.constant_edges_in: 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 @@ -791,6 +963,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, @@ -837,7 +1130,16 @@ def update_slice_copy(self, node: torch.fx.Node, start_permute: list[int]) -> No dim = get_arg(node, "dim", int) set_arg(node, "dim", start_permute[dim]) - def update_pad(self, node: torch.fx.Node, start_permute: list[int]) -> None: + def update_dim(self, node: torch.fx.Node, start_permute: list[int]) -> None: + dim = get_arg(node, "dim", int) % len(start_permute) + set_arg(node, "dim", start_permute[dim]) + + def update_pad( + self, + node: torch.fx.Node, + start_permute: list[int], + layout_region: bool, + ) -> None: pad = list(cast(list[int], node.args[1])) rank = len(start_permute) pad_pairs = [[0, 0] for _ in range(rank)] @@ -854,6 +1156,13 @@ def update_pad(self, node: torch.fx.Node, start_permute: list[int]) -> None: remapped_pad = remapped_pad[:-2] node.update_arg(1, remapped_pad) + if ( + layout_region + and len(start_permute) == 4 + and node.target == exir_ops.edge.aten.constant_pad_nd.default + and self.layout_pad_target is not None + ): + node.target = self.layout_pad_target def update_view_copy(self, node: torch.fx.Node, start_permute: list[int]) -> None: """Adjust view_copy shape arg after permute removal. @@ -892,7 +1201,7 @@ def update_view_copy(self, node: torch.fx.Node, start_permute: list[int]) -> Non node.update_arg(1, new_shape) def get_permutation(self, permute_node: torch.fx.Node) -> list[int] | None: - assert permute_node.target == exir_ops.edge.aten.permute_copy.default + assert permute_node.target in PERMUTE_COPY_TARGETS raw_permute: list[int] if len(permute_node.args) >= 2: raw_permute = list(cast(list[int], permute_node.args[1])) diff --git a/backends/transforms/replace_nop_transpose_or_permute_with_view.py b/backends/transforms/replace_nop_transpose_or_permute_with_view.py index ccfb4ebe8b9..aeff987954b 100644 --- a/backends/transforms/replace_nop_transpose_or_permute_with_view.py +++ b/backends/transforms/replace_nop_transpose_or_permute_with_view.py @@ -10,6 +10,7 @@ import torch import torch.fx +from executorch.backends.transforms.channels_last_layout import PERMUTE_COPY_TARGETS from executorch.backends.transforms.permute_pass_utils import ( RemoveOrReplacePassInterface, ) @@ -28,7 +29,7 @@ class ReplaceNopTransposeOrPermuteWithViewPass(RemoveOrReplacePassInterface): def targets(self) -> list[EdgeOpOverload]: return [ exir_ops.edge.aten.transpose_copy.int, - exir_ops.edge.aten.permute_copy.default, + *PERMUTE_COPY_TARGETS, ] def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: @@ -61,7 +62,7 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: node.replace_all_uses_with(new_node) return True - elif node.target == exir_ops.edge.aten.permute_copy.default: + elif node.target in PERMUTE_COPY_TARGETS: old_dims = list(range(len(in_shape))) new_dims = cast(Sequence[int], node.args[1]) # If the permute does not change anything, return the input as output. diff --git a/backends/transforms/replace_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py index 5e8a4612163..641c3925b65 100644 --- a/backends/transforms/replace_ops_with_channels_last_variants.py +++ b/backends/transforms/replace_ops_with_channels_last_variants.py @@ -11,6 +11,7 @@ import torch +from executorch.backends.transforms.channels_last_layout import LAYOUT_PERMUTE_COPY from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops @@ -166,7 +167,7 @@ def _permute_node_input( res = graph.create_node( "call_function", - target=exir_ops.edge.channels_last.permute_copy.default, + target=LAYOUT_PERMUTE_COPY, args=(node_input, _NCHW_TO_NHWC_PERM), ) res.meta = {} @@ -182,7 +183,7 @@ def _permute_node_output( ): output = graph.create_node( "call_function", - target=exir_ops.edge.channels_last.permute_copy.default, + target=LAYOUT_PERMUTE_COPY, args=(node_output, _NHWC_TO_NCHW_PERM), ) output.meta = {} diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index f72d847a8ef..4283c949e13 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -221,6 +221,21 @@ def define_common_targets(): ], ) + runtime.python_library( + name = "channels_last_layout", + srcs = [ + "channels_last_layout.py", + ], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + ":channels_last_ops", + "//executorch/exir/dialects:lib", + ], + ) + runtime.python_library( name = "decompose_channels_last_pass", srcs = [ @@ -375,6 +390,7 @@ def define_common_targets(): "//executorch/backends/...", ], deps = [ + ":channels_last_layout", "//caffe2:torch", "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", @@ -388,6 +404,7 @@ def define_common_targets(): "//executorch/backends/...", ], deps = [ + ":channels_last_layout", "//caffe2:torch", "//executorch/exir/dialects:lib", ":permute_pass_utils", @@ -414,6 +431,7 @@ def define_common_targets(): "//executorch/backends/...", ], deps = [ + ":channels_last_layout", "//caffe2:torch", "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", @@ -429,8 +447,10 @@ def define_common_targets(): "@EXECUTORCH_CLIENTS", ], deps = [ + ":channels_last_layout", ":permute_pass_utils", "//caffe2:torch", + "//executorch/exir:lib", "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", ], @@ -443,6 +463,7 @@ def define_common_targets(): "//executorch/backends/...", ], deps = [ + ":channels_last_layout", "//caffe2:torch", "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", @@ -457,6 +478,7 @@ def define_common_targets(): "//executorch/backends/...", ], deps = [ + ":channels_last_layout", "//caffe2:torch", "//executorch/exir/dialects:lib", ":permute_pass_utils", @@ -508,6 +530,7 @@ def define_common_targets(): ], deps = [ "//caffe2:torch", + ":channels_last_layout", ":channels_last_ops", "//executorch/exir:pass_base", "//executorch/exir:lib", diff --git a/backends/transforms/test/test_channels_last_ops.py b/backends/transforms/test/test_channels_last_ops.py index db5e42c317d..e53ab101382 100644 --- a/backends/transforms/test/test_channels_last_ops.py +++ b/backends/transforms/test/test_channels_last_ops.py @@ -207,6 +207,16 @@ def test_max_pool2d_matches_aten(): assert torch.equal(actual, expected) +def test_constant_pad_nd_uses_nhwc_axis_order(): + nchw = torch.arange(2 * 3 * 4 * 5, dtype=torch.float32).reshape(2, 3, 4, 5) + nhwc = _to_nhwc(nchw) + + expected = _to_nhwc(torch.ops.aten.constant_pad_nd(nchw, [1, 2, 3, 0], -0.5)) + actual = torch.ops.channels_last.constant_pad_nd(nhwc, [0, 0, 1, 2, 3, 0], -0.5) + + torch.testing.assert_close(actual, expected) + + def test_grid_sampler_2d_matches_aten(): torch.manual_seed(0) nchw = torch.randn(2, 3, 8, 8) diff --git a/backends/transforms/test/test_decompose_channels_last_pass.py b/backends/transforms/test/test_decompose_channels_last_pass.py index 9df0567c99f..1971bf7b9a1 100644 --- a/backends/transforms/test/test_decompose_channels_last_pass.py +++ b/backends/transforms/test/test_decompose_channels_last_pass.py @@ -88,6 +88,11 @@ def forward(self, x): ) +class _ConstantPadModule(torch.nn.Module): + def forward(self, x): + return torch.ops.channels_last.constant_pad_nd(x, [0, 0, 1, 2, 3, 0], -0.5) + + class _PermuteModule(torch.nn.Module): def forward(self, x): return torch.ops.channels_last.permute_copy(x, [0, 3, 1, 2]) @@ -157,6 +162,13 @@ def forward(self, x): exir_ops.edge.aten.max_pool2d_with_indices.default, 2, ), + ( + _ConstantPadModule(), + (torch.randn(2, 4, 5, 3),), + exir_ops.edge.channels_last.constant_pad_nd.default, + exir_ops.edge.aten.constant_pad_nd.default, + 0, + ), ( _PermuteModule(), (torch.randn(2, 8, 8, 3),), diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index cf4583357d7..fc0968c9740 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -99,6 +99,112 @@ 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_channels_last_input_normalization_pair_is_preserved(self) -> None: + builder = GraphBuilder() + x_data = torch.randn(1, 2, 3, 4).to(memory_format=torch.channels_last) + x = builder.placeholder("x", x_data) + first = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + second = builder.call_operator( + op=exir_ops.edge.channels_last.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.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 1, + ) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), + 1, + ) + + 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)) @@ -555,6 +661,40 @@ def test_negative_not_squeeze_like(self) -> None: class FuseTransposeOrPermuteOpPairsTest(unittest.TestCase): + def test_channels_last_input_normalization_pair_is_preserved(self) -> None: + builder = GraphBuilder() + x_data = torch.randn(1, 2, 3, 4).to(memory_format=torch.channels_last) + x = builder.placeholder("x", x_data) + to_nhwc = builder.call_operator( + op=exir_ops.edge.aten.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()(graph_module)) + + self.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 1, + ) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), + 1, + ) + def test_per_tensor_qdq_is_bypassed(self) -> None: for op, x_data in ( ( @@ -740,10 +880,627 @@ 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_retarget_is_opt_in(self) -> None: + graph_module, x_data = self._layout_pad_graph( + (1, 8, 8, 3), + [0, 3, 1, 2], + [0, 2, 3, 1], + [0, 0, 0, 0, 0, 1], + ) + 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, + ) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.constant_pad_nd.default, + ), + 0, + ) + validate_numerics( + before, + result.graph_module, + [x_data], + "RemovePermutesAroundElementwiseOps", + ) + + def test_layout_pad_retarget_requires_rank_four(self) -> None: + for shape, to_inner, to_outer, pad, expected_target in ( + ( + (1, 8, 8, 3), + [0, 3, 1, 2], + [0, 2, 3, 1], + [0, 0, 0, 0, 0, 1], + exir_ops.edge.channels_last.constant_pad_nd.default, + ), + ( + (2, 8, 3), + [0, 2, 1], + [0, 2, 1], + [0, 0, 0, 1], + exir_ops.edge.aten.constant_pad_nd.default, + ), + ): + 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, + layout_pad_target=( + exir_ops.edge.channels_last.constant_pad_nd.default + ), + )(graph_module), + ) + + self.assertTrue(result.modified) + self.assertEqual(count_node(result.graph_module, expected_target), 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.channels_last.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], + layout_region=True, + ) + + 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)