From 6a5005d2f38e1e6daf4fe5123b1c093d0aee07a7 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 21 Aug 2026 13:55:46 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../absorb_boundary_layout_copies.py | 208 +++++++++++++++++ backends/transforms/targets.bzl | 31 +++ .../test_absorb_boundary_layout_copies.py | 214 ++++++++++++++++++ 3 files changed, 453 insertions(+) create mode 100644 backends/transforms/absorb_boundary_layout_copies.py create mode 100644 backends/transforms/test/test_absorb_boundary_layout_copies.py 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/targets.bzl b/backends/transforms/targets.bzl index 4283c949e13..913ea377687 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -537,6 +537,37 @@ def define_common_targets(): ], ) + runtime.python_library( + name = "absorb_boundary_layout_copies", + srcs = [ + "absorb_boundary_layout_copies.py", + ], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + ":channels_last_layout", + "//executorch/exir:lib", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_absorb_boundary_layout_copies", + srcs = [ + "test/test_absorb_boundary_layout_copies.py", + ], + deps = [ + "//caffe2:torch", + ":absorb_boundary_layout_copies", + "//executorch/exir:lib", + "//executorch/exir/dialects:lib", + "fbsource//third-party/pypi/pytest:pytest", + ], + ) + runtime.python_test( name = "test_replace_ops_with_channels_last_variants", srcs = [ 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()