NXP backend: Enable Amax with new Neutron flow

Author: roman-janik-nxp

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,

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

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

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

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:
                 #  <reduce_op>(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor
                 keep_dim = try_get_arg(node, 2) or False

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

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

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.