diff --git a/backends/transforms/channels_last_ops.py b/backends/transforms/channels_last_ops.py index 0a90f0d1889..d7bfd76fb50 100644 --- a/backends/transforms/channels_last_ops.py +++ b/backends/transforms/channels_last_ops.py @@ -175,6 +175,7 @@ def _permute_copy(input, dims): lib.impl("max_pool2d", _max_pool2d, "CompositeExplicitAutograd") register_fake("channels_last::max_pool2d", _max_pool2d, lib=lib) + lib.define( "grid_sampler_2d(Tensor input, Tensor grid, int interpolation_mode, " "int padding_mode, bool align_corners) -> Tensor" diff --git a/backends/transforms/decompose_channels_last_pass.py b/backends/transforms/decompose_channels_last_pass.py index 05ec247c30a..2a28600eed8 100644 --- a/backends/transforms/decompose_channels_last_pass.py +++ b/backends/transforms/decompose_channels_last_pass.py @@ -27,6 +27,10 @@ exir_ops.edge.channels_last.grid_sampler_2d.default: exir_ops.edge.aten.grid_sampler_2d.default, } +_DIRECT_DECOMPOSITIONS = { + exir_ops.edge.channels_last.permute_copy.default: exir_ops.edge.aten.permute_copy.default, +} + class DecomposeChannelsLastPass(ExportPass): """Decompose channels_last dialect ops into permute + aten op + permute. @@ -39,6 +43,10 @@ class DecomposeChannelsLastPass(ExportPass): """ def call_operator(self, op, args, kwargs, meta): + direct_op = _DIRECT_DECOMPOSITIONS.get(op) + if direct_op is not None: + return super().call_operator(direct_op, args, kwargs, meta) + aten_op = _DECOMPOSITIONS.get(op) if aten_op is not None: nchw_in = super().call_operator( @@ -90,8 +98,4 @@ def call_operator(self, op, args, kwargs, meta): meta, ) return values, indices - if op == exir_ops.edge.channels_last.permute_copy.default: - return super().call_operator( - exir_ops.edge.aten.permute_copy.default, args, kwargs, meta - ) return super().call_operator(op, args, kwargs, meta) diff --git a/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py b/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py index 83b33533995..62183da2ebd 100644 --- a/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py +++ b/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py @@ -6,6 +6,7 @@ # pyre-unsafe +from collections import deque from typing import Any, Callable, cast import torch @@ -41,6 +42,42 @@ class FuseTransposeOrPermuteOpPairsPass(FuseOpPairsAcrossBranchesPass): exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, } + def __init__( + self, + can_propagate: Callable[[torch.fx.Node], bool] | None = None, + ) -> None: + super().__init__() + self.can_propagate = can_propagate + + def get_fuse_candidates( + self, + producer: torch.fx.Node, + consumer_op_packets: set[EdgeOpOverloadPacket], + bypass_ops: set[EdgeOpOverload], + ) -> list[torch.fx.Node]: + if self.can_propagate is None: + return super().get_fuse_candidates( + producer, consumer_op_packets, bypass_ops + ) + + users = deque(producer.users) + visited: set[torch.fx.Node] = set() + removal_candidates = [] + while users: + user = users.popleft() + if user in visited: + continue + visited.add(user) + if user.target in bypass_ops: + if not self.can_propagate(user): + return [] + users.extend(user.users) + elif self.can_fuse_for_chain(producer, user, consumer_op_packets): + removal_candidates.append(user) + else: + return [] + return removal_candidates + def can_fuse_for_chain( self, producer: torch.fx.Node, diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 34068e97ecd..0344ee3a58f 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -7,12 +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 ( + ATEN_PERMUTE_COPY, is_permute_copy, PERMUTE_COPY_TARGETS, ) @@ -29,6 +31,8 @@ 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. """ @dataclass() @@ -45,6 +49,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) @@ -56,8 +65,14 @@ 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, + ) -> None: super().__init__() + self.can_propagate = can_propagate self._permutable_ops = { exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.mul.Tensor, @@ -329,9 +344,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 self._interleave_cache.clear() subgraphs_found: list[RemovePermutesAroundElementwiseOps.Subgraph] = [] processed_nodes: set[torch.fx.Node] = set() - for node in graph_module.graph.nodes: - if not is_permute_copy(node): - continue + permute_nodes = [ + node for node in graph_module.graph.nodes if is_permute_copy(node) + ] + for node in permute_nodes: start_permute = self.get_permutation(node) if start_permute is None: continue @@ -516,7 +532,9 @@ def visit( # noqa: C901 continue return False elif user.op == "output": - return False + subgraph.output_boundaries.add( + (users_source, user, tuple(downstream_start)) + ) elif self._is_permutation_sink_view(user): # The permutation dies at this reshape (see # _is_permutation_sink_view), so terminate the region here with @@ -525,6 +543,13 @@ def visit( # noqa: C901 # terminates cleanly, whereas crossing it would leave the region # hunting for an end permute that layout-invariance made moot. continue + elif self.can_propagate is not None and not self.can_propagate(user): + # A backend barrier. 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. + subgraph.output_boundaries.add( + (users_source, user, tuple(downstream_start)) + ) elif not self.visit( user, subgraph, processed_nodes, downstream_end, downstream_start ): @@ -618,6 +643,8 @@ def _is_pointwise(target) -> bool: return False def is_node_permutable(self, node: torch.fx.Node) -> bool: + if self.can_propagate is not None and not self.can_propagate(node): + return False if node.target in self._PAD_OPS and not self._is_constant_pad(node): return False if node.target in self._permutable_ops: @@ -654,6 +681,8 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 # Ensure that the subgraph's edges have not been modified by an earlier rewrite before applying changes. if not self._subgraph_edges_are_current(subgraph): return False + if not self._constant_edges_are_free(subgraph): + return False # Nodes belonging to a repeat_interleave triple are rewritten as a unit # below, so they must skip the per-node dim handling and the view rank @@ -736,9 +765,10 @@ 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, + ATEN_PERMUTE_COPY, args=(const_node, node_end_perm), ) + new_node.meta = {} elif ( const_rank is not None and const_rank < permute_rank @@ -757,6 +787,8 @@ 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 PERMUTE_COPY_TARGETS @@ -764,7 +796,7 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 return True - def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: + def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: # noqa: C901 """Return false if an earlier rewrite invalidated this candidate.""" for inp, out in subgraph.edges_in: if inp.target not in PERMUTE_COPY_TARGETS or inp not in out.all_input_nodes: @@ -778,6 +810,18 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: if const_node not in user_node.all_input_nodes: return False + for producer, output_node, _ in subgraph.output_boundaries: + if producer not in output_node.all_input_nodes: + return False + future_occurrences = self._node_argument_count(output_node, producer) + future_occurrences += sum( + self._node_argument_count(output_node, permute) + for source, permute in subgraph.edges_out + if source is producer + ) + if future_occurrences != 1: + return False + for head, (_, _, expand_node, view_node) in subgraph.interleaves.items(): if ( len(head.users) != 1 @@ -789,6 +833,69 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: return True + def _insert_output_boundary_permutations(self, subgraph: Subgraph) -> None: + 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( + ATEN_PERMUTE_COPY, + args=(producer, list(permutation)), + ) + new_permute.meta = dict(producer.meta) + for output in outputs: + output.replace_input_with(producer, new_permute) + + @staticmethod + def _node_argument_count(node: torch.fx.Node, target: torch.fx.Node) -> int: + count = 0 + + def visit(argument): + nonlocal count + if argument is target: + count += 1 + return argument + + torch.fx.map_arg((node.args, node.kwargs), visit) + return count + + def _constant_edges_are_free(self, subgraph: Subgraph) -> bool: + """Reject a rewrite that would make a constant need a real permute.""" + for const_node, user_node in subgraph.constant_edges_in: + node_end_perm = subgraph.node_end_permute.get( + user_node, subgraph.end_permute + ) + if self._constant_transform_requires_data_copy(const_node, node_end_perm): + return False + return True + + def _constant_transform_requires_data_copy( + self, node: torch.fx.Node, permutation: list[int] + ) -> bool: + val = node.meta.get("val") + if not isinstance(val, torch.Tensor): + return True + shape = list(val.shape) + if len(shape) > len(permutation) or not all( + isinstance(dim, int) for dim in shape + ): + return True + rank_difference = len(permutation) - len(shape) + padded_shape = [1] * rank_difference + shape + output_shape = [padded_shape[dim] for dim in permutation] + if any(size != 1 for size in output_shape[:rank_difference]): + return True + old_order = [dim for dim, size in enumerate(padded_shape) if size != 1] + new_order = [dim for dim, size in zip(permutation, output_shape) if size != 1] + return old_order != new_order + def update_interleave( self, head: torch.fx.Node, diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index ed5578f7fb0..7884ac70d10 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -101,6 +101,38 @@ def get_compute_nodes( class FuseCascadedTransposeOrPermuteOpsTest(unittest.TestCase): + def test_ordinary_mixed_inverse_permutes_are_fused(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 2, 3, 4)) + first = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + second = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(first, [0, 3, 1, 2]), + ) + builder.output([second]) + graph_module = builder.get_graph_module() + + result = cast( + PassResult, + FuseCascadedTransposeOrPermuteOps()(graph_module), + ) + + self.assertTrue(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 0, + ) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), + 0, + ) + def test_permute_transpose_fusion(self) -> None: builder = GraphBuilder() x = builder.placeholder("x", torch.randn(3, 1, 3, 1, 4)) @@ -593,6 +625,7 @@ def test_negative_not_squeeze_like(self) -> None: class FuseTransposeOrPermuteOpPairsTest(unittest.TestCase): + def test_per_tensor_qdq_is_bypassed(self) -> None: for op, x_data in ( ( @@ -778,10 +811,282 @@ def test_per_channel_branch_blocks_shared_permute_fusion(self) -> None: # ────────────────────────────────────────────────────────────────────── -# Tests for ReplaceNopTransposeOrPermuteWithViewPass +# 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_respects_backend_propagation_barrier(self) -> None: + builder = GraphBuilder() + x_data = torch.randn(1, 2, 3, 4) + x = builder.placeholder("x", x_data) + to_nhwc = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + quantize = builder.call_operator( + op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(to_nhwc, 0.25, 0, -128, 127, torch.int8), + ) + to_nchw = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(quantize, [0, 3, 1, 2]), + ) + builder.output([to_nchw]) + graph_module = builder.get_graph_module() + + result = cast( + PassResult, + FuseTransposeOrPermuteOpPairsPass( + can_propagate=lambda node: node.target + != exir_ops.edge.quantized_decomposed.quantize_per_tensor.default + )(graph_module), + ) + + self.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 2, + ) + + def test_pair_fusion_does_not_bypass_structural_per_channel_qdq(self) -> None: + for op, x_data in ( + ( + exir_ops.edge.quantized_decomposed.quantize_per_channel.default, + torch.randn(1, 2, 3, 4), + ), + ( + exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, + torch.randint(-128, 127, (1, 2, 3, 4), dtype=torch.int8), + ), + ): + with self.subTest(op=op): + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + scales = builder.placeholder("scales", torch.tensor([0.25, 0.5])) + zero_points = builder.placeholder( + "zero_points", torch.tensor([0, 0], dtype=torch.int64) + ) + to_nhwc = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + qdq = builder.call_operator( + op=op, + args=(to_nhwc, scales, zero_points, 3, -128, 127, torch.int8), + ) + to_nchw = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(qdq, [0, 3, 1, 2]), + ) + builder.output([to_nchw]) + graph_module = builder.get_graph_module() + + result = cast( + PassResult, FuseTransposeOrPermuteOpPairsPass()(graph_module) + ) + + self.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 2, + ) + + def test_layout_copy_rejects_spatial_constant_reordering(self) -> None: + bias_data = torch.randn(4, 8, 8) + graph_module, x_data = self._layout_add_graph("b_bias", bias_data) + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps()(graph_module), + ) + + self.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ), + 2, + ) + validate_numerics( + before, + result.graph_module, + [x_data, bias_data], + "RemovePermutesAroundElementwiseOps", + ) + + def test_layout_copy_reshapes_channel_constant_without_copy(self) -> None: + bias_data = torch.randn(4, 1, 1) + graph_module, x_data = self._layout_add_graph("b_bias", bias_data) + before = copy.deepcopy(graph_module) + + result = cast( + PassResult, + RemovePermutesAroundElementwiseOps()(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)