diff --git a/backends/nxp/backend/edge_program_converter.py b/backends/nxp/backend/edge_program_converter.py index 2a76a558329..dfdc054fe41 100644 --- a/backends/nxp/backend/edge_program_converter.py +++ b/backends/nxp/backend/edge_program_converter.py @@ -45,6 +45,7 @@ exir_ops.edge.aten.leaky_relu.default: LeakyReluConverter, # noqa F405 exir_ops.edge.aten.log.default: LogConverter, # noqa F405 exir_ops.edge.aten.max_pool2d_with_indices.default: MaxPool2DWithIndicesConverter, # noqa F405 + exir_ops.edge.aten.maximum.default: MaximumConverter, # noqa F405 exir_ops.edge.aten.mean.dim: MeanDimConverter, # noqa F405 exir_ops.edge.aten.mm.default: MMConverter, # noqa F405 exir_ops.edge.aten.mul.Tensor: MulTensorConverter, # noqa F405 diff --git a/backends/nxp/backend/ir/converter/node_converter.py b/backends/nxp/backend/ir/converter/node_converter.py index 4c71dab1285..ad95794ccdb 100755 --- a/backends/nxp/backend/ir/converter/node_converter.py +++ b/backends/nxp/backend/ir/converter/node_converter.py @@ -226,7 +226,12 @@ def supports_partitioning_result( @staticmethod def _has_shared_q_params_if_quantized(node: Node) -> bool: - """Check if node has shared quantization parameters if it's quantized.""" + """Check if node has shared quantization parameters for first input and output if it's quantized. + + :param node: PyTorch fx Node with 'all_input_nodes' initialized. + :return: True, if first input and output parameters have the same quantization parameters. + False otherwise. + """ if len(node.users) < 1 or len(node.all_input_nodes) < 1: # Some exotic operator (only consumer or only produces) return True @@ -493,3 +498,38 @@ def inputs_satisfy_broadcast_condition(node: Node) -> bool: prod(input_tensor.meta["val"].shape) == num_elements for input_tensor in node.all_input_nodes ) + + @staticmethod + def _all_io_shares_quantization_parameters(node: Node) -> bool: + """Determine if given PyTorch fx Node has all inputs and output with same quantization parameters. + + :param node: PyTorch fx Node with 'all_input_nodes' initialized. + :return: True, if all input and output parameters have the same quantization parameters. + False otherwise. + """ + post_node = list(node.users.keys())[0] + if not _is_quant_node(post_node): + return False + output_zp, output_scale, output_type = ( + post_node.args[1], + post_node.args[2], + post_node.args[5], + ) + + for input_node in node.all_input_nodes: + if not _is_dequant_node(input_node): + return False + + input_zp, input_scale, input_type = ( + input_node.args[1], + input_node.args[2], + input_node.args[5], + ) + if (input_zp, input_scale, input_type) != ( + output_zp, + output_scale, + output_type, + ): + return False + + return True 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 ece9d5aa672..09168f75d3e 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 @@ -52,6 +52,9 @@ from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.max_pool2d_with_indices_converter import ( MaxPool2DWithIndicesConverter, ) +from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.maximum_converter import ( + MaximumConverter, +) from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.mean_dim_converter import ( MeanDimConverter, ) @@ -127,6 +130,7 @@ "LeakyReluConverter", "LogConverter", "MaxPool2DWithIndicesConverter", + "MaximumConverter", "MeanDimConverter", "MMConverter", "MulTensorConverter", diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/maximum_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/maximum_converter.py new file mode 100644 index 00000000000..8d1d05ba3d7 --- /dev/null +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/maximum_converter.py @@ -0,0 +1,68 @@ +# 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.tflite_generator.builtin_options import ( + maximum_options, +) +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec +from torch.fx import Node +from torch.nn import Parameter + + +class MaximumConverter(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.inputs_satisfy_broadcast_condition(node): + return False + + supported_types = [torch.int8, torch.uint8] + if not NodeConverter.uses_quantization_type_for_io( + node, supported_types, [0, 1], [0] + ): + return False + + return True + + @staticmethod + def _is_supported_in_IR( + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + if len(node.args) != 2: + return False + + if not NodeConverter._all_io_shares_quantization_parameters(node): + # The IR requires all inputs to have the same quantization parameters as the output. + # The quantizer should quantize the operator so that this case does not happen. + return False + + return True + + def convert(self, node: Node): + """Convert 'maximum' operator to NeutronIR 'Maximum'. + The ExecuTorch schema is: + maximum(Tensor input, Tensor other, *, Tensor out=None) + """ + self.assert_convertible(node) + t_op = self._create_tflite_op_with_io_tensors(node) + t_op.builtin_options = maximum_options.Maximum() + + ops = OpsList(middle_op=t_op) + # Create additional ops in case of shape broadcasting + ops.add_pre(self.builder.ensure_correct_broadcasting(t_op, t_op.tmp_outputs[0])) + self.builder.append_operators(ops.flatten()) diff --git a/backends/nxp/neutron_partitioner.py b/backends/nxp/neutron_partitioner.py index 3c056ce239e..46bb8e6d0a7 100644 --- a/backends/nxp/neutron_partitioner.py +++ b/backends/nxp/neutron_partitioner.py @@ -218,6 +218,7 @@ def tag_qdq_clusters(self, nodes: list[torch.fx.Node]): exir_ops.edge.aten.leaky_relu.default: LeakyReluConverter, # noqa F405 exir_ops.edge.aten.log.default: LogConverter, # noqa F405 exir_ops.edge.aten.max_pool2d_with_indices.default: MaxPool2DWithIndicesConverter, # noqa F405 + exir_ops.edge.aten.maximum.default: MaximumConverter, # noqa F405 exir_ops.edge.aten.mean.dim: MeanDimConverter, # noqa F405 exir_ops.edge.aten.mm.default: MMConverter, # noqa F405 exir_ops.edge.aten.mul.Tensor: MulTensorConverter, # noqa F405 diff --git a/backends/nxp/quantizer/neutron_quantizer.py b/backends/nxp/quantizer/neutron_quantizer.py index 9262d1a5814..97204b609f4 100644 --- a/backends/nxp/quantizer/neutron_quantizer.py +++ b/backends/nxp/quantizer/neutron_quantizer.py @@ -34,6 +34,7 @@ LeakyReluPattern, LinearPattern, LogPattern, + MaximumPattern, MaxPool1DPattern, MaxPool2DPattern, MeanDimPattern, @@ -282,6 +283,7 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec, is_qat: bool = False) OpQuantizer(LeakyReluInPlacePattern(is_qat=is_qat), static_fc_qconfig), OpQuantizer(LinearPattern(self, is_qat=is_qat), static_fc_qconfig), OpQuantizer(LogPattern(is_qat=is_qat), static_qconfig), + OpQuantizer(MaximumPattern(is_qat=is_qat), static_qconfig), OpQuantizer(MaxPool1DPattern(is_qat=is_qat), static_qconfig), OpQuantizer(MaxPool2DPattern(is_qat=is_qat), static_qconfig), OpQuantizer(MeanDimPattern(is_qat=is_qat), static_qconfig), diff --git a/backends/nxp/quantizer/patterns.py b/backends/nxp/quantizer/patterns.py index 5a4d86b65d7..c4ef35b4a5b 100644 --- a/backends/nxp/quantizer/patterns.py +++ b/backends/nxp/quantizer/patterns.py @@ -120,7 +120,7 @@ class SharedSpecPattern(QuantizationPattern): """ @abstractmethod - def partition_types(self) -> list[torch.nn.Module]: + def partition_types(self) -> list[OpOverload]: pass @staticmethod @@ -152,6 +152,52 @@ def get_anchors( return self.get_shared_spec_anchors(gm, fused_partition) +class SharedSpecMultipleInputPattern(QuantizationPattern): + """ + Quantization pattern for shared quantization with multiple inputs. + + The quantization is derived from the previous node quantization and inputs and the output shares the same + quantization parameters (scale and zero-point). + """ + + @abstractmethod + def partition_types(self) -> list[OpOverload]: + pass + + @staticmethod + def get_shared_spec_anchors( + gm: fx.GraphModule, fused_partition: list[fx.GraphModule] + ) -> PartitionAnchors | None: + node = fused_partition[0].nodes[-1] + + quantized_input = None + for prev_node in node.args: + if Q_ANNOTATION_KEY in prev_node.meta: + quantized_input = prev_node + break + + if quantized_input is not None: + inputs = [] + for idx, _ in enumerate(node.args): + inputs.append( + (node, NodeArgsIdx(idx), SharedQuantizationSpec(quantized_input)) + ) + else: + return None + + return PartitionAnchors( + inputs=inputs, + weights=[], + biases=[], + output=[(node, SharedQuantizationSpec(quantized_input))], + ) + + def get_anchors( + self, gm: fx.GraphModule, fused_partition: list[fx.GraphModule] + ) -> PartitionAnchors | None: + return self.get_shared_spec_anchors(gm, fused_partition) + + class SingleInputBasicPattern(QuantizationPattern): @abstractmethod def partition_types(self) -> list[OpOverload]: @@ -300,7 +346,7 @@ class AddTensorPattern(QuantizationPattern): Basic quantization for all inputs and output. """ - def partition_types(self) -> list[torch.nn.Module]: + def partition_types(self) -> list[OpOverload]: return [torch.ops.aten.add.Tensor] def get_anchors( @@ -333,7 +379,7 @@ class BMMPattern(QuantizationPattern): Quantizer for BatchMatMul operator. """ - def partition_types(self) -> list[torch.nn.Module]: + def partition_types(self) -> list[OpOverload]: return [torch.ops.aten.bmm.default] def get_anchors( @@ -372,7 +418,7 @@ class SubTensorPattern(QuantizationPattern): Basic quantization for all inputs and output. """ - def partition_types(self) -> list[torch.nn.Module]: + def partition_types(self) -> list[OpOverload]: return [torch.ops.aten.sub.Tensor] def get_anchors( @@ -426,7 +472,7 @@ def get_anchors( quantized_input = None for prev_node in node.args[0]: - if "quantization_annotation" in prev_node.meta: + if Q_ANNOTATION_KEY in prev_node.meta: quantized_input = prev_node break @@ -876,6 +922,17 @@ def partition_types(self): return [torch.ops.aten.log.default] +class MaximumPattern(SharedSpecMultipleInputPattern): + """ + Quantization pattern for Maximum quantization. Accepts 1 or 2 input nodes. + + Basic quantization for all inputs and output. + """ + + def partition_types(self) -> list[OpOverload]: + return [torch.ops.aten.maximum.default] + + class MaxPool1DPattern(SharedSpecPattern): """Quantizer for the MaxPool1D operator.""" diff --git a/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py new file mode 100644 index 00000000000..da1be530e3f --- /dev/null +++ b/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py @@ -0,0 +1,269 @@ +# 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.tests.dataset_creator import RandomDatasetCreator +from executorch.backends.nxp.tests.executorch_pipeline import ( + ModelInputSpec, + 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.models import MaximumModule, MaxPoolMaximumModule +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.ops_aliases import ( + ExecutorchDelegateCall, + GetItem, + Maximum, + 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 TestMaximum: + @pytest.mark.parametrize( + "x_input_shape", + [ + pytest.param((1,), id="1D."), + pytest.param((6, 5), id="2D."), + pytest.param((6, 82), id="2D alt."), + pytest.param((1, 4, 7), id="3D."), + pytest.param((1, 68, 7), id="3D alt."), + pytest.param((2, 4, 3, 15), id="4D."), + pytest.param((1, 4, 9, 11, 4), id="5D."), + ], + ) + def test__basic_nsys_inference(self, mocker, request, x_input_shape): + x_input_spec = ModelInputSpec(x_input_shape) + model = MaximumModule() + graph_verifier = DetailedGraphVerifier( + mocker, expected_delegated_ops={Maximum: 1}, expected_non_delegated_ops={} + ) + dataset_creator = RandomDatasetCreator(low=-1.0, high=1.0) + + # Quantize the dataset and allow a single bit error. + remove_quant_io_ops = True + comparator = AllCloseOutputComparator(atol=1) + + lower_run_compare( + model, + [x_input_spec, x_input_spec], + graph_verifier, + request, + dataset_creator, + comparator, + remove_quant_io_ops=remove_quant_io_ops, + ) + + def test__basic_nsys_inference_qat(self, mocker, request): + x_input_spec = ModelInputSpec((1, 4, 7)) + model = MaximumModule() + graph_verifier = DetailedGraphVerifier( + mocker, expected_delegated_ops={Maximum: 1}, expected_non_delegated_ops={} + ) + dataset_creator = RandomDatasetCreator(low=-1.0, high=1.0) + + # Quantize the dataset and allow a single bit error. + remove_quant_io_ops = True + comparator = AllCloseOutputComparator(atol=1) + + lower_run_compare( + model, + [x_input_spec, x_input_spec], + graph_verifier, + request, + dataset_creator, + comparator, + remove_quant_io_ops=remove_quant_io_ops, + use_qat=True, + ) + + @pytest.mark.parametrize( + "input_spec", + [ + pytest.param( + [ModelInputSpec((4, 6)), ModelInputSpec((1, 6))], id="2 inputs 2D." + ), + pytest.param( + [ModelInputSpec((69, 73)), ModelInputSpec((1, 73))], + id="2 inputs 2D alt.", + ), + pytest.param( + [ModelInputSpec((5, 3, 4)), ModelInputSpec((1, 3, 1))], + id="2 inputs 3D.", + ), + pytest.param( + [ModelInputSpec((4,)), ModelInputSpec((4, 4))], + id="2 inputs 1D + 2D.", + marks=pytest.mark.xfail( + reason="AIR-14864: conversion fails", strict=True + ), + ), + pytest.param( + [ModelInputSpec((1, 4, 8, 8)), ModelInputSpec((8, 8))], + id="2 inputs 4D + 2D.", + ), + pytest.param( + [ModelInputSpec((10,)), ModelInputSpec((1, 1))], + id="2 inputs 1D + 2D, num_elems of input == num_elems of output", + ), + ], + ) + def test__broadcast(self, mocker, request, input_spec): + model = MaximumModule() + graph_verifier = DetailedGraphVerifier( + mocker, expected_delegated_ops={Maximum: 1}, expected_non_delegated_ops={} + ) + dataset_creator = RandomDatasetCreator(low=-1.0, high=1.0) + + # Quantize the dataset and allow a single bit error. + remove_quant_io_ops = True + comparator = AllCloseOutputComparator(atol=1) + + lower_run_compare( + model, + input_spec, + graph_verifier, + request, + dataset_creator, + comparator, + remove_quant_io_ops=remove_quant_io_ops, + ) + + @pytest.mark.parametrize( + "input_spec", + [ + pytest.param( + [ModelInputSpec((4, 1)), ModelInputSpec((1, 6))], id="2 inputs 2D." + ), + pytest.param( + [ModelInputSpec((1, 3, 4)), ModelInputSpec((5, 3, 1))], + id="2 inputs 3D.", + ), + pytest.param( + [ModelInputSpec((6, 4)), ModelInputSpec((6, 6, 1))], + id="2 inputs 2D + 3D.", + ), + ], + ) + def test__broadcast_unsupported(self, input_spec): + # Broadcast where at least one of the inputs is not equal to output is not supported + model = MaximumModule() + + delegated_ep = to_quantized_edge_program(model, input_spec).exported_program() + + # Make sure the `Maximum` was NOT delegated. + assert not graph_contains_any_of_ops( + delegated_ep.graph, [ExecutorchDelegateCall] + ) + assert graph_contains_any_of_ops(delegated_ep.graph, [Maximum]) + + @pytest.mark.parametrize( + "input_spec", + [ + pytest.param( + ModelInputSpec((1, 4, 5, 5)), + id="4D, product of dims is not a multiple of 8.", + ), + ], + ) + def test__channels_first_input(self, mocker, request, input_spec): + model = MaxPoolMaximumModule() + + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={Maximum: 1, MaxPool2DWithIndices: 1, GetItem: 1}, + expected_non_delegated_ops={}, + ) + dataset_creator = RandomDatasetCreator(low=-1.0, high=1.0) + + lower_run_compare( + model, + [input_spec, input_spec], + graph_verifier, + request, + dataset_creator, + ) + + @pytest.mark.parametrize( + "input_spec", + [ + pytest.param( + [ModelInputSpec((1, 8, 5, 5)), ModelInputSpec((1, 8, 5, 1))], + id="2 inputs 4D + 4D.", + ), + pytest.param( + [ModelInputSpec((1, 8, 1, 67)), ModelInputSpec((1, 8, 5, 67))], + id="2 inputs 4D + 4D same width.", + marks=pytest.mark.xfail( + reason="AIR-14864: conversion fails", strict=True + ), + ), + pytest.param( + [ModelInputSpec((1, 8, 5, 5)), ModelInputSpec((1, 1))], + id="2 inputs 4D + 2D ones tensor.", + ), + pytest.param( + [ModelInputSpec((1, 8, 8, 8)), ModelInputSpec((8, 8))], + id="2 inputs 4D + 2D both dims 8.", + ), + pytest.param( + [ModelInputSpec((1, 8, 5, 5)), ModelInputSpec((1, 5))], + id="2 inputs 4D + 2D one dim 5.", + ), + pytest.param( + [ModelInputSpec((1, 8, 12, 10)), ModelInputSpec((8, 1, 10))], + id="2 inputs 4D + 3D channels dim 1.", + ), + pytest.param( + [ModelInputSpec((1, 8, 4, 10)), ModelInputSpec((1, 4, 1))], + id="2 inputs 4D + 3D channels dim 4.", + ), + pytest.param( + [ModelInputSpec((1, 8, 25, 18)), ModelInputSpec((4, 1, 8, 25, 18))], + id="2 inputs 4D + 5D conversion fails.", + marks=pytest.mark.xfail( + reason="AIR-14864: conversion fails", strict=True + ), + ), + ], + ) + def test__broadcast__channels_first_input(self, mocker, request, input_spec): + model = MaxPoolMaximumModule() + + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={Maximum: 1, MaxPool2DWithIndices: 1, GetItem: 1}, + expected_non_delegated_ops={}, + ) + dataset_creator = RandomDatasetCreator(low=-1.0, high=1.0) + + # Quantize the dataset and allow a single bit error. + remove_quant_io_ops = True + comparator = AllCloseOutputComparator(atol=1) + + lower_run_compare( + model, + input_spec, + graph_verifier, + request, + dataset_creator, + comparator, + remove_quant_io_ops=remove_quant_io_ops, + ) diff --git a/backends/nxp/tests/models.py b/backends/nxp/tests/models.py index a00ef602395..b7f73ed396e 100644 --- a/backends/nxp/tests/models.py +++ b/backends/nxp/tests/models.py @@ -1008,3 +1008,24 @@ def forward(self, x, y): x = torch.bmm(x, y) x = self.noop_max_pool_1d(x) return x + + +class MaximumModule(torch.nn.Module): + def __init__(self): + super().__init__() + + @staticmethod + def forward(x, y): + return torch.maximum(x, y) + + +class MaxPoolMaximumModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.max_pool2d = torch.nn.MaxPool2d( + kernel_size=1 + ) # No-op, but it enforces the channels first format. + + def forward(self, x, y): + x = self.max_pool2d(x) + return torch.maximum(x, y) diff --git a/backends/nxp/tests/ops_aliases.py b/backends/nxp/tests/ops_aliases.py index 17833edb14b..9aa7fa393d8 100644 --- a/backends/nxp/tests/ops_aliases.py +++ b/backends/nxp/tests/ops_aliases.py @@ -35,6 +35,7 @@ LeakyRelu = exir_ops.edge.aten.leaky_relu.default Log = exir_ops.edge.aten.log.default MM = exir_ops.edge.aten.mm.default +Maximum = exir_ops.edge.aten.maximum.default MaxPool2D = exir_ops.edge.aten.max_pool2d.default MaxPool2DWithIndices = exir_ops.edge.aten.max_pool2d_with_indices.default MeanDim = exir_ops.edge.aten.mean.dim diff --git a/docs/source/backends/nxp/op-support.csv b/docs/source/backends/nxp/op-support.csv index 5f97cf5d8b9..08793d46c47 100644 --- a/docs/source/backends/nxp/op-support.csv +++ b/docs/source/backends/nxp/op-support.csv @@ -21,6 +21,7 @@ aten.log.default,int8,static int8, aten.max_pool1d.default,int8,static int8,"dilation=1, ceil_mode=False, channels%8=0, batch_size=1, stride_h=1 or 2" aten.max_pool2d.default,int8,static int8,"dilation=1, ceil_mode=False, channels%8=0, batch_size=1, stride_h=1 or 2" aten.max_pool2d_with_indices.default,int8,static int8,"dilation=1, ceil_mode=False, channels%8=0, batch_size=1, stride_h=1 or 2" +aten.maximum.default,int8,static int8, aten.mean.dim,int8,static int8,"4D tensor only, dims = [-1,-2] or [-2,-1]" aten.mul.Tensor,int8,static int8,"tensor-size % 8 = 0" aten.mm.default,int8,static int8,"2D tensor only"