NXP backend: Add support for the aten.cat operator. (#13505)

### Summary
Add delegation of `aten.cat` to Neutron, and a
`CustomDelegationOptions` class allowing delegation rules
to be overridden.

The CustomDelegationOption is introduced to allow force delegation of aten.cat operator, opportunistically, in cases when the constaint on number of channels cannot be determined automatically.

### Test plan
Unit tests provided in
`backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py`.


cc @digantdesai @JakeStevens @robert-kalmar

Author: MartinPavella

backends/nxp/backend/custom_delegation_options.py

@@ -0,0 +1,19 @@
+# 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.
+
+
+from dataclasses import dataclass
+
+
+@dataclass
+class CustomDelegationOptions:
+    """The class allows the user to specify details which affect which nodes will be delegated."""
+
+    # Neutron requires the channel dimension to be multiple of `num_macs` for concatenation (cat op).
+    #  Due to different dim ordering in torch (channel_first) and Neutron IR (channel last), dim of the channel is
+    #  ambiguous. Cat converter will defensively require both possible dimension index for the channels to be multiple
+    #  of `num_macs`. The `force_delegate_cat` allows the user to turn off the defensive check if from the model design
+    #  it is known this constraint will be satisfied.
+    force_delegate_cat: bool = False

backends/nxp/backend/edge_program_converter.py

@@ -1,4 +1,4 @@
-# Copyright 2024 NXP
+# Copyright 2024-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.
@@ -10,6 +10,9 @@
 from executorch.backends.nxp.backend.ir.converter.builder.aten_model_builder_director import (
     AtenModelBuilderDirector,
 )
+from executorch.backends.nxp.backend.ir.converter.node_converter import (
+    CustomDelegationOptions,
+)
 from torch.export import ExportedProgram
 from torch.export.graph_signature import InputKind
 from torch.fx import Node
@@ -28,6 +31,7 @@
     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.cat.default: CatConverter,  # 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
@@ -49,24 +53,30 @@ class EdgeProgramToIRConverter:
     """
 
     _default_conversion_config = ConversionConfig()
+    _default_delegation_options = CustomDelegationOptions()
 
     def convert_program(
         self,
         edge_program: ExportedProgram,
         conversion_config=_default_conversion_config,
+        custom_delegation_options: CustomDelegationOptions = _default_delegation_options,
     ) -> (bytes, dict):
         """
         Convert ExportedProgram in Edge dialect to IR (TFLite flatbuffers) as bytes.
 
         :param edge_program: Converter ExportedProgram.
         :param conversion_config: ConversionConfig instance.
+        :param custom_delegation_options: Custom user options which affect node delegation.
         :return: TFLite flatbuffers as bytes.
         """
         node_formats = NodeFormatInference(edge_program).identify_node_formats()
         parameters_mapping = self.map_inputs_to_parameters(edge_program)
 
         cc = self.build_conversion_context(
-            parameters_mapping, node_formats, conversion_config
+            parameters_mapping,
+            node_formats,
+            conversion_config,
+            custom_delegation_options,
         )
 
         # Program conversion
@@ -162,6 +172,7 @@ def build_conversion_context(
         parameters_mapping: dict,
         node_formats: dict[Node, NodeFormat],
         conversion_config: ConversionConfig = _default_conversion_config,
+        custom_delegation_options: CustomDelegationOptions = _default_delegation_options,
     ) -> ConversionContext:
         tflite_builder = AtenModelBuilderDirector(
             3, "TFLite from EdgeProgram", conversion_config
@@ -171,7 +182,11 @@ def build_conversion_context(
         tflite_builder.build_empty_buffer()
 
         context = ConversionContext(
-            tflite_builder, conversion_config, parameters_mapping, node_formats
+            tflite_builder,
+            conversion_config,
+            parameters_mapping,
+            node_formats,
+            custom_delegation_options,
         )
 
         return context

backends/nxp/backend/ir/conversion_context.py

@@ -1,8 +1,11 @@
-# Copyright 2024 NXP
+# Copyright 2024-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.
 
+from executorch.backends.nxp.backend.custom_delegation_options import (
+    CustomDelegationOptions,
+)
 from executorch.backends.nxp.backend.ir.conversion_config import ConversionConfig
 from executorch.backends.nxp.backend.ir.converter.builder.aten_model_builder_director import (
     AtenModelBuilderDirector,
@@ -17,13 +20,15 @@ class ConversionContext:
     conversion_config: ConversionConfig
     parameters_mapping: dict[str, Parameter]
     node_formats: dict[Node, NodeFormat]
+    custom_delegation_options: CustomDelegationOptions
 
     def __init__(
         self,
         tflite_builder: AtenModelBuilderDirector,
         conversion_config: ConversionConfig,
         parameters_mapping: dict,
         node_formats: dict[Node, NodeFormat],
+        custom_delegation_options: CustomDelegationOptions,
     ):
         """
         Context with data related to current conversion.
@@ -35,3 +40,4 @@ def __init__(
         self.conversion_config = conversion_config
         self.parameters_mapping = parameters_mapping
         self.node_formats = node_formats
+        self.custom_delegation_options = custom_delegation_options

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

@@ -8,6 +8,9 @@
 
 import torch
 
+from executorch.backends.nxp.backend.custom_delegation_options import (
+    CustomDelegationOptions,
+)
 from executorch.backends.nxp.backend.ir.conversion_context import ConversionContext
 from executorch.backends.nxp.backend.ir.converter.builder.aten_model_builder_director import (
     AtenModelBuilderDirector,
@@ -70,19 +73,25 @@ def convert(self, node: Node):
     @staticmethod
     @abstractmethod
     def _is_supported_in_IR(
-        node: Node, parameters_mapping: dict[str, Parameter]
+        node: Node,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
         """Check if the `node` can be converted to the intermediate representation.
             Classes which implement conversion for individual operators must overwrite this method.
 
         :param node: torch.Node to check.
         :param parameters_mapping: Dictionary mapping tensor names to their static data (if they have it).
+        :param custom_delegation_options: Custom options which affect delegation.
         """
         pass
 
     @staticmethod
     def _is_supported_on_target(
-        node: Node, target: Target, parameters_mapping: dict[str, Parameter]
+        node: Node,
+        target: Target,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
         """Check if the node is supported on the target platform.
             Child classes should overwrite this method to implement specific target checks. The default implementation
@@ -91,22 +100,30 @@ def _is_supported_on_target(
         :param node: The node (edge operator) to check.
         :param target: Value of the `Target` enum representing the target platform to check for.
         :param parameters_mapping: Dictionary mapping tensor names to their static data (if they have it).
+        :param custom_delegation_options: Custom options which affect delegation.
         """
         return target == Target.RT700
 
     @classmethod
     def is_supported(
-        cls, node: Node, target: Target, parameters_mapping: dict[str, Parameter]
+        cls,
+        node: Node,
+        target: Target,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
         """Check if the given `node` is supported in the IR and on the given `target` platform.
 
         :param node: torch.Node to check.
         :param target: Value of the `Target` enum representing the target platform to check for.
         :param parameters_mapping: Dict mapping tensor names to their data.
+        :param custom_delegation_options: Custom user options which affect node delegation.
         """
         return cls._is_supported_in_IR(
-            node, parameters_mapping
-        ) and cls._is_supported_on_target(node, target, parameters_mapping)
+            node, parameters_mapping, custom_delegation_options
+        ) and cls._is_supported_on_target(
+            node, target, parameters_mapping, custom_delegation_options
+        )
 
     @staticmethod
     def _has_shared_q_params_if_quantized(node: Node) -> bool:
@@ -145,7 +162,11 @@ def assert_convertible(self, node):
         """Assert that the call `_is_supported_in_IR()` returns `True`. Otherwise, raise an exception and print an
         error message.
         """
-        assert self._is_supported_in_IR(node, self.context.parameters_mapping), (
+        assert self._is_supported_in_IR(
+            node,
+            self.context.parameters_mapping,
+            self.context.custom_delegation_options,
+        ), (
             f"Node `{node}` is not convertible to the intermediate representation. "
             "There is an error in the partitioner."
         )
@@ -169,7 +190,15 @@ def _create_tflite_op_with_io_tensors(self, node: Node) -> tflite_model.Operator
 
         # Initialize node's inputs
         t_operator.inputs = tflite_model.OperatorInputs()
-        input_nodes = [arg for arg in node.args if isinstance(arg, Node)]
+
+        input_nodes = []
+        for arg in node.args:
+            match arg:
+                case Node():
+                    input_nodes.append(arg)
+                case list() if all(isinstance(node_, Node) for node_ in arg):
+                    input_nodes.extend(arg)
+
         for ancestor_node in input_nodes:
             assert self.context.tflite_builder.tensor_exists(ancestor_node.name)
             t_operator.tmp_inputs.append(

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

@@ -13,6 +13,9 @@
 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.cat_converter import (
+    CatConverter,
+)
 from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.clone_converter import (
     CloneConverter,
 )
@@ -58,6 +61,7 @@
 
 __all__ = [
     "AddMMConverter",
+    "CatConverter",
     "ConvolutionConverter",
     "MMConverter",
     "PermuteCopyConverter",

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

@@ -4,7 +4,10 @@
 # LICENSE file in the root directory of this source tree.
 
 
-from executorch.backends.nxp.backend.ir.converter.node_converter import NodeConverter
+from executorch.backends.nxp.backend.ir.converter.node_converter import (
+    CustomDelegationOptions,
+    NodeConverter,
+)
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
     abs_options,
 )
@@ -16,7 +19,9 @@ class AbsConverter(NodeConverter):
 
     @staticmethod
     def _is_supported_in_IR(
-        node: Node, parameters_mapping: dict[str, Parameter]
+        node: Node,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
         return True
 

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

@@ -5,7 +5,10 @@
 
 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
+from executorch.backends.nxp.backend.ir.converter.node_converter import (
+    CustomDelegationOptions,
+    NodeConverter,
+)
 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,
@@ -19,7 +22,9 @@ class AdaptiveAvgPool2dConverter(NodeConverter):
 
     @staticmethod
     def _is_supported_in_IR(
-        node: Node, parameters_mapping: dict[str, Parameter]
+        node: Node,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
         input_size = node.args[0].meta["val"].shape
         output_size = node.args[1]

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

@@ -7,6 +7,7 @@
     node_uses_shape_broadcasting,
 )
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
+    CustomDelegationOptions,
     NodeConverter,
     Target,
 )
@@ -20,7 +21,10 @@
 class AddTensorConverter(NodeConverter):
     @staticmethod
     def _is_supported_on_target(
-        node: Node, target: Target, parameters_mapping: dict[str, Parameter]
+        node: Node,
+        target: Target,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
         match target:
             case Target.RT700:
@@ -35,7 +39,9 @@ def _is_supported_on_target(
 
     @staticmethod
     def _is_supported_in_IR(
-        node: Node, parameters_mapping: dict[str, Parameter]
+        node: Node,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
         if len(node.args) != 2:
             return False

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

@@ -5,7 +5,10 @@
 
 from executorch.backends.nxp.backend.edge_helper import input_rank
 from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList
-from executorch.backends.nxp.backend.ir.converter.node_converter import NodeConverter
+from executorch.backends.nxp.backend.ir.converter.node_converter import (
+    CustomDelegationOptions,
+    NodeConverter,
+)
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
     fully_connected_options,
 )
@@ -18,7 +21,9 @@ class AddMMConverter(NodeConverter):
 
     @staticmethod
     def _is_supported_in_IR(
-        node: Node, parameters_mapping: dict[str, Parameter]
+        node: Node,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
         if len(node.all_input_nodes) != 3:
             return False

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

@@ -8,7 +8,10 @@
     common,
 )
 from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList
-from executorch.backends.nxp.backend.ir.converter.node_converter import NodeConverter
+from executorch.backends.nxp.backend.ir.converter.node_converter import (
+    CustomDelegationOptions,
+    NodeConverter,
+)
 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,
@@ -21,7 +24,9 @@ class AvgPool2dConverter(NodeConverter):
 
     @staticmethod
     def _is_supported_in_IR(
-        node: Node, parameters_mapping: dict[str, Parameter]
+        node: Node,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
         n_args = len(node.args)