NXP backend: Enable sum with new Neutron flow

Author: roman-janik-nxp

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

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

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

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

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:
                 #  <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

@@ -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

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

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.