NXP backend: Add maximum support using new neutron flow

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

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

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",

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())

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

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),

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."""