diff --git a/backends/nxp/backend/edge_helper.py b/backends/nxp/backend/edge_helper.py index c2fd7c4f220..849f103520a 100644 --- a/backends/nxp/backend/edge_helper.py +++ b/backends/nxp/backend/edge_helper.py @@ -10,6 +10,7 @@ from executorch.backends.nxp.tests.ops_aliases import ( AddTensor, + Amax, Amin, Cat, Clone, @@ -46,6 +47,7 @@ no_op_candidates = { AddTensor, Amin, + Amax, MulTensor, PermuteCopy, SubTensor, diff --git a/backends/nxp/backend/edge_program_converter.py b/backends/nxp/backend/edge_program_converter.py index e1cdee4fc28..8a98a321a2a 100644 --- a/backends/nxp/backend/edge_program_converter.py +++ b/backends/nxp/backend/edge_program_converter.py @@ -31,6 +31,7 @@ exir_ops.edge.aten._adaptive_avg_pool2d.default: AdaptiveAvgPool2dConverter, # noqa F405 exir_ops.edge.aten.addmm.default: AddMMConverter, # noqa F405 exir_ops.edge.aten.add.Tensor: AddTensorConverter, # noqa F405 + exir_ops.edge.aten.amax.default: AmaxConverter, # noqa F405 exir_ops.edge.aten.amin.default: AminConverter, # noqa F405 exir_ops.edge.aten.avg_pool2d.default: AvgPool2dConverter, # noqa F405 exir_ops.edge.aten.bmm.default: BMMConverter, # 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..4d63cdc2123 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 @@ -10,6 +10,9 @@ from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.addmm_converter import ( AddMMConverter, ) +from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.amax_converter import ( + AmaxConverter, +) from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.amin_converter import ( AminConverter, ) @@ -110,6 +113,7 @@ "AdaptiveAvgPool2dConverter", "AddMMConverter", "AddTensorConverter", + "AmaxConverter", "AminConverter", "AvgPool2dConverter", "BMMConverter", diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/amax_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/amax_converter.py new file mode 100644 index 00000000000..9172d5341f5 --- /dev/null +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/amax_converter.py @@ -0,0 +1,76 @@ +# 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 ( + reduce_max_options, +) +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec +from torch.fx import Node +from torch.nn import Parameter + + +class AmaxConverter(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 'amax' operator to NeutronIR 'ReduceMax'. + The ExecuTorch schema is: + amax( + Tensor self, + int[1]? dim, + bool keepdim=False, + ) -> 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 = reduce_max_options.ReduceMax(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..e091ad23a33 100644 --- a/backends/nxp/backend/node_format_inference.py +++ b/backends/nxp/backend/node_format_inference.py @@ -16,6 +16,7 @@ from executorch.backends.nxp.backend.edge_program_converter import functions_converters from executorch.backends.nxp.tests.ops_aliases import ( AdaptiveAvgPool2D, + Amax, Amin, AvgPool2D, Convolution, @@ -59,6 +60,7 @@ class NodeFormatInference: ViewCopy, PermuteCopy, MeanDim, + Amax, Amin, } @@ -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, Amax, Amin]: # 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..e11c7c836e9 100644 --- a/backends/nxp/neutron_partitioner.py +++ b/backends/nxp/neutron_partitioner.py @@ -204,6 +204,7 @@ def tag_qdq_clusters(self, nodes: list[torch.fx.Node]): exir_ops.edge.aten._adaptive_avg_pool2d.default: AdaptiveAvgPool2dConverter, # noqa F405 exir_ops.edge.aten.addmm.default: AddMMConverter, # noqa F405 exir_ops.edge.aten.add.Tensor: AddTensorConverter, # noqa F405 + exir_ops.edge.aten.amax.default: AmaxConverter, # noqa F405 exir_ops.edge.aten.amin.default: AminConverter, # noqa F405 exir_ops.edge.aten.avg_pool2d.default: AvgPool2dConverter, # noqa F405 exir_ops.edge.aten.bmm.default: BMMConverter, # noqa F405 diff --git a/backends/nxp/quantizer/neutron_quantizer.py b/backends/nxp/quantizer/neutron_quantizer.py index 66f92777194..7d9e319fa72 100644 --- a/backends/nxp/quantizer/neutron_quantizer.py +++ b/backends/nxp/quantizer/neutron_quantizer.py @@ -16,6 +16,7 @@ AdaptiveAvgPoolPattern, AddmmPattern, AddTensorPattern, + AmaxPattern, AminPattern, AvgPool1DPattern, AvgPool2DPattern, @@ -261,6 +262,7 @@ 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(AmaxPattern(is_qat=is_qat), static_qconfig), OpQuantizer(AminPattern(is_qat=is_qat), static_qconfig), OpQuantizer(AvgPool1DPattern(is_qat=is_qat), static_qconfig), OpQuantizer(AvgPool2DPattern(is_qat=is_qat), static_qconfig), diff --git a/backends/nxp/quantizer/patterns.py b/backends/nxp/quantizer/patterns.py index b594e0bc663..7b7239099eb 100644 --- a/backends/nxp/quantizer/patterns.py +++ b/backends/nxp/quantizer/patterns.py @@ -319,6 +319,15 @@ def get_anchors( ) +class AmaxPattern(SharedSpecPattern): + """ + Quantizer for Amax operator. + """ + + def partition_types(self): + return [torch.ops.aten.amax.default] + + class AminPattern(SharedSpecPattern): """ Quantizer for Amin operator. diff --git a/backends/nxp/tests/ir/converter/node_converter/test_amax_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_amax_converter.py new file mode 100644 index 00000000000..37505c3e478 --- /dev/null +++ b/backends/nxp/tests/ir/converter/node_converter/test_amax_converter.py @@ -0,0 +1,412 @@ +# 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.reduce_max_options import ( + ReduceMax, +) +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, + Amax, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, +) +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 AmaxModule(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.amax(x, dim=self.dim, keepdim=self.keepdim) + + +class AmaxAddModule(AmaxModule): + def forward(self, x): + x = super().forward(x) + return x + x + + +class MaxPoolAmaxModule(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.amax(x, dim=self.dim, keepdim=self.keepdim) + return x + + +class AmaxMaxPoolModule(MaxPoolAmaxModule): + def forward(self, x): + x = torch.amax(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 = {Amax: 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 `amax` was NOT delegated. + assert not graph_contains_any_of_ops(delegated_ep.graph, [ExecutorchDelegateCall]) + assert graph_contains_any_of_ops(delegated_ep.graph, [Amax]) + + +class TestAmaxConverter: + # 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 = AmaxModule(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 = AmaxModule(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 = AmaxModule(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 = AmaxModule(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 = AmaxAddModule(dim, keep_dim) + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={Amax: 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 = AmaxModule(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 = MaxPoolAmaxModule(dim, True) + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={MaxPool2DWithIndices: 1, GetItem: 1, Amax: 1}, + ) + + class TestKeepDimFalseFormatHandling: + """When `keep_dim = False`, the `amax` 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 `amax` output will always be equal in channels first + # and channels last, so no `Transpose` ops are added. + input_shape = (1, 7, 3, 3) + model = MaxPoolAmaxModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + Amax: 1, + }, + ) + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Transpose, + MaxPool2D, + ReduceMax, + ], + ) + + @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 `amax` output will always be equal in channels + # first and channels last, so no `Transpose` ops are added. + input_shape = (1, 7, 3, 3) + model = MaxPoolAmaxModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + Amax: 1, + }, + ) + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Transpose, + MaxPool2D, + ReduceMax, + ], + ) + + @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 = MaxPoolAmaxModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + Amax: 1, + }, + ) + + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + Transpose, + MaxPool2D, + Transpose, # The necessary `Transpose` operator. + ReduceMax, + ], + ) + + @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 = AmaxMaxPoolModule(dim, False) + + model_builder_finish_spy = mocker.spy(ModelBuilder, "finish") + assert_delegated( + model, + input_shape, + mocker, + request, + expected_delegated_ops={ + MaxPool2DWithIndices: 1, + GetItem: 1, + Amax: 1, + }, + ) + + self._assert_neutron_ir_model_has_ops( + model_builder_finish_spy, + expected_ops=[ + ReduceMax, + 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..07c909febeb 100644 --- a/backends/nxp/tests/ops_aliases.py +++ b/backends/nxp/tests/ops_aliases.py @@ -15,6 +15,7 @@ AdaptiveAvgPool2D = exir_ops.edge.aten._adaptive_avg_pool2d.default AddMM = exir_ops.edge.aten.addmm.default AddTensor = exir_ops.edge.aten.add.Tensor +Amax = exir_ops.edge.aten.amax.default Amin = exir_ops.edge.aten.amin.default AvgPool2D = exir_ops.edge.aten.avg_pool2d.default BMM = exir_ops.edge.aten.bmm.default diff --git a/docs/source/backends/nxp/op-support.csv b/docs/source/backends/nxp/op-support.csv index 228fbec7ed4..f845948d729 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 tensors of equal shape" +aten.amax.default,int8,static int8, aten.amin.default,int8,static int8, 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"