diff --git a/backends/arm/quantizer/arm_quantizer_utils.py b/backends/arm/quantizer/arm_quantizer_utils.py index 58f953d8795..41b798d4347 100644 --- a/backends/arm/quantizer/arm_quantizer_utils.py +++ b/backends/arm/quantizer/arm_quantizer_utils.py @@ -411,6 +411,20 @@ def validate(self, model: torch.fx.GraphModule) -> bool: # type: ignore[overrid return True +# Ops whose second operand contributes only its shape, never its values, so it +# must not be tied into the shared qparams of the first. +_SHAPE_AS_OPS = { + torch.ops.aten.expand_as.default, + torch.ops.aten.view_as.default, +} + + +# Whether ``node`` shares qparams with its input ``input_node``: always for an +# ordinary shared op, but for a shape-as op only through its data operand. +def _shares_qparams_with(node: Node, input_node: Node) -> bool: + return node.target not in _SHAPE_AS_OPS or input_node is node.args[0] + + class SharedQspecQuantizer(Quantizer, QuantizerReporterUser): """Assures that specific ops share quantization parameters on all inputs/outputs. @@ -442,6 +456,7 @@ class SharedQspecQuantizer(Quantizer, QuantizerReporterUser): torch.ops.aten.repeat_interleave.self_int, torch.ops.aten.expand_copy.default, torch.ops.aten.expand.default, + torch.ops.aten.expand_as.default, torch.ops.aten.select.int, torch.ops.aten.select_copy.int, torch.ops.aten.slice.Tensor, @@ -543,7 +558,11 @@ def _is_annotated(self, node: Node) -> bool: return Q_ANNOTATION_KEY in node.meta def _get_input_nodes_with_float_output(self, node: Node) -> list[Node]: - return [n for n in node.all_input_nodes if has_float_output(n)] + return [ + n + for n in node.all_input_nodes + if has_float_output(n) and _shares_qparams_with(node, n) + ] def _get_user_nodes_with_float_input(self, node: Node) -> list[Node]: return [n for n in node.users.keys() if has_float_output(node)] @@ -640,6 +659,8 @@ def _get_shared_clique( shared_nodes.add(node) for input_node in node.all_input_nodes: + if not _shares_qparams_with(node, input_node): + continue self._maybe_enqueue_shared_node(input_node, shared_nodes, bfs_queue) self._append_output_qspec(input_node, adjacent_qspecs) touches_quantized_io |= self._is_quantized_io_boundary(input_node) @@ -648,6 +669,8 @@ def _get_shared_clique( ) for output_node in node.users.keys(): + if not _shares_qparams_with(output_node, node): + continue self._maybe_enqueue_shared_node(output_node, shared_nodes, bfs_queue) self._append_input_qspec(output_node, node, adjacent_qspecs) touches_quantized_io |= self._is_quantized_io_boundary(output_node) diff --git a/backends/arm/test/misc/test_shared_qspecs.py b/backends/arm/test/misc/test_shared_qspecs.py index 7287aefb8af..6174fe267e4 100644 --- a/backends/arm/test/misc/test_shared_qspecs.py +++ b/backends/arm/test/misc/test_shared_qspecs.py @@ -16,6 +16,8 @@ from executorch.backends.arm.test.tester.test_pipeline import QuantizationPipeline from executorch.backends.arm.tosa import TosaSpecification from executorch.backends.cortex_m.test.tester import McuTestCase, ramp_tensor +from torch.export import export +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e _QUANT_CONFIG_INT8 = get_symmetric_quantization_config() @@ -734,3 +736,36 @@ def test_silu_does_not_share_input_qspec(): ) pipeline.run() _check_quant_params(pipeline, model.quant_params) + + +class ShapeAsOperand(torch.nn.Module): + """``op(x + x, other)`` reads only the shape of ``other``, so its qparams + must stay ``other``'s own rather than join those of ``x + x``. + """ + + def __init__(self, op: str) -> None: + super().__init__() + self.op = op + + def forward(self, x, y): + other = y + y + return getattr(x + x, self.op)(other) + other + + +@parametrize( + "test_data", + {"view_as": ("view_as", (2, 16)), "expand_as": ("expand_as", (1, 8))}, +) +def test_shape_operand_does_not_share_qspec(test_data): + op, x_shape = test_data + inputs = (ramp_tensor(-1, 1, x_shape), ramp_tensor(-50, 50, (4, 8))) + model = ShapeAsOperand(op).eval() + prepared = prepare_pt2e(export(model, inputs).module(), _get_quantizer()) + prepared(*inputs) + quantized = convert_pt2e(prepared) + + (node,) = quantized.graph.find_nodes( + op="call_function", target=getattr(torch.ops.aten, op).default + ) + data_dequantize, other_dequantize = node.args + assert data_dequantize.args[1] != other_dequantize.args[1] diff --git a/backends/arm/test/models/test_swin2sr_arm.py b/backends/arm/test/models/test_swin2sr_arm.py index d863c7016a4..596ada74636 100644 --- a/backends/arm/test/models/test_swin2sr_arm.py +++ b/backends/arm/test/models/test_swin2sr_arm.py @@ -29,7 +29,7 @@ ] # TODO/MLETORCH-2163: Investigate Swin2SR delegation gaps around index/view -# in FP and Q/DQ, clamp, and expand_copy in INT. +# in FP and Q/DQ and clamp in INT. swin2sr_fp_lowered_outer_graph_ops = { "torch.ops.higher_order.executorch_call_delegate": 2, "executorch_exir_dialects_edge__ops_aten_index_Tensor": 2, @@ -38,7 +38,6 @@ swin2sr_int_lowered_outer_graph_ops = { "torch.ops.higher_order.executorch_call_delegate": 3, "executorch_exir_dialects_edge__ops_aten_clamp_default": 4, - "executorch_exir_dialects_edge__ops_aten_expand_copy_default": 4, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 5, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 5, } diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index 2bf0892a7e9..462e967c9dd 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -31,6 +31,7 @@ fbcode_target(_kind = runtime.python_library, "activation_fusion_pass.py", "aten_to_cortex_m_pass.py", "clamp_hardswish_pass.py", + "compute_quantized_constant_ops_pass.py", "cortex_m_pass.py", "cortex_m_pass_manager.py", "decompose_hardswish_pass.py", @@ -56,6 +57,7 @@ fbcode_target(_kind = runtime.python_library, "//executorch/backends/transforms:channels_last_layout", "//executorch/backends/transforms:channels_last_ops", "//executorch/backends/transforms:convert_conv1d_to_conv2d_pass", + "//executorch/backends/transforms:postpone_permute_below_squeeze_view", "//executorch/backends/transforms:remove_getitem_op", "//executorch/backends/transforms:remove_permutes_around_elementwise_ops", "//executorch/backends/transforms:remove_unused_constants_pass", diff --git a/backends/cortex_m/passes/compute_quantized_constant_ops_pass.py b/backends/cortex_m/passes/compute_quantized_constant_ops_pass.py new file mode 100644 index 00000000000..b7d19812a5e --- /dev/null +++ b/backends/cortex_m/passes/compute_quantized_constant_ops_pass.py @@ -0,0 +1,34 @@ +# 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 torch +from executorch.backends.arm._passes import ComputeConstantOpsAOTPass +from executorch.exir.dialects._ops import ops as exir_ops + + +class ComputeQuantizedConstantOpsPass(ComputeConstantOpsAOTPass): + """Give the quantized constant ops their int8 values. + + Folding removes the Q node after these ops but leaves them producing float, + so the int8 kernels they feed would read float data. A ``full`` stays a + runtime fill with its fill value quantized, so it adds no constant data; the + other ops are computed in float and stored as quantized constants. Runs + after ``ConvertFullLikeToFullPass``, so a ``full_like`` arrives as ``full``. + """ + + def compute_node_aot(self, node: torch.fx.Node) -> bool: + output_qparams = node.meta.get("output_qparams") + if not output_qparams: + return False + if node.target == exir_ops.edge.aten.full.default: + qparams = output_qparams[0] + fill_value = node.args[1] + assert isinstance(fill_value, (int, float)) + fill = qparams.quantize_value(torch.tensor(float(fill_value))) + node.args = (node.args[0], int(fill.item())) + node.kwargs = {**node.kwargs, "dtype": qparams.dtype} + return True + return super().compute_node_aot(node) diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index c5e475c58a7..7475e9d57bf 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -10,6 +10,7 @@ import torch from executorch.backends.arm._passes import ( + ConvertFullLikeToFullPass, DeduplicateGetAttrPass, FoldAndAnnotateQParamsPass, ScalarsToAttributePass, @@ -48,11 +49,13 @@ from .activation_fusion_pass import ActivationFusionPass from .aten_to_cortex_m_pass import AtenToCortexMPass from .clamp_hardswish_pass import ClampHardswishPass +from .compute_quantized_constant_ops_pass import ComputeQuantizedConstantOpsPass from .decompose_hardswish_pass import DecomposeHardswishPass from .decompose_mean_pass import DecomposeMeanPass from .decompose_sdpa_pass import DecomposeSDPAPass from .explicit_layout_pass import ( CortexMCanonicalizeViewCopyPermutePass, + CortexMPostponeLayoutCopyBelowSqueezePass, CortexMReplaceOpsWithChannelsLastVariants, ValidateCortexMExplicitLayoutPass, ) @@ -149,6 +152,8 @@ class CortexMPassManager(ExportedProgramPassManager): FoldAndAnnotateQParamsPass, ActivationFusionPass, QuantizedClampActivationPass, + ConvertFullLikeToFullPass, + ComputeQuantizedConstantOpsPass, DecomposeHardswishPass, AtenToCortexMPass, # Retrace only after quantized ATen ops have become Cortex-M kernels. @@ -165,10 +170,14 @@ class CortexMPassManager(ExportedProgramPassManager): ReplaceScalarWithTensorArgPass, ActivationFusionPass, QuantizedClampActivationPass, + ConvertFullLikeToFullPass, + ComputeQuantizedConstantOpsPass, DecomposeHardswishPass, ConvertConv1dToConv2dPass, CortexMReplaceOpsWithChannelsLastVariants, ReplaceSqueezeAndUnsqueezeWithViewPass, + # Let elementwise regions meet layout copies of their own rank. + CortexMPostponeLayoutCopyBelowSqueezePass, # Move layout copies across pads before singleton permutations become views. RemovePermutesAroundElementwiseOps, CortexMCanonicalizeViewCopyPermutePass, diff --git a/backends/cortex_m/passes/explicit_layout_pass.py b/backends/cortex_m/passes/explicit_layout_pass.py index 659b67d729e..c90cdc973ca 100644 --- a/backends/cortex_m/passes/explicit_layout_pass.py +++ b/backends/cortex_m/passes/explicit_layout_pass.py @@ -19,12 +19,16 @@ LAYOUT_PERMUTE_COPY, PERMUTE_COPY_TARGETS, ) +from executorch.backends.transforms.postpone_permute_below_squeeze_view import ( + PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView, +) from executorch.backends.transforms.replace_ops_with_channels_last_variants import ( ChannelsLastOpSpec, ReplaceOpsWithChannelsLastVariants, ) from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.pass_base import ExportPass, PassResult from torch.fx import GraphModule, Node from torch.fx.node import Target @@ -117,6 +121,33 @@ def _set_node_op(self, node, target, input_node, arg) -> None: ) +class CortexMPostponeLayoutCopyBelowSqueezePass( + PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView +): + """Sink a layout copy below a following squeeze, never below an unsqueeze. + + Only the layout copies this pipeline inserts are moved; user permutes keep + their place. Sinking below an unsqueeze would raise the copy's rank, up to + a rank-5 layout copy that Cortex-M cannot run. + """ + + @property + def targets(self) -> list[EdgeOpOverload]: + return [LAYOUT_PERMUTE_COPY] + + def maybe_remove_or_replace(self, node: Node) -> bool: + users = list(node.users) + user_value = users[0].meta.get("val") if len(users) == 1 else None + value = node.meta.get("val") + if ( + isinstance(user_value, torch.Tensor) + and isinstance(value, torch.Tensor) + and user_value.dim() > value.dim() + ): + return False + return super().maybe_remove_or_replace(node) + + class ValidateCortexMExplicitLayoutPass(ExportPass): def call(self, graph_module: GraphModule) -> PassResult: for node in graph_module.graph.nodes: diff --git a/backends/cortex_m/test/misc/test_portable_int8.py b/backends/cortex_m/test/misc/test_portable_int8.py index ce0aeeb1de6..3c2b125f21a 100644 --- a/backends/cortex_m/test/misc/test_portable_int8.py +++ b/backends/cortex_m/test/misc/test_portable_int8.py @@ -221,6 +221,12 @@ def _quantize_and_export( (torch.randn(1, 3, 4, 5), torch.randn(1, 3, 4, 5)), None, ), + "expand_as": OpCase( + torch.ops.aten.expand_as.default, + _build_module(lambda x, y: torch.ops.aten.expand_as.default(x, y)), + (torch.randn(1, 3, 4, 5), torch.randn(2, 3, 4, 5)), + None, + ), "select": OpCase( torch.ops.aten.select.int, _build_module(lambda x, y: torch.ops.aten.select.int(x, 1, 0)), diff --git a/backends/cortex_m/test/misc/test_quantization.py b/backends/cortex_m/test/misc/test_quantization.py index 6f27d1a6c28..4b3c8e971ae 100644 --- a/backends/cortex_m/test/misc/test_quantization.py +++ b/backends/cortex_m/test/misc/test_quantization.py @@ -6,7 +6,7 @@ import torch from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor -from executorch.backends.arm.test.common import parametrize, xfail_type +from executorch.backends.arm.test.common import parametrize from executorch.backends.cortex_m.test.tester import ( CortexMTester, McuTestCase, @@ -358,12 +358,8 @@ def forward(self, x): ), } -xfails: dict[str, xfail_type] = { - "surrounded_quantized_op_constant": "Numerical error since the add is forced to have non-correct qparams.", -} - -@parametrize("test_case", test_cases, xfails=xfails) +@parametrize("test_case", test_cases) def test_shared_qspec_quantizer(test_case): """ Test that ops which does not change dynamic range are able to use int8 portable kernels. diff --git a/backends/cortex_m/test/targets.bzl b/backends/cortex_m/test/targets.bzl index ac48ee03869..011424b8d4f 100644 --- a/backends/cortex_m/test/targets.bzl +++ b/backends/cortex_m/test/targets.bzl @@ -83,6 +83,8 @@ def define_common_targets(is_fbcode = False): "//executorch/backends/cortex_m/passes:cortex_passes", "//executorch/backends/cortex_m/quantizer:quantizer", "//executorch/backends/test/harness:tester", + "//executorch/backends/test:graph_builder", + "//executorch/backends/transforms:channels_last_layout", "//executorch/backends/transforms:remove_unused_constants_pass", "//executorch/exir:lib", "//executorch/exir/dialects:lib", diff --git a/backends/cortex_m/test/test_explicit_layout_pipeline.py b/backends/cortex_m/test/test_explicit_layout_pipeline.py index 689d117e7b0..2b9edca2743 100644 --- a/backends/cortex_m/test/test_explicit_layout_pipeline.py +++ b/backends/cortex_m/test/test_explicit_layout_pipeline.py @@ -13,6 +13,9 @@ CortexMPassManager, LiftConstantTensorsPass, ) +from executorch.backends.cortex_m.passes.explicit_layout_pass import ( + CortexMPostponeLayoutCopyBelowSqueezePass, +) from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig from executorch.backends.cortex_m.test.tester import ( @@ -20,7 +23,9 @@ CortexMRunPasses, CortexMTester, ) +from executorch.backends.test.graph_builder import GraphBuilder from executorch.backends.test.harness.stages import Quantize, StageType +from executorch.backends.transforms.channels_last_layout import LAYOUT_PERMUTE_COPY from executorch.backends.transforms.remove_unused_constants_pass import ( RemoveUnusedConstantsPass, ) @@ -88,6 +93,18 @@ def forward(self, x): return self.conv2(self.conv1(torch.tanh(x)).view(1, 4, 32)) +class Conv1dChain(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = torch.nn.Conv1d(8, 8, 3, padding=1) + self.conv2 = torch.nn.Conv1d(8, 8, 3, padding=1) + self.conv3 = torch.nn.Conv1d(8, 8, 1) + + def forward(self, x): + y = torch.tanh(self.conv1(x)) + return self.conv3(torch.relu(self.conv2(y) + y)) + + class Apply(torch.nn.Module): def __init__(self, fn) -> None: super().__init__() @@ -546,14 +563,185 @@ def test_convs_around_a_reshape_behind_lut_activation_lower_to_cmsis_nn(): tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) -def test_explicit_layout_rejects_float_input_to_int8_conv(): +def test_explicit_layout_lowers_conv_on_a_quantized_constant(): + inputs = (torch.randn(1, 8, 16),) model = ThenConv1d(lambda x: torch.full_like(x, 0.5, dtype=torch.float32)) - tester = CortexMTester(model.eval(), (torch.randn(1, 8, 16),)) + tester = _run_explicit_layout_passes(CortexMTester(model.eval(), inputs)) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() - with pytest.raises(Exception) as caught: - _run_explicit_layout_passes(tester) + assert _count(program, exir_ops.edge.aten.full_like.default) == 0 + assert _count(program, exir_ops.edge.aten.full.default) == 1 + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) - error = caught.value - while error.__cause__ is not None: - error = error.__cause__ - assert "must move an int8 tensor" in str(error) + +@pytest.mark.parametrize("use_explicit_layout", [False, True]) +@pytest.mark.parametrize( + "constant,constant_op,shape", + [ + ( + lambda: torch.arange(-1.0, 1.0, 0.125), + exir_ops.edge.aten.arange.start_step, + (16,), + ), + ( + lambda: torch.scalar_tensor(0.5), + torch.ops.aten.scalar_tensor.default, + (), + ), + ], + ids=["arange", "scalar_tensor"], +) +def test_quantized_constant_is_computed_before_quantizing( + use_explicit_layout, constant, constant_op, shape +): + torch.manual_seed(0) + inputs = (torch.randn(shape),) + tester = CortexMTester(Apply(lambda x: x + constant()).eval(), inputs) + tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout))) + tester.export().to_edge() + program = _run_passes(tester, use_explicit_layout) + + assert _count(program, constant_op) == 0 + assert _count(program, exir_ops.edge.cortex_m.quantized_add.default) == 1 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize("use_explicit_layout", [False, True]) +@pytest.mark.parametrize( + "fill", + [lambda x: torch.full((1, 8, 16), 0.5), lambda x: torch.full_like(x, 0.5)], + ids=["full", "full_like"], +) +def test_quantized_fill_stays_a_runtime_int8_full(use_explicit_layout, fill): + torch.manual_seed(0) + inputs = (torch.randn(1, 8, 16),) + tester = CortexMTester(Apply(lambda x: x + fill(x)).eval(), inputs) + tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout))) + tester.export().to_edge() + program = _run_passes(tester, use_explicit_layout) + + (full,) = program.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.full.default + ) + assert full.meta["val"].dtype == torch.int8 + assert _count(program, exir_ops.edge.aten.full_like.default) == 0 + assert _count(program, exir_ops.edge.cortex_m.quantized_add.default) == 1 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize("use_explicit_layout", [False, True]) +def test_clamp_with_constant_tensor_bounds_quantizes_them_once(use_explicit_layout): + torch.manual_seed(0) + inputs = (torch.randn(2, 3, 4, 5),) + tester = CortexMTester( + Apply( + lambda x: torch.clamp(x, torch.full_like(x, -1.0), torch.full_like(x, 1.0)) + ).eval(), + inputs, + ) + tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout))) + tester.export().to_edge() + program = _run_passes(tester, use_explicit_layout) + + (clamp,) = program.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.clamp.default + ) + qparams = clamp.meta["output_qparams"][0] + expected = tuple( + int(qparams.quantize_value(torch.tensor(bound)).item()) for bound in (-1.0, 1.0) + ) + assert tuple(clamp.args[1:]) == expected + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize("use_explicit_layout", [False, True]) +def test_unquantized_full_like_keeps_int8_after_float_promoting_op(use_explicit_layout): + inputs = (torch.randn(1, 8, 16),) + tester = CortexMTester( + Apply( + lambda x: ( + torch.sigmoid(x).permute(0, 2, 1), + torch.full_like(x, 2, dtype=torch.int32), + ) + ).eval(), + inputs, + ) + tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout))) + tester.export().to_edge() + program = _run_passes(tester, use_explicit_layout) + + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 1 + assert _count(program, exir_ops.edge.aten.permute_copy.default) == 0 + + +@pytest.mark.parametrize("use_explicit_layout", [False, True]) +def test_unquantized_constant_is_not_precomputed(use_explicit_layout): + inputs = (torch.randn(1, 8, 16),) + tester = CortexMTester( + Apply(lambda x: x + torch.full((1, 8, 16), 0.5)).eval(), inputs + ) + tester.export().to_edge() + program = _run_passes(tester, use_explicit_layout) + + assert _count(program, exir_ops.edge.aten.full.default) == 1 + tester.run_method_and_compare_outputs(inputs=inputs) + + +def test_conv1d_chain_keeps_only_the_boundary_layout_copies(): + torch.manual_seed(0) + inputs = (torch.randn(1, 8, 16),) + tester = _run_explicit_layout_passes(CortexMTester(Conv1dChain().eval(), inputs)) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 3 + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + assert _count(program, exir_ops.edge.aten.permute_copy.default) == 0 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +def _permute_then_view(permute, shape, view_shape) -> torch.fx.GraphModule: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(shape)) + permuted = builder.call_operator(op=permute, args=(x, [0, 3, 1, 2])) + view = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(permuted, view_shape) + ) + builder.output([view]) + return builder.get_graph_module() + + +def test_layout_copy_is_sunk_below_a_squeeze(): + x = torch.randn(1, 1, 4, 6) + graph_module = _permute_then_view(LAYOUT_PERMUTE_COPY, x.shape, [1, 6, 4]) + + result = CortexMPostponeLayoutCopyBelowSqueezePass()(graph_module) + + assert result.modified + assert [ + (node.target, node.args[1]) + for node in result.graph_module.graph.nodes + if node.op == "call_function" + ] == [ + (exir_ops.edge.aten.view_copy.default, [1, 4, 6]), + (LAYOUT_PERMUTE_COPY, [0, 2, 1]), + ] + torch.testing.assert_close( + result.graph_module(x)[0], x.permute(0, 3, 1, 2).reshape(1, 6, 4) + ) + + +@pytest.mark.parametrize( + "permute,shape,view_shape", + [ + (LAYOUT_PERMUTE_COPY, (1, 5, 4, 6), [1, 1, 6, 5, 4]), + (exir_ops.edge.aten.permute_copy.default, (1, 1, 4, 6), [1, 6, 4]), + ], + ids=["layout_copy_unsqueeze", "user_permute_squeeze"], +) +def test_other_permutes_are_not_sunk(permute, shape, view_shape): + graph_module = _permute_then_view(permute, shape, view_shape) + + result = CortexMPostponeLayoutCopyBelowSqueezePass()(graph_module) + + assert not result.modified