From 5df001086f664efaf226c813ce97ac4a663f6310 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Mon, 31 Aug 2026 09:15:00 -0700 Subject: [PATCH 1/5] Transforms: support backend-defined layout copies The existing view and permute optimizers assumed that every layout change used the ATen permute target. Backends with a graph-visible layout dialect therefore could not reuse propagation or region cancellation without changing the operator family mid-rewrite. Make permute targets, emitted targets, fusion scope, propagation barriers, and output-boundary compensation explicit extension points while preserving existing defaults. Keep synthesized nodes in their source dialect, refresh their metadata, reject mixed-dialect chains conservatively, and cover reconvergent and rank-changing regions. AI-assisted: Codex. --- .../propagate_view_copy_permute_pass.py | 3 + ...est_canonicalize_view_copy_permute_pass.py | 58 +++ .../canonicalize_view_copy_permute_pass.py | 21 +- .../fuse_identical_input_transforms_pass.py | 2 +- ...fuse_transpose_or_permute_op_pairs_pass.py | 4 + backends/transforms/permute_view_meta.py | 26 +- .../propagate_view_copy_permute_pass.py | 29 +- .../remove_permutes_around_elementwise_ops.py | 150 +++++- ...replace_ops_with_channels_last_variants.py | 11 +- .../test/test_permute_optimization_passes.py | 447 ++++++++++++++++++ .../test_propagate_view_copy_permute_pass.py | 20 +- ...replace_ops_with_channels_last_variants.py | 25 + 12 files changed, 725 insertions(+), 71 deletions(-) diff --git a/backends/arm/_passes/propagate_view_copy_permute_pass.py b/backends/arm/_passes/propagate_view_copy_permute_pass.py index 374f3857f9d..9022412a092 100644 --- a/backends/arm/_passes/propagate_view_copy_permute_pass.py +++ b/backends/arm/_passes/propagate_view_copy_permute_pass.py @@ -55,6 +55,9 @@ def duplicate_user_fusion_key(self, node: torch.fx.Node) -> Any: exir_ops.edge.aten.slice_copy.Tensor, } + def duplicate_user_fusion_key(self, node: torch.fx.Node) -> Any: + return quantization_metadata_key(node) + def blocks_moving( self, moving_node: torch.fx.Node, diff --git a/backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py b/backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py index 38a55c8ba10..c6203a75919 100644 --- a/backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py +++ b/backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py @@ -571,3 +571,61 @@ def test_canonicalize_unsupported_end_view_does_not_block_prefix() -> None: ) assert _count_node(result.graph_module, exir_ops.edge.aten.view_copy.default) == 1 _validate_numerics(gm_before, result.graph_module, (x_data,)) + + +def test_canonicalize_mixed_permute_dialects_is_idempotent() -> None: + builder = GraphBuilder() + x_data = torch.randn(1, 2, 3, 4) + x = builder.placeholder("x", x_data) + layout_permute = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + view = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(layout_permute, [1, 12, 2]), + ) + permute = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(view, [0, 2, 1]), + ) + builder.output([permute]) + graph_module = builder.get_graph_module() + + pass_instance = CanonicalizeViewCopyPermutePass( + permute_targets={ + exir_ops.edge.channels_last.permute_copy.default, + exir_ops.edge.aten.permute_copy.default, + } + ) + first = cast(PassResult, pass_instance.call(graph_module)) + second = cast(PassResult, pass_instance.call(first.graph_module)) + + assert not first.modified + assert not second.modified + assert _compute_nodes(second.graph_module) == [ + exir_ops.edge.channels_last.permute_copy.default, + exir_ops.edge.aten.view_copy.default, + exir_ops.edge.aten.permute_copy.default, + ] + + +def test_canonicalize_refreshes_backend_permute_metadata() -> 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]), + ) + + pass_instance = CanonicalizeViewCopyPermutePass( + permute_targets={exir_ops.edge.channels_last.permute_copy.default} + ) + pass_instance._set_node_op( + permute.node, + exir_ops.edge.channels_last.permute_copy.default, + x.node, + [0, 3, 1, 2], + ) + + assert permute.node.meta["val"].shape == torch.Size([1, 4, 2, 3]) diff --git a/backends/transforms/canonicalize_view_copy_permute_pass.py b/backends/transforms/canonicalize_view_copy_permute_pass.py index 0a76f10011d..701238da62a 100644 --- a/backends/transforms/canonicalize_view_copy_permute_pass.py +++ b/backends/transforms/canonicalize_view_copy_permute_pass.py @@ -43,8 +43,7 @@ class CanonicalizeViewCopyPermutePass(ExportPass): def __init__(self, permute_targets: Iterable[Any] | None = None) -> None: super().__init__() # Which targets count as a permute. A backend carrying its own layout - # dialect passes them here; the pass still emits _PERMUTE_TARGET when it - # has to create one. + # dialect passes them here. Mixed-dialect chains are left unchanged. self._permute_targets = frozenset(permute_targets or (self._PERMUTE_TARGET,)) self._targets = {self._VIEW_TARGET} | self._permute_targets @@ -55,6 +54,12 @@ def call(self, graph_module: GraphModule) -> PassResult: modified = False for chain in self._collect_chains(graph_module): + permute_targets = { + node.target for node in chain if node.target in self._permute_targets + } + if len(permute_targets) > 1: + continue + updated_chain = chain while True: @@ -169,7 +174,7 @@ def _fuse_sequential_ops( _normalize_dim(dim, len(self._shape(input_node))) for dim in dims ] - self._set_node_op(node, self._PERMUTE_TARGET, input_node, dims) + self._set_node_op(node, node.target, input_node, dims) changed = True any_changed = True @@ -214,13 +219,17 @@ def _fuse_sequential_ops( any_changed = True continue - if self._is_permute(node) and self._is_permute(next_node): + if ( + self._is_permute(node) + and self._is_permute(next_node) + and node.target == next_node.target + ): # Fuse consecutive permutes dims = self._permute_dims(node) next_dims = self._permute_dims(next_node) self._set_node_op( node, - self._PERMUTE_TARGET, + node.target, input_node, [dims[dim] for dim in next_dims], ) @@ -372,7 +381,7 @@ def _set_node_op( ) -> None: node.target = target node.args = (input_node, list(arg)) - refresh_permute_view_meta(node) + refresh_permute_view_meta(node, self._permute_targets) def _permute_dims(self, node: Node) -> list[int]: assert self._is_permute(node), "Expected permute node" diff --git a/backends/transforms/fuse_identical_input_transforms_pass.py b/backends/transforms/fuse_identical_input_transforms_pass.py index b82da62c1d3..dcdf09e4f94 100644 --- a/backends/transforms/fuse_identical_input_transforms_pass.py +++ b/backends/transforms/fuse_identical_input_transforms_pass.py @@ -259,7 +259,7 @@ def _sink_identical_input_transforms(self, node: Node) -> bool: kwargs=dict(transform.kwargs), ) new_node.meta = self._new_transform_meta(node, transform) - refresh_permute_view_meta(new_node) + refresh_permute_view_meta(new_node, self._permute_targets) for user in list(node.users): if user is not new_node: 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 b8ae98d9335..166908638bb 100644 --- a/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py +++ b/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py @@ -49,6 +49,10 @@ def can_fuse_for_chain( ) -> bool: if not super().can_fuse_for_chain(producer, consumer, consumer_op_packets): return False + if (producer.target == LAYOUT_PERMUTE_COPY) != ( + consumer.target == LAYOUT_PERMUTE_COPY + ): + return False # checking that permut2(permut1(identity)) == identity, modulo unitary dimensions producer_input = cast(torch.fx.Node, producer.args[0]) diff --git a/backends/transforms/permute_view_meta.py b/backends/transforms/permute_view_meta.py index 18972329176..d0d739ac075 100644 --- a/backends/transforms/permute_view_meta.py +++ b/backends/transforms/permute_view_meta.py @@ -6,8 +6,8 @@ # pyre-unsafe -from collections.abc import Sequence -from typing import cast +from collections.abc import Iterable, Sequence +from typing import Any, cast import torch import torch.fx @@ -19,16 +19,26 @@ from executorch.exir.dialects._ops import ops as exir_ops -def refresh_permute_view_meta(node: torch.fx.Node) -> None: +def refresh_permute_view_meta( + node: torch.fx.Node, permute_targets: Iterable[Any] | None = None +) -> None: """Compute new meta-vals, specifically preserving SymInts for view/permute nodes. """ input_node = node.all_input_nodes[0] input_val = input_node.meta.get("val") - if input_val is None or node.target not in { - exir_ops.edge.aten.view_copy.default, - exir_ops.edge.aten.permute_copy.default, - }: + if input_val is None: + return + + permute_targets = frozenset( + (exir_ops.edge.aten.permute_copy.default,) + if permute_targets is None + else permute_targets + ) + if ( + node.target != exir_ops.edge.aten.view_copy.default + and node.target not in permute_targets + ): return if not isinstance(input_val, torch.Tensor): @@ -45,7 +55,7 @@ def refresh_permute_view_meta(node: torch.fx.Node) -> None: ) ) ) - case exir_ops.edge.aten.permute_copy.default: + case target if target in permute_targets: dims = _normalize_dims( cast(Sequence[int], node.args[1]), len(input_val.shape) ) diff --git a/backends/transforms/propagate_view_copy_permute_pass.py b/backends/transforms/propagate_view_copy_permute_pass.py index dc851b011e0..0da14b7248b 100644 --- a/backends/transforms/propagate_view_copy_permute_pass.py +++ b/backends/transforms/propagate_view_copy_permute_pass.py @@ -47,9 +47,8 @@ class PropagateViewCopyPermutePass(ExportPass, ABC): _passes_required_after: Set[Type[ExportPass]] = set() _VIEW_TARGET = exir_ops.edge.aten.view_copy.default - _VIEW_DEFAULT_TARGET = exir_ops.edge.aten.view.default _PERMUTE_TARGET = exir_ops.edge.aten.permute_copy.default - _TARGETS = {_VIEW_TARGET, _VIEW_DEFAULT_TARGET, _PERMUTE_TARGET} + _TARGETS = {_VIEW_TARGET, _PERMUTE_TARGET} _TRANSPARENT_TARGETS = { exir_ops.edge.dim_order_ops._clone_dim_order.default, exir_ops.edge.dim_order_ops._to_dim_order_copy.default, @@ -77,13 +76,9 @@ def __init__( self.exported_program = exported_program self.compile_spec = compile_spec # Which targets count as a permute. A backend carrying its own layout - # dialect passes them here; the pass still emits self._PERMUTE_TARGET - # when it has to create one. + # dialect passes them here. self._permute_targets = frozenset(permute_targets or (self._PERMUTE_TARGET,)) - self._targets = { - self._VIEW_TARGET, - self._VIEW_DEFAULT_TARGET, - } | self._permute_targets + self._targets = {self._VIEW_TARGET} | self._permute_targets @staticmethod def _dim_arg(arg: Any) -> int | Sequence[int] | None: @@ -137,6 +132,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: if modified: graph_module = self._retrace(graph_module) + graph_module.recompile() return PassResult(graph_module, modified) @@ -240,6 +236,11 @@ def _propagate(self, node: torch.fx.Node, stale_nodes: set[torch.fx.Node]) -> bo moved = True continue + if self._maybe_distribute_upwards_permute_over_elementwise( + node, frontier, next_node + ): + return True + # Concats are a special case since they branch the graph. # Perform the swap directly in this case and return. # Otherwise break and move the node before the concat @@ -337,6 +338,14 @@ def _maybe_split_upwards_cat_fanout( """ return False + def _maybe_distribute_upwards_permute_over_elementwise( + self, + node: torch.fx.Node, + frontier: torch.fx.Node, + next_node: torch.fx.Node, + ) -> bool: + return False + def _maybe_split_fork( self, node: torch.fx.Node, @@ -357,7 +366,7 @@ def _maybe_swap_args( """ if node.target in self._permute_targets: return self._maybe_swap_permute_args(node, next_node) - elif node.target in {self._VIEW_TARGET, self._VIEW_DEFAULT_TARGET}: + elif node.target == self._VIEW_TARGET: return self._maybe_swap_view_args(node, next_node) else: raise ValueError( @@ -764,7 +773,7 @@ def _maybe_split_upwards_cat_fanout( output_shape = [input_val.shape[dim] for dim in permute_args[0]] with next_node.graph.inserting_before(next_node): permute = next_node.graph.call_function( - self._PERMUTE_TARGET, + cast(Any, node.target), args=(input_node, permute_args[0]), ) permute.meta = dict(input_node.meta) diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 1a3a6ada995..4dac35ea46e 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -7,13 +7,14 @@ # 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 ( - is_permute_copy, + ATEN_PERMUTE_COPY, PERMUTE_COPY_TARGETS, ) from executorch.backends.transforms.permute_pass_utils import get_arg, set_arg @@ -29,6 +30,15 @@ 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. ``can_propagate`` is the backend's statement that a layout copy + must not move across a node. + + ``compensate_at_output`` lets a region end at a graph output instead of + being abandoned there. It is opt-in because it relocates a layout copy + toward the outputs, and a backend that also runs the propagation passes has + already chosen where its copies sit. """ @dataclass() @@ -51,6 +61,11 @@ class Subgraph: constant_edges_in: set[tuple[torch.fx.Node, torch.fx.Node]] = field( default_factory=set ) + # Region values that are also returned. The permute cannot simply be + # dropped there, so it is re-inserted on the output edge instead. + output_boundaries: set[tuple[torch.fx.Node, torch.fx.Node, tuple[int, ...]]] = ( + field(default_factory=set) + ) # 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) @@ -62,8 +77,22 @@ 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, + *, + can_propagate: Callable[[torch.fx.Node], bool] | None = None, + compensate_at_output: bool = False, + permute_targets: set | frozenset | None = None, + permute_target=None, + ) -> None: super().__init__() + self.can_propagate = can_propagate + self.compensate_at_output = compensate_at_output + self._permute_targets = frozenset(permute_targets or PERMUTE_COPY_TARGETS) + self._permute_target = permute_target or ATEN_PERMUTE_COPY + if self._permute_target not in self._permute_targets: + raise ValueError("permute_target must be included in permute_targets") self._permutable_ops = { exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.mul.Tensor, @@ -94,10 +123,7 @@ def __init__(self, extra_permutable_ops: set | None = None) -> None: tuple[int, int, torch.fx.Node, torch.fx.Node] | None, ] = {} - _VIEW_OPS = ( - exir_ops.edge.aten.view_copy.default, - exir_ops.edge.aten.view.default, - ) + _VIEW_OPS = (exir_ops.edge.aten.view_copy.default,) @staticmethod def _concrete_shape(node: torch.fx.Node) -> list[int] | None: @@ -198,7 +224,7 @@ def _sink_users_are_layout_invariant(self, sink: torch.fx.Node) -> bool: if node.op == "output": continue - if node.target == exir_ops.edge.aten.permute_copy.default: + if node.target in self._permute_targets: # This explicit transform re-establishes the downstream layout, # so consumers beyond it do not depend on the sink's layout. continue @@ -370,7 +396,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 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): + if node.target not in self._permute_targets: continue start_permute = self.get_permutation(node) if start_permute is None: @@ -482,7 +508,7 @@ def visit( # noqa: C901 # Traverse downstream: for user in users_source.users: - if user.target in PERMUTE_COPY_TARGETS: + if user.target in self._permute_targets: user_perm = self.get_permutation(user) if user_perm == downstream_end: subgraph.edges_out.add((users_source, user)) @@ -499,7 +525,20 @@ def visit( # noqa: C901 ) ) elif user.op == "output": - return False + if not self.compensate_at_output: + return False + subgraph.output_boundaries.add( + (users_source, user, tuple(downstream_start)) + ) + elif self.can_propagate is not None and not self.can_propagate(user): + # A backend barrier. With compensation the region ends here + # rather than being abandoned: the permute is re-inserted on + # this edge so the barrier still sees the layout it expects. + if not self.compensate_at_output: + return False + subgraph.output_boundaries.add( + (users_source, user, tuple(downstream_start)) + ) elif self._is_permutation_sink_view(user): # The tensor's element order is invariant at this reshape, but # its output shape can still carry broadcast-axis meaning. @@ -513,7 +552,7 @@ def visit( # noqa: C901 # Traverse upstream: for inp in node.all_input_nodes: - if inp.target in PERMUTE_COPY_TARGETS: + if inp.target in self._permute_targets: if self.get_permutation(inp) != current_start_permute: return False subgraph.edges_in.add((inp, node)) @@ -599,6 +638,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: @@ -713,7 +754,7 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 # Skip incoming permutes. for inp, out in subgraph.edges_in: - assert inp.target in PERMUTE_COPY_TARGETS + assert inp.target in self._permute_targets if len(inp.args) >= 1: out.replace_input_with(inp, cast(torch.fx.Node, inp.args[0])) else: @@ -733,7 +774,7 @@ 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, + self._permute_target, args=(const_node, node_end_perm), ) elif ( @@ -741,22 +782,53 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 and const_rank < permute_rank and const_node.meta.get("val") is not None ): + # Broadcasting widens the constant to the region's rank + # before the permutation applies. original_shape = list(const_node.meta["val"].shape) padded = [1] * (permute_rank - const_rank) + original_shape - target_shape = [padded[d] for d in node_end_perm] - target_shape = target_shape[permute_rank - const_rank :] - new_node = graph.create_node( - "call_function", - exir_ops.edge.aten.view_copy.default, - args=(const_node, target_shape), - ) + target_shape = [padded[dim] for dim in node_end_perm] + + # Where each non-unit axis ends up. Unit axes carry no + # elements, so only the order of these decides whether the + # permutation rearranges data or merely reshapes. + destinations = [ + node_end_perm.index(axis) + for axis, size in enumerate(padded) + if size != 1 + ] + if destinations == sorted(destinations): + # Only unit extents moved, so this is a pure reshape and + # a view says it exactly -- and says it for free, since + # view_copy later becomes a memory.view alias. + new_node = graph.create_node( + "call_function", + exir_ops.edge.aten.view_copy.default, + args=(const_node, target_shape), + ) + else: + # Reordering a non-unit extent moves data. A view would + # reinterpret the strides and read different elements, + # so widen with a view and permute at full rank. + widened = graph.create_node( + "call_function", + exir_ops.edge.aten.view_copy.default, + args=(const_node, padded), + ) + with graph.inserting_after(widened): + new_node = graph.create_node( + "call_function", + self._permute_target, + args=(widened, node_end_perm), + ) else: continue user_node.replace_input_with(const_node, new_node) + self._insert_output_boundary_permutations(subgraph) + # Skip outgoing permutes. for inp, out in subgraph.edges_out: - assert out.target in PERMUTE_COPY_TARGETS + assert out.target in self._permute_targets out.replace_all_uses_with(inp) # Update outgoing permutes that can't be eliminated. @@ -769,7 +841,10 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: """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: + if ( + inp.target not in self._permute_targets + or inp not in out.all_input_nodes + ): return False # edges_out_to_update can rewrite a permute in place, leaving it wired. @@ -779,7 +854,7 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: return False for inp, out in subgraph.edges_out: - if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users: + if out.target not in self._permute_targets or out not in inp.users: return False for inp, out, _ in subgraph.edges_out_to_update: @@ -801,6 +876,33 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: return True + def _insert_output_boundary_permutations(self, subgraph: Subgraph) -> None: + """Put the region's permutation back on the edges that leave it. + + A returned value still has to be in the layout the caller was promised, + so the permute the region cancelled everywhere else is re-inserted here + rather than the whole region being abandoned. + + """ + if not subgraph.output_boundaries: + return + 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( + self._permute_target, + 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, @@ -902,7 +1004,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 in PERMUTE_COPY_TARGETS + assert permute_node.target in self._permute_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_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py index 641c3925b65..3e36761bd4e 100644 --- a/backends/transforms/replace_ops_with_channels_last_variants.py +++ b/backends/transforms/replace_ops_with_channels_last_variants.py @@ -145,12 +145,14 @@ def __init__( self, exported_program: ExportedProgram, op_map: dict[Target, ChannelsLastOpSpec] | None = None, + require_contiguous_output: bool = True, ) -> None: super().__init__() self.exported_program = exported_program self.op_map: dict[Target, ChannelsLastOpSpec] = ( op_map if op_map is not None else dict(_DEFAULT_OP_MAP) ) + self.require_contiguous_output = require_contiguous_output @staticmethod def _permute_node_input( @@ -208,12 +210,15 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: continue if (spec := self.op_map.get(node.target)) is None: continue + if spec.filter_fn is not None and not spec.filter_fn(node): + continue val = node.meta["val"] val = val[0] if isinstance(val, (list, tuple)) else val contiguous_dim_order = tuple(range(val.dim())) - if val.dim_order() != contiguous_dim_order: - continue - if spec.filter_fn is not None and not spec.filter_fn(node): + if ( + self.require_contiguous_output + and val.dim_order() != contiguous_dim_order + ): continue # In case of implicit batch size, insert also `unsqueeze_copy.default` and `squeeze_copy.dims` operators. diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index 1f357deb171..34478374de5 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -782,6 +782,250 @@ def test_per_channel_branch_blocks_shared_permute_fusion(self) -> None: # ────────────────────────────────────────────────────────────────────── +# ────────────────────────────────────────────────────────────────────── +# Tests for structural layout boundary propagation +# ────────────────────────────────────────────────────────────────────── + + +class LayoutDialectHandlingTest(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()(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_preserves_layout_dialect_across_aten_transpose(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(2, 3)) + transpose = builder.call_operator( + op=exir_ops.edge.aten.transpose_copy.int, + args=(x, 0, 1), + ) + layout_permute = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(transpose, [1, 0]), + ) + builder.output([layout_permute]) + 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.aten.transpose_copy.int), 1 + ) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 1, + ) + + 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_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()(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) @@ -2354,6 +2598,209 @@ def test_chained_regions_absorb_into_last_permute(self) -> None: "chained_regions_absorb_into_last_permute", ) + def _assert_region_cancels( + self, module: torch.nn.Module, inputs: tuple[torch.Tensor, ...] + ) -> None: + """The region's boundary permutes go away and the values do not change.""" + module = module.eval() + expected = module(*inputs) + with torch.no_grad(): + exported = torch.export.export(module, inputs) + edge = to_edge( + exported, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, _skip_dim_order=True + ), + ) + before = count_node( + edge.exported_program().graph_module, + exir_ops.edge.aten.permute_copy.default, + ) + transformed = edge.transform([RemovePermutesAroundElementwiseOps()]) + actual = transformed.exported_program().module()(*inputs) + + after = count_node( + transformed.exported_program().graph_module, + exir_ops.edge.aten.permute_copy.default, + ) + self.assertLess(after, before, "the boundary permutes should have cancelled") + torch.testing.assert_close(actual, expected) + + def test_lower_rank_constant_reorder_preserves_values(self) -> None: + """A broadcast constant is widened and permuted, never reinterpreted. + + A view cannot express the reorder -- it reinterprets strides rather than + moving elements -- so each rank below the region's needs its own case. + """ + + class Rank3(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer( + "bias", + torch.arange(4 * 8 * 8, dtype=torch.float32).reshape(4, 8, 8), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x.permute(0, 3, 1, 2) + self.bias).permute(0, 2, 3, 1) + + class Rank2(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer( + "bias", torch.arange(8, dtype=torch.float32).reshape(1, 8) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x.permute(2, 0, 1) + self.bias).permute(1, 2, 0) + + class Rank1(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("scale", torch.arange(1.0, 9.0)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x.permute(0, 3, 1, 2) * self.scale).permute(0, 2, 3, 1) + + class EqualExtents(torch.nn.Module): + """Two non-unit axes of the same size swap. + + The extents read the same before and after, so only their order + distinguishes a reshape from a reorder. + """ + + def __init__(self) -> None: + super().__init__() + self.register_buffer( + "bias", torch.arange(4, dtype=torch.float32).reshape(2, 2) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x.permute(0, 1, 3, 2) + self.bias).permute(0, 1, 3, 2) + + self._assert_region_cancels(Rank3(), (torch.randn(1, 8, 8, 4),)) + self._assert_region_cancels(Rank2(), (torch.randn(8, 8, 8),)) + self._assert_region_cancels(Rank1(), (torch.randn(1, 8, 8, 8),)) + self._assert_region_cancels(EqualExtents(), (torch.randn(1, 4, 2, 2),)) + + def test_region_returning_a_value_reinserts_the_permute(self) -> None: + """A returned region value keeps its layout, and the region still cancels. + + Reaching an output used to abandon the whole region, so a graph that + returned an intermediate paid for every permute in it. + """ + + class ReturnsIntermediate(torch.nn.Module): + def forward(self, x: torch.Tensor): + permuted = x.permute(0, 3, 1, 2) + shifted = permuted + 1.0 + scaled = shifted * 2.0 + return scaled.permute(0, 2, 3, 1), shifted + + inputs = (torch.randn(1, 8, 8, 4),) + module = ReturnsIntermediate().eval() + expected = module(*inputs) + with torch.no_grad(): + edge = to_edge( + torch.export.export(module, inputs), + compile_config=EdgeCompileConfig( + _check_ir_validity=False, _skip_dim_order=True + ), + ) + before = count_node( + edge.exported_program().graph_module, + exir_ops.edge.aten.permute_copy.default, + ) + transformed = edge.transform( + [RemovePermutesAroundElementwiseOps(compensate_at_output=True)] + ) + actual = transformed.exported_program().module()(*inputs) + + after = count_node( + transformed.exported_program().graph_module, + exir_ops.edge.aten.permute_copy.default, + ) + self.assertEqual(before, 2) + # One survives, on the edge that returns the intermediate. + self.assertEqual(after, 1) + torch.testing.assert_close(actual, expected) + + def test_output_boundary_compensation_is_off_by_default(self) -> None: + """Without the opt-in the pass still gives up at a graph output. + + Compensating relocates a layout copy toward the outputs, and a backend + that also runs the propagation passes has already chosen where its + copies sit -- Arm parks them near the inputs. Changing that under it + would have the two pulling the same copy apart. + """ + + class SinkOnly(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + permuted = x.permute(0, 2, 3, 1) + return permuted + permuted * 2.0 + + inputs = (torch.randn(1, 2, 3, 4),) + module = SinkOnly().eval() + expected = module(*inputs) + with torch.no_grad(): + edge = to_edge( + torch.export.export(module, inputs), + compile_config=EdgeCompileConfig( + _check_ir_validity=False, _skip_dim_order=True + ), + ) + transformed = edge.transform([RemovePermutesAroundElementwiseOps()]) + graph = transformed.exported_program().graph_module.graph + actual = transformed.exported_program().module()(*inputs) + + targets = [n.target for n in graph.nodes if n.op == "call_function"] + self.assertEqual( + targets.index(exir_ops.edge.aten.permute_copy.default), + 0, + "the permute should not have moved to the far side of the region", + ) + torch.testing.assert_close(actual, expected) + + def test_value_returned_both_raw_and_through_the_region_exit(self) -> None: + """One value returned in both layouts still leaves exactly one permute. + + The boundary permute is inserted before the exit permute is skipped, so + the two output slots stay distinct: the slot that went through the exit + takes the region's own value and the raw slot takes the new permute. + Inserting after the skip would collapse both onto one layout. + """ + + class BothLayouts(torch.nn.Module): + def forward(self, x: torch.Tensor): + shifted = x.permute(0, 3, 1, 2) + 1.0 + return shifted.permute(0, 2, 3, 1), shifted + + inputs = (torch.randn(1, 8, 8, 4),) + module = BothLayouts().eval() + expected = module(*inputs) + with torch.no_grad(): + edge = to_edge( + torch.export.export(module, inputs), + compile_config=EdgeCompileConfig( + _check_ir_validity=False, _skip_dim_order=True + ), + ) + transformed = edge.transform( + [RemovePermutesAroundElementwiseOps(compensate_at_output=True)] + ) + actual = transformed.exported_program().module()(*inputs) + + self.assertEqual( + count_node( + transformed.exported_program().graph_module, + exir_ops.edge.aten.permute_copy.default, + ), + 1, + ) + for index, (want, got) in enumerate(zip(expected, actual)): + self.assertEqual(want.shape, got.shape, f"output {index} changed layout") + torch.testing.assert_close(got, want) + class LayoutPermuteVisibilityTest(unittest.TestCase): """The data-movement passes must see both permute dialects. diff --git a/backends/transforms/test/test_propagate_view_copy_permute_pass.py b/backends/transforms/test/test_propagate_view_copy_permute_pass.py index daf5a5b7ff6..c6c6c47959a 100644 --- a/backends/transforms/test/test_propagate_view_copy_permute_pass.py +++ b/backends/transforms/test/test_propagate_view_copy_permute_pass.py @@ -4,7 +4,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import pytest import torch from executorch.backends.transforms.propagate_view_copy_permute_pass import ( PropagateViewCopyPermuteDownPass, @@ -13,6 +12,7 @@ from executorch.exir.dialects._ops import ops as exir_ops PERMUTE = exir_ops.edge.aten.permute_copy.default +VIEW = exir_ops.edge.aten.view_copy.default ABS = exir_ops.edge.aten.abs.default NEG = exir_ops.edge.aten.neg.default SIGMOID = exir_ops.edge.aten.sigmoid.default @@ -58,24 +58,6 @@ def count() -> int: return before, after_down, count() -@pytest.mark.xfail( - strict=True, - reason="Splitting a fork where some branches rejoin and others do not leaves " - "one copy below the meeting node and one at the source, and the up pass has " - "no fork split of its own to hoist the first above the rejoin. No model in a " - "15-model sweep produces this shape, so the driver does not special-case it; " - "the general fix is to stop propagation increasing the copy count at all.", -) -def test_mixed_reconvergence_fork_does_not_strand_a_permute() -> None: - graph = torch.fx.Graph() - left, right, diverging = _forked_permute(graph, branches=3) - rejoin = graph.call_function(ADD, args=(left, right)) - rejoin.meta["val"] = torch.empty(PERMUTED_SHAPE) - graph.output((rejoin, diverging)) - - assert _permute_counts(graph) == (1, 1, 1) - - def test_fork_split_still_applies_when_every_branch_rejoins() -> None: graph = torch.fx.Graph() left, right = _forked_permute(graph, branches=2) 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 3c8fdd8309f..15e7d11ec4b 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 @@ -464,6 +464,31 @@ def test_modified_false_when_no_matching_ops(self): _, modified = _run_pass(ep) assert not modified + def test_noncontiguous_output_can_be_normalized_by_opt_in_backend(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] + conv.meta["val"] = conv.meta["val"].to(memory_format=torch.channels_last) + + result = ReplaceOpsWithChannelsLastVariants(ep)(ep.graph_module) + + assert not result.modified + + ep = _export_to_edge(Conv2dModule(), (torch.randn(1, 4, 8, 8),)) + conv = _find_nodes(ep.graph_module, exir_ops.edge.aten.convolution.default)[0] + conv.meta["val"] = conv.meta["val"].to(memory_format=torch.channels_last) + result = ReplaceOpsWithChannelsLastVariants( + ep, require_contiguous_output=False + )(ep.graph_module) + + assert result.modified + assert ( + _count( + result.graph_module, + exir_ops.edge.channels_last.convolution.default, + ) + == 1 + ) + def test_empty_op_map_leaves_graph_unchanged(self): ep = _export_to_edge(Conv2dModule(), (torch.randn(1, 4, 8, 8),)) From c8c32d40d35314a20a1b2174e58cdd01fcd4a66f Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Tue, 1 Sep 2026 08:26:52 -0700 Subject: [PATCH 2/5] Transforms: document opt-in reconvergence handling Restore the shared mixed-reconvergence xfail and document that upward distribution is a backend opt-in because it can increase layout-copy count. AI-assisted: Codex. --- .../propagate_view_copy_permute_pass.py | 6 ++++++ .../test_propagate_view_copy_permute_pass.py | 16 ++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/backends/transforms/propagate_view_copy_permute_pass.py b/backends/transforms/propagate_view_copy_permute_pass.py index 0da14b7248b..6370097f72d 100644 --- a/backends/transforms/propagate_view_copy_permute_pass.py +++ b/backends/transforms/propagate_view_copy_permute_pass.py @@ -344,6 +344,12 @@ def _maybe_distribute_upwards_permute_over_elementwise( frontier: torch.fx.Node, next_node: torch.fx.Node, ) -> bool: + """Optionally distribute an upward-moving permute over multiple inputs. + + The shared driver leaves this disabled because distribution can increase + the number of layout copies. Backends may opt in when their layout + strategy requires crossing a reconvergent elementwise node. + """ return False def _maybe_split_fork( diff --git a/backends/transforms/test/test_propagate_view_copy_permute_pass.py b/backends/transforms/test/test_propagate_view_copy_permute_pass.py index c6c6c47959a..39e9a7b9bfe 100644 --- a/backends/transforms/test/test_propagate_view_copy_permute_pass.py +++ b/backends/transforms/test/test_propagate_view_copy_permute_pass.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import pytest import torch from executorch.backends.transforms.propagate_view_copy_permute_pass import ( PropagateViewCopyPermuteDownPass, @@ -58,6 +59,21 @@ def count() -> int: return before, after_down, count() +@pytest.mark.xfail( + strict=True, + reason="The shared driver does not distribute an upward permute across a " + "multi-input elementwise node unless a backend explicitly opts in.", +) +def test_mixed_reconvergence_fork_does_not_strand_a_permute() -> None: + graph = torch.fx.Graph() + left, right, diverging = _forked_permute(graph, branches=3) + rejoin = graph.call_function(ADD, args=(left, right)) + rejoin.meta["val"] = torch.empty(PERMUTED_SHAPE) + graph.output((rejoin, diverging)) + + assert _permute_counts(graph) == (1, 1, 1) + + def test_fork_split_still_applies_when_every_branch_rejoins() -> None: graph = torch.fx.Graph() left, right = _forked_permute(graph, branches=2) From 9c0b58d6a6cbe6a8bbfe3d46c58e815e1fb18f71 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Tue, 1 Sep 2026 08:47:43 -0700 Subject: [PATCH 3/5] Transforms: preserve layout-copy dialect in regions Track each cancellation region's source permute target so constant and output compensation cannot silently switch dialects. Document why propagation excludes aliasing view.default and why backend policy hooks are opt-in. AI-assisted: Codex. --- .../propagate_view_copy_permute_pass.py | 7 +- .../remove_permutes_around_elementwise_ops.py | 33 +++--- ...replace_ops_with_channels_last_variants.py | 2 + .../test/test_permute_optimization_passes.py | 109 ++++++++++++++++++ 4 files changed, 135 insertions(+), 16 deletions(-) diff --git a/backends/transforms/propagate_view_copy_permute_pass.py b/backends/transforms/propagate_view_copy_permute_pass.py index 6370097f72d..a776d32edc9 100644 --- a/backends/transforms/propagate_view_copy_permute_pass.py +++ b/backends/transforms/propagate_view_copy_permute_pass.py @@ -40,12 +40,17 @@ class PropagateViewCopyPermutePass(ExportPass, ABC): """Abstract implementation of a permute/view_copy propagation pass. To be used for upwards/downwards propagation by implementing the abstract - methods for the direction of propagation. + methods for the direction of propagation. Backends may supply a closed set + of equivalent permute targets and override the policy hooks when their + layout contract permits more movement. Every hook preserves the existing + behavior by default. """ _passes_required_after: Set[Type[ExportPass]] = set() + # Moving an aliasing aten.view.default requires alias-aware reasoning that + # this pass does not provide. Restrict propagation to copy semantics. _VIEW_TARGET = exir_ops.edge.aten.view_copy.default _PERMUTE_TARGET = exir_ops.edge.aten.permute_copy.default _TARGETS = {_VIEW_TARGET, _PERMUTE_TARGET} diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 4dac35ea46e..dc8b6618ee1 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -9,14 +9,11 @@ from collections.abc import Callable from dataclasses import dataclass, field -from typing import cast +from typing import Any, cast import torch import torch.fx -from executorch.backends.transforms.channels_last_layout import ( - ATEN_PERMUTE_COPY, - PERMUTE_COPY_TARGETS, -) +from executorch.backends.transforms.channels_last_layout import PERMUTE_COPY_TARGETS from executorch.backends.transforms.permute_pass_utils import get_arg, set_arg from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult @@ -33,18 +30,23 @@ class RemovePermutesAroundElementwiseOps(ExportPass): ``extra_permutable_ops`` must be layout-equivariant without argument remapping. ``can_propagate`` is the backend's statement that a layout copy - must not move across a node. + may move across a node. ``compensate_at_output`` lets a region end at a graph output instead of being abandoned there. It is opt-in because it relocates a layout copy toward the outputs, and a backend that also runs the propagation passes has already chosen where its copies sit. + + ``permute_targets`` is the closed family of layout-copy operators the + backend considers equivalent. A rewritten region must use one member of + that family consistently, and any synthesized copies retain that member. """ @dataclass() class Subgraph: start_permute: list[int] end_permute: list[int] + permute_target: Any # Nodes in the subgraph, does not include permutes. nodes: set[torch.fx.Node] = field(default_factory=set) # Incoming edges to the subgraph from permute nodes. @@ -84,15 +86,11 @@ def __init__( can_propagate: Callable[[torch.fx.Node], bool] | None = None, compensate_at_output: bool = False, permute_targets: set | frozenset | None = None, - permute_target=None, ) -> None: super().__init__() self.can_propagate = can_propagate self.compensate_at_output = compensate_at_output self._permute_targets = frozenset(permute_targets or PERMUTE_COPY_TARGETS) - self._permute_target = permute_target or ATEN_PERMUTE_COPY - if self._permute_target not in self._permute_targets: - raise ValueError("permute_target must be included in permute_targets") self._permutable_ops = { exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.mul.Tensor, @@ -410,7 +408,7 @@ 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, node.target) if self.visit(user, subgraph, processed_nodes): subgraphs_found.append(subgraph) for n in subgraph.nodes: @@ -509,6 +507,8 @@ def visit( # noqa: C901 # Traverse downstream: for user in users_source.users: if user.target in self._permute_targets: + if user.target != subgraph.permute_target: + return False user_perm = self.get_permutation(user) if user_perm == downstream_end: subgraph.edges_out.add((users_source, user)) @@ -553,7 +553,10 @@ def visit( # noqa: C901 # Traverse upstream: for inp in node.all_input_nodes: if inp.target in self._permute_targets: - if self.get_permutation(inp) != current_start_permute: + if ( + inp.target != subgraph.permute_target + or self.get_permutation(inp) != current_start_permute + ): return False subgraph.edges_in.add((inp, node)) elif (inp_val := inp.meta.get("val")) is not None and inp_val.numel() == 1: @@ -774,7 +777,7 @@ 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", - self._permute_target, + subgraph.permute_target, args=(const_node, node_end_perm), ) elif ( @@ -817,7 +820,7 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 with graph.inserting_after(widened): new_node = graph.create_node( "call_function", - self._permute_target, + subgraph.permute_target, args=(widened, node_end_perm), ) else: @@ -896,7 +899,7 @@ def _insert_output_boundary_permutations(self, subgraph: Subgraph) -> None: first_output = min(outputs, key=node_order.__getitem__) with producer.graph.inserting_before(first_output): new_permute = producer.graph.call_function( - self._permute_target, + subgraph.permute_target, args=(producer, list(permutation)), ) new_permute.meta = dict(producer.meta) diff --git a/backends/transforms/replace_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py index 3e36761bd4e..23cdbf94b53 100644 --- a/backends/transforms/replace_ops_with_channels_last_variants.py +++ b/backends/transforms/replace_ops_with_channels_last_variants.py @@ -134,6 +134,8 @@ 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. + ``require_contiguous_output`` preserves the legacy eligibility rule by + default; a backend with an explicit logical-layout contract may disable it. Metadata from each replaced operator is preserved so provenance and backend annotations survive the rewrite. ExportPass recomputes shape metadata after diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index 34478374de5..23c85aa739f 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -1020,6 +1020,115 @@ def test_layout_copy_reshapes_channel_constant_without_copy(self) -> None: "RemovePermutesAroundElementwiseOps", ) + def test_layout_copy_constant_compensation_preserves_dialect(self) -> None: + bias_data = torch.randn(1, 4, 8, 8) + graph_module, x_data = self._layout_add_graph("b_bias", bias_data) + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps()(graph_module), + ) + + self.assertTrue(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), + 0, + ) + validate_numerics( + before, + result.graph_module, + [x_data, bias_data], + "RemovePermutesAroundElementwiseOps", + ) + + def test_layout_copy_output_compensation_preserves_dialect(self) -> None: + builder = GraphBuilder() + x_data = torch.randn(1, 8, 8, 4) + x = builder.placeholder("x", x_data) + to_nchw = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 3, 1, 2]), + ) + shifted = builder.call_operator( + op=exir_ops.edge.aten.add.Tensor, + args=(to_nchw, 1.0), + ) + to_nhwc = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(shifted, [0, 2, 3, 1]), + ) + builder.output((to_nhwc, shifted)) + graph_module = builder.get_graph_module() + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps(compensate_at_output=True)(graph_module), + ) + + self.assertTrue(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), + 0, + ) + validate_numerics( + before, + result.graph_module, + [x_data], + "RemovePermutesAroundElementwiseOps", + ) + + def test_region_cancellation_rejects_mixed_permute_dialects(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 8, 8, 4)) + to_nchw = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 3, 1, 2]), + ) + shifted = builder.call_operator( + op=exir_ops.edge.aten.add.Tensor, + args=(to_nchw, 1.0), + ) + to_nhwc = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(shifted, [0, 2, 3, 1]), + ) + builder.output([to_nhwc]) + graph_module = builder.get_graph_module() + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps()(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, + ) + # ───────────────────────────────────── # Tests for ReplaceNopTransposeOrPermuteWithViewPass From ad179c71962d012d91daec7b41a3926652ff0f83 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Tue, 1 Sep 2026 16:11:40 -0700 Subject: [PATCH 4/5] Transforms: narrow layout-copy policy surface Keep this stack focused on preserving configured layout-copy dialects and stopping cancellation at backend barriers. Remove upward distribution and output-boundary compensation because either can increase the number of copies, and drop the unused noncontiguous-formation override. Preserve the existing aten.view.default behavior so alias handling can be reviewed separately with a dedicated reproducer. AI-assisted: Codex. --- .../propagate_view_copy_permute_pass.py | 3 - .../propagate_view_copy_permute_pass.py | 37 +- .../remove_permutes_around_elementwise_ops.py | 65 +--- ...replace_ops_with_channels_last_variants.py | 13 +- .../test/test_permute_optimization_passes.py | 349 ++---------------- .../test_propagate_view_copy_permute_pass.py | 8 +- ...replace_ops_with_channels_last_variants.py | 25 -- 7 files changed, 60 insertions(+), 440 deletions(-) diff --git a/backends/arm/_passes/propagate_view_copy_permute_pass.py b/backends/arm/_passes/propagate_view_copy_permute_pass.py index 9022412a092..374f3857f9d 100644 --- a/backends/arm/_passes/propagate_view_copy_permute_pass.py +++ b/backends/arm/_passes/propagate_view_copy_permute_pass.py @@ -55,9 +55,6 @@ def duplicate_user_fusion_key(self, node: torch.fx.Node) -> Any: exir_ops.edge.aten.slice_copy.Tensor, } - def duplicate_user_fusion_key(self, node: torch.fx.Node) -> Any: - return quantization_metadata_key(node) - def blocks_moving( self, moving_node: torch.fx.Node, diff --git a/backends/transforms/propagate_view_copy_permute_pass.py b/backends/transforms/propagate_view_copy_permute_pass.py index a776d32edc9..69df4fc69c5 100644 --- a/backends/transforms/propagate_view_copy_permute_pass.py +++ b/backends/transforms/propagate_view_copy_permute_pass.py @@ -41,19 +41,17 @@ class PropagateViewCopyPermutePass(ExportPass, ABC): To be used for upwards/downwards propagation by implementing the abstract methods for the direction of propagation. Backends may supply a closed set - of equivalent permute targets and override the policy hooks when their - layout contract permits more movement. Every hook preserves the existing - behavior by default. + of equivalent permute targets; synthesized permutes retain their source + target. """ _passes_required_after: Set[Type[ExportPass]] = set() - # Moving an aliasing aten.view.default requires alias-aware reasoning that - # this pass does not provide. Restrict propagation to copy semantics. _VIEW_TARGET = exir_ops.edge.aten.view_copy.default + _VIEW_DEFAULT_TARGET = exir_ops.edge.aten.view.default _PERMUTE_TARGET = exir_ops.edge.aten.permute_copy.default - _TARGETS = {_VIEW_TARGET, _PERMUTE_TARGET} + _TARGETS = {_VIEW_TARGET, _VIEW_DEFAULT_TARGET, _PERMUTE_TARGET} _TRANSPARENT_TARGETS = { exir_ops.edge.dim_order_ops._clone_dim_order.default, exir_ops.edge.dim_order_ops._to_dim_order_copy.default, @@ -83,7 +81,10 @@ def __init__( # Which targets count as a permute. A backend carrying its own layout # dialect passes them here. self._permute_targets = frozenset(permute_targets or (self._PERMUTE_TARGET,)) - self._targets = {self._VIEW_TARGET} | self._permute_targets + self._targets = { + self._VIEW_TARGET, + self._VIEW_DEFAULT_TARGET, + } | self._permute_targets @staticmethod def _dim_arg(arg: Any) -> int | Sequence[int] | None: @@ -137,7 +138,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: if modified: graph_module = self._retrace(graph_module) - graph_module.recompile() return PassResult(graph_module, modified) @@ -241,11 +241,6 @@ def _propagate(self, node: torch.fx.Node, stale_nodes: set[torch.fx.Node]) -> bo moved = True continue - if self._maybe_distribute_upwards_permute_over_elementwise( - node, frontier, next_node - ): - return True - # Concats are a special case since they branch the graph. # Perform the swap directly in this case and return. # Otherwise break and move the node before the concat @@ -343,20 +338,6 @@ def _maybe_split_upwards_cat_fanout( """ return False - def _maybe_distribute_upwards_permute_over_elementwise( - self, - node: torch.fx.Node, - frontier: torch.fx.Node, - next_node: torch.fx.Node, - ) -> bool: - """Optionally distribute an upward-moving permute over multiple inputs. - - The shared driver leaves this disabled because distribution can increase - the number of layout copies. Backends may opt in when their layout - strategy requires crossing a reconvergent elementwise node. - """ - return False - def _maybe_split_fork( self, node: torch.fx.Node, @@ -377,7 +358,7 @@ def _maybe_swap_args( """ if node.target in self._permute_targets: return self._maybe_swap_permute_args(node, next_node) - elif node.target == self._VIEW_TARGET: + elif node.target in {self._VIEW_TARGET, self._VIEW_DEFAULT_TARGET}: return self._maybe_swap_view_args(node, next_node) else: raise ValueError( diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index dc8b6618ee1..9b7dfa49d84 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -29,13 +29,8 @@ class RemovePermutesAroundElementwiseOps(ExportPass): recognised as a single rank-preserving unit; see _interleave_triple. ``extra_permutable_ops`` must be layout-equivariant without argument - remapping. ``can_propagate`` is the backend's statement that a layout copy - may move across a node. - - ``compensate_at_output`` lets a region end at a graph output instead of - being abandoned there. It is opt-in because it relocates a layout copy - toward the outputs, and a backend that also runs the propagation passes has - already chosen where its copies sit. + remapping. ``can_propagate`` lets a backend reject nodes that the shared + pass would otherwise treat as layout-equivariant. ``permute_targets`` is the closed family of layout-copy operators the backend considers equivalent. A rewritten region must use one member of @@ -63,11 +58,6 @@ class Subgraph: constant_edges_in: set[tuple[torch.fx.Node, torch.fx.Node]] = field( default_factory=set ) - # Region values that are also returned. The permute cannot simply be - # dropped there, so it is re-inserted on the output edge instead. - output_boundaries: set[tuple[torch.fx.Node, torch.fx.Node, tuple[int, ...]]] = ( - field(default_factory=set) - ) # 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) @@ -84,12 +74,10 @@ def __init__( extra_permutable_ops: set | None = None, *, can_propagate: Callable[[torch.fx.Node], bool] | None = None, - compensate_at_output: bool = False, permute_targets: set | frozenset | None = None, ) -> None: super().__init__() self.can_propagate = can_propagate - self.compensate_at_output = compensate_at_output self._permute_targets = frozenset(permute_targets or PERMUTE_COPY_TARGETS) self._permutable_ops = { exir_ops.edge.aten.add.Tensor, @@ -121,7 +109,10 @@ def __init__( tuple[int, int, torch.fx.Node, torch.fx.Node] | None, ] = {} - _VIEW_OPS = (exir_ops.edge.aten.view_copy.default,) + _VIEW_OPS = ( + exir_ops.edge.aten.view_copy.default, + exir_ops.edge.aten.view.default, + ) @staticmethod def _concrete_shape(node: torch.fx.Node) -> list[int] | None: @@ -525,20 +516,9 @@ def visit( # noqa: C901 ) ) elif user.op == "output": - if not self.compensate_at_output: - return False - subgraph.output_boundaries.add( - (users_source, user, tuple(downstream_start)) - ) + return False elif self.can_propagate is not None and not self.can_propagate(user): - # A backend barrier. With compensation the region ends here - # rather than being abandoned: the permute is re-inserted on - # this edge so the barrier still sees the layout it expects. - if not self.compensate_at_output: - return False - subgraph.output_boundaries.add( - (users_source, user, tuple(downstream_start)) - ) + return False elif self._is_permutation_sink_view(user): # The tensor's element order is invariant at this reshape, but # its output shape can still carry broadcast-axis meaning. @@ -827,8 +807,6 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 continue user_node.replace_input_with(const_node, new_node) - self._insert_output_boundary_permutations(subgraph) - # Skip outgoing permutes. for inp, out in subgraph.edges_out: assert out.target in self._permute_targets @@ -879,33 +857,6 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: return True - def _insert_output_boundary_permutations(self, subgraph: Subgraph) -> None: - """Put the region's permutation back on the edges that leave it. - - A returned value still has to be in the layout the caller was promised, - so the permute the region cancelled everywhere else is re-inserted here - rather than the whole region being abandoned. - - """ - if not subgraph.output_boundaries: - return - 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( - subgraph.permute_target, - 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/replace_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py index 23cdbf94b53..641c3925b65 100644 --- a/backends/transforms/replace_ops_with_channels_last_variants.py +++ b/backends/transforms/replace_ops_with_channels_last_variants.py @@ -134,8 +134,6 @@ 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. - ``require_contiguous_output`` preserves the legacy eligibility rule by - default; a backend with an explicit logical-layout contract may disable it. Metadata from each replaced operator is preserved so provenance and backend annotations survive the rewrite. ExportPass recomputes shape metadata after @@ -147,14 +145,12 @@ def __init__( self, exported_program: ExportedProgram, op_map: dict[Target, ChannelsLastOpSpec] | None = None, - require_contiguous_output: bool = True, ) -> None: super().__init__() self.exported_program = exported_program self.op_map: dict[Target, ChannelsLastOpSpec] = ( op_map if op_map is not None else dict(_DEFAULT_OP_MAP) ) - self.require_contiguous_output = require_contiguous_output @staticmethod def _permute_node_input( @@ -212,15 +208,12 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: continue if (spec := self.op_map.get(node.target)) is None: continue - if spec.filter_fn is not None and not spec.filter_fn(node): - continue val = node.meta["val"] val = val[0] if isinstance(val, (list, tuple)) else val contiguous_dim_order = tuple(range(val.dim())) - if ( - self.require_contiguous_output - and val.dim_order() != contiguous_dim_order - ): + if val.dim_order() != contiguous_dim_order: + continue + if spec.filter_fn is not None and not spec.filter_fn(node): continue # In case of implicit batch size, insert also `unsqueeze_copy.default` and `squeeze_copy.dims` operators. diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index 23c85aa739f..0271e1ba783 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -811,112 +811,6 @@ def _layout_add_graph( 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()(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_preserves_layout_dialect_across_aten_transpose(self) -> None: builder = GraphBuilder() x = builder.placeholder("x", torch.randn(2, 3)) @@ -945,52 +839,6 @@ def test_pair_fusion_preserves_layout_dialect_across_aten_transpose(self) -> Non 1, ) - 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_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) @@ -1049,50 +897,6 @@ def test_layout_copy_constant_compensation_preserves_dialect(self) -> None: "RemovePermutesAroundElementwiseOps", ) - def test_layout_copy_output_compensation_preserves_dialect(self) -> None: - builder = GraphBuilder() - x_data = torch.randn(1, 8, 8, 4) - x = builder.placeholder("x", x_data) - to_nchw = builder.call_operator( - op=exir_ops.edge.channels_last.permute_copy.default, - args=(x, [0, 3, 1, 2]), - ) - shifted = builder.call_operator( - op=exir_ops.edge.aten.add.Tensor, - args=(to_nchw, 1.0), - ) - to_nhwc = builder.call_operator( - op=exir_ops.edge.channels_last.permute_copy.default, - args=(shifted, [0, 2, 3, 1]), - ) - builder.output((to_nhwc, shifted)) - graph_module = builder.get_graph_module() - before = copy.deepcopy(graph_module) - - result = cast( - PassResult, - RemovePermutesAroundElementwiseOps(compensate_at_output=True)(graph_module), - ) - - self.assertTrue(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), - 0, - ) - validate_numerics( - before, - result.graph_module, - [x_data], - "RemovePermutesAroundElementwiseOps", - ) - def test_region_cancellation_rejects_mixed_permute_dialects(self) -> None: builder = GraphBuilder() x = builder.placeholder("x", torch.randn(1, 8, 8, 4)) @@ -2707,6 +2511,41 @@ def test_chained_regions_absorb_into_last_permute(self) -> None: "chained_regions_absorb_into_last_permute", ) + def test_backend_policy_can_stop_region_cancellation(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 2, 3, 4)) + to_nchw = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(x, [0, 3, 1, 2]), + ) + shifted = builder.call_operator( + op=exir_ops.edge.aten.add.Tensor, + args=(to_nchw, 1.0), + ) + barrier = builder.call_operator( + op=exir_ops.edge.aten.sigmoid.default, + args=(shifted,), + ) + to_nhwc = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(barrier, [0, 2, 3, 1]), + ) + builder.output([to_nhwc]) + graph_module = builder.get_graph_module() + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps( + can_propagate=lambda node: node.target + != exir_ops.edge.aten.sigmoid.default + )(graph_module), + ) + + self.assertFalse(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2 + ) + def _assert_region_cancels( self, module: torch.nn.Module, inputs: tuple[torch.Tensor, ...] ) -> None: @@ -2792,124 +2631,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self._assert_region_cancels(Rank1(), (torch.randn(1, 8, 8, 8),)) self._assert_region_cancels(EqualExtents(), (torch.randn(1, 4, 2, 2),)) - def test_region_returning_a_value_reinserts_the_permute(self) -> None: - """A returned region value keeps its layout, and the region still cancels. - - Reaching an output used to abandon the whole region, so a graph that - returned an intermediate paid for every permute in it. - """ - - class ReturnsIntermediate(torch.nn.Module): - def forward(self, x: torch.Tensor): - permuted = x.permute(0, 3, 1, 2) - shifted = permuted + 1.0 - scaled = shifted * 2.0 - return scaled.permute(0, 2, 3, 1), shifted - - inputs = (torch.randn(1, 8, 8, 4),) - module = ReturnsIntermediate().eval() - expected = module(*inputs) - with torch.no_grad(): - edge = to_edge( - torch.export.export(module, inputs), - compile_config=EdgeCompileConfig( - _check_ir_validity=False, _skip_dim_order=True - ), - ) - before = count_node( - edge.exported_program().graph_module, - exir_ops.edge.aten.permute_copy.default, - ) - transformed = edge.transform( - [RemovePermutesAroundElementwiseOps(compensate_at_output=True)] - ) - actual = transformed.exported_program().module()(*inputs) - - after = count_node( - transformed.exported_program().graph_module, - exir_ops.edge.aten.permute_copy.default, - ) - self.assertEqual(before, 2) - # One survives, on the edge that returns the intermediate. - self.assertEqual(after, 1) - torch.testing.assert_close(actual, expected) - - def test_output_boundary_compensation_is_off_by_default(self) -> None: - """Without the opt-in the pass still gives up at a graph output. - - Compensating relocates a layout copy toward the outputs, and a backend - that also runs the propagation passes has already chosen where its - copies sit -- Arm parks them near the inputs. Changing that under it - would have the two pulling the same copy apart. - """ - - class SinkOnly(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - permuted = x.permute(0, 2, 3, 1) - return permuted + permuted * 2.0 - - inputs = (torch.randn(1, 2, 3, 4),) - module = SinkOnly().eval() - expected = module(*inputs) - with torch.no_grad(): - edge = to_edge( - torch.export.export(module, inputs), - compile_config=EdgeCompileConfig( - _check_ir_validity=False, _skip_dim_order=True - ), - ) - transformed = edge.transform([RemovePermutesAroundElementwiseOps()]) - graph = transformed.exported_program().graph_module.graph - actual = transformed.exported_program().module()(*inputs) - - targets = [n.target for n in graph.nodes if n.op == "call_function"] - self.assertEqual( - targets.index(exir_ops.edge.aten.permute_copy.default), - 0, - "the permute should not have moved to the far side of the region", - ) - torch.testing.assert_close(actual, expected) - - def test_value_returned_both_raw_and_through_the_region_exit(self) -> None: - """One value returned in both layouts still leaves exactly one permute. - - The boundary permute is inserted before the exit permute is skipped, so - the two output slots stay distinct: the slot that went through the exit - takes the region's own value and the raw slot takes the new permute. - Inserting after the skip would collapse both onto one layout. - """ - - class BothLayouts(torch.nn.Module): - def forward(self, x: torch.Tensor): - shifted = x.permute(0, 3, 1, 2) + 1.0 - return shifted.permute(0, 2, 3, 1), shifted - - inputs = (torch.randn(1, 8, 8, 4),) - module = BothLayouts().eval() - expected = module(*inputs) - with torch.no_grad(): - edge = to_edge( - torch.export.export(module, inputs), - compile_config=EdgeCompileConfig( - _check_ir_validity=False, _skip_dim_order=True - ), - ) - transformed = edge.transform( - [RemovePermutesAroundElementwiseOps(compensate_at_output=True)] - ) - actual = transformed.exported_program().module()(*inputs) - - self.assertEqual( - count_node( - transformed.exported_program().graph_module, - exir_ops.edge.aten.permute_copy.default, - ), - 1, - ) - for index, (want, got) in enumerate(zip(expected, actual)): - self.assertEqual(want.shape, got.shape, f"output {index} changed layout") - torch.testing.assert_close(got, want) - class LayoutPermuteVisibilityTest(unittest.TestCase): """The data-movement passes must see both permute dialects. diff --git a/backends/transforms/test/test_propagate_view_copy_permute_pass.py b/backends/transforms/test/test_propagate_view_copy_permute_pass.py index 39e9a7b9bfe..daf5a5b7ff6 100644 --- a/backends/transforms/test/test_propagate_view_copy_permute_pass.py +++ b/backends/transforms/test/test_propagate_view_copy_permute_pass.py @@ -13,7 +13,6 @@ from executorch.exir.dialects._ops import ops as exir_ops PERMUTE = exir_ops.edge.aten.permute_copy.default -VIEW = exir_ops.edge.aten.view_copy.default ABS = exir_ops.edge.aten.abs.default NEG = exir_ops.edge.aten.neg.default SIGMOID = exir_ops.edge.aten.sigmoid.default @@ -61,8 +60,11 @@ def count() -> int: @pytest.mark.xfail( strict=True, - reason="The shared driver does not distribute an upward permute across a " - "multi-input elementwise node unless a backend explicitly opts in.", + reason="Splitting a fork where some branches rejoin and others do not leaves " + "one copy below the meeting node and one at the source, and the up pass has " + "no fork split of its own to hoist the first above the rejoin. No model in a " + "15-model sweep produces this shape, so the driver does not special-case it; " + "the general fix is to stop propagation increasing the copy count at all.", ) def test_mixed_reconvergence_fork_does_not_strand_a_permute() -> None: graph = torch.fx.Graph() 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 15e7d11ec4b..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 @@ -464,31 +464,6 @@ def test_modified_false_when_no_matching_ops(self): _, modified = _run_pass(ep) assert not modified - def test_noncontiguous_output_can_be_normalized_by_opt_in_backend(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] - conv.meta["val"] = conv.meta["val"].to(memory_format=torch.channels_last) - - result = ReplaceOpsWithChannelsLastVariants(ep)(ep.graph_module) - - assert not result.modified - - ep = _export_to_edge(Conv2dModule(), (torch.randn(1, 4, 8, 8),)) - conv = _find_nodes(ep.graph_module, exir_ops.edge.aten.convolution.default)[0] - conv.meta["val"] = conv.meta["val"].to(memory_format=torch.channels_last) - result = ReplaceOpsWithChannelsLastVariants( - ep, require_contiguous_output=False - )(ep.graph_module) - - assert result.modified - assert ( - _count( - result.graph_module, - exir_ops.edge.channels_last.convolution.default, - ) - == 1 - ) - def test_empty_op_map_leaves_graph_unchanged(self): ep = _export_to_edge(Conv2dModule(), (torch.randn(1, 4, 8, 8),)) From 57e481631439be15425c36e7189917183af34a03 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 2 Sep 2026 16:03:17 -0700 Subject: [PATCH 5/5] Transforms: keep layout-copy support narrowly scoped Drop the unrelated transpose-pair guard and the general lower-rank constant reorder fix. Retain only the dialect, metadata, barrier, and region-cancellation behavior consumed by Cortex-M explicit layout. AI-assisted: Codex. --- ...fuse_transpose_or_permute_op_pairs_pass.py | 4 - .../remove_permutes_around_elementwise_ops.py | 43 ++----- .../test/test_permute_optimization_passes.py | 113 ------------------ 3 files changed, 7 insertions(+), 153 deletions(-) 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 166908638bb..b8ae98d9335 100644 --- a/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py +++ b/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py @@ -49,10 +49,6 @@ def can_fuse_for_chain( ) -> bool: if not super().can_fuse_for_chain(producer, consumer, consumer_op_packets): return False - if (producer.target == LAYOUT_PERMUTE_COPY) != ( - consumer.target == LAYOUT_PERMUTE_COPY - ): - return False # checking that permut2(permut1(identity)) == identity, modulo unitary dimensions producer_input = cast(torch.fx.Node, producer.args[0]) diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 9b7dfa49d84..9d66c2384a7 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -765,44 +765,15 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 and const_rank < permute_rank and const_node.meta.get("val") is not None ): - # Broadcasting widens the constant to the region's rank - # before the permutation applies. original_shape = list(const_node.meta["val"].shape) padded = [1] * (permute_rank - const_rank) + original_shape - target_shape = [padded[dim] for dim in node_end_perm] - - # Where each non-unit axis ends up. Unit axes carry no - # elements, so only the order of these decides whether the - # permutation rearranges data or merely reshapes. - destinations = [ - node_end_perm.index(axis) - for axis, size in enumerate(padded) - if size != 1 - ] - if destinations == sorted(destinations): - # Only unit extents moved, so this is a pure reshape and - # a view says it exactly -- and says it for free, since - # view_copy later becomes a memory.view alias. - new_node = graph.create_node( - "call_function", - exir_ops.edge.aten.view_copy.default, - args=(const_node, target_shape), - ) - else: - # Reordering a non-unit extent moves data. A view would - # reinterpret the strides and read different elements, - # so widen with a view and permute at full rank. - widened = graph.create_node( - "call_function", - exir_ops.edge.aten.view_copy.default, - args=(const_node, padded), - ) - with graph.inserting_after(widened): - new_node = graph.create_node( - "call_function", - subgraph.permute_target, - args=(widened, node_end_perm), - ) + target_shape = [padded[d] for d in node_end_perm] + target_shape = target_shape[permute_rank - const_rank :] + new_node = graph.create_node( + "call_function", + exir_ops.edge.aten.view_copy.default, + args=(const_node, target_shape), + ) else: continue user_node.replace_input_with(const_node, new_node) diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index 0271e1ba783..f0ade7139c2 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -811,34 +811,6 @@ def _layout_add_graph( builder.output([to_nhwc]) return builder.get_graph_module(), x_data - def test_pair_fusion_preserves_layout_dialect_across_aten_transpose(self) -> None: - builder = GraphBuilder() - x = builder.placeholder("x", torch.randn(2, 3)) - transpose = builder.call_operator( - op=exir_ops.edge.aten.transpose_copy.int, - args=(x, 0, 1), - ) - layout_permute = builder.call_operator( - op=exir_ops.edge.channels_last.permute_copy.default, - args=(transpose, [1, 0]), - ) - builder.output([layout_permute]) - 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.aten.transpose_copy.int), 1 - ) - self.assertEqual( - count_node( - result.graph_module, - exir_ops.edge.channels_last.permute_copy.default, - ), - 1, - ) - 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) @@ -2546,91 +2518,6 @@ def test_backend_policy_can_stop_region_cancellation(self) -> None: count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2 ) - def _assert_region_cancels( - self, module: torch.nn.Module, inputs: tuple[torch.Tensor, ...] - ) -> None: - """The region's boundary permutes go away and the values do not change.""" - module = module.eval() - expected = module(*inputs) - with torch.no_grad(): - exported = torch.export.export(module, inputs) - edge = to_edge( - exported, - compile_config=EdgeCompileConfig( - _check_ir_validity=False, _skip_dim_order=True - ), - ) - before = count_node( - edge.exported_program().graph_module, - exir_ops.edge.aten.permute_copy.default, - ) - transformed = edge.transform([RemovePermutesAroundElementwiseOps()]) - actual = transformed.exported_program().module()(*inputs) - - after = count_node( - transformed.exported_program().graph_module, - exir_ops.edge.aten.permute_copy.default, - ) - self.assertLess(after, before, "the boundary permutes should have cancelled") - torch.testing.assert_close(actual, expected) - - def test_lower_rank_constant_reorder_preserves_values(self) -> None: - """A broadcast constant is widened and permuted, never reinterpreted. - - A view cannot express the reorder -- it reinterprets strides rather than - moving elements -- so each rank below the region's needs its own case. - """ - - class Rank3(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.register_buffer( - "bias", - torch.arange(4 * 8 * 8, dtype=torch.float32).reshape(4, 8, 8), - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return (x.permute(0, 3, 1, 2) + self.bias).permute(0, 2, 3, 1) - - class Rank2(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.register_buffer( - "bias", torch.arange(8, dtype=torch.float32).reshape(1, 8) - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return (x.permute(2, 0, 1) + self.bias).permute(1, 2, 0) - - class Rank1(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.register_buffer("scale", torch.arange(1.0, 9.0)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return (x.permute(0, 3, 1, 2) * self.scale).permute(0, 2, 3, 1) - - class EqualExtents(torch.nn.Module): - """Two non-unit axes of the same size swap. - - The extents read the same before and after, so only their order - distinguishes a reshape from a reorder. - """ - - def __init__(self) -> None: - super().__init__() - self.register_buffer( - "bias", torch.arange(4, dtype=torch.float32).reshape(2, 2) - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return (x.permute(0, 1, 3, 2) + self.bias).permute(0, 1, 3, 2) - - self._assert_region_cancels(Rank3(), (torch.randn(1, 8, 8, 4),)) - self._assert_region_cancels(Rank2(), (torch.randn(8, 8, 8),)) - self._assert_region_cancels(Rank1(), (torch.randn(1, 8, 8, 8),)) - self._assert_region_cancels(EqualExtents(), (torch.randn(1, 4, 2, 2),)) - class LayoutPermuteVisibilityTest(unittest.TestCase): """The data-movement passes must see both permute dialects.