NXP backend: Add support for aten.add.Tensor, aten.mean.dim, aten._adaptive_avg_pool2d, aten.clone and aten.abs operators (pytorch#12585)

### Summary

Add quantization and conversion support for the following operators: 
- aten.add.Tensor
- aten.mean.dim
- aten._adaptive_avg_pool2d
- aten.clone
- aten.abs

### Test plan
All newly supported operators should be covered by unit tests.

---------

Co-authored-by: Robert Kalmar <[email protected]>
Co-authored-by: Roman Janik <[email protected]>

Author: 3 people

backends/nxp/backend/edge_program_converter.py

@@ -23,15 +23,20 @@
 
 # noinspection PyProtectedMember
 functions_converters = {
+    exir_ops.edge.aten.abs.default: AbsConverter,  # noqa F405
+    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.avg_pool2d.default: AvgPool2dConverter,  # noqa F405
+    exir_ops.edge.aten.clone.default: CloneConverter,  # noqa F405
     exir_ops.edge.aten.constant_pad_nd.default: ConstantPadNDConverter,  # noqa F405
     exir_ops.edge.aten.convolution.default: ConvolutionConverter,  # noqa F405
+    exir_ops.edge.aten.hardtanh.default: HardTanhConverter,  # noqa F405
     exir_ops.edge.aten.max_pool2d.default: MaxPool2dConverter,  # noqa F405
+    exir_ops.edge.aten.mean.dim: MeanDimConverter,  # noqa F405
     exir_ops.edge.aten.mm.default: MMConverter,  # noqa F405
     exir_ops.edge.aten.permute_copy.default: PermuteCopyConverter,  # noqa F405
     exir_ops.edge.aten.relu.default: ReLUConverter,  # noqa F405
-    exir_ops.edge.aten.hardtanh.default: HardTanhConverter,  # noqa F405
     exir_ops.edge.aten._softmax.default: SoftmaxConverter,  # noqa F405
     exir_ops.edge.aten.view_copy.default: ViewCopyConverter,  # noqa F405
 }

backends/nxp/backend/ir/converter/conversion/common.py

@@ -22,6 +22,7 @@
     max_pool_2d_options,
     transpose_conv_options,
 )
+from torch.fx import Node
 
 
 def exactly_one_is_none(obj1: Optional, obj2: Optional) -> bool:
@@ -166,6 +167,34 @@ def uses_shape_broadcasting(t_op: tflite_model.Operator) -> bool:
     )
 
 
+def node_uses_shape_broadcasting(node: Node) -> bool:
+    """Determine if given PyTorch fx Node uses shape broadcasting for it's input nodes or not.
+
+    :param node: PyTorch fx Node with 'all_input_nodes' initialized.
+    :return: True, if the node uses shape broadcasting for it's input nodes.
+             False otherwise.
+    """
+
+    if node.all_input_nodes is None:
+        logger.e(
+            logger.Code.INTERNAL_ERROR,
+            "common.node_uses_shape_broadcasting(): 'all_input_nodes' are None!",
+        )
+
+    if len(node.all_input_nodes) == 0:
+        logger.e(
+            logger.Code.INTERNAL_ERROR,
+            "common.node_uses_shape_broadcasting(): Operator has no inputs!",
+        )
+
+    first_input_shape = node.all_input_nodes[0].meta["val"].shape
+
+    return any(
+        input_tensor.meta["val"].shape != first_input_shape
+        for input_tensor in node.all_input_nodes[1:]
+    )
+
+
 def uses_multiple_input_types(t_op: tflite_model.Operator) -> bool:
     """Determine if the input tensors of given TFLite operator use different data types or not.
 

backends/nxp/backend/ir/converter/node_converter.py

@@ -173,7 +173,8 @@ def _create_tflite_op_with_io_tensors(self, node: Node) -> tflite_model.Operator
 
         # Initialize node's inputs
         t_operator.inputs = tflite_model.OperatorInputs()
-        for ancestor_node in node.all_input_nodes:
+        input_nodes = [arg for arg in node.args if isinstance(arg, Node)]
+        for ancestor_node in input_nodes:
             assert self.context.tflite_builder.tensor_exists(ancestor_node.name)
             t_operator.tmp_inputs.append(
                 self.context.tflite_builder.tensor_for_name(ancestor_node.name)

backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py

@@ -1,9 +1,21 @@
+from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.abs_converter import (
+    AbsConverter,
+)
+from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.adaptive_avg_pool_2d_converter import (
+    AdaptiveAvgPool2dConverter,
+)
+from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.add_tensor_converter import (
+    AddTensorConverter,
+)
 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.avg_pool_2d_converter import (
     AvgPool2dConverter,
 )
+from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.clone_converter import (
+    CloneConverter,
+)
 from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.constant_pad_nd_converter import (
     ConstantPadNDConverter,
 )
@@ -16,6 +28,9 @@
 from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.max_pool_2d_converter import (
     MaxPool2dConverter,
 )
+from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.mean_dim_converter import (
+    MeanDimConverter,
+)
 from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.mm_converter import (
     MMConverter,
 )
@@ -49,7 +64,12 @@
     "QDQQuantizeConverter",
     "ConstantPadNDConverter",
     "ReLUConverter",
+    "MeanDimConverter",
     "MaxPool2dConverter",
     "AvgPool2dConverter",
+    "AddTensorConverter",
+    "CloneConverter",
+    "AbsConverter",
+    "AdaptiveAvgPool2dConverter",
     "HardTanhConverter",
 ]

backends/nxp/backend/ir/converter/node_converters/ops_converters/abs_converter.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2025 NXP
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+from executorch.backends.nxp.backend.ir.converter.node_converter import (
+    NodeConverter,
+    Target,
+)
+from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
+    abs_options,
+)
+from torch.fx import Node
+from torch.nn import Parameter
+
+
+class AbsConverter(NodeConverter):
+    supported_targets = [Target.RT700]
+
+    @staticmethod
+    def _is_supported_in_IR(
+        node: Node, parameters_mapping: dict[str, Parameter]
+    ) -> bool:
+        return True
+
+    def convert(self, node: Node):
+        """Convert 'aten::abs' operator to TFLite 'Abs'."""
+        self.assert_convertible(node)
+
+        t_op = self._create_tflite_op_with_io_tensors(node)
+
+        t_op.builtin_options = abs_options.Abs()
+        self.builder.append_operators([t_op])

backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2025 NXP
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import executorch.backends.nxp.backend.ir.lib.tflite.Padding as tflPadding
+from executorch.backends.nxp.backend.ir.converter.conversion import common
+from executorch.backends.nxp.backend.ir.converter.node_converter import (
+    NodeConverter,
+    Target,
+)
+from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model
+from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
+    average_pool_2d_options,
+)
+from torch import Size
+from torch.fx import Node
+from torch.nn import Parameter
+
+
+class AdaptiveAvgPool2dConverter(NodeConverter):
+    supported_targets = [Target.RT700]
+
+    @staticmethod
+    def _is_supported_in_IR(
+        node: Node, parameters_mapping: dict[str, Parameter]
+    ) -> bool:
+        input_size = node.args[0].meta["val"].shape
+        output_size = node.args[1]
+
+        if (input_size[-1] % output_size[-1] != 0) or (
+            input_size[-2] % output_size[-2] != 0
+        ):
+            return False
+
+        if not NodeConverter._has_shared_q_params_if_quantized(node):
+            return False
+
+        return True
+
+    # noinspection PyMethodMayBeStatic
+    def _convert_adaptive_avg_pool_2d(
+        self, input_size: Size, output_size: list[int], t_op: tflite_model.Operator
+    ):
+        t_op.builtin_options = average_pool_2d_options.AveragePool2D()
+        stride = [input_size[-2] // output_size[-2], input_size[-1] // output_size[-1]]
+        common.assign_2d_strides(t_op.builtin_options, stride)
+        t_op.builtin_options.filter_h = (
+            input_size[-2] - (output_size[-2] - 1) * stride[-2]
+        )
+        t_op.builtin_options.filter_w = (
+            input_size[-1] - (output_size[-1] - 1) * stride[-1]
+        )
+        t_op.builtin_options.padding = tflPadding.Padding.VALID
+
+    # AdaptiveAvgPool2d Node format: (Tensor self, SymInt[2] output_size)
+    def convert(self, node: Node):
+        """Convert '_adaptive_avg_pool2d' operator to TFLite 'AveragePool2D'."""
+        self.assert_convertible(node)
+
+        input_size = node.args[0].meta["val"].shape
+        output_size = node.args[1]
+
+        t_op = self._create_tflite_op_with_io_tensors(node)
+
+        self._convert_adaptive_avg_pool_2d(input_size, output_size, t_op)
+        self.builder.append_operators([t_op])

backends/nxp/backend/ir/converter/node_converters/ops_converters/add_tensor_converter.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2025 NXP
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from executorch.backends.nxp.backend.ir.converter.conversion.common import (
+    node_uses_shape_broadcasting,
+)
+from executorch.backends.nxp.backend.ir.converter.node_converter import (
+    NodeConverter,
+    Target,
+)
+from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
+    add_options,
+)
+from torch.fx import Node
+from torch.nn import Parameter
+
+
+class AddTensorConverter(NodeConverter):
+    supported_targets = [Target.RT700]
+
+    @staticmethod
+    def _is_supported_in_IR(
+        node: Node, parameters_mapping: dict[str, Parameter]
+    ) -> bool:
+        if len(node.args) != 2:
+            return False
+
+        if hasattr(node.kwargs, "alpha"):
+            return False
+
+        # Don't convert if broadcasting input tensors
+        if node_uses_shape_broadcasting(node):
+            return False
+
+        return True
+
+    # add.Tensor Node format: (Tensor self, Tensor other, *, Scalar alpha=1)
+    def convert(self, node: Node):
+        """Convert 'add_tensor' operator to TFLite 'add'."""
+        self.assert_convertible(node)
+
+        t_op = self._create_tflite_op_with_io_tensors(node)
+
+        t_op.builtin_options = add_options.Add()
+        self.builder.append_operators([t_op])

backends/nxp/backend/ir/converter/node_converters/ops_converters/clone_converter.py

@@ -0,0 +1,38 @@
+# Copyright 2025 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.node_converter import (
+    NodeConverter,
+    Target,
+)
+from torch.fx import Node
+from torch.nn import Parameter
+
+
+def _has_supported_memory_format(node: Node) -> bool:
+    if "memory_format" in node.kwargs.keys():
+        return node.kwargs["memory_format"] == torch.preserve_format
+
+    return True
+
+
+class CloneConverter(NodeConverter):
+    supported_targets = [Target.RT700]
+
+    @staticmethod
+    def _is_supported_in_IR(
+        node: Node, parameters_mapping: dict[str, Parameter]
+    ) -> bool:
+        return _has_supported_memory_format(node)
+
+    def convert(self, node: Node):
+        """Skip `aten.clone` operator if it has no `memory_format` specified."""
+        self.assert_convertible(node)
+
+        t_op = self._create_tflite_op_with_io_tensors(node)
+
+        self.builder.turn_operator_to_identity(t_op)
+        self.builder.append_operators([t_op])