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/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..69df4fc69c5 100644 --- a/backends/transforms/propagate_view_copy_permute_pass.py +++ b/backends/transforms/propagate_view_copy_permute_pass.py @@ -40,7 +40,9 @@ 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; synthesized permutes retain their source + target. """ @@ -77,8 +79,7 @@ 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, @@ -764,7 +765,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..9d66c2384a7 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -7,15 +7,13 @@ # pyre-unsafe +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 ( - is_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 @@ -29,12 +27,21 @@ 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`` 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 + 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. @@ -62,8 +69,16 @@ 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, + permute_targets: set | frozenset | None = None, + ) -> None: super().__init__() + self.can_propagate = can_propagate + self._permute_targets = frozenset(permute_targets or PERMUTE_COPY_TARGETS) self._permutable_ops = { exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.mul.Tensor, @@ -198,7 +213,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 +385,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: @@ -384,7 +399,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: @@ -482,7 +497,9 @@ 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: + 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)) @@ -500,6 +517,8 @@ def visit( # noqa: C901 ) elif user.op == "output": return False + elif self.can_propagate is not None and not self.can_propagate(user): + 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. @@ -513,8 +532,11 @@ def visit( # noqa: C901 # Traverse upstream: for inp in node.all_input_nodes: - if inp.target in PERMUTE_COPY_TARGETS: - if self.get_permutation(inp) != current_start_permute: + if inp.target in self._permute_targets: + 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: @@ -599,6 +621,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 +737,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 +757,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, + subgraph.permute_target, args=(const_node, node_end_perm), ) elif ( @@ -756,7 +780,7 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 # 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 +793,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 +806,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: @@ -902,7 +929,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/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index 1f357deb171..f0ade7139c2 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -782,6 +782,135 @@ 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 + + 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", + ) + + 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_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 +# ───────────────────────────────────── + + class ReplaceNopTransposeOrPermuteWithViewTest(unittest.TestCase): def test_replace_nop_transpose_with_view_float(self) -> None: x = torch.randn(2, 1, 3, 1) @@ -2354,6 +2483,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 + ) + class LayoutPermuteVisibilityTest(unittest.TestCase): """The data-movement passes must see both permute dialects.