diff --git a/backends/nxp/backend/edge_helper.py b/backends/nxp/backend/edge_helper.py index c2fd7c4f220..2cc5856c807 100644 --- a/backends/nxp/backend/edge_helper.py +++ b/backends/nxp/backend/edge_helper.py @@ -21,6 +21,7 @@ QuantizePerChannel, QuantizePerTensor, SubTensor, + SumDimIntList, ViewCopy, ) from torch.fx import GraphModule, Node @@ -49,6 +50,7 @@ MulTensor, PermuteCopy, SubTensor, + SumDimIntList, } diff --git a/backends/nxp/backend/edge_program_converter.py b/backends/nxp/backend/edge_program_converter.py index e1cdee4fc28..2a76a558329 100644 --- a/backends/nxp/backend/edge_program_converter.py +++ b/backends/nxp/backend/edge_program_converter.py @@ -56,6 +56,7 @@ exir_ops.edge.aten.slice_copy.Tensor: SliceTensorConverter, # noqa F405 exir_ops.edge.aten._softmax.default: SoftmaxConverter, # noqa F405 exir_ops.edge.aten.sub.Tensor: SubTensorConverter, # noqa F405 + exir_ops.edge.aten.sum.dim_IntList: SumDimIntListConverter, # noqa F405 exir_ops.edge.aten.tanh.default: TanhConverter, # noqa F405 exir_ops.edge.aten.upsample_bilinear2d.vec: UpsampleBilinear2DConverter, # noqa F405 exir_ops.edge.aten.upsample_nearest2d.vec: UpsampleNearest2DConverter, # noqa F405 diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py index cc648b9fef8..ece9d5aa672 100755 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py @@ -92,6 +92,9 @@ from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.sub_tensor_converter import ( SubTensorConverter, ) +from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.sum_dim_int_list_converter import ( + SumDimIntListConverter, +) from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.tanh_converter import ( TanhConverter, ) @@ -138,6 +141,7 @@ "SliceTensorConverter", "SoftmaxConverter", "SubTensorConverter", + "SumDimIntListConverter", "TanhConverter", "UpsampleBilinear2DConverter", "UpsampleNearest2DConverter", diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/sum_dim_int_list_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/sum_dim_int_list_converter.py new file mode 100644 index 00000000000..9e54e661f31 --- /dev/null +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/sum_dim_int_list_converter.py @@ -0,0 +1,78 @@ +# 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 torch + +from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList +from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, + NodeConverter, +) +from executorch.backends.nxp.backend.ir.converter.node_converters.shared.reduce_utils import ( + convert_axes_from_attribute, + get_dim_and_handle_io_formats, + get_reduce_node_attrs, +) +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import ( + sum_options, +) +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec +from torch.fx import Node +from torch.nn import Parameter + + +class SumDimIntListConverter(NodeConverter): + + @staticmethod + def _is_supported_on_target( + node: Node, + neutron_target_spec: NeutronTargetSpec, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + if not NodeConverter.uses_quantization_type_for_io( + node, + supported_types=[torch.int8, torch.uint8], + input_indices=[0], + output_indices=[0], + ): + return False + + return True + + @staticmethod + def _is_supported_in_IR( + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + if not NodeConverter._has_shared_q_params_if_quantized(node): + return False + + return True + + def convert(self, node: Node): + """Convert the 'sum.dim_IntList' operator to NeutronIR 'Sum'. + The ExecuTorch schema is: + sum.dim_IntList( + Tensor self, + int[1]? dim, + bool keepdim=False, + *, + dtype=None, + ) -> Tensor + """ + self.assert_convertible(node) + + dim, keepdim = get_reduce_node_attrs(node) + + t_op = self._create_tflite_op_with_io_tensors(node) + t_op.builtin_options = sum_options.Sum(keepdim) + + ops = OpsList(middle_op=t_op) + dim = get_dim_and_handle_io_formats(self.builder, ops, dim, keepdim) + + convert_axes_from_attribute(t_op, self.builder, dim) + self.builder.append_operators(ops.flatten()) diff --git a/backends/nxp/backend/node_format_inference.py b/backends/nxp/backend/node_format_inference.py index 80d24a5cd92..44e28ac748b 100644 --- a/backends/nxp/backend/node_format_inference.py +++ b/backends/nxp/backend/node_format_inference.py @@ -27,6 +27,7 @@ MeanDim, PermuteCopy, QuantizePerTensor, + SumDimIntList, UpsampleBilinear2D, UpsampleNearest2D, ViewCopy, @@ -60,6 +61,7 @@ class NodeFormatInference: PermuteCopy, MeanDim, Amin, + SumDimIntList, } _type_changed_during_last_run: bool @@ -136,7 +138,7 @@ def _infer_format_of_nodes(self, node: Node): self._node_inputs[node][0], DataFormat.FORMATLESS ) - elif op_type in [MeanDim, Amin]: + elif op_type in [MeanDim, Amin, SumDimIntList]: # The operator schema is: # (Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor keep_dim = try_get_arg(node, 2) or False diff --git a/backends/nxp/neutron_partitioner.py b/backends/nxp/neutron_partitioner.py index 8b81eb00505..3c056ce239e 100644 --- a/backends/nxp/neutron_partitioner.py +++ b/backends/nxp/neutron_partitioner.py @@ -229,6 +229,7 @@ def tag_qdq_clusters(self, nodes: list[torch.fx.Node]): exir_ops.edge.aten.slice_copy.Tensor: SliceTensorConverter, # noqa F405 exir_ops.edge.aten._softmax.default: SoftmaxConverter, # noqa F405 exir_ops.edge.aten.sub.Tensor: SubTensorConverter, # noqa F405 + exir_ops.edge.aten.sum.dim_IntList: SumDimIntListConverter, # noqa F405 exir_ops.edge.aten.tanh.default: TanhConverter, # noqa F405 exir_ops.edge.aten.upsample_bilinear2d.vec: UpsampleBilinear2DConverter, # noqa F405 exir_ops.edge.aten.upsample_nearest2d.vec: UpsampleNearest2DConverter, # noqa F405 diff --git a/backends/nxp/quantizer/neutron_quantizer.py b/backends/nxp/quantizer/neutron_quantizer.py index 66f92777194..9262d1a5814 100644 --- a/backends/nxp/quantizer/neutron_quantizer.py +++ b/backends/nxp/quantizer/neutron_quantizer.py @@ -56,6 +56,7 @@ SqueezeDimsPattern, SqueezePattern, SubTensorPattern, + SumDimIntListPattern, TanhInPlacePattern, TanhPattern, TransposeIntPattern, @@ -300,6 +301,7 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec, is_qat: bool = False) 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(SumDimIntListPattern(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), diff --git a/backends/nxp/quantizer/patterns.py b/backends/nxp/quantizer/patterns.py index b594e0bc663..ae373b14d64 100644 --- a/backends/nxp/quantizer/patterns.py +++ b/backends/nxp/quantizer/patterns.py @@ -1141,6 +1141,15 @@ def partition_types(self): return [torch.ops.aten.squeeze.dims] +class SumDimIntListPattern(SharedSpecPattern): + """ + Quantizer for the `aten.sum.dim_IntList` operator. + """ + + def partition_types(self): + return [torch.ops.aten.sum.dim_IntList] + + class TanhPattern(QuantizationPattern): """ Quantizer for Tanh operator. diff --git a/backends/nxp/tests/ir/converter/node_converter/test_sum_dim_int_list_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_sum_dim_int_list_converter.py new file mode 100644 index 00000000000..6ad9fd8539c --- /dev/null +++ b/backends/nxp/tests/ir/converter/node_converter/test_sum_dim_int_list_converter.py @@ -0,0 +1,416 @@ +# 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 + +# noinspection PyUnusedImports +import pytest +import torch + +from executorch.backends.nxp.backend.ir.converter.builder.model_builder import ( + ModelBuilder, +) +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.max_pool_2d_options import ( + MaxPool2D, +) +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.sum_options import ( + Sum, +) +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import ( + Transpose, +) +from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator +from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program +from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.model_output_comparator import ( + AllCloseOutputComparator, +) +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.ops_aliases import ( + AddTensor, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, + SumDimIntList, +) +from executorch.backends.nxp.tests.use_qat import * # noqa F403 + + +@pytest.fixture(autouse=True) +def reseed_model_per_test_run(): + torch.manual_seed(23) + np.random.seed(23) + + +class SumModule(torch.nn.Module): + def __init__( + self, dim: int | torch.Size | list[int] | tuple[int, ...], keepdim: bool + ): + super().__init__() + self.dim = dim + self.keepdim = keepdim + + def forward(self, x): + return torch.sum(x, dim=self.dim, keepdim=self.keepdim) + + +class SumAddModule(SumModule): + def forward(self, x): + x = super().forward(x) + return x + x + + +class MaxPoolSumModule(torch.nn.Module): + @staticmethod + def noop_max_pool_2d(x): + """Call `torch.max_pool2d` that is a NoOp, but it enforces the ChannelsFirst format in the `NodeFormatInference`.""" + return torch.max_pool2d(x, kernel_size=1) + + def __init__( + self, dim: int | torch.Size | list[int] | tuple[int, ...], keepdim: bool + ): + super().__init__() + self.dim, self.keepdim = dim, keepdim + + def forward(self, x): + x = self.noop_max_pool_2d(x) + x = torch.sum(x, dim=self.dim, keepdim=self.keepdim) + return x + + +class SumDimIntListMaxPoolModule(MaxPoolSumModule): + def forward(self, x): + x = torch.sum(x, dim=self.dim, keepdim=self.keepdim) + x = self.noop_max_pool_2d(x) + return x + + +def assert_delegated( + model, + input_shape, + mocker, + request, + use_qat=False, + expected_delegated_ops=None, +): + if expected_delegated_ops is None: + expected_delegated_ops = {SumDimIntList: 1} + + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops=expected_delegated_ops, + expected_non_delegated_ops={}, + ) + + # Cover also negative values to thoroughly test the operator. + dataset_creator = RandomDatasetCreator(low=-2, high=2) + + remove_quant_io_ops = True # Use quantized dataset. + output_comparator = AllCloseOutputComparator(atol=1) # Allow single bit error. + + lower_run_compare( + model, + input_shape, + graph_verifier, + request, + dataset_creator, + output_comparator, + use_qat=use_qat, + remove_quant_io_ops=remove_quant_io_ops, + ) + + +def assert_not_delegated(model, input_shape): + delegated_ep = to_quantized_edge_program(model, input_shape).exported_program() + + # Make sure the `sum` was NOT delegated. + assert not graph_contains_any_of_ops(delegated_ep.graph, [ExecutorchDelegateCall]) + assert graph_contains_any_of_ops(delegated_ep.graph, [SumDimIntList]) + + +class TestSumDimIntListConverter: + # noinspection PyMethodMayBeStatic + @pytest.fixture(params=[True, False], ids=lambda keep_dim: f"keep_dim = {keep_dim}") + def keep_dim(self, request): + return request.param + + def test__basic_nsys_inference__qat(self, mocker, request, use_qat, keep_dim): + input_shape = (23,) + model = SumModule(0, keep_dim) + assert_delegated(model, input_shape, mocker, request, use_qat=use_qat) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((5,), 0, id="1D, dim = 0."), + pytest.param((4, 2), 0, id="2D, dim = 0."), + pytest.param((4, 2), -1, id="2D, dim = -1."), + pytest.param((3, 1, 4), 2, id="3D, dim = 2."), + pytest.param((1, 3, 3, 7), 3, id="4D, dim = 3."), + pytest.param((3, 1, 4, 1, 5), -1, id="5D, dim = -1."), + pytest.param((3, 1, 4, 1, 5), 0, id="5D, dim = 0."), + ], + ) + def test__single_dims(self, mocker, request, input_shape, dim, keep_dim): + model = SumModule(dim, keep_dim) + assert_delegated(model, input_shape, mocker, request) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((4, 2), (-2,), id="2D, dim = (-2,)."), + pytest.param((2, 3, 4), (0, 2), id="3D, dim = (0, 2,)."), + pytest.param((1, 3, 3, 7), (2, -3), id="4D, dim = (2, -3)."), + pytest.param((1, 3, 3, 7), -2, id="4D, dim = -2."), + pytest.param((3, 1, 4, 1, 5), (3, -5, -4), id="5D, dim = (3, -5 ,-4)."), + ], + ) + def test__tuple_dims(self, mocker, request, input_shape, dim, keep_dim): + model = SumModule(dim, keep_dim) + assert_delegated(model, input_shape, mocker, request) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((3, 1, 4), 1, id="3D, dim = 1."), + pytest.param((3, 1, 4, 1, 5), -2, id="5D, dim = -2."), + ], + ) + def test__noop__only_node__not_delegated(self, input_shape, dim): + keep_dim = True # Reduction over a dimension of size `1` with `keep_dim=True` is a no-op. + model = SumModule(dim, keep_dim) + assert_not_delegated(model, input_shape) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((3, 1, 4), 1, id="3D, dim = 1."), + pytest.param((3, 1, 4, 1, 5), -2, id="5D, dim = -2."), + ], + ) + def test__noop__not_only_node__delegated(self, mocker, request, input_shape, dim): + keep_dim = True # Reduction over a dimension of size `1` with `keep_dim=True` is a no-op. + model = SumAddModule(dim, keep_dim) + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={SumDimIntList: 1, AddTensor: 1}, + ) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((3, 1, 4), 1, id="3D, dim = 1."), + pytest.param((3, 1, 4, 1, 5), -2, id="5D, dim = -2."), + pytest.param((1, 7, 3, 3), [0], id="4D, dim = [0]."), + ], + ) + def test__no_reduction__keepdim_false__delegated( + self, mocker, request, input_shape, dim + ): + # These cases reduce over a dimension of size 1. + # When `keep_dim=True` the node is a noop, and it's not delegated (see `test__noop__only_node__not_delegated`), + # but with `keep_dim=False` it changes the shape so it's not a noop and is therefore delegated successfully. + keep_dim = False + model = SumModule(dim, keep_dim) + assert_delegated(model, input_shape, mocker, request) + + def test__channels_first__keep_dim__true(self, mocker, request): + # Just 1 test case to verify correct handling of the `dim`. + # Most cases fall into the single bit error case, and since this test uses 2 operators, the error accumulates + # and the final error is larger. We cannot with 100% certainty say that the error is only caused by the single + # bit errors and not related to the format. That's why only this 1 case with no errors is used. + input_shape, dim = (1, 7, 3, 3), 1 + model = MaxPoolSumModule(dim, True) + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + SumDimIntList: 1, + }, + ) + + class TestKeepDimFalseFormatHandling: + """When `keep_dim = False`, the `sum` operator changes the rank, so the format have to be explicitly + handled. The tests in this class focus on the related edge cases. + """ + + def _assert_neutron_ir_model_has_ops( + self, model_builder_finish_spy, expected_ops + ): + assert ( + model_builder_finish_spy.call_count == 1 + ), "Conversion to Neutron IR happened multiple times." + + neutron_ir_ops = model_builder_finish_spy.spy_return.sub_graphs[ + 0 + ].operators.vector + assert len(neutron_ir_ops) == len( + expected_ops + ), "Neutron IR model doesn't have the expected number of ops." + + for op, expected_op in zip(neutron_ir_ops, expected_ops, strict=True): + assert isinstance( + op.builtin_options, expected_op + ), f"Expected {expected_op}, got {op}." + + @pytest.mark.parametrize( + "dim", + [ + 1, + [0, -3], + (-4, 1, 2), + [-3, 3], + [1, 2, 3], + ], + ids=lambda dim: f"dim={dim}", + ) + def test__channels_first_input__reducing_channels(self, mocker, request, dim): + # If the channels dimension is reduced (removed), the `sum` output will always be equal in channels first + # and channels last, so no `Transpose` ops are added. + input_shape = (1, 7, 3, 3) + model = MaxPoolSumModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + SumDimIntList: 1, + }, + ) + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Transpose, + MaxPool2D, + Sum, + ], + ) + + @pytest.mark.parametrize( + "dim", + [ + (2, 3), + [1, -2, 3], + [-1, -2, 0], + ], + ids=lambda dim: f"dim={dim}", + ) + def test__channels_first_input__reducing_all_spatial_dims( + self, mocker, request, dim + ): + # If the spatial dimensions are reduced (removed), the `sum` output will always be equal in channels + # first and channels last, so no `Transpose` ops are added. + input_shape = (1, 7, 3, 3) + model = MaxPoolSumModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + SumDimIntList: 1, + }, + ) + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Transpose, + MaxPool2D, + Sum, + ], + ) + + @pytest.mark.parametrize( + "dim", + [ + 0, + (2,), + [-1, 0], + ], + ids=lambda dim: f"dim={dim}", + ) + def test__channels_first_input__not_reducing_channels_or_all_spatial_dims( + self, mocker, request, dim + ): + # If the channels dimension is not reduced, a `Transpose` operator must be added to make the input channels + # first in Neutron IR. + + input_shape = (1, 7, 3, 3) + model = MaxPoolSumModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + SumDimIntList: 1, + }, + ) + + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Transpose, + MaxPool2D, + Transpose, # The necessary `Transpose` operator. + Sum, + ], + ) + + @pytest.mark.parametrize( + "input_shape, dim", + [ + pytest.param((2, 3, 4, 5, 6), 0, id="dim=0, 5D->4D"), + pytest.param((2, 3, 4, 5, 6), [-3], id="dim=[-3], 5D->4D"), + pytest.param((1, 2, 3, 4, 5, 6), (1, -1), id="dim=(1, -1), 6D->4D"), + ], + ids=lambda dim: f"dim={dim}", + ) + def test__channels_first_output(self, mocker, request, input_shape, dim): + model = SumDimIntListMaxPoolModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + SumDimIntList: 1, + }, + ) + + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Sum, + Transpose, # The necessary `Transpose` operator. + MaxPool2D, + Transpose, + ], + ) diff --git a/backends/nxp/tests/ops_aliases.py b/backends/nxp/tests/ops_aliases.py index 5c87635b7d7..be5fc5e0be0 100644 --- a/backends/nxp/tests/ops_aliases.py +++ b/backends/nxp/tests/ops_aliases.py @@ -51,6 +51,7 @@ SqueezeDim = exir_ops.edge.aten.squeeze.dim SqueezeDims = exir_ops.edge.aten.squeeze.dims SubTensor = exir_ops.edge.aten.sub.Tensor +SumDimIntList = exir_ops.edge.aten.sum.dim_IntList Tanh = exir_ops.edge.aten.tanh.default Tanh_ = exir_ops.edge.aten.tanh_.default Unsqueeze = exir_ops.edge.aten.unsqueeze.default diff --git a/docs/source/backends/nxp/op-support.csv b/docs/source/backends/nxp/op-support.csv index 228fbec7ed4..5f97cf5d8b9 100644 --- a/docs/source/backends/nxp/op-support.csv +++ b/docs/source/backends/nxp/op-support.csv @@ -37,6 +37,7 @@ aten.split_with_sizes.default,N/A, N/A, "transforming split -> getitem to slice, aten.squeeze.default,int8,static int8, aten.squeeze.dim,int8,static int8, aten.squeeze.dims,int8,static int8, +aten.sum.dim_IntList,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"