From 07f2fd373113fcda514319cdd07078c141d38d11 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 21 Aug 2026 12:12:54 -0700 Subject: [PATCH 1/2] Update [ghstack-poisoned] --- .../absorb_boundary_layout_copies.py | 208 ++++++++++++++++ backends/transforms/targets.bzl | 34 +++ .../test_absorb_boundary_layout_copies.py | 228 ++++++++++++++++++ 3 files changed, 470 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 e15e089b051..fb7b3f54838 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -562,6 +562,40 @@ 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", + # The permute-count matrix the absorption totals are measured on. + "test/test_to_contiguous_channels_last_pass.py", + ], + deps = [ + "//caffe2:torch", + ":absorb_boundary_layout_copies", + ":to_contiguous_channels_last_pass", + "//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..e3675d51d05 --- /dev/null +++ b/backends/transforms/test/test_absorb_boundary_layout_copies.py @@ -0,0 +1,228 @@ +# 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.backends.transforms.to_contiguous_channels_last_pass import ( + ToContiguousChannelsLastPass, +) +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 + + +class Conv(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 ResidualConvPool(torch.nn.Module): + """One input feeding two branches, so the region brackets it twice.""" + + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + self.pool = torch.nn.MaxPool2d(3, stride=1, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + self.pool(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 _lower(module, inputs): + 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 + ), + ) + edge = edge.transform([ToContiguousChannelsLastPass(edge.exported_program())]) + return edge + + +def _absorb(edge): + layout_pass = AbsorbBoundaryLayoutCopies(edge.exported_program()) + return edge.transform([layout_pass]), layout_pass.contract + + +def _run(edge, contract, inputs): + 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", [Conv(), ResidualConvPool()]) +def test_boundary_copies_are_absorbed_and_numerics_hold(module) -> None: + inputs = (torch.randn(1, 4, 8, 8),) + expected = module.eval()(*inputs) + edge = _lower(module, inputs) + assert _count(edge.exported_program().graph_module, _LAYOUT_COPY) > 0 + + 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: + """Both branches of a residual share the input, so one entry covers them.""" + module = ResidualConvPool() + inputs = (torch.randn(1, 4, 8, 8),) + edge = _lower(module, inputs) + copies_on_input = [ + node + for node in edge.exported_program().graph_module.graph.nodes + if node.op == "placeholder" + and node.name in edge.exported_program().graph_signature.user_inputs + for _ in node.users + ] + assert len(copies_on_input) > 1 + + _, contract = _absorb(edge) + + assert list(contract.inputs) == [0] + + +def test_mixed_users_are_left_alone() -> None: + """An input consumed both by a layout region and directly is not a boundary.""" + + class MixedUse(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x): + return self.conv(x) + x + + inputs = (torch.randn(1, 4, 8, 8),) + module = MixedUse() + expected = module.eval()(*inputs) + edge = _lower(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_absorbing_is_idempotent() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + edge = _lower(Conv(), 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 = _lower(Conv(), inputs) + edge, contract = _absorb(edge) + + assert contract + edge.exported_program()._validate() + + +# The layout pass is measured against these 36 models in +# test_to_contiguous_channels_last_pass.py, which pins its own per-case counts +# but stops before absorption. These are the totals across that matrix. +_MATRIX_BASELINE_PERMUTES = 138 +_MATRIX_LAYOUT_ONLY_PERMUTES = 156 +_MATRIX_ABSORBED_PERMUTES = 137 + +# Absorption does not destroy 19 permutes, it moves them across the method +# boundary: 17 become obligations on the caller. Pinning both numbers keeps the +# in-graph count honest, since it could otherwise be driven to zero by handing +# the caller unlimited work. The gap is the genuine saving, and it comes from +# fan-out — one placeholder feeding several branches needs several copies but +# only one contract entry. +_MATRIX_CONTRACT_ENTRIES = 17 + +# Cases still above baseline after absorbing. Every one of these is an internal +# copy left by a region that a non-permutable op (batch_norm, linear) cut in +# two, which is region-merging work rather than boundary work. +_MATRIX_RESIDUAL_REGRESSIONS = { + "conv2d_rank3", + "model_1_conv_maxpool_residual_linear", + "model_8_conv_batchnorm_maxpool_residual", + "model_9_dilated_conv_batchnorm_avgpool_residual", + "views", +} + +_PERMUTE_TARGETS = { + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.channels_last.permute_copy.default, +} + + +def _permutes(edge) -> int: + return sum( + node.op == "call_function" and node.target in _PERMUTE_TARGETS + for node in edge.exported_program().graph.nodes + ) + + +def test_absorption_pays_for_the_layout_pass_across_the_model_matrix() -> None: + from executorch.backends.transforms.test.test_to_contiguous_channels_last_pass import ( + cases, + ) + + baseline = layout_only = absorbed = contract_entries = 0 + regressions = set() + for name, case in cases.items(): + case.module.eval() + with torch.no_grad(): + exported = torch.export.export(case.module, case.inputs) + config = EdgeCompileConfig(_check_ir_validity=False, _skip_dim_order=True) + case_baseline = _permutes(to_edge(exported, compile_config=config)) + + edge = to_edge(exported, compile_config=config) + edge = edge.transform( + [ToContiguousChannelsLastPass(edge.exported_program())] + ) + case_layout = _permutes(edge) + + absorb = AbsorbBoundaryLayoutCopies(edge.exported_program()) + edge = edge.transform([absorb]) + case_absorbed = _permutes(edge) + + baseline += case_baseline + layout_only += case_layout + absorbed += case_absorbed + contract_entries += len(absorb.contract.inputs) + len(absorb.contract.outputs) + if case_absorbed > case_baseline: + regressions.add(name) + + assert baseline == _MATRIX_BASELINE_PERMUTES + assert layout_only == _MATRIX_LAYOUT_ONLY_PERMUTES + assert absorbed == _MATRIX_ABSORBED_PERMUTES + assert contract_entries == _MATRIX_CONTRACT_ENTRIES + assert regressions == _MATRIX_RESIDUAL_REGRESSIONS From 88da265ff4d53b9765d0f7f7088d159821a89455 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 21 Aug 2026 13:52:46 -0700 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- ...ve_permutes_around_elementwise_tosa_ops.py | 7 +- .../remove_permutes_around_elementwise_ops.py | 70 ++++++++----------- .../test/test_permute_optimization_passes.py | 38 +++------- 3 files changed, 42 insertions(+), 73 deletions(-) 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/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/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)