NXP backend: Enable cat with new Neutron flow.

Author: MartinPavella

backends/nxp/backend/custom_delegation_options.py

@@ -11,13 +11,6 @@
 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
-
     # Proposed partitions which only contain Neutron no-ops are normally not delegated, as the NeutronConverter would
     #  not create any NeutronGraph that can be called. This is done by the partitioner itself, and is not handled by
     #  the individual node converters.

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

@@ -8,13 +8,7 @@
 from executorch.backends.nxp.backend.custom_delegation_options import (
     CustomDelegationOptions,
 )
-from executorch.backends.nxp.backend.data_format import NXP_NODE_FORMAT
-from executorch.backends.nxp.backend.edge_helper import previous_non_qdq_node
 from executorch.backends.nxp.backend.ir.converter.conversion import translator
-from executorch.backends.nxp.backend.ir.converter.conversion.translator import (
-    apply_permutation_to,
-    create_channels_first_to_channels_last_permutation,
-)
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     _is_dequant_node,
     _is_quant_node,
@@ -25,7 +19,6 @@
 )
 from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
 from torch.fx import Node
-from torch.fx.passes.infra.partitioner import Partition
 from torch.nn import Parameter
 
 
@@ -83,56 +76,12 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        if custom_delegation_options.force_delegate_cat:
-            return True
-
-        dim = CatConverter._get_normalized_dim(node)
-
-        # Neutron requires the channels to be a multiple of `num_macs`. The channels could either be the second or the
-        #  last dimension, depending on the formats of the node.
-        if node.meta[NXP_NODE_FORMAT].is_channels_first():
-            # During conversion to IR, the shape will be permuted to channels last, and the dimension on index
-            #  `1` will end up being the channels (last dim in NHWC).
-            channels_index = 1
-            to_nhwc_perm = create_channels_first_to_channels_last_permutation(
-                len(node.meta["val"].shape), True
-            )
-            dim = to_nhwc_perm.index(
-                dim
-            )  # Make sure the dim points to the NHWC dimension.
-        else:
-            # The shape will not be permuted during conversion, so the channels will remain the last dimension.
-            channels_index = -1
-
-        input_channels = [
-            _get_shape(input_)[channels_index] for input_ in node.all_input_nodes
-        ]
-        output_channels = _get_shape(node)[channels_index]
-
-        num_macs = neutron_target_spec.get_num_macs()
-        input_shapes = [_get_shape(input_) for input_ in node.all_input_nodes]
-        if any((input_channel % num_macs) != 0 for input_channel in input_channels):
-            # neutron-library/src/utils/NeutronLibraryInterrogation.cpp#1492
-
-            # If all input shapes are equal, the neutron is able to pad the last dimension of the inputs.
-            if not (
-                input_shapes.count(input_shapes[0]) == len(input_shapes)
-                and dim == len(input_shapes[0]) - 1
-            ):
-                return False
-
-        if (output_channels % num_macs) != 0:
-            # neutron-library/src/utils/NeutronLibraryInterrogation.cpp#1493
-
-            # If all input shapes are equal, the neutron is able to pad the last dimension of the output.
-            if not (
-                input_shapes.count(input_shapes[0]) == len(input_shapes)
-                and dim == len(input_shapes[0]) - 1
-            ):
-                return False
-
-        if len(node.all_input_nodes) < 2:  # Not supported on Neutron
-            # TODO Try to skip the operator if this case is realistic.
+        # `cat` uses a list of inputs as its first argument, so the indices are tuples of (0, i).
+        input_indices = [(0, i) for i in range(len(node.args[0]))]
+        supported_types = [torch.int8, torch.uint8]
+        if not NodeConverter.uses_quantization_type_for_io(
+            node, supported_types, input_indices=input_indices, output_indices=[0]
+        ):
             return False
 
         return True
@@ -150,50 +99,14 @@ def _is_supported_in_IR(
 
         return True
 
-    @classmethod
-    def supports_partitioning_result(
-        cls,
-        node: Node,
-        partition_list: list[Partition],
-        custom_delegation_options: CustomDelegationOptions,
-        neutron_target_spec: NeutronTargetSpec,
-        parameters_mapping: dict[str, Parameter],
-    ) -> bool:
-        # There is a bug in the NeutronConverter, where if none of the input dimensions before the one referenced by
-        #  `dim` are `!= 1`, the `Concat` is not delegated.
-        # This only happens when the inputs to the `Concat` are model inputs, and not outputs of other
-        #  operators.
-        cat_partition = [p for p in partition_list if node in p.nodes][0]
-        cat_inputs = map(previous_non_qdq_node, node.args[0])
-
-        if not all(
-            input_.op == "call_function" and input_ in cat_partition.nodes
-            for input_ in cat_inputs
-        ):
-            # Some inputs of the `cat` are NOT in the same partition as `cat`.
-            dim = CatConverter._get_normalized_dim(node)
-            input_shapes = [list(n.meta["val"].shape) for n in node.args[0]]
-            if node.meta[NXP_NODE_FORMAT].is_channels_first():
-                # Transform the shapes to channels last.
-                to_nhwc_perm = create_channels_first_to_channels_last_permutation(
-                    len(node.meta["val"].shape), True
-                )
-                input_shapes = [
-                    apply_permutation_to(shape, to_nhwc_perm) for shape in input_shapes
-                ]
-
-                # Transform the `dim` to refer to a channels last dimension.
-                dim = to_nhwc_perm.index(dim)
-
-            for input_shape in input_shapes:
-                if not any(d != 1 for d in input_shape[:dim]):
-                    # Do not delegate if there are no "non-1" dimensions in the shape before the `dim` dimension.
-                    return False
-
-        return True
-
     def convert(self, node: Node):
-        """Convert the 'aten.cat' operator to TFLite 'Concatenation'."""
+        """Convert the 'aten.cat' operator to NeutronIR 'Concatenation'.
+        The ExecuTorch schema is:
+            cat(
+                Tensor[] tensors,
+                int dim=0
+            ) -> Tensor
+        """
         self.assert_convertible(node)
 
         t_op = self._create_tflite_op_with_io_tensors(node)
@@ -205,5 +118,5 @@ def convert(self, node: Node):
                 t_op.tmp_inputs[0].rank
             )[dim]
 
-        t_op.builtin_options = Concatenation(dim)
+        t_op.builtin_options = Concatenation(int(dim))
         self.builder.append_operators([t_op])

backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py

@@ -120,7 +120,15 @@ def test_noop_partitions__concatenate_one_tensor_and_add_zeros():
     )
 
 
+@pytest.mark.xfail(
+    strict=True,
+    reason="Neutron Converter currently supports these 2 noops in sequence.",
+)
 def test_noop_partitions__concatenate_one_tensor_and_add_zeros__forced_delegation():
+    # When the noop `Concatenate` and noop `Add` are in sequence, Neutron Converter supports them. This edge case is
+    #  not reflected in our logic. But as this edge case is extremely rare (and even if it ever happened in a real
+    #  model, the consequences would be minimal), fixing it is not a priority.
+
     input_shape = (1, 2, 3, 4)
     module = ConcatAddNoOpModel(input_shape)