diff --git a/backends/nxp/aten_passes/convert_unsqueeze_to_view.py b/backends/nxp/aten_passes/convert_unsqueeze_to_view.py deleted file mode 100644 index 613c0a3b9b1..00000000000 --- a/backends/nxp/aten_passes/convert_unsqueeze_to_view.py +++ /dev/null @@ -1,79 +0,0 @@ -# Copyright 2026 NXP -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from typing import Optional - -import torch -from torch._subclasses import FakeTensor, FakeTensorMode -from torch.fx import GraphModule, Node -from torch.fx.passes.infra.pass_base import PassBase, PassResult - - -class ConvertUnsqueezeToViewPass(PassBase): - """Replace 'aten.unsqueeze.default' with 'aten.view.default'. - - x x - │ │ - ┌─────────────▼─────────────┐ replace with ┌─────────────▼─────────────┐ - │ aten.unsqueeze(x, dim) │ ──────────────► │ aten.view.default(x, S) │ - └─────────────┬─────────────┘ └─────────────┬─────────────┘ - │ │ - ▼ ▼ - out out - """ - - @staticmethod - def _is_unsqueeze(node_: Node) -> bool: - return ( - node_.op == "call_function" - and node_.target == torch.ops.aten.unsqueeze.default - ) - - def _create_view_node(self, *view_args) -> Node: - view_target = torch.ops.aten.view.default - view_node = self.graph_module.graph.call_function(view_target, view_args) - - view_node.meta["source_fn_stack"] = [ - (view_node.name, torch.ops.aten.view.default) - ] - - x_val = view_args[0].meta["val"] - with FakeTensorMode() as mode: - fake_input = FakeTensor.from_tensor( - torch.empty(x_val.shape, dtype=x_val.dtype), mode - ) - output_shape = view_target(fake_input, *view_args[1:]).shape - view_node.meta["val"] = FakeTensor.from_tensor( - torch.empty(output_shape, dtype=x_val.dtype), mode - ) - - return view_node - - def call(self, graph_module: GraphModule) -> Optional[PassResult]: - self.graph_module = graph_module - made_changes = False - - if not any(self._is_unsqueeze(n) for n in graph_module.graph.nodes): - return PassResult(graph_module, made_changes) - - for node in list(graph_module.graph.nodes): - if not self._is_unsqueeze(node): - continue - - input_node = node.all_input_nodes[0] - target_size = node.meta["val"].shape - - with self.graph_module.graph.inserting_after(node): - view_node = self._create_view_node(input_node, target_size) - - node.replace_all_uses_with(view_node) - self.graph_module.graph.erase_node(node) - - made_changes = True - - self.graph_module.recompile() - self.graph_module.graph.eliminate_dead_code() - - return PassResult(graph_module, made_changes) diff --git a/backends/nxp/aten_passes/neutron_aten_pass_manager.py b/backends/nxp/aten_passes/neutron_aten_pass_manager.py index 578d287dfa7..ac5a4d47b30 100644 --- a/backends/nxp/aten_passes/neutron_aten_pass_manager.py +++ b/backends/nxp/aten_passes/neutron_aten_pass_manager.py @@ -7,9 +7,6 @@ import torch -from executorch.backends.nxp.aten_passes.convert_unsqueeze_to_view import ( - ConvertUnsqueezeToViewPass, -) from executorch.backends.nxp.aten_passes.decompose_split_to_slices_pass import ( DecomposeSplitToSlicesPass, ) @@ -50,7 +47,6 @@ def _get_default_passes(neutron_target_spec, qat_mode: bool = False) -> list[Pas RemoveNodesWithKnownOutputs(), FuseLinearAndAddPass(), MoveActivationBeforeConcat(neutron_target_spec), - ConvertUnsqueezeToViewPass(), ] if not qat_mode: diff --git a/backends/nxp/edge_passes/convert_reshaping_nodes_to_view.py b/backends/nxp/edge_passes/convert_reshaping_nodes_to_view.py new file mode 100644 index 00000000000..9a0478f6894 --- /dev/null +++ b/backends/nxp/edge_passes/convert_reshaping_nodes_to_view.py @@ -0,0 +1,103 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Optional + +import torch + +from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass +from executorch.exir.dialects._ops import ops as exir_ops +from torch._subclasses import FakeTensor, FakeTensorMode +from torch.fx import GraphModule, Node +from torch.fx.passes.infra.pass_base import PassResult + + +class ConvertReshapingNodesToViewPass(NeutronEdgePass): + """Replaces: + - 'aten.squeeze.default', 'aten.squeeze.dims' and 'aten.squeeze.dim' with 'aten.view_copy.default'. + + x x + │ │ + ┌──────────────▼──────────────┐ replace with ┌───────────────▼────────────────┐ + │ aten.[un]squeeze(x, dim) │ ──────────────► │ aten.view_copy.default(x, S) │ + └──────────────┬──────────────┘ └───────────────┬────────────────┘ + │ │ + ▼ ▼ + out out + + - 'aten.unsqueeze.default' with 'aten.view_copy.default'. + + x x + │ │ + ┌─────────────▼─────────────┐ replace with ┌───────────────▼────────────────┐ + │ aten.unsqueeze(x, dim) │ ──────────────► │ aten.view_copy.default(x, S) │ + └─────────────┬─────────────┘ └───────────────┬────────────────┘ + │ │ + ▼ ▼ + out out + """ + + graph_module: GraphModule + + @staticmethod + def _is_squeeze(node_: Node) -> bool: + return node_.op == "call_function" and ( + node_.target == exir_ops.edge.aten.squeeze_copy.dim + or node_.target == exir_ops.edge.aten.squeeze_copy.dims + or node_.target == exir_ops.edge.aten.squeeze_copy.default + ) + + @staticmethod + def _is_unsqueeze(node_: Node) -> bool: + return ( + node_.op == "call_function" + and node_.target == exir_ops.edge.aten.unsqueeze_copy.default + ) + + def _create_view_copy_node(self, *view_args) -> Node: + view_target = exir_ops.edge.aten.view_copy.default + view_node = self.graph_module.graph.call_function(view_target, view_args) + + view_node.meta["source_fn_stack"] = [ + (view_node.name, exir_ops.edge.aten.view_copy.default) + ] + + x_val = view_args[0].meta["val"] + with FakeTensorMode() as mode: + fake_input = FakeTensor.from_tensor( + torch.empty(x_val.shape, dtype=x_val.dtype), mode + ) + output_shape = view_target(fake_input, *view_args[1:]).shape + view_node.meta["val"] = FakeTensor.from_tensor( + torch.empty(output_shape, dtype=x_val.dtype), mode + ) + + return view_node + + def run(self, graph_module: GraphModule) -> Optional[PassResult]: + self.graph_module = graph_module + + for node in list(graph_module.graph.nodes): + if not (self._is_squeeze(node) or self._is_unsqueeze(node)): + continue + + input_node = node.all_input_nodes[0] + target_shape = node.meta["val"].shape + + with self.graph_module.graph.inserting_after(node): + view_copy_node = self._create_view_copy_node(input_node, target_shape) + + node.replace_all_uses_with(view_copy_node) + self.graph_module.graph.erase_node(node) + + self.graph_module.recompile() + self.graph_module.graph.eliminate_dead_code() + + # Return immediately to avoid traversing a modified graph. + # The parent class will call this pass again. + return PassResult(graph_module, True) + + # The graph was not modified. + return PassResult(graph_module, False) diff --git a/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py b/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py index 489bb4a4999..cc927b8f1c7 100644 --- a/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py +++ b/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py @@ -1,4 +1,4 @@ -# Copyright 2025 NXP +# Copyright 2025-2026 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -13,15 +13,18 @@ # Operator aliases for better readability. AddMM = exir_ops.edge.aten.addmm.default -ViewCopy = exir_ops.edge.aten.view_copy.default -MM = exir_ops.edge.aten.mm.default +AvgPool2D = exir_ops.edge.aten.avg_pool2d.default Conv = exir_ops.edge.aten.convolution.default +Clone = exir_ops.edge.aten.clone.default +CloneDimOrder = exir_ops.edge.dim_order_ops._clone_dim_order.default HardTanh = exir_ops.edge.aten.hardtanh.default +MM = exir_ops.edge.aten.mm.default Relu = exir_ops.edge.aten.relu.default Sigmoid = exir_ops.edge.aten.sigmoid.default +SqueezeCopy = exir_ops.edge.aten.squeeze_copy.dims Tanh = exir_ops.edge.aten.tanh.default -Clone = exir_ops.edge.aten.clone.default -CloneDimOrder = exir_ops.edge.dim_order_ops._clone_dim_order.default +UnsqueezeCopy = exir_ops.edge.aten.unsqueeze_copy.default +ViewCopy = exir_ops.edge.aten.view_copy.default def insert_qdq_pair_after_node( @@ -105,6 +108,13 @@ class MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass): ViewCopy, ], ViewCopy: [Clone, CloneDimOrder], + Conv: [ + ViewCopy, # For 1D conv + ], + AvgPool2D: [ + ViewCopy, # For 1D AvgPool + UnsqueezeCopy, + ], } def run(self, graph_module: torch.fx.GraphModule) -> PassResult: @@ -200,8 +210,13 @@ class MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass): Relu, Sigmoid, Tanh, + ViewCopy, # For 1D conv. ], ViewCopy: [Clone, CloneDimOrder], + AvgPool2D: [ + ViewCopy, + SqueezeCopy, + ], } def run(self, graph_module: torch.fx.GraphModule) -> PassResult: diff --git a/backends/nxp/edge_passes/neutron_edge_pass_manager.py b/backends/nxp/edge_passes/neutron_edge_pass_manager.py index 44a887b7039..2252ff05a21 100644 --- a/backends/nxp/edge_passes/neutron_edge_pass_manager.py +++ b/backends/nxp/edge_passes/neutron_edge_pass_manager.py @@ -3,6 +3,9 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from executorch.backends.nxp.edge_passes.convert_reshaping_nodes_to_view import ( + ConvertReshapingNodesToViewPass, +) from executorch.backends.nxp.edge_passes.move_auxiliary_operator_into_separate_qdq_cluster_pass import ( MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass, MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass, @@ -21,6 +24,7 @@ def __init__(self, passes: list[NeutronEdgePass] = None): MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass(), MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(), RemoveUselessAsStridedCopyNodes(), + ConvertReshapingNodesToViewPass(), ] super().__init__( diff --git a/backends/nxp/neutron_partitioner.py b/backends/nxp/neutron_partitioner.py index 0723e2f23a0..bc994c5a1d8 100644 --- a/backends/nxp/neutron_partitioner.py +++ b/backends/nxp/neutron_partitioner.py @@ -75,13 +75,15 @@ class QDQCluster: AUXILIARY_OPS = [ operator.getitem, - exir_ops.edge.aten.view_copy.default, - exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.clone.default, exir_ops.edge.aten.hardtanh.default, + exir_ops.edge.aten.permute_copy.default, exir_ops.edge.aten.relu.default, exir_ops.edge.aten.sigmoid.default, + exir_ops.edge.aten.squeeze_copy.dims, exir_ops.edge.aten.tanh.default, - exir_ops.edge.aten.clone.default, + exir_ops.edge.aten.unsqueeze_copy.default, + exir_ops.edge.aten.view_copy.default, exir_ops.edge.dim_order_ops._clone_dim_order.default, ] diff --git a/backends/nxp/quantizer/neutron_quantizer.py b/backends/nxp/quantizer/neutron_quantizer.py index 274839bfb24..8dc6290dbaf 100644 --- a/backends/nxp/quantizer/neutron_quantizer.py +++ b/backends/nxp/quantizer/neutron_quantizer.py @@ -17,7 +17,8 @@ AdaptiveAvgPoolPattern, AddmmPattern, AddTensorPattern, - AvgPoolPattern, + AvgPool1DPattern, + AvgPool2DPattern, BatchNormPattern, CatPattern, Conv1dPattern, @@ -43,10 +44,14 @@ SigmoidPattern, SliceTensorPattern, SoftMaxPattern, + SqueezeDimPattern, + SqueezeDimsPattern, + SqueezePattern, SubTensorPattern, TanhInPlacePattern, TanhPattern, TransposeIntPattern, + UnsqueezePattern, UpsampleBilinear2DPattern, UpsampleNearest2DPattern, ViewPattern, @@ -248,7 +253,8 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec, is_qat: bool = False) OpQuantizer(AdaptiveAvgPoolPattern(is_qat=is_qat), static_qconfig), OpQuantizer(AddTensorPattern(is_qat=is_qat), static_qconfig), OpQuantizer(AddmmPattern(self, is_qat=is_qat), static_fc_qconfig), - OpQuantizer(AvgPoolPattern(is_qat=is_qat), static_qconfig), + OpQuantizer(AvgPool1DPattern(is_qat=is_qat), static_qconfig), + OpQuantizer(AvgPool2DPattern(is_qat=is_qat), static_qconfig), OpQuantizer(BatchNormPattern(is_qat=is_qat), static_qconfig), OpQuantizer(CatPattern(is_qat=is_qat), static_qconfig), OpQuantizer(Conv1dPattern(is_qat=is_qat), static_qconfig), @@ -271,10 +277,14 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec, is_qat: bool = False) OpQuantizer(SigmoidPattern(is_qat=is_qat), static_qconfig), OpQuantizer(SliceTensorPattern(is_qat=is_qat), static_qconfig), OpQuantizer(SoftMaxPattern(is_qat=is_qat), static_qconfig), + OpQuantizer(SqueezeDimPattern(is_qat=is_qat), static_qconfig), + OpQuantizer(SqueezeDimsPattern(is_qat=is_qat), static_qconfig), + OpQuantizer(SqueezePattern(is_qat=is_qat), static_qconfig), OpQuantizer(SubTensorPattern(is_qat=is_qat), static_qconfig), OpQuantizer(TanhPattern(is_qat=is_qat), static_qconfig), OpQuantizer(TanhInPlacePattern(is_qat=is_qat), static_qconfig), OpQuantizer(TransposeIntPattern(is_qat=is_qat), static_qconfig), + OpQuantizer(UnsqueezePattern(is_qat=is_qat), static_qconfig), OpQuantizer(UpsampleBilinear2DPattern(is_qat=is_qat), static_qconfig), OpQuantizer(UpsampleNearest2DPattern(is_qat=is_qat), static_qconfig), OpQuantizer(ViewPattern(is_qat=is_qat), static_qconfig), diff --git a/backends/nxp/quantizer/patterns.py b/backends/nxp/quantizer/patterns.py index d34f010ffff..f5f6fcac5be 100644 --- a/backends/nxp/quantizer/patterns.py +++ b/backends/nxp/quantizer/patterns.py @@ -25,7 +25,6 @@ QuantizationSpec, SharedQuantizationSpec, ) - from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY @@ -325,7 +324,16 @@ def get_anchors( ) -class AvgPoolPattern(SharedSpecPattern): +class AvgPool1DPattern(SharedSpecPattern): + """ + Quantizer for AvgPool1D operator. + """ + + def partition_types(self): + return [torch.ops.aten.avg_pool1d.default] + + +class AvgPool2DPattern(SharedSpecPattern): """ Quantizer for AvgPool2D operator. """ @@ -908,6 +916,33 @@ def get_anchors( ) +class SqueezePattern(SharedSpecPattern): + """ + Quantizer for the `aten.squeeze.default` operator. + """ + + def partition_types(self): + return [torch.ops.aten.squeeze.default] + + +class SqueezeDimPattern(SharedSpecPattern): + """ + Quantizer for the `aten.squeeze.dim` operator. + """ + + def partition_types(self): + return [torch.ops.aten.squeeze.dim] + + +class SqueezeDimsPattern(SharedSpecPattern): + """ + Quantizer for the `aten.squeeze.dims` operator. + """ + + def partition_types(self): + return [torch.ops.aten.squeeze.dims] + + class TanhPattern(QuantizationPattern): """ Quantizer for Tanh operator. @@ -944,6 +979,13 @@ def get_anchors( ) +class UnsqueezePattern(SharedSpecPattern): + """Quantizer for the `aten.unsqueeze.default` operator.""" + + def partition_types(self): + return [torch.ops.aten.unsqueeze.default] + + class UpsampleBilinear2DPattern(SharedSpecPattern): """ Quantizer for `aten.upsample_bilinear2d.vec` operator. diff --git a/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py index b6083d1e816..b0c2608a52d 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py @@ -9,7 +9,6 @@ from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) - from executorch.backends.nxp.backend.ir.conversion_config import ConversionConfig from executorch.backends.nxp.backend.ir.converter.builder.model_builder import ( ModelBuilder, @@ -23,12 +22,25 @@ ) from executorch.backends.nxp.tests.executors import ( convert_run_compare, + graph_contains_any_of_ops, + ToChannelFirstPreprocess, + ToChannelLastPreprocess, ToNCHWPreprocess, ToNHWCPreprocess, ) from executorch.backends.nxp.tests.models import AvgPool2dConvModule, AvgPool2dModule from torch.export import ExportedProgram from executorch.backends.nxp.tests.use_qat import * # noqa F403 +from executorch.exir.dialects._ops import ops as exir_ops + +# noinspection PyProtectedMember +AvgPool2D = exir_ops.edge.aten.avg_pool2d.default +ExecutorchDelegateCall = torch._higher_order_ops.executorch_call_delegate +Squeeze = exir_ops.edge.aten.squeeze.default +SqueezeDim = exir_ops.edge.aten.squeeze.dim +SqueezeDims = exir_ops.edge.aten.squeeze.dims +Unsqueeze = exir_ops.edge.aten.unsqueeze.default +ViewCopy = exir_ops.edge.aten.view_copy.default @pytest.fixture(autouse=True) @@ -219,3 +231,59 @@ def test_avg_pool_2d_quant_conversion__padded(mocker, use_qat): assert ( pad_value == ops[1].tmp_inputs[0].quantization.zero_point[0] ) # `AvgPool` input zp. + + +class AvgPool1DModule(torch.nn.Module): + def __init__(self): + super().__init__() + + self.avg_pool = torch.nn.AvgPool1d( + kernel_size=3, + ) + + def forward(self, x): + return self.avg_pool(x) + + +def test_from_avg_pool_1d(mocker): + model = AvgPool1DModule() + input_shape = ( + 1, + 3, + 12, + ) # Don't use multiples of `num_macs` so the `view_copy` nodes will NOT be deleagted. + extended_shape = (1, 3, 1, 12) + + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + delegated_ep = to_quantized_edge_program( + model, input_shape, use_neutron_for_format_conversion=False + ).exported_program() + + # Make sure the `avg_pool` was delegated. + assert graph_contains_any_of_ops(delegated_ep.graph, [ExecutorchDelegateCall]) + assert not graph_contains_any_of_ops(delegated_ep.graph, [AvgPool2D]) + + # Make sure both `view_copy` nodes were added, and there is no `squeeze` or `unsqueeze`. + assert len([n for n in delegated_ep.graph.nodes if n.target == ViewCopy]) == 2 + assert not graph_contains_any_of_ops( + delegated_ep.graph, [Unsqueeze, Squeeze, SqueezeDim, SqueezeDims] + ) + + # Verify correct behavior of the converted NeutronIR model. + intermediate_ep = converter_spy.call_args.args[1] + neutron_ir_model, _ = converter_spy.spy_return + + input_data = ( + np.random.random(extended_shape).astype(np.float32) * 256.0 - 128.0 + ).astype(np.int8) + + # Make sure the tested program contains the `avg_pool`. + assert graph_contains_any_of_ops(intermediate_ep.graph, [AvgPool2D]) + + convert_run_compare( + intermediate_ep, + tfl_model=neutron_ir_model, + input_data=input_data, + tflite_input_preprocess=ToChannelLastPreprocess(), + tflite_output_preprocess=ToChannelFirstPreprocess(), + ) diff --git a/backends/nxp/tests/models.py b/backends/nxp/tests/models.py index d8ca09a1bd9..98dc93a8732 100644 --- a/backends/nxp/tests/models.py +++ b/backends/nxp/tests/models.py @@ -735,3 +735,15 @@ def __init__(self, dim): def forward(self, x, y): return torch.unsqueeze(x + y, self.dim) + + +class SqueezeAddModel(torch.nn.Module): + def __init__(self, dim=None): + super().__init__() + self.dim = dim + + def forward(self, x, y): + if self.dim is None: + return torch.squeeze(x + y) + else: + return torch.squeeze(x + y, self.dim) diff --git a/backends/nxp/tests/test_batch_norm_fusion.py b/backends/nxp/tests/test_batch_norm_fusion.py index eeb4b03d7a6..9a879534d5f 100644 --- a/backends/nxp/tests/test_batch_norm_fusion.py +++ b/backends/nxp/tests/test_batch_norm_fusion.py @@ -1,4 +1,4 @@ -# Copyright 2025 NXP +# Copyright 2025-2026 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -172,9 +172,7 @@ def test_batch_norm_conv_fusing__full_pipeline__1d(bias: bool): ).exported_program() nodes = list(edge_program.graph.nodes) - assert ( - len(nodes) == 17 - ) # 1D Conv currently isn't delegated, because it doesn't get quantized. + assert len(nodes) == 13 assert not any( node.op == "call_function" and "batch_norm" in node.target.__name__ for node in nodes diff --git a/backends/nxp/tests/test_convert_reshaping_nodes_to_view.py b/backends/nxp/tests/test_convert_reshaping_nodes_to_view.py new file mode 100644 index 00000000000..d0c5aeb886a --- /dev/null +++ b/backends/nxp/tests/test_convert_reshaping_nodes_to_view.py @@ -0,0 +1,151 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import pytest +import torch + +from executorch.backends.nxp.backend.edge_program_converter import ( + EdgeProgramToIRConverter, +) +from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program +from executorch.backends.nxp.tests.executors import ( + convert_run_compare, + graph_contains_any_of_ops, +) +from executorch.backends.nxp.tests.models import SqueezeAddModel, UnsqueezeAddModel +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import ExportedProgram + + +def _create_example_inputs(input_shape): + """Helper function to create random int8 example inputs.""" + example_input_1 = (np.random.random(input_shape).astype(np.float32) * 50).astype( + np.int8 + ) + example_input_2 = (np.random.random(input_shape).astype(np.float32) * 50).astype( + np.int8 + ) + return {0: example_input_1, 1: example_input_2} + + +@pytest.fixture(autouse=True) +def reseed_model_per_test_run(): + torch.manual_seed(42) + np.random.seed(23) + + +@pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((2,), 0, id="1D."), + pytest.param((8, 4, 6), 2, id="3D."), + pytest.param((8, 4, 6, 8), -2, id="4D, negative dim."), + pytest.param((8, 4, 6), 3, id="3D, dim arg is clipped."), + pytest.param((8, 4, 6), -4, id="3D, dim arg is clipped."), + ], +) +def test_convert_unsqueeze_to_view_full_pipeline(mocker, input_shape, dim): + model = UnsqueezeAddModel(dim) + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + + # Run conversion + full_delegated_program = to_quantized_edge_program( + model, + [input_shape, input_shape], + ).exported_program() + + # Make sure no "aten.unsqueeze_copy.default" is in the model. + assert not graph_contains_any_of_ops( + full_delegated_program.graph, + [ + exir_ops.edge.aten.unsqueeze_copy.default, + ], + ) + + # Capture generated model + neutron_ir_partition = converter_spy.spy_return[0] + exported_program_partition: ExportedProgram = converter_spy.call_args.args[1] + + # Make sure "edge.aten.view_copy.default" is in the model that was converted to NeutronIR and delegated. + assert graph_contains_any_of_ops( + exported_program_partition.graph, + [ + exir_ops.edge.aten.view_copy.default, + ], + ) + + example_input = _create_example_inputs(input_shape) + + convert_run_compare( + exported_program_partition, + input_data=example_input, + tfl_model=neutron_ir_partition, + ) + + +@pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((8, 1, 1), None, id="3D, dim = None."), + pytest.param((8, 4, 1), 2, id="3D, dim hit."), + pytest.param((8, 4, 1), 1, id="3D, dim miss."), + pytest.param((8, 4, 1), -1, id="3D, negative dim hit."), + pytest.param((8, 1, 1, 8), [1, 2], id="4D, full dims overlap."), + pytest.param((8, 1, 4, 8), [1, 2], id="4D, partial dims overlap."), + pytest.param((1, 8, 4, 8), [1, 2], id="4D, no dims overlap."), + pytest.param((8, 1, 1, 8), [-2, -3], id="4D, negative full dims overlap."), + pytest.param((8, 1, 4, 8), [-2, -3], id="4D, negative partial dims overlap."), + pytest.param((1, 8, 4, 8), [-2, -3], id="4D, negative no dims overlap."), + pytest.param( + (8, 1, 1, 8), (1, 2), id="4D, tuple instead of list, full dims overlap." + ), + pytest.param( + (8, 1, 4, 8), (1, 2), id="4D, tuple instead of list, partial dims overlap." + ), + pytest.param( + (1, 8, 4, 8), (1, 2), id="4D, tuple instead of list, no dims overlap." + ), + ], +) +def test_convert_squeeze_to_view_full_pipeline(mocker, input_shape, dim): + model = SqueezeAddModel(dim) + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + + # Run conversion + edge_program = to_quantized_edge_program( + model, + [input_shape, input_shape], + ).exported_program() + + # Check that `Squeeze` is no longer present in the model + assert not graph_contains_any_of_ops( + edge_program.graph, + [ + torch.ops.aten.squeeze.dim, + torch.ops.aten.squeeze.dims, + torch.ops.aten.squeeze.default, + ], + ) + + # Capture generated model + neutron_ir_model = converter_spy.spy_return[0] + exported_program: ExportedProgram = converter_spy.call_args.args[1] + + # Make sure `edge.aten.view_copy.default` is in the model. + assert graph_contains_any_of_ops( + exported_program.graph, + [ + exir_ops.edge.aten.view_copy.default, + ], + ) + + example_input = _create_example_inputs(input_shape) + + convert_run_compare( + exported_program, + input_data=example_input, + tfl_model=neutron_ir_model, + ) diff --git a/backends/nxp/tests/test_convert_unsqueeze_to_view.py b/backends/nxp/tests/test_convert_unsqueeze_to_view.py deleted file mode 100644 index 1d2555e5809..00000000000 --- a/backends/nxp/tests/test_convert_unsqueeze_to_view.py +++ /dev/null @@ -1,143 +0,0 @@ -# Copyright 2026 NXP -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -import numpy as np -import pytest -import torch -from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import ( - ConvertUnsqueezeToViewPass, - NeutronAtenPassManager, -) -from executorch.backends.nxp.backend.edge_program_converter import ( - EdgeProgramToIRConverter, -) -from executorch.backends.nxp.tests.executorch_pipeline import ( - neutron_target_spec, - to_quantized_edge_program, -) -from executorch.backends.nxp.tests.executors import ( - convert_run_compare, - graph_contains_any_of_ops, -) - -from executorch.backends.nxp.tests.models import UnsqueezeAddModel -from executorch.exir.dialects._ops import ops as exir_ops -from torch.export import ExportedProgram - - -@pytest.fixture(autouse=True) -def reseed_model_per_test_run(): - torch.manual_seed(42) - np.random.seed(23) - - -@pytest.mark.parametrize( - "input_shape, dim", - [ - pytest.param((2,), 0, id="1D."), - pytest.param((8, 4, 6), 2, id="3D."), - pytest.param((8, 4, 6, 8), -2, id="4D, negative dim."), - pytest.param((8, 4, 6), 3, id="3D, dim arg is clipped."), - pytest.param((8, 4, 6), -4, id="3D, dim arg is clipped."), - ], -) -def test_convert_unsqueeze_to_view_simple(mocker, input_shape, dim): - model = UnsqueezeAddModel(dim) - - example_input_1 = torch.rand(input_shape) - example_input_2 = torch.rand(input_shape) - - exir_program_aten = torch.export.export( - model, - (example_input_1, example_input_2), - ).module() - - # Check "aten.unsqueeze.default" is present - assert graph_contains_any_of_ops( - exir_program_aten.graph, [torch.ops.aten.unsqueeze.default] - ) - - example_input = (example_input_1, example_input_2) - outputs_before = [o.detach().numpy() for o in exir_program_aten(*example_input)] - - # Apply the optimization. - NeutronAtenPassManager(neutron_target_spec, [ConvertUnsqueezeToViewPass()])( - exir_program_aten - ) - - # Make sure no "aten.unsqueeze.default" is in the model. - assert not graph_contains_any_of_ops( - exir_program_aten.graph, - [torch.ops.aten.unsqueeze.default], - ) - - # Make sure there is "aten.view.default" in the model. - assert graph_contains_any_of_ops( - exir_program_aten.graph, - [torch.ops.aten.view.default], - ) - - outputs_after = [o.detach().numpy() for o in exir_program_aten(*example_input)] - - # Make sure the model still produces the exact same output. - assert len(outputs_before) == len(outputs_after) - - for i in range(len(outputs_before)): - assert np.allclose(outputs_before[i], outputs_after[i]) - - -@pytest.mark.parametrize( - "input_shape, dim", - [ - pytest.param((2,), 0, id="1D."), - pytest.param((8, 4, 6), 2, id="3D."), - pytest.param((8, 4, 6, 8), -2, id="4D, negative dim."), - pytest.param((8, 4, 6), 3, id="3D, dim arg is clipped."), - pytest.param((8, 4, 6), -4, id="3D, dim arg is clipped."), - ], -) -def test_convert_unsqueeze_to_view_full_pipeline(mocker, input_shape, dim): - model = UnsqueezeAddModel(dim) - converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") - - # Run conversion - edge_program = to_quantized_edge_program( - model, - [input_shape, input_shape], - ).exported_program() - - # Make sure no "aten.unsqueeze.default" is in the model. - assert not graph_contains_any_of_ops( - edge_program.graph, - [ - torch.ops.aten.unsqueeze.default, - ], - ) - - # Capture generated model - neutron_ir_model = converter_spy.spy_return[0] - exported_program: ExportedProgram = converter_spy.call_args.args[1] - - # Make sure "edge.aten.view_copy.default" is in the model. - assert graph_contains_any_of_ops( - exported_program.graph, - [ - exir_ops.edge.aten.view_copy.default, - ], - ) - - example_input_1 = (np.random.random(input_shape).astype(np.float32) * 50).astype( - np.int8 - ) - example_input_2 = (np.random.random(input_shape).astype(np.float32) * 50).astype( - np.int8 - ) - example_input = {0: example_input_1, 1: example_input_2} - - convert_run_compare( - exported_program, - input_data=example_input, - tfl_model=neutron_ir_model, - ) diff --git a/docs/source/backends/nxp/op-support.csv b/docs/source/backends/nxp/op-support.csv index 5f1956a4905..e1457329946 100644 --- a/docs/source/backends/nxp/op-support.csv +++ b/docs/source/backends/nxp/op-support.csv @@ -3,6 +3,7 @@ aten.abs.default,int8,static int8, aten._adaptive_avg_pool2d.default,int8,static int8,"ceil_mode=False, count_include_pad=False, divisor_override=False" aten.addmm.default,int8,static int8,2D tensor only aten.add.Tensor,int8,static int8,"alpha = 1, input tensor of rame rank" +aten.avg_pool1d.default,int8,static int8,"ceil_mode=False, count_include_pad=False, divisor_override=False" aten.avg_pool2d.default,int8,static int8,"ceil_mode=False, count_include_pad=False, divisor_override=False" aten.cat.default,int8,static int8,"input_channels % 8 = 0, output_channels %8 = 0" aten.clone.default,int8,static int8, @@ -17,7 +18,11 @@ aten.mm.default,int8,static int8,"2D tensor only" aten.relu.default,int8,static int8, aten.sigmoid.default,int8,static int8, aten.slice_copy.Tensor, int8, static int8 +aten.squeeze.default,int8,static int8, +aten.squeeze.dim,int8,static int8, +aten.squeeze.dims,int8,static int8, aten.tanh.default,int8,static int8, +aten.unsqueeze.default,int8,static int8, aten.upsample_bilinear2d.vec,int8,static int8,"channels % 8 = 0, H_scale = W_scale = 2 or 4" aten.upsample_nearest2d.vec,int8,static int8,"channels % 8 = 0, H_scale = W_scale = 2 or 4" aten.view_copy.default,int8,static int8,