diff --git a/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py b/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py index e84343f8cac..9e909e1a338 100644 --- a/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py +++ b/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py @@ -15,6 +15,9 @@ class RemovePermutesAroundElementwiseTosaOps(RemovePermutesAroundElementwiseOps): + # The base takes an optional program; this pass always has one. + exported_program: ExportedProgram + def __init__(self, exported_program: ExportedProgram) -> None: super().__init__( extra_permutable_ops={ @@ -28,9 +31,7 @@ def __init__(self, exported_program: ExportedProgram) -> None: def _is_constant(self, node: torch.fx.Node) -> bool: # Override fragile string match check with exported program check - exported_program = self.exported_program - assert exported_program is not None - return super()._is_constant(node) or is_param_node(exported_program, node) + return super()._is_constant(node) or is_param_node(self.exported_program, node) def permute_subgraph(self, subgraph) -> bool: # TABLE lookup inputs are already tied to the table layout. diff --git a/backends/transforms/absorb_boundary_layout_copies.py b/backends/transforms/absorb_boundary_layout_copies.py new file mode 100644 index 00000000000..d97ef62e900 --- /dev/null +++ b/backends/transforms/absorb_boundary_layout_copies.py @@ -0,0 +1,208 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from dataclasses import dataclass, field + +import torch + +from executorch.backends.transforms.channels_last_layout import is_layout_copy +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + +# Elementwise per-tensor quantization is layout-agnostic, so a layout copy on +# the far side of one is still a boundary copy. Per-channel quantization is not: +# its axis is dimension-dependent. +_LAYOUT_AGNOSTIC_TARGETS = frozenset( + { + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + } +) + + +def _inverse(dims: tuple[int, ...]) -> list[int]: + inverse = [0] * len(dims) + for position, dim in enumerate(dims): + inverse[dim] = position + return inverse + + +@dataclass(frozen=True) +class BoundaryLayoutContract: + """Which method inputs and outputs changed layout, and to what. + + An entry ``{0: (0, 2, 3, 1)}`` in ``inputs`` means argument 0 must now be + passed as ``argument.permute(0, 2, 3, 1)``. An entry in ``outputs`` means + the returned tensor needs the same permutation applied to recover what the + method used to return. + """ + + inputs: dict[int, tuple[int, ...]] = field(default_factory=dict) + outputs: dict[int, tuple[int, ...]] = field(default_factory=dict) + + def __bool__(self) -> bool: + return bool(self.inputs or self.outputs) + + +class AbsorbBoundaryLayoutCopies(ExportPass): + """Move layout copies that sit on the method boundary into the signature. + + A layout region formed by ``ToContiguousChannelsLastPass`` is bracketed by + ``channels_last.permute_copy``. Those in the interior cancel against each + other; the ones on the boundary have nothing to cancel against and survive. + Deleting them and declaring the corresponding method input or output to be + channels-last moves the transpose to the caller, which is free whenever the + caller already has the data in that layout. + + Run this *after* region formation. Permuting the boundary first and hoping + the copies cancel is measurably worse: it inserts copies into graphs that + have no anchors at all, where nothing can cancel them. + + A copy need not touch the boundary directly. Quantized graphs put a + per-tensor ``quantize``/``dequantize`` in between, which reorders nothing, + so the search walks through those and relabels them on the way. + + Changing a method's layout is caller-visible, so the applied changes are + reported in ``contract`` rather than assumed. + """ + + def __init__(self, exported_program: ExportedProgram) -> None: + super().__init__() + self.exported_program = exported_program + self.contract = BoundaryLayoutContract() + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + if graph_module is not self.exported_program.graph_module: + raise RuntimeError( + "AbsorbBoundaryLayoutCopies rewrites the ExportedProgram's graph " + "and signature together; run it as its own transform rather than " + "after a pass that replaces the graph module." + ) + + inputs = self._absorb_inputs(graph_module) + outputs = self._absorb_outputs(graph_module) + self.contract = BoundaryLayoutContract(inputs=inputs, outputs=outputs) + + modified = bool(self.contract) + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) + + def _forward_to_copies(self, node: torch.fx.Node): + """Follow every path out of ``node`` until it reaches a layout copy. + + Returns the layout-agnostic nodes crossed, the copies terminating the + paths, and the permutation they agree on; ``None`` if any path ends + somewhere else or the copies disagree. + """ + interior: list[torch.fx.Node] = [] + copies: list[torch.fx.Node] = [] + dims: tuple[int, ...] | None = None + frontier = [node] + + while frontier: + users = list(frontier.pop().users) + if not users: + return None + for user in users: + if is_layout_copy(user): + user_dims = tuple(user.args[1]) + if dims is not None and user_dims != dims: + # A residual block feeds one value to several branches, + # each with its own copy. They collapse to a single + # contract entry only if they agree. + return None + dims = user_dims + copies.append(user) + elif user.target in _LAYOUT_AGNOSTIC_TARGETS: + if user in interior: + continue + interior.append(user) + frontier.append(user) + else: + return None + + return None if dims is None else (interior, copies, dims) + + def _backward_to_copy(self, result: torch.fx.Node): + """Walk back from a returned value to the layout copy that produced it.""" + interior: list[torch.fx.Node] = [] + current = result + + while True: + if is_layout_copy(current): + if len(current.users) != 1: + return None + source = current.args[0] + if not isinstance(source, torch.fx.Node): + return None + return interior, current, source, tuple(current.args[1]) + if ( + current.target not in _LAYOUT_AGNOSTIC_TARGETS + or len(current.users) != 1 + ): + return None + interior.append(current) + current = current.args[0] + if not isinstance(current, torch.fx.Node): + return None + + def _absorb_inputs(self, graph_module) -> dict[int, tuple[int, ...]]: + user_inputs = list(self.exported_program.graph_signature.user_inputs) + absorbed: dict[int, tuple[int, ...]] = {} + + for node in list(graph_module.graph.nodes): + if node.op != "placeholder" or node.name not in user_inputs: + continue + found = self._forward_to_copies(node) + if found is None: + continue + interior, copies, dims = found + + for member in (node, *interior): + member.meta["val"] = member.meta["val"].permute(dims) + for copy in copies: + copy.replace_all_uses_with(copy.args[0]) + graph_module.graph.erase_node(copy) + absorbed[user_inputs.index(node.name)] = dims + + return absorbed + + def _absorb_outputs(self, graph_module) -> dict[int, tuple[int, ...]]: + output_node = graph_module.graph.output_node() + results = list(output_node.args[0]) + specs = self.exported_program.graph_signature.output_specs + absorbed: dict[int, tuple[int, ...]] = {} + + for index, result in enumerate(results): + if not isinstance(result, torch.fx.Node): + continue + found = self._backward_to_copy(result) + if found is None: + continue + interior, copy, source, dims = found + + for member in interior: + member.meta["val"] = member.meta["val"].permute(_inverse(dims)) + if interior: + interior[-1].replace_input_with(copy, source) + else: + results[index] = source + # The manager re-derives the signature, but the direct + # exported_program= path does not. + if index < len(specs) and getattr(specs[index].arg, "name", None) == ( + result.name + ): + specs[index].arg.name = source.name + absorbed[index] = dims + + if absorbed: + output_node.args = (results,) + return absorbed diff --git a/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py b/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py index 24ed19a9310..6b08f1ef415 100644 --- a/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py +++ b/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py @@ -6,10 +6,6 @@ # pyre-unsafe -from collections.abc import Callable - -import torch - from executorch.backends.transforms.channels_last_layout import ( composed_permute_target, PERMUTE_COPY_TARGETS, @@ -42,13 +38,6 @@ class FuseCascadedTransposeOrPermuteOps(RemoveOrReplacePassInterface): exir_ops.edge.aten.view.default, } - def __init__( - self, - can_propagate: Callable[[torch.fx.Node], bool] | None = None, - ) -> None: - super().__init__() - self.can_propagate = can_propagate - @property def targets(self) -> list[EdgeOpOverload]: return list(self.transpose_or_permute_target) @@ -119,8 +108,6 @@ def _apply_view_to_dims( def _fuse_across_view(self, node: Node, view_node: Node) -> bool: # noqa: C901 """Fuse permute -> view(squeeze/unsqueeze) -> permute into a view_copy.""" - if self.can_propagate is not None and not self.can_propagate(view_node): - return False # view_node must have exactly one user (this permute node) if len(view_node.users) != 1: return False diff --git a/backends/transforms/postpone_permute_below_squeeze_view.py b/backends/transforms/postpone_permute_below_squeeze_view.py index fe6ea618d91..226924cbea7 100644 --- a/backends/transforms/postpone_permute_below_squeeze_view.py +++ b/backends/transforms/postpone_permute_below_squeeze_view.py @@ -7,7 +7,6 @@ # pyre-unsafe -from collections.abc import Callable from typing import cast, List import torch @@ -36,13 +35,6 @@ class PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView(RemoveOrReplacePassInterf mean the view_copy is normalized from squeeze or unsqueeze. """ - def __init__( - self, - can_propagate: Callable[[torch.fx.Node], bool] | None = None, - ) -> None: - super().__init__() - self.can_propagate = can_propagate - @property def targets(self) -> list[EdgeOpOverload]: return list(PERMUTE_COPY_TARGETS) @@ -77,9 +69,6 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: exir_ops.edge.aten.view.default, ): return False - view_node = users[0] - if self.can_propagate is not None and not self.can_propagate(view_node): - return False # If the permute_node/view_node was newly added to the # graph, it may not have the meta["val"] FakeTensor. @@ -90,6 +79,8 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: permute_node_shape = [*cast(list, get_shape(node.graph.owning_module, node))] permute_dims = cast(list, node.args[1]) + view_node = users[0] + if view_node.meta.get("val") is None: return False diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index a6d55deb4f1..8b6b9ba461f 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -77,12 +77,10 @@ def __init__( extra_permutable_ops: set | None = None, *, exported_program: ExportedProgram | None = None, - allow_layout_boundary_propagation: bool = False, can_propagate: Callable[[torch.fx.Node], bool] | None = None, ) -> None: super().__init__() self.exported_program = exported_program - self.allow_layout_boundary_propagation = allow_layout_boundary_propagation self.can_propagate = can_propagate self._permutable_ops = { exir_ops.edge.aten.add.Tensor, @@ -362,9 +360,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 start_permute = self.get_permutation(node) if start_permute is None: continue - layout_region = self.allow_layout_boundary_propagation and is_layout_copy( - node - ) + layout_region = is_layout_copy(node) # Expected end permutation for the subgraph. end_permute = [start_permute.index(i) for i in range(len(start_permute))] @@ -453,30 +449,29 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 for n in subgraph.nodes: processed_nodes.add(n) - if self.allow_layout_boundary_propagation: - for node in permute_nodes: - end_permute = self.get_permutation(node) - if end_permute is None: - continue - if not is_layout_copy(node): - continue - producer = node.args[0] if node.args else None - if not isinstance(producer, torch.fx.Node): - continue - if ( - not self.is_node_permutable(producer) - and self._interleave_triple(producer) is None - ): - continue - start_permute = [end_permute.index(i) for i in range(len(end_permute))] - subgraph = self.Subgraph( - start_permute, - end_permute, - layout_region=True, - ) - if self.visit(producer, subgraph, processed_nodes): - subgraphs_found.append(subgraph) - processed_nodes.update(subgraph.nodes) + for node in permute_nodes: + end_permute = self.get_permutation(node) + if end_permute is None: + continue + if not is_layout_copy(node): + continue + producer = node.args[0] if node.args else None + if not isinstance(producer, torch.fx.Node): + continue + if ( + not self.is_node_permutable(producer) + and self._interleave_triple(producer) is None + ): + continue + start_permute = [end_permute.index(i) for i in range(len(end_permute))] + subgraph = self.Subgraph( + start_permute, + end_permute, + layout_region=True, + ) + if self.visit(producer, subgraph, processed_nodes): + subgraphs_found.append(subgraph) + processed_nodes.update(subgraph.nodes) modified = False for subgraph in subgraphs_found: @@ -599,10 +594,7 @@ def visit( # noqa: C901 continue return False elif user.op == "output": - if ( - not self.allow_layout_boundary_propagation - or not subgraph.layout_region - ): + if not subgraph.layout_region: return False subgraph.output_boundaries.add( (users_source, user, tuple(downstream_start)) @@ -616,8 +608,7 @@ def visit( # noqa: C901 # hunting for an end permute that layout-invariance made moot. continue elif ( - self.allow_layout_boundary_propagation - and subgraph.layout_region + subgraph.layout_region and self.can_propagate is not None and not self.can_propagate(user) ): @@ -653,16 +644,13 @@ def visit( # noqa: C901 return False subgraph.constant_edges_in.add((inp, node)) elif self._is_user_input(inp): - if ( - not self.allow_layout_boundary_propagation - or not subgraph.layout_region - or self._get_node_rank(inp) != len(current_end_permute) + if not subgraph.layout_region or self._get_node_rank(inp) != len( + current_end_permute ): return False subgraph.input_boundaries.add((inp, node, tuple(current_end_permute))) elif ( - self.allow_layout_boundary_propagation - and subgraph.layout_region + subgraph.layout_region and self.can_propagate is not None and not self.can_propagate(inp) and self._get_node_rank(inp) == len(current_end_permute) diff --git a/backends/transforms/replace_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py index 40f38df574b..641c3925b65 100644 --- a/backends/transforms/replace_ops_with_channels_last_variants.py +++ b/backends/transforms/replace_ops_with_channels_last_variants.py @@ -151,8 +151,6 @@ def __init__( self.op_map: dict[Target, ChannelsLastOpSpec] = ( op_map if op_map is not None else dict(_DEFAULT_OP_MAP) ) - self.candidate_count = 0 - self.replacement_count = 0 @staticmethod def _permute_node_input( @@ -202,8 +200,6 @@ def _permute_node_output( original_node_output.replace_all_uses_with(output) def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - self.candidate_count = 0 - self.replacement_count = 0 modified = False graph = graph_module.graph @@ -212,14 +208,13 @@ 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 - self.candidate_count += 1 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): + continue # In case of implicit batch size, insert also `unsqueeze_copy.default` and `squeeze_copy.dims` operators. # With `convolution`, this already happens during lowering to edge. But it doesn't happen for example with @@ -279,7 +274,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: graph.erase_node(node) modified = True - self.replacement_count += 1 if modified: graph.eliminate_dead_code() diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index e15e089b051..913ea377687 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -538,9 +538,9 @@ def define_common_targets(): ) runtime.python_library( - name = "to_contiguous_channels_last_pass", + name = "absorb_boundary_layout_copies", srcs = [ - "to_contiguous_channels_last_pass.py", + "absorb_boundary_layout_copies.py", ], visibility = [ "//executorch/backends/...", @@ -548,14 +548,6 @@ def define_common_targets(): deps = [ "//caffe2:torch", ":channels_last_layout", - ":fuse_cascaded_transpose_or_permute_ops", - ":fuse_cascaded_view_ops", - ":fuse_transpose_or_permute_op_pairs_pass", - ":postpone_permute_below_squeeze_view", - ":remove_permutes_around_elementwise_ops", - ":replace_nop_transpose_or_permute_with_view", - ":replace_ops_with_channels_last_variants", - ":replace_squeeze_unsqueeze_with_view", "//executorch/exir:lib", "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", @@ -563,30 +555,30 @@ def define_common_targets(): ) runtime.python_test( - name = "test_replace_ops_with_channels_last_variants", + name = "test_absorb_boundary_layout_copies", srcs = [ - "test/test_replace_ops_with_channels_last_variants.py", + "test/test_absorb_boundary_layout_copies.py", ], deps = [ "//caffe2:torch", - ":channels_last_ops", - ":remove_getitem_op", - ":replace_ops_with_channels_last_variants", + ":absorb_boundary_layout_copies", "//executorch/exir:lib", + "//executorch/exir/dialects:lib", "fbsource//third-party/pypi/pytest:pytest", ], ) runtime.python_test( - name = "test_to_contiguous_channels_last_pass", + name = "test_replace_ops_with_channels_last_variants", srcs = [ - "test/test_to_contiguous_channels_last_pass.py", + "test/test_replace_ops_with_channels_last_variants.py", ], deps = [ "//caffe2:torch", - ":to_contiguous_channels_last_pass", + ":channels_last_ops", + ":remove_getitem_op", + ":replace_ops_with_channels_last_variants", "//executorch/exir:lib", - "//executorch/exir/dialects:lib", "fbsource//third-party/pypi/pytest:pytest", ], ) diff --git a/backends/transforms/test/test_absorb_boundary_layout_copies.py b/backends/transforms/test/test_absorb_boundary_layout_copies.py new file mode 100644 index 00000000000..a3bab1b7ece --- /dev/null +++ b/backends/transforms/test/test_absorb_boundary_layout_copies.py @@ -0,0 +1,214 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import pytest +import torch + +from executorch.backends.transforms.absorb_boundary_layout_copies import ( + AbsorbBoundaryLayoutCopies, +) +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops + +_LAYOUT_COPY = exir_ops.edge.channels_last.permute_copy.default +_ATEN_PERMUTE = exir_ops.edge.aten.permute_copy.default +_QUANTIZE = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default +_DEQUANTIZE = exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default + +_TO_NHWC = (0, 2, 3, 1) +_TO_NCHW = (0, 3, 1, 2) + + +class Region(torch.nn.Module): + """A layout region: a body bracketed by a permute and its inverse.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(x.permute(*_TO_NHWC)).permute(*_TO_NCHW) + + +class Fork(torch.nn.Module): + """A region whose entry copy feeds two consumers.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = x.permute(*_TO_NHWC) + return (torch.relu(y) + torch.sigmoid(y)).permute(*_TO_NCHW) + + +class MixedUse(torch.nn.Module): + """An input read both through a region and directly.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(x.permute(*_TO_NHWC)).permute(*_TO_NCHW) + x + + +def _count(graph_module, target) -> int: + return sum( + node.op == "call_function" and node.target == target + for node in graph_module.graph.nodes + ) + + +def _build_region(module, inputs): + """Export ``module`` and retarget its permutes to the layout dialect. + + ``ToContiguousChannelsLastPass`` emits exactly these nodes, but building + them here keeps this suite independent of it: absorption is defined against + the dialect operator, not against whoever produced it. + """ + module.eval() + 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 + ), + ) + graph_module = edge.exported_program().graph_module + for node in graph_module.graph.nodes: + if node.op == "call_function" and node.target == _ATEN_PERMUTE: + node.target = _LAYOUT_COPY + graph_module.recompile() + assert _count(graph_module, _LAYOUT_COPY) > 0 + return edge + + +def _split_shared_copy(edge): + """Give each consumer of the entry copy its own copy. + + Region formation inserts one copy per anchor; export would have collapsed + identical ones, so the fan-out shape is constructed directly. + """ + graph_module = edge.exported_program().graph_module + graph = graph_module.graph + copy = next( + node + for node in graph.nodes + if node.op == "call_function" + and node.target == _LAYOUT_COPY + and node.args[0].op == "placeholder" + ) + users = list(copy.users) + assert len(users) > 1 + for user in users[1:]: + with graph.inserting_before(user): + clone = graph.call_function(_LAYOUT_COPY, copy.args, copy.kwargs) + clone.meta.update(copy.meta) + user.replace_input_with(copy, clone) + graph_module.recompile() + return edge + + +def _absorb(edge): + layout_pass = AbsorbBoundaryLayoutCopies(edge.exported_program()) + return edge.transform([layout_pass]), layout_pass.contract + + +def _run(edge, contract, inputs): + """Invoke the method through its (possibly rewritten) layout contract.""" + args = list(inputs) + for index, dims in contract.inputs.items(): + args[index] = args[index].permute(list(dims)).contiguous() + result = edge.exported_program().module()(*args) + results = list(result) if isinstance(result, (tuple, list)) else [result] + for index, dims in contract.outputs.items(): + results[index] = results[index].permute(list(dims)) + return results[0] if len(results) == 1 else results + + +@pytest.mark.parametrize("module", [Region(), Fork()]) +def test_boundary_copies_are_absorbed_and_numerics_hold(module) -> None: + inputs = (torch.randn(1, 4, 8, 8),) + expected = module.eval()(*inputs) + edge = _build_region(module, inputs) + + edge, contract = _absorb(edge) + + assert contract.inputs and contract.outputs + assert _count(edge.exported_program().graph_module, _LAYOUT_COPY) == 0 + assert torch.allclose(_run(edge, contract, inputs), expected, atol=1e-6) + + +def test_fan_out_collapses_to_one_contract_entry() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + module = Fork() + expected = module.eval()(*inputs) + edge = _split_shared_copy(_build_region(module, inputs)) + assert _count(edge.exported_program().graph_module, _LAYOUT_COPY) == 3 + + edge, contract = _absorb(edge) + + assert list(contract.inputs) == [0] + assert _count(edge.exported_program().graph_module, _LAYOUT_COPY) == 0 + assert torch.allclose(_run(edge, contract, inputs), expected, atol=1e-6) + + +def test_mixed_users_are_left_alone() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + module = MixedUse() + expected = module.eval()(*inputs) + edge = _build_region(module, inputs) + + edge, contract = _absorb(edge) + + assert 0 not in contract.inputs + assert torch.allclose(_run(edge, contract, inputs), expected, atol=1e-6) + + +def test_per_tensor_quantization_is_traversed() -> None: + """A copy behind a per-tensor quantize is still a boundary copy. + + Quantized graphs interpose q/dq between the placeholder and the region; + those reorder nothing, so absorption has to see through them. + """ + inputs = (torch.randn(1, 4, 8, 8),) + edge = _build_region(Region(), inputs) + graph_module = edge.exported_program().graph_module + graph = graph_module.graph + entry = next( + node + for node in graph.nodes + if node.op == "call_function" + and node.target == _LAYOUT_COPY + and node.args[0].op == "placeholder" + ) + placeholder = entry.args[0] + with graph.inserting_after(placeholder): + quantize = graph.call_function( + _QUANTIZE, (placeholder, 1.0, 0, -128, 127, torch.int8) + ) + with graph.inserting_after(quantize): + dequantize = graph.call_function( + _DEQUANTIZE, (quantize, 1.0, 0, -128, 127, torch.int8) + ) + quantize.meta.update(placeholder.meta) + dequantize.meta.update(placeholder.meta) + entry.replace_input_with(placeholder, dequantize) + graph_module.recompile() + + _, contract = _absorb(edge) + + assert contract.inputs == {0: _TO_NHWC} + + +def test_absorbing_is_idempotent() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + edge = _build_region(Region(), inputs) + edge, first = _absorb(edge) + edge, second = _absorb(edge) + + assert first + assert not second + + +def test_signature_stays_valid() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + edge = _build_region(Region(), inputs) + + edge, contract = _absorb(edge) + + assert contract + edge.exported_program()._validate() diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index fc6aec148d5..9db4bd8511f 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -882,9 +882,7 @@ def test_layout_pad_argument_is_remapped(self) -> None: result = cast( PassResult, - RemovePermutesAroundElementwiseOps( - allow_layout_boundary_propagation=True, - )(graph_module), + RemovePermutesAroundElementwiseOps()(graph_module), ) self.assertTrue(result.modified) @@ -1052,9 +1050,7 @@ def test_layout_copy_moves_to_static_output_boundary(self) -> None: result = cast( PassResult, - RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( - graph_module - ), + RemovePermutesAroundElementwiseOps()(graph_module), ) self.assertTrue(result.modified) @@ -1100,7 +1096,6 @@ def test_layout_region_terminates_at_backend_barrier(self) -> None: result = cast( PassResult, RemovePermutesAroundElementwiseOps( - allow_layout_boundary_propagation=True, can_propagate=lambda node: node.target != exir_ops.edge.aten.relu.default, )(graph_module), @@ -1141,9 +1136,7 @@ def test_layout_copy_does_not_fork_to_more_boundaries(self) -> None: result = cast( PassResult, - RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( - graph_module - ), + RemovePermutesAroundElementwiseOps()(graph_module), ) self.assertFalse(result.modified) @@ -1175,9 +1168,7 @@ def test_shared_incoming_layout_copy_is_not_credited_as_removed(self) -> None: result = cast( PassResult, - RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( - graph_module - ), + RemovePermutesAroundElementwiseOps()(graph_module), ) self.assertFalse(result.modified) @@ -1207,9 +1198,7 @@ def test_input_boundary_does_not_alias_permuted_and_unpermuted_input(self) -> No result = cast( PassResult, - RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( - graph_module - ), + RemovePermutesAroundElementwiseOps()(graph_module), ) self.assertFalse(result.modified) @@ -1242,9 +1231,7 @@ def test_output_boundary_does_not_alias_two_output_edges(self) -> None: result = cast( PassResult, - RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( - graph_module - ), + RemovePermutesAroundElementwiseOps()(graph_module), ) self.assertFalse(result.modified) @@ -1282,7 +1269,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: PassResult, RemovePermutesAroundElementwiseOps( exported_program=edge.exported_program(), - allow_layout_boundary_propagation=True, )(graph_module), ) @@ -1305,9 +1291,7 @@ def test_layout_copy_rejects_rank_mismatched_runtime_input(self) -> None: result = cast( PassResult, - RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( - graph_module - ), + RemovePermutesAroundElementwiseOps()(graph_module), ) self.assertFalse(result.modified) @@ -1332,9 +1316,7 @@ def test_layout_copy_rejects_spatial_constant_reordering(self) -> None: result = cast( PassResult, - RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( - graph_module - ), + RemovePermutesAroundElementwiseOps()(graph_module), ) self.assertFalse(result.modified) @@ -1359,9 +1341,7 @@ def test_layout_copy_reshapes_channel_constant_without_copy(self) -> None: result = cast( PassResult, - RemovePermutesAroundElementwiseOps(allow_layout_boundary_propagation=True)( - graph_module - ), + RemovePermutesAroundElementwiseOps()(graph_module), ) self.assertTrue(result.modified) diff --git a/backends/transforms/test/test_to_contiguous_channels_last_pass.py b/backends/transforms/test/test_to_contiguous_channels_last_pass.py index c02f068dcc1..65dd39f09b1 100644 --- a/backends/transforms/test/test_to_contiguous_channels_last_pass.py +++ b/backends/transforms/test/test_to_contiguous_channels_last_pass.py @@ -3,18 +3,17 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import unittest from dataclasses import dataclass from typing import Any, Tuple import pytest - import torch -from executorch.backends.transforms.to_contiguous_channels_last_pass import ( - ToContiguousChannelsLastPass, -) -from executorch.exir import EdgeCompileConfig, to_edge +from executorch.backends.transforms.test import common +from executorch.exir import to_edge_transform_and_lower from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass +from torch.fx import GraphModule +from torch.fx.passes.infra.pass_base import PassResult InputT = Tuple[Any, ...] @@ -354,11 +353,9 @@ def forward(self, x: torch.Tensor): ), "conv1d_rank3": PermuteCountTestCase(Conv1dModule(), (torch.randn(1, 2, 8),), 0), "conv2d_rank3": PermuteCountTestCase( - Conv2dModule(), (torch.randn(2, 8, 8),), 0, 2, 2, 2 - ), - "conv2d_rank4": PermuteCountTestCase( - Conv2dModule(), (torch.randn(1, 2, 8, 8),), 0, 0, 2, 0 + Conv2dModule(), (torch.randn(2, 8, 8),), 0, 2, 0, 2 ), + "conv2d_rank4": PermuteCountTestCase(Conv2dModule(), (torch.randn(1, 2, 8, 8),), 0), "conv3d_rank4": PermuteCountTestCase( Conv3dModule(), (torch.randn(2, 6, 6, 6),), 0, 2, 0, 2 ), @@ -417,34 +414,25 @@ def forward(self, x: torch.Tensor): GroupedConvModule(), (torch.randn(1, 4, 8, 8),), 0, - 0, - 2, - 0, ), "transpose_conv": PermuteCountTestCase( TransposeConvModule(), (torch.randn(1, 2, 8, 8),), 0, - 0, - 2, - 0, ), - "views": PermuteCountTestCase(ViewsModule(), (torch.rand(1, 2, 2, 4),), 0, 2, 4, 2), + "views": PermuteCountTestCase(ViewsModule(), (torch.rand(1, 2, 2, 4),), 0, 2, 0, 2), "transposes": PermuteCountTestCase( TransposesModule(), (torch.randn(1, 2, 3, 4),), 2, 0, - 1, + 2, 0, ), "maxpool2d_dilation": PermuteCountTestCase( MaxPool2dDilatedModule(), (torch.randn(1, 2, 8, 8),), 0, - 0, - 2, - 0, ), "lstm": PermuteCountTestCase( LstmModule(), @@ -452,7 +440,7 @@ def forward(self, x: torch.Tensor): 7, 19, 7, - 16, + 19, ), "groupnorm": PermuteCountTestCase( GroupNormModule(), @@ -464,16 +452,16 @@ def forward(self, x: torch.Tensor): (torch.randn(4, 8),), 11, 24, - 8, - 14, + 11, + 24, ), "multihead_attention_rank3": PermuteCountTestCase( MultiheadAttentionModule(), (torch.randn(2, 4, 8),), 12, 20, - 10, - 18, + 12, + 20, ), "cumsum_rank3_dim0": PermuteCountTestCase( CumsumModule(), @@ -486,41 +474,31 @@ def forward(self, x: torch.Tensor): 0, ), "model_1_conv_maxpool_residual_linear": PermuteCountTestCase( - Model1ConvMaxPoolResidualLinear(), (torch.randn(2, 8, 64),), 2, 7, 6, 7 + Model1ConvMaxPoolResidualLinear(), (torch.randn(2, 8, 64),), 2, 7, 2, 7 ), "model_2_conv_mha_linear_layernorm": PermuteCountTestCase( - Model2ConvMhaLinearLayerNorm(), (torch.randn(2, 8, 32),), 14, 23, 11, 21 + Model2ConvMhaLinearLayerNorm(), (torch.randn(2, 8, 32),), 14, 23, 14, 23 ), "model_3_lstm_linear": PermuteCountTestCase( - Model3LstmLinear(), (torch.randn(2, 16, 8),), 20, 58, 20, 55 + Model3LstmLinear(), (torch.randn(2, 16, 8),), 20, 58, 20, 58 ), "model_4_conv_lstm_linear_layernorm": PermuteCountTestCase( - Model4ConvLstmLinearLayerNorm(), (torch.randn(2, 8, 32),), 37, 106, 36, 103 + Model4ConvLstmLinearLayerNorm(), (torch.randn(2, 8, 32),), 37, 106, 37, 106 ), "model_5_dwconv_gelu_layernorm_avgpool": PermuteCountTestCase( - Model5DwConvGeluLayerNormAvgPool(), (torch.randn(1, 8, 16, 16),), 2, 0, 4, 0 + Model5DwConvGeluLayerNormAvgPool(), (torch.randn(1, 8, 16, 16),), 2, 0, 2, 0 ), "model_6_gru_linear": PermuteCountTestCase( - Model6GruLinear(), (torch.randn(2, 16, 8),), 20, 56, 20, 55 + Model6GruLinear(), (torch.randn(2, 16, 8),), 20, 56, 20, 56 ), "model_7_dwconv_batchnorm_linear": PermuteCountTestCase( Model7DwConvBatchNormLinear(), (torch.randn(2, 8, 64),), 2, 3, 2, 3 ), "model_8_conv_batchnorm_maxpool_residual": PermuteCountTestCase( - Model8ConvBatchNormMaxPoolResidual(), - (torch.randn(1, 8, 16, 16),), - 0, - 0, - 5, - 0, + Model8ConvBatchNormMaxPoolResidual(), (torch.randn(1, 8, 16, 16),), 0 ), "model_9_dilated_conv_batchnorm_avgpool_residual": PermuteCountTestCase( - Model9DilatedConvBatchNormAvgPoolResidual(), - (torch.randn(1, 8, 16, 16),), - 0, - 0, - 5, - 0, + Model9DilatedConvBatchNormAvgPoolResidual(), (torch.randn(1, 8, 16, 16),), 0 ), "model_10_dwconv_batchnorm_linear_cat": PermuteCountTestCase( Model10DwConvBatchNormLinearCat(), (torch.randn(2, 8, 64),), 3, 6, 3, 6 @@ -530,14 +508,12 @@ def forward(self, x: torch.Tensor): (torch.randn(1, 2, 3, 4),), 2, 0, - 0, + 2, 0, ), } -# Channels-last inputs are left alone: the replacement pass only converts -# contiguous anchors, so these pin the fallback rather than a conversion. cases_channels_last = { "conv2d_rank4_channels_last": PermuteCountTestCase( Conv2dModule(), @@ -604,7 +580,7 @@ def forward(self, x: torch.Tensor): (torch.randn(1, 2, 3, 4).to(memory_format=torch.channels_last),), 2, 0, - 1, + 2, 0, ), "maxpool2d_dilation_channels_last": PermuteCountTestCase( @@ -624,475 +600,101 @@ def forward(self, x: torch.Tensor): ), } -_CHANNELS_LAST_XFAILS = { - "views_channels_last": "Views are not supported for channels last tensors", -} -_PERMUTE_TARGETS = { - exir_ops.edge.aten.permute.default, - exir_ops.edge.aten.permute_copy.default, - exir_ops.edge.aten.transpose.int, - exir_ops.edge.aten.transpose_copy.int, - exir_ops.edge.channels_last.permute_copy.default, -} -_VIEW_TARGETS = { - exir_ops.edge.aten._unsafe_view.default, - exir_ops.edge.aten.reshape.default, - exir_ops.edge.aten.squeeze.default, - exir_ops.edge.aten.squeeze.dim, - exir_ops.edge.aten.squeeze.dims, - exir_ops.edge.aten.squeeze_copy.default, - exir_ops.edge.aten.squeeze_copy.dim, - exir_ops.edge.aten.squeeze_copy.dims, - exir_ops.edge.aten.unsqueeze.default, - exir_ops.edge.aten.unsqueeze_copy.default, - exir_ops.edge.aten.view.default, - exir_ops.edge.aten.view_copy.default, -} +class ToContiguousChannelsLastPassTestPass(ExportPass): + """ + A test pass which runs the pass pipeline intended to and verifies that permutes and + views are fused as expected. + + TODO: Currently no permute-view passes are implemented, proof of concept only. + """ + + _PERMUTE_TARGETS = { + exir_ops.edge.aten.permute.default, + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.transpose.int, + exir_ops.edge.aten.transpose_copy.int, + } + _VIEW_TARGETS = { + exir_ops.edge.aten._unsafe_view.default, + exir_ops.edge.aten.reshape.default, + exir_ops.edge.aten.squeeze.default, + exir_ops.edge.aten.squeeze.dim, + exir_ops.edge.aten.squeeze.dims, + exir_ops.edge.aten.squeeze_copy.default, + exir_ops.edge.aten.squeeze_copy.dim, + exir_ops.edge.aten.squeeze_copy.dims, + exir_ops.edge.aten.unsqueeze.default, + exir_ops.edge.aten.unsqueeze_copy.default, + exir_ops.edge.aten.view.default, + exir_ops.edge.aten.view_copy.default, + } + + def __init__(self): + super().__init__() + self.initial_permutes = 0 + self.initial_views = 0 + self.final_permutes = 0 + self.final_views = 0 + def count_ops(self, graph_module: GraphModule, targets: set) -> int: + return sum( + 1 + for node in graph_module.graph.nodes + if node.op == "call_function" and node.target in targets + ) -def _count_ops(graph_module: torch.fx.GraphModule, targets: set) -> int: - return sum( - node.op == "call_function" and node.target in targets - for node in graph_module.graph.nodes - ) + def call(self, graph_module: GraphModule) -> PassResult: + self.initial_permutes = self.count_ops(graph_module, self._PERMUTE_TARGETS) + self.initial_views = self.count_ops(graph_module, self._VIEW_TARGETS) + result = super().call(graph_module) + self.final_permutes = self.count_ops(result.graph_module, self._PERMUTE_TARGETS) + self.final_views = self.count_ops(result.graph_module, self._VIEW_TARGETS) + return result def run_test(case: PermuteCountTestCase) -> None: case.module.eval() with torch.no_grad(): exported_program = torch.export.export(case.module, case.inputs) - edge_program = to_edge( - exported_program, - compile_config=EdgeCompileConfig( - _check_ir_validity=False, - _skip_dim_order=True, - ), - ) - initial_graph = edge_program.exported_program().graph_module - initial_permutes = _count_ops(initial_graph, _PERMUTE_TARGETS) - initial_views = _count_ops(initial_graph, _VIEW_TARGETS) - - layout_pass = ToContiguousChannelsLastPass(edge_program.exported_program()) - transformed = edge_program.transform([layout_pass]) - final_graph = transformed.exported_program().graph_module - final_permutes = _count_ops(final_graph, _PERMUTE_TARGETS) - final_views = _count_ops(final_graph, _VIEW_TARGETS) - - assert initial_permutes == case.expected_initial_permutes - assert initial_views == case.expected_initial_views - assert final_permutes == case.expected_final_permutes - assert final_views == case.expected_final_views - ref_result = exported_program.module()(*case.inputs) - edge_result = transformed.exported_program().module()(*case.inputs) - assert torch.allclose(ref_result, edge_result, atol=1e-6) - - -class TestToContiguousChannelsLastPass(unittest.TestCase): - def test_permute_view_counts(self) -> None: - for name, case in cases.items(): - with self.subTest(name=name): - run_test(case) - - def test_permute_view_counts_channels_last(self) -> None: - for name, case in cases_channels_last.items(): - if name in _CHANNELS_LAST_XFAILS: - continue - with self.subTest(name=name): - run_test(case) - - -class ConvChain(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv1 = torch.nn.Conv2d(4, 4, 3, padding=1) - self.conv2 = torch.nn.Conv2d(4, 4, 3, padding=1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.conv2(torch.relu(self.conv1(x))) - - -class ConvOnly(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.conv(x) - - -class DynamicViewConv(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.conv(x.view(x.shape[0], x.shape[1], x.shape[2], x.shape[3])) - - -class ConvThenLinear(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) - self.linear = torch.nn.Linear(4 * 8 * 8, 3) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.linear(self.conv(x).flatten(1)) - - -class UserPermute(torch.nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return x.permute(0, 2, 3, 1) - - -class ConvChannelBias(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) - self.register_buffer("channel_bias", torch.randn(1, 4, 1, 1)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.conv(x) + self.channel_bias - - -class PadConv(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv = torch.nn.Conv2d(4, 4, 3) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.conv(torch.nn.functional.pad(x, (1, 1, 1, 1))) - - -class ConvSoftmax(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.softmax(self.conv(x), dim=-1) - - -class ConvRuntimeBiasConv(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv1 = torch.nn.Conv2d(4, 4, 3, padding=1) - self.conv2 = torch.nn.Conv2d(4, 4, 3, padding=1) - - def forward(self, x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: - return self.conv2(self.conv1(x) + bias) - - -class ConvSpatialBufferConv(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv1 = torch.nn.Conv2d(4, 4, 3, padding=1) - self.conv2 = torch.nn.Conv2d(4, 4, 3, padding=1) - self.register_buffer("bias", torch.randn(4, 8, 8)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.conv2(self.conv1(x) + self.bias) - - -def _edge(module: torch.nn.Module, inputs: tuple[torch.Tensor, ...]): - exported = torch.export.export(module.eval(), inputs) - return to_edge( - exported, - compile_config=EdgeCompileConfig( - _check_ir_validity=False, - _skip_dim_order=True, - ), - ) - - -def _count(graph_module: torch.fx.GraphModule, target: object) -> int: - return sum( - node.op == "call_function" and node.target == target - for node in graph_module.graph.nodes - ) - - -def test_conv_chain_folds_to_boundary_copies() -> None: - torch.manual_seed(0) - module = ConvChain().eval() - inputs = (torch.randn(1, 4, 8, 8),) - expected = module(*inputs) - edge = _edge(module, inputs) - layout_pass = ToContiguousChannelsLastPass(edge.exported_program()) - - transformed = edge.transform([layout_pass]) - graph_module = transformed.exported_program().graph_module - - assert _count(graph_module, exir_ops.edge.channels_last.convolution.default) == 2 - assert _count(graph_module, exir_ops.edge.channels_last.permute_copy.default) == 2 - assert layout_pass.report.candidate_anchor_count == 2 - assert layout_pass.report.converted_anchor_count == 2 - assert layout_pass.report.inserted_copy_count == 4 - assert layout_pass.report.eliminated_copy_count == 2 - assert layout_pass.report.boundary_copy_count == 2 - assert layout_pass.report.internal_copy_count == 0 - assert layout_pass.report.unknown_copy_count == 0 - assert layout_pass.report.boundary_copy_bytes == 2048 - assert layout_pass.report.internal_copy_bytes == 0 - assert layout_pass.report.unknown_copy_bytes == 0 - assert layout_pass.report.copies_with_unknown_size == 0 - actual = transformed.exported_program().module()(*inputs) - assert torch.allclose(actual, expected, atol=1e-6) - - -def test_strict_rejects_internal_layout_copy() -> None: - inputs = (torch.randn(1, 4, 8, 8),) - edge = _edge(ConvThenLinear().eval(), inputs) - - with pytest.raises(RuntimeError, match="left .* internal"): - ToContiguousChannelsLastPass(edge.exported_program(), strict=True).call( - edge.exported_program().graph_module + test_pass = ToContiguousChannelsLastPassTestPass() + edge_program = to_edge_transform_and_lower( + exported_program, transform_passes=[test_pass] ) - -def test_strict_rejects_supported_anchor_with_noncontiguous_dim_order() -> None: - module = ConvOnly().eval().to(memory_format=torch.channels_last) - inputs = (torch.randn(1, 4, 8, 8).to(memory_format=torch.channels_last),) - exported = torch.export.export(module, inputs) - edge = to_edge( - exported, - compile_config=EdgeCompileConfig(_check_ir_validity=False), - ) - layout_pass = ToContiguousChannelsLastPass( - edge.exported_program(), - strict=True, - ) - - with pytest.raises(RuntimeError, match="0 converted of 1 candidate anchors"): - layout_pass.call(edge.exported_program().graph_module) - - assert layout_pass.report.candidate_anchor_count == 1 - assert layout_pass.report.converted_anchor_count == 0 - - -def test_strict_rejects_unknown_boundary_copy_size() -> None: - inputs = (torch.randn(1, 4, 8, 8),) - exported = torch.export.export( - ConvOnly().eval(), - inputs, - dynamic_shapes={"x": {2: torch.export.Dim("height", min=4, max=16)}}, - ) - edge = to_edge( - exported, - compile_config=EdgeCompileConfig( - _check_ir_validity=False, - _skip_dim_order=True, - ), - ) - layout_pass = ToContiguousChannelsLastPass(edge.exported_program(), strict=True) - - with pytest.raises(RuntimeError, match="unknown sizes"): - layout_pass.call(edge.exported_program().graph_module) - - assert layout_pass.report.boundary_copy_count == 2 - assert layout_pass.report.boundary_copy_bytes == 0 - assert layout_pass.report.copies_with_unknown_size == 2 - - -def test_dynamic_conv_chain_eliminates_internal_layout_copies() -> None: - inputs = (torch.randn(1, 4, 8, 8),) - exported = torch.export.export( - ConvChain().eval(), - inputs, - dynamic_shapes={"x": {2: torch.export.Dim("height", min=4, max=16)}}, - ) - edge = to_edge( - exported, - compile_config=EdgeCompileConfig( - _check_ir_validity=False, - _skip_dim_order=True, - ), - ) - layout_pass = ToContiguousChannelsLastPass(edge.exported_program()) - - edge.transform([layout_pass]) - - assert layout_pass.report.boundary_copy_count == 2 - assert layout_pass.report.internal_copy_count == 0 - assert layout_pass.report.copies_with_unknown_size == 2 - - -def test_dynamic_view_does_not_hide_input_boundary_copy() -> None: - inputs = (torch.randn(1, 4, 8, 8),) - exported = torch.export.export( - DynamicViewConv().eval(), - inputs, - dynamic_shapes={"x": {2: torch.export.Dim("height", min=4, max=16)}}, - ) - edge = to_edge( - exported, - compile_config=EdgeCompileConfig( - _check_ir_validity=False, - _skip_dim_order=True, - ), - ) - layout_pass = ToContiguousChannelsLastPass(edge.exported_program()) - - edge.transform([layout_pass]) - - assert layout_pass.report.boundary_copy_count == 2 - assert layout_pass.report.internal_copy_count == 0 - assert layout_pass.report.copies_with_unknown_size == 2 - - -def test_backend_can_block_layout_propagation_at_a_node() -> None: - torch.manual_seed(0) - module = ConvChain().eval() - inputs = (torch.randn(1, 4, 8, 8),) - expected = module(*inputs) - edge = _edge(module, inputs) - layout_pass = ToContiguousChannelsLastPass( - edge.exported_program(), - can_propagate=lambda node: node.target != exir_ops.edge.aten.relu.default, - ) - - transformed = edge.transform([layout_pass]) - - assert layout_pass.report.boundary_copy_count == 2 - assert layout_pass.report.internal_copy_count == 2 - actual = transformed.exported_program().module()(*inputs) - torch.testing.assert_close(actual, expected) - - -def test_backend_barrier_blocks_view_reordering() -> None: - module = DynamicViewConv().eval() - inputs = (torch.randn(1, 4, 8, 8),) - expected = module(*inputs) - edge = _edge(module, inputs) - layout_pass = ToContiguousChannelsLastPass( - edge.exported_program(), - can_propagate=lambda node: node.target - not in ( - exir_ops.edge.aten.view.default, - exir_ops.edge.aten.view_copy.default, - ), - ) - - transformed = edge.transform([layout_pass]) - - assert layout_pass.report.boundary_copy_count == 1 - assert layout_pass.report.internal_copy_count == 1 - actual = transformed.exported_program().module()(*inputs) - torch.testing.assert_close(actual, expected) - - -@pytest.mark.parametrize("module", [ConvChannelBias(), PadConv()]) -def test_one_sided_propagation_reaches_graph_boundary( - module: torch.nn.Module, -) -> None: - torch.manual_seed(0) - module.eval() - inputs = (torch.randn(1, 4, 8, 8),) - expected = module(*inputs) - edge = _edge(module, inputs) - layout_pass = ToContiguousChannelsLastPass(edge.exported_program()) - - transformed = edge.transform([layout_pass]) - - assert layout_pass.report.boundary_copy_count == 2 - assert layout_pass.report.internal_copy_count == 0 - assert layout_pass.report.unknown_copy_count == 0 - actual = transformed.exported_program().module()(*inputs) - assert torch.allclose(actual, expected, atol=1e-6) - - if isinstance(module, PadConv): - assert ( - _count( - transformed.exported_program().graph_module, - exir_ops.edge.aten.constant_pad_nd.default, + if not ( + (test_pass.initial_permutes == case.expected_initial_permutes) + and (test_pass.initial_views == case.expected_initial_views) + and (test_pass.final_permutes == case.expected_final_permutes) + and (test_pass.final_views == case.expected_final_views) + ): + raise AssertionError( + f"Operator counts do not match for case {case.module.__class__.__name__}\n" + f"Expected initial permutes: {case.expected_initial_permutes}, got: {test_pass.initial_permutes}\n" + f"Expected initial views: {case.expected_initial_views}, got: {test_pass.initial_views}\n" + f"Expected final permutes: {case.expected_final_permutes}, got: {test_pass.final_permutes}\n" + f"Expected final views: {case.expected_final_views}, got: {test_pass.final_views}\n" ) - == 1 - ) + + ref_result = exported_program.module()(*case.inputs) + edge_result = edge_program.exported_program().module()(*case.inputs) + assert torch.allclose(ref_result, edge_result, atol=1e-6) -def test_softmax_blocks_layout_propagation() -> None: - module = ConvSoftmax().eval() - inputs = (torch.randn(1, 4, 8, 8),) - expected = module(*inputs) - edge = _edge(module, inputs) - layout_pass = ToContiguousChannelsLastPass(edge.exported_program()) +@pytest.mark.skip( + reason="Proof of concept - currently no permute-view passes implemented." +) +@common.parametrize("case", cases) +def test_permute_view_counts(case: PermuteCountTestCase) -> None: + run_test(case) - transformed = edge.transform([layout_pass]) - assert layout_pass.report.boundary_copy_count == 1 - assert layout_pass.report.internal_copy_count == 1 - softmax = next( - node - for node in transformed.exported_program().graph.nodes - if node.target - in (exir_ops.edge.aten._softmax.default, exir_ops.edge.aten.softmax.int) - ) - assert softmax.args[1] in (-1, 3) - actual = transformed.exported_program().module()(*inputs) - torch.testing.assert_close(actual, expected) +xfails = {"views_channels_last": "Views are not supported for channels last tensors"} -@pytest.mark.parametrize( - "module, inputs", - [ - ( - ConvRuntimeBiasConv(), - (torch.randn(1, 4, 8, 8), torch.randn(4, 1, 1)), - ), - (ConvSpatialBufferConv(), (torch.randn(1, 4, 8, 8),)), - ], +@pytest.mark.skip( + reason="Proof of concept - currently no permute-view passes implemented." ) -def test_boundary_propagation_rejects_unsafe_broadcast_rewrites( - module: torch.nn.Module, inputs: tuple[torch.Tensor, ...] -) -> None: - module.eval() - expected = module(*inputs) - edge = _edge(module, inputs) - layout_pass = ToContiguousChannelsLastPass(edge.exported_program()) - - transformed = edge.transform([layout_pass]) - - assert layout_pass.report.boundary_copy_count == 2 - assert layout_pass.report.internal_copy_count == 2 - actual = transformed.exported_program().module()(*inputs) - torch.testing.assert_close(actual, expected) - - -def test_layout_copy_report_is_idempotent() -> None: - inputs = (torch.randn(1, 4, 8, 8),) - edge = _edge(ConvChain().eval(), inputs) - first_pass = ToContiguousChannelsLastPass(edge.exported_program()) - transformed = edge.transform([first_pass]) - second_pass = ToContiguousChannelsLastPass(transformed.exported_program()) - - transformed.transform([second_pass]) - - assert second_pass.report.inserted_copy_count == 0 - assert second_pass.report.eliminated_copy_count == 0 - assert second_pass.report.candidate_anchor_count == 0 - assert second_pass.report.converted_anchor_count == 0 - - -def test_user_permute_is_not_reported_as_layout_copy() -> None: - inputs = (torch.randn(1, 4, 8, 8),) - edge = _edge(UserPermute(), inputs) - layout_pass = ToContiguousChannelsLastPass(edge.exported_program(), op_map={}) - - transformed = edge.transform([layout_pass]) - permutes = [ - node - for node in transformed.exported_program().graph.nodes - if node.target == exir_ops.edge.aten.permute_copy.default - ] - - assert len(permutes) == 1 - assert ( - _count( - transformed.exported_program().graph_module, - exir_ops.edge.channels_last.permute_copy.default, - ) - == 0 - ) - assert layout_pass.report.inserted_copy_count == 0 - assert layout_pass.report.boundary_copy_count == 0 - assert layout_pass.report.internal_copy_count == 0 +@common.parametrize("case", cases_channels_last, xfails=xfails) +def test_permute_view_counts_channels_last(case: PermuteCountTestCase) -> None: + run_test(case) diff --git a/backends/transforms/to_contiguous_channels_last_pass.py b/backends/transforms/to_contiguous_channels_last_pass.py deleted file mode 100644 index 0ab78cd36de..00000000000 --- a/backends/transforms/to_contiguous_channels_last_pass.py +++ /dev/null @@ -1,298 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-unsafe - -from collections.abc import Callable -from dataclasses import dataclass -from typing import Any - -import torch - -from executorch.backends.transforms.channels_last_layout import is_layout_copy -from executorch.backends.transforms.fuse_cascaded_transpose_or_permute_ops import ( - FuseCascadedTransposeOrPermuteOps, -) -from executorch.backends.transforms.fuse_cascaded_view_ops import FuseCascadedViewOps -from executorch.backends.transforms.fuse_transpose_or_permute_op_pairs_pass import ( - FuseTransposeOrPermuteOpPairsPass, -) -from executorch.backends.transforms.postpone_permute_below_squeeze_view import ( - PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView, -) -from executorch.backends.transforms.remove_permutes_around_elementwise_ops import ( - RemovePermutesAroundElementwiseOps, -) -from executorch.backends.transforms.replace_nop_transpose_or_permute_with_view import ( - ReplaceNopTransposeOrPermuteWithViewPass, -) -from executorch.backends.transforms.replace_ops_with_channels_last_variants import ( - ChannelsLastOpSpec, - ReplaceOpsWithChannelsLastVariants, -) -from executorch.backends.transforms.replace_squeeze_unsqueeze_with_view import ( - ReplaceSqueezeAndUnsqueezeWithViewPass, -) -from executorch.exir import ExportedProgram -from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.pass_base import ExportPass, PassResult -from torch.fx.node import Target - -_BOUNDARY_TRANSPARENT_TARGETS = { - exir_ops.edge.aten.squeeze_copy.default, - exir_ops.edge.aten.squeeze_copy.dim, - exir_ops.edge.aten.squeeze_copy.dims, - exir_ops.edge.aten.unsqueeze_copy.default, - exir_ops.edge.aten.view.default, - exir_ops.edge.aten.view_copy.default, -} -_BOUNDARY_TRANSPARENT_TARGETS.update( - { - exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, - } -) - - -@dataclass(frozen=True) -class ChannelsLastLayoutReport: - candidate_anchor_count: int = 0 - converted_anchor_count: int = 0 - inserted_copy_count: int = 0 - eliminated_copy_count: int = 0 - boundary_copy_count: int = 0 - internal_copy_count: int = 0 - unknown_copy_count: int = 0 - boundary_copy_bytes: int = 0 - internal_copy_bytes: int = 0 - unknown_copy_bytes: int = 0 - copies_with_unknown_size: int = 0 - internal_copy_nodes: tuple[str, ...] = () - unknown_copy_nodes: tuple[str, ...] = () - - -class ToContiguousChannelsLastPass(ExportPass): - """Build and optimize explicit contiguous-NHWC regions. - - The pass replaces selected NCHW operators with channels-last dialect - anchors surrounded by ``channels_last.permute_copy`` nodes. It then runs - the common data-movement optimizers to a fixed point and reports only the - surviving layout copies. Strict mode rejects structurally unsafe copies and - supported source anchors that were not converted. It does not reject - user-authored permutes or estimate peak arena usage after memory planning. - ``can_propagate`` is consulted by every transform that moves a layout copy - across a graph node. - """ - - # The matrix in test_to_contiguous_channels_last_pass reaches its fixed - # point in at most two rounds; this only catches a pass that reports - # progress it did not make. - _MAX_OPTIMIZATION_ITERATIONS = 4 - - def __init__( - self, - exported_program: ExportedProgram, - op_map: dict[Target, ChannelsLastOpSpec] | None = None, - can_propagate: Callable[[torch.fx.Node], bool] | None = None, - strict: bool = False, - ) -> None: - super().__init__() - self.exported_program = exported_program - self.op_map = op_map - self.can_propagate = can_propagate - self.strict = strict - self.report = ChannelsLastLayoutReport() - - def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - existing_copy_count = len(self._layout_copy_nodes(graph_module)) - replacement_pass = ReplaceOpsWithChannelsLastVariants( - self.exported_program, - op_map=self.op_map, - ) - replacement = replacement_pass.call(graph_module) - graph_module = replacement.graph_module - modified = replacement.modified - preoptimization_copy_count = len(self._layout_copy_nodes(graph_module)) - inserted_copy_count = max(0, preoptimization_copy_count - existing_copy_count) - - for iteration in range(self._MAX_OPTIMIZATION_ITERATIONS): - iteration_modified = False - for transform in self._optimization_passes(): - result = transform.call(graph_module) - graph_module = result.graph_module - iteration_modified |= result.modified - - modified |= iteration_modified - if not iteration_modified: - break - if iteration == self._MAX_OPTIMIZATION_ITERATIONS - 1: - raise RuntimeError( - "Channels-last layout optimization did not converge after " - f"{self._MAX_OPTIMIZATION_ITERATIONS} iterations." - ) - - self.report = self._build_report( - graph_module, - replacement_pass.candidate_count, - replacement_pass.replacement_count, - inserted_copy_count, - preoptimization_copy_count, - ) - if self.strict and ( - self.report.candidate_anchor_count != self.report.converted_anchor_count - or self.report.internal_copy_count - or self.report.unknown_copy_count - or self.report.copies_with_unknown_size - ): - raise RuntimeError( - "Channels-last layout optimization left " - f"{self.report.converted_anchor_count} converted of " - f"{self.report.candidate_anchor_count} candidate anchors, " - f"{self.report.internal_copy_count} internal and " - f"{self.report.unknown_copy_count} unknown copies, with " - f"{self.report.copies_with_unknown_size} unknown sizes. " - f"Internal nodes: {self.report.internal_copy_nodes}; " - f"unknown nodes: {self.report.unknown_copy_nodes}." - ) - - return PassResult(graph_module, modified) - - def _optimization_passes(self) -> tuple[ExportPass, ...]: - return ( - ReplaceSqueezeAndUnsqueezeWithViewPass(), - ReplaceNopTransposeOrPermuteWithViewPass(), - PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView( - can_propagate=self.can_propagate - ), - FuseCascadedViewOps(), - FuseCascadedTransposeOrPermuteOps(can_propagate=self.can_propagate), - RemovePermutesAroundElementwiseOps( - exported_program=self.exported_program, - allow_layout_boundary_propagation=True, - can_propagate=self.can_propagate, - ), - FuseTransposeOrPermuteOpPairsPass(can_propagate=self.can_propagate), - FuseCascadedViewOps(), - FuseCascadedTransposeOrPermuteOps(can_propagate=self.can_propagate), - ) - - @staticmethod - def _layout_copy_nodes( - graph_module: torch.fx.GraphModule, - ) -> list[torch.fx.Node]: - return [node for node in graph_module.graph.nodes if is_layout_copy(node)] - - def _build_report( - self, - graph_module: torch.fx.GraphModule, - candidate_anchor_count: int, - converted_anchor_count: int, - inserted_copy_count: int, - preoptimization_copy_count: int, - ) -> ChannelsLastLayoutReport: - boundary_nodes: list[torch.fx.Node] = [] - internal_nodes: list[torch.fx.Node] = [] - unknown_nodes: list[torch.fx.Node] = [] - - for node in self._layout_copy_nodes(graph_module): - dims = self._normalized_dims(node) - if dims is None: - unknown_nodes.append(node) - elif self._reaches_user_input(node) or self._reaches_graph_output(node): - boundary_nodes.append(node) - else: - internal_nodes.append(node) - - boundary_bytes, boundary_unknown = self._copy_bytes(boundary_nodes) - internal_bytes, internal_unknown = self._copy_bytes(internal_nodes) - unknown_bytes, unknown_unknown = self._copy_bytes(unknown_nodes) - surviving_count = len(boundary_nodes) + len(internal_nodes) + len(unknown_nodes) - return ChannelsLastLayoutReport( - candidate_anchor_count=candidate_anchor_count, - converted_anchor_count=converted_anchor_count, - inserted_copy_count=inserted_copy_count, - eliminated_copy_count=max(0, preoptimization_copy_count - surviving_count), - boundary_copy_count=len(boundary_nodes), - internal_copy_count=len(internal_nodes), - unknown_copy_count=len(unknown_nodes), - boundary_copy_bytes=boundary_bytes, - internal_copy_bytes=internal_bytes, - unknown_copy_bytes=unknown_bytes, - copies_with_unknown_size=( - boundary_unknown + internal_unknown + unknown_unknown - ), - internal_copy_nodes=tuple(node.name for node in internal_nodes), - unknown_copy_nodes=tuple(node.name for node in unknown_nodes), - ) - - def _reaches_user_input(self, node: torch.fx.Node) -> bool: - current = node.args[0] if node.args else None - visited: set[torch.fx.Node] = set() - while isinstance(current, torch.fx.Node) and current not in visited: - visited.add(current) - if current.op == "placeholder": - return current.name in self.exported_program.graph_signature.user_inputs - if self.can_propagate is not None and not self.can_propagate(current): - return False - if ( - current.op != "call_function" - or current.target not in _BOUNDARY_TRANSPARENT_TARGETS - or not current.args - or not isinstance(current.args[0], torch.fx.Node) - ): - return False - current = current.args[0] - return False - - def _reaches_graph_output(self, node: torch.fx.Node) -> bool: - pending = list(node.users) - visited: set[torch.fx.Node] = set() - reached_output = False - while pending: - current = pending.pop() - if current in visited: - continue - visited.add(current) - if current.op == "output": - reached_output = True - continue - if self.can_propagate is not None and not self.can_propagate(current): - return False - if ( - current.op != "call_function" - or current.target not in _BOUNDARY_TRANSPARENT_TARGETS - or not current.users - ): - return False - pending.extend(current.users) - return reached_output - - @staticmethod - def _normalized_dims(node: torch.fx.Node) -> list[int] | None: - if len(node.args) < 2 or not isinstance(node.args[1], (list, tuple)): - return None - dims = list(node.args[1]) - if not all(isinstance(dim, int) for dim in dims): - return None - rank = len(dims) - normalized = [dim + rank if dim < 0 else dim for dim in dims] - if sorted(normalized) != list(range(rank)): - return None - return normalized - - @staticmethod - def _copy_bytes(nodes: list[torch.fx.Node]) -> tuple[int, int]: - known_bytes = 0 - unknown_count = 0 - for node in nodes: - val: Any = node.meta.get("val") - if not isinstance(val, torch.Tensor) or not all( - isinstance(dim, int) for dim in val.shape - ): - unknown_count += 1 - continue - known_bytes += val.numel() * val.element_size() - return known_bytes, unknown_count