From 3a14e9e769c4c3ce95e8b413dc71e74a4b17a327 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 19 Aug 2026 21:50:46 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/transforms/BUCK | 16 + .../fuse_cascaded_transpose_or_permute_ops.py | 13 + .../postpone_permute_below_squeeze_view.py | 13 +- ...replace_ops_with_channels_last_variants.py | 10 +- backends/transforms/targets.bzl | 38 ++ .../test_to_contiguous_channels_last_pass.py | 297 +++++-------- ...st_to_contiguous_channels_last_pipeline.py | 417 ++++++++++++++++++ .../to_contiguous_channels_last_pass.py | 307 +++++++++++++ 8 files changed, 916 insertions(+), 195 deletions(-) create mode 100644 backends/transforms/test/test_to_contiguous_channels_last_pipeline.py create mode 100644 backends/transforms/to_contiguous_channels_last_pass.py diff --git a/backends/transforms/BUCK b/backends/transforms/BUCK index f5029903c21..8bcff0e0563 100644 --- a/backends/transforms/BUCK +++ b/backends/transforms/BUCK @@ -1,6 +1,22 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +load("@fbcode_macros//build_defs:python_pytest.bzl", "python_pytest") load(":targets.bzl", "define_common_targets") oncall("executorch") fbcode_target(_kind = define_common_targets,) + +fbcode_target( + _kind = python_pytest, + name = "test_to_contiguous_channels_last_pipeline", + srcs = ["test/test_to_contiguous_channels_last_pipeline.py"], + compile = "with-source", + typing = False, + deps = [ + "//caffe2:torch", + ":to_contiguous_channels_last_pass", + "//executorch/exir:lib", + "//executorch/exir/dialects:lib", + "fbsource//third-party/pypi/pytest:pytest", + ], +) diff --git a/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py b/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py index ab75a6e0b5f..5c76b078036 100644 --- a/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py +++ b/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py @@ -6,6 +6,10 @@ # pyre-unsafe +from collections.abc import Callable + +import torch + from executorch.backends.transforms.channels_last_layout import ( composed_permute_target, is_channels_last_input_normalization_pair, @@ -39,6 +43,13 @@ 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) @@ -111,6 +122,8 @@ 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 226924cbea7..fe6ea618d91 100644 --- a/backends/transforms/postpone_permute_below_squeeze_view.py +++ b/backends/transforms/postpone_permute_below_squeeze_view.py @@ -7,6 +7,7 @@ # pyre-unsafe +from collections.abc import Callable from typing import cast, List import torch @@ -35,6 +36,13 @@ 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) @@ -69,6 +77,9 @@ 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. @@ -79,8 +90,6 @@ 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/replace_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py index 641c3925b65..40f38df574b 100644 --- a/backends/transforms/replace_ops_with_channels_last_variants.py +++ b/backends/transforms/replace_ops_with_channels_last_variants.py @@ -151,6 +151,8 @@ 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( @@ -200,6 +202,8 @@ 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 @@ -208,13 +212,14 @@ 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 @@ -274,6 +279,7 @@ 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 4283c949e13..426c96174f6 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -537,6 +537,31 @@ def define_common_targets(): ], ) + runtime.python_library( + name = "to_contiguous_channels_last_pass", + srcs = [ + "to_contiguous_channels_last_pass.py", + ], + visibility = [ + "//executorch/backends/...", + ], + 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", + ], + ) + runtime.python_test( name = "test_replace_ops_with_channels_last_variants", srcs = [ @@ -552,6 +577,19 @@ def define_common_targets(): ], ) + runtime.python_test( + name = "test_to_contiguous_channels_last_pass", + srcs = [ + "test/test_to_contiguous_channels_last_pass.py", + ], + deps = [ + "//caffe2:torch", + ":to_contiguous_channels_last_pass", + "//executorch/exir:lib", + "//executorch/exir/dialects:lib", + ], + ) + runtime.python_test( name = "test_convert_conv1d_to_conv2d_pass", srcs = [ 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 65dd39f09b1..54df4f51f6c 100644 --- a/backends/transforms/test/test_to_contiguous_channels_last_pass.py +++ b/backends/transforms/test/test_to_contiguous_channels_last_pass.py @@ -3,17 +3,16 @@ # 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.test import common -from executorch.exir import to_edge_transform_and_lower +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 -from executorch.exir.pass_base import ExportPass -from torch.fx import GraphModule -from torch.fx.passes.infra.pass_base import PassResult InputT = Tuple[Any, ...] @@ -353,9 +352,11 @@ 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, 0, 2 + Conv2dModule(), (torch.randn(2, 8, 8),), 0, 2, 2, 2 + ), + "conv2d_rank4": PermuteCountTestCase( + Conv2dModule(), (torch.randn(1, 2, 8, 8),), 0, 0, 2, 0 ), - "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 ), @@ -414,25 +415,34 @@ 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, 0, 2), + "views": PermuteCountTestCase(ViewsModule(), (torch.rand(1, 2, 2, 4),), 0, 2, 4, 2), "transposes": PermuteCountTestCase( TransposesModule(), (torch.randn(1, 2, 3, 4),), 2, 0, - 2, + 1, 0, ), "maxpool2d_dilation": PermuteCountTestCase( MaxPool2dDilatedModule(), (torch.randn(1, 2, 8, 8),), 0, + 0, + 2, + 0, ), "lstm": PermuteCountTestCase( LstmModule(), @@ -440,7 +450,7 @@ def forward(self, x: torch.Tensor): 7, 19, 7, - 19, + 16, ), "groupnorm": PermuteCountTestCase( GroupNormModule(), @@ -452,16 +462,16 @@ def forward(self, x: torch.Tensor): (torch.randn(4, 8),), 11, 24, - 11, - 24, + 8, + 14, ), "multihead_attention_rank3": PermuteCountTestCase( MultiheadAttentionModule(), (torch.randn(2, 4, 8),), 12, 20, - 12, - 20, + 10, + 18, ), "cumsum_rank3_dim0": PermuteCountTestCase( CumsumModule(), @@ -474,227 +484,132 @@ def forward(self, x: torch.Tensor): 0, ), "model_1_conv_maxpool_residual_linear": PermuteCountTestCase( - Model1ConvMaxPoolResidualLinear(), (torch.randn(2, 8, 64),), 2, 7, 2, 7 + Model1ConvMaxPoolResidualLinear(), (torch.randn(2, 8, 64),), 2, 7, 6, 7 ), "model_2_conv_mha_linear_layernorm": PermuteCountTestCase( - Model2ConvMhaLinearLayerNorm(), (torch.randn(2, 8, 32),), 14, 23, 14, 23 + Model2ConvMhaLinearLayerNorm(), (torch.randn(2, 8, 32),), 14, 23, 11, 21 ), "model_3_lstm_linear": PermuteCountTestCase( - Model3LstmLinear(), (torch.randn(2, 16, 8),), 20, 58, 20, 58 + Model3LstmLinear(), (torch.randn(2, 16, 8),), 20, 58, 20, 55 ), "model_4_conv_lstm_linear_layernorm": PermuteCountTestCase( - Model4ConvLstmLinearLayerNorm(), (torch.randn(2, 8, 32),), 37, 106, 37, 106 + Model4ConvLstmLinearLayerNorm(), (torch.randn(2, 8, 32),), 37, 106, 36, 103 ), "model_5_dwconv_gelu_layernorm_avgpool": PermuteCountTestCase( - Model5DwConvGeluLayerNormAvgPool(), (torch.randn(1, 8, 16, 16),), 2, 0, 2, 0 + Model5DwConvGeluLayerNormAvgPool(), (torch.randn(1, 8, 16, 16),), 2, 0, 4, 0 ), "model_6_gru_linear": PermuteCountTestCase( - Model6GruLinear(), (torch.randn(2, 16, 8),), 20, 56, 20, 56 + Model6GruLinear(), (torch.randn(2, 16, 8),), 20, 56, 20, 55 ), "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 - ), - "model_9_dilated_conv_batchnorm_avgpool_residual": PermuteCountTestCase( - 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 - ), - "permute_silu_permute": PermuteCountTestCase( - PermuteSiluPermute(), - (torch.randn(1, 2, 3, 4),), - 2, + Model8ConvBatchNormMaxPoolResidual(), + (torch.randn(1, 8, 16, 16),), 0, - 2, 0, - ), -} - - -cases_channels_last = { - "conv2d_rank4_channels_last": PermuteCountTestCase( - Conv2dModule(), - (torch.randn(1, 2, 8, 8).to(memory_format=torch.channels_last),), + 5, 0, ), - "conv3d_rank4_channels_last": PermuteCountTestCase( - Conv3dModule(), - (torch.randn(2, 6, 6, 6).to(memory_format=torch.channels_last),), - 0, - 2, - 0, - 2, - ), - "conv3d_rank5_channels_last": PermuteCountTestCase( - Conv3dModule(), - (torch.randn(1, 2, 6, 6, 6).to(memory_format=torch.channels_last_3d),), - 0, - ), - "linear_rank4_channels_last": PermuteCountTestCase( - LinearModule(), - (torch.randn(1, 2, 2, 8).to(memory_format=torch.channels_last),), - 1, - 3, - 1, - 3, - ), - "matmul_rank4_channels_last": PermuteCountTestCase( - MatmulModule(), - ( - torch.randn(2, 2, 2, 3).to(memory_format=torch.channels_last), - torch.randn(2, 2, 3, 4).to(memory_format=torch.channels_last), - ), + "model_9_dilated_conv_batchnorm_avgpool_residual": PermuteCountTestCase( + Model9DilatedConvBatchNormAvgPoolResidual(), + (torch.randn(1, 8, 16, 16),), 0, - 3, 0, - 3, - ), - "pixel_shuffle_channels_last": PermuteCountTestCase( - PixelShuffleModule(), - (torch.randn(1, 8, 2, 2).to(memory_format=torch.channels_last),), - 1, - 2, - 1, - 2, - ), - "grouped_conv_channels_last": PermuteCountTestCase( - GroupedConvModule(), - (torch.randn(1, 4, 8, 8).to(memory_format=torch.channels_last),), - 0, - ), - "transpose_conv_channels_last": PermuteCountTestCase( - TransposeConvModule(), - (torch.randn(1, 2, 8, 8).to(memory_format=torch.channels_last),), + 5, 0, ), - "views_channels_last": PermuteCountTestCase( - ViewsModule(), - (torch.rand(1, 2, 2, 4).to(memory_format=torch.channels_last),), - -1, # The test crashes before reaching the transpose count + "model_10_dwconv_batchnorm_linear_cat": PermuteCountTestCase( + Model10DwConvBatchNormLinearCat(), (torch.randn(2, 8, 64),), 3, 6, 3, 6 ), - "transposes_channels_last": PermuteCountTestCase( - TransposesModule(), - (torch.randn(1, 2, 3, 4).to(memory_format=torch.channels_last),), - 2, - 0, + "permute_silu_permute": PermuteCountTestCase( + PermuteSiluPermute(), + (torch.randn(1, 2, 3, 4),), 2, 0, - ), - "maxpool2d_dilation_channels_last": PermuteCountTestCase( - MaxPool2dDilatedModule(), - (torch.randn(1, 2, 8, 8).to(memory_format=torch.channels_last),), 0, - ), - "groupnorm_channels_last": PermuteCountTestCase( - GroupNormModule(), - (torch.randn(1, 4, 4, 4).to(memory_format=torch.channels_last),), - 0, - ), - "cumsum_rank4_dim3_channels_last": PermuteCountTestCase( - CumsumModule(), - (torch.randn(1, 2, 3, 4).to(memory_format=torch.channels_last), 3), 0, ), } -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, - } +_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, +} - 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 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 _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 run_test(case: PermuteCountTestCase) -> None: case.module.eval() with torch.no_grad(): exported_program = torch.export.export(case.module, case.inputs) - test_pass = ToContiguousChannelsLastPassTestPass() - edge_program = to_edge_transform_and_lower( - exported_program, transform_passes=[test_pass] + edge_program = to_edge( + exported_program, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=True, + ), ) - - 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" - ) - + 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 = edge_program.exported_program().module()(*case.inputs) + edge_result = transformed.exported_program().module()(*case.inputs) assert torch.allclose(ref_result, edge_result, atol=1e-6) -@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) +_EXPECTED_AGGREGATE_COUNTS = (138, 339, 156, 315) -xfails = {"views_channels_last": "Views are not supported for channels last tensors"} - +class TestToContiguousChannelsLastPass(unittest.TestCase): + def test_aggregate_permute_view_counts(self) -> None: + self.assertEqual( + ( + sum(case.expected_initial_permutes for case in cases.values()), + sum(case.expected_initial_views for case in cases.values()), + sum(case.expected_final_permutes for case in cases.values()), + sum(case.expected_final_views for case in cases.values()), + ), + _EXPECTED_AGGREGATE_COUNTS, + ) -@pytest.mark.skip( - reason="Proof of concept - currently no permute-view passes implemented." -) -@common.parametrize("case", cases_channels_last, xfails=xfails) -def test_permute_view_counts_channels_last(case: PermuteCountTestCase) -> None: - run_test(case) + def test_permute_view_counts(self) -> None: + for name, case in cases.items(): + with self.subTest(name=name): + run_test(case) diff --git a/backends/transforms/test/test_to_contiguous_channels_last_pipeline.py b/backends/transforms/test/test_to_contiguous_channels_last_pipeline.py new file mode 100644 index 00000000000..e9dcece4293 --- /dev/null +++ b/backends/transforms/test/test_to_contiguous_channels_last_pipeline.py @@ -0,0 +1,417 @@ +# 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.to_contiguous_channels_last_pass import ( + ToContiguousChannelsLastPass, +) +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops + + +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 + ) + + +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(), + layout_pad_target=( + exir_ops.edge.channels_last.constant_pad_nd.default + if isinstance(module, PadConv) + else None + ), + ) + + 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.channels_last.constant_pad_nd.default, + ) + == 1 + ) + + +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()) + + 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) + + +@pytest.mark.parametrize( + "module, inputs", + [ + ( + ConvRuntimeBiasConv(), + (torch.randn(1, 4, 8, 8), torch.randn(4, 1, 1)), + ), + (ConvSpatialBufferConv(), (torch.randn(1, 4, 8, 8),)), + ], +) +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 diff --git a/backends/transforms/to_contiguous_channels_last_pass.py b/backends/transforms/to_contiguous_channels_last_pass.py new file mode 100644 index 00000000000..1df7dcd49d1 --- /dev/null +++ b/backends/transforms/to_contiguous_channels_last_pass.py @@ -0,0 +1,307 @@ +# 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, +} +try: + _BOUNDARY_TRANSPARENT_TARGETS.update( + { + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + } + ) +except AttributeError: + pass + + +@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. + Additional permutable ops must be layout-equivariant without argument + remapping. ``can_propagate`` is consulted by every transform that moves a + layout copy across a graph node. + """ + + _MAX_OPTIMIZATION_ITERATIONS = 8 + + def __init__( + self, + exported_program: ExportedProgram, + op_map: dict[Target, ChannelsLastOpSpec] | None = None, + extra_permutable_ops: set[Target] | None = None, + can_propagate: Callable[[torch.fx.Node], bool] | None = None, + layout_pad_target: Target | None = None, + strict: bool = False, + ) -> None: + super().__init__() + self.exported_program = exported_program + self.op_map = op_map + self.extra_permutable_ops: set[Target] = set() + if extra_permutable_ops: + self.extra_permutable_ops |= extra_permutable_ops + self.can_propagate = can_propagate + self.layout_pad_target = layout_pad_target + 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( + self.extra_permutable_ops, + exported_program=self.exported_program, + allow_layout_boundary_propagation=True, + layout_pad_target=self.layout_pad_target, + 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