NXP backend: Add mean.dim support with new Neutron flow. (#19740)

### Summary
Add `mean.dim` support with new Neutron flow.

### Test plan
Unit tests provided.

cc @robert-kalmar @JakeStevens @digantdesai @rascani

Author: MartinPavella

backends/nxp/backend/edge_helper.py

@@ -318,7 +318,7 @@ def is_no_op_on_neutron(node: Node, parameters_mapping: dict[str, Parameter]) ->
                         input_data = torch.rand(val.shape, dtype=val.dtype) * 10 - 5
                         args_with_random_data.append(input_data)
 
-                case list():
+                case list() if any(isinstance(a, Node) for a in arg):
                     # Lists of input nodes are not supported to keep the code simple. It is not crucial to support this
                     #  case as the affected operators are either not supported on Neutron, or are extremely unlikely to
                     #  be no-ops (e.g. GRU). One exception is `aten.cat`, which is explicitly supported above.

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

@@ -152,9 +152,7 @@ def _get_node_args(
         :return: Tuple of (kernel_size, stride, padding, dilation, ceil_mode).
         """
         kernel_size = node.args[1]
-        stride = node.args[
-            2
-        ]  # The default value is equal to the kernel_size, so it is never empty here.
+        stride = try_get_arg(node, 2) or kernel_size
         padding = try_get_arg(node, 3) or (0, 0)
         dilation = try_get_arg(node, 4) or (1, 1)
         ceil_mode = try_get_arg(node, 5) or False

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

@@ -4,13 +4,15 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
+
 from executorch.backends.nxp.backend.data_format import NXP_NODE_FORMAT
 
 from executorch.backends.nxp.backend.ir.converter.conversion.translator import (
     create_channels_last_to_channels_first_permutation,
 )
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
+    is_not_qdq_node,
     NodeConverter,
 )
 from executorch.backends.nxp.backend.ir.converter.node_converters.shared.reduce_utils import (
@@ -21,45 +23,90 @@
 )
 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
 
 
 class MeanDimConverter(NodeConverter):
+
+    @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:
+        if custom_delegation_options.use_new_flow_neutron_c:
+            dim, keepdim = MeanDimConverter._get_attrs(node)
+            input_shape = node.args[0].meta["val"].shape
+
+            is_alone_in_partition = cls.is_node_alone_in_partition(
+                node, partition_list, filter_fn=is_not_qdq_node
+            )
+
+            if (
+                is_alone_in_partition
+                and keepdim
+                and all(input_shape[d] == 1 for d in dim)
+            ):
+                # The operator is a no-op, so the Neutron Converter will skip it. If it's the only node in the
+                #  partition, the graph would end up empty.
+                return False
+
+        return True
+
     @staticmethod
     def _is_supported_on_target(
         node: Node,
         neutron_target_spec: NeutronTargetSpec,
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        keepdim = node.args[2] if len(node.args) >= 3 else False
-        rank = len(node.args[0].meta["val"].shape)
-        dim = [MeanDimConverter._to_pos_dim(d, rank) for d in node.args[1]]
+        if custom_delegation_options.use_new_flow_neutron_c:
+            # Requirements specified by the new Neutron flow documentation.
+
+            if not NodeConverter.uses_quantization_type_for_io(
+                node,
+                supported_types=[torch.int8, torch.uint8],
+                input_indices=[0],
+                output_indices=[0],
+            ):
+                return False
 
-        if rank != 4 or not keepdim:
-            # neutron-converter/src/OperatorC/GlobalAvgPoolPlugin.cpp#74-77
-            return False
+            return True
 
-        # The `mean.dim` gets converted to AveragePool by the NeutronConverter, so the channels must be a
-        #  multiple of `num_macs`.
-        # neutron-converter/src/OperatorC/GlobalAvgPoolPlugin.cpp#59-85
-        num_macs = neutron_target_spec.get_num_macs()
-        channels_dim = 1 if node.meta[NXP_NODE_FORMAT].is_channels_first() else -1
-        if (node.meta["val"].shape[channels_dim] % num_macs) != 0:
-            return False
+        else:
+            # Requirements of the old Neutron flow.
+            rank = len(node.args[0].meta["val"].shape)
+            dim, keepdim = MeanDimConverter._get_attrs(node)
+            dim = [MeanDimConverter._to_pos_dim(d, rank) for d in dim]
 
-        # Neutron only supports reduction over the spatial dimensions H, W.
-        if node.meta[NXP_NODE_FORMAT].is_channels_first():
-            # The input is NCHW. H and W are at indices 2 and 3.
-            if dim not in [[2, 3], [3, 2]]:
+            if rank != 4 or not keepdim:
+                # neutron-converter/src/OperatorC/GlobalAvgPoolPlugin.cpp#74-77
                 return False
-        else:
-            # The input is formatless. It can be considered as NHWC, as this is the way Neutron will look at
-            #  the dimensions. So H and W are the middle dimensions.
-            if dim not in [[1, 2], [2, 1]]:
+
+            # The `mean.dim` gets converted to AveragePool by the NeutronConverter, so the channels must be a
+            #  multiple of `num_macs`.
+            # neutron-converter/src/OperatorC/GlobalAvgPoolPlugin.cpp#59-85
+            num_macs = neutron_target_spec.get_num_macs()
+            channels_dim = 1 if node.meta[NXP_NODE_FORMAT].is_channels_first() else -1
+            if (node.meta["val"].shape[channels_dim] % num_macs) != 0:
                 return False
 
-        return True
+            # Neutron only supports reduction over the spatial dimensions H, W.
+            if node.meta[NXP_NODE_FORMAT].is_channels_first():
+                # The input is NCHW. H and W are at indices 2 and 3.
+                if dim not in [[2, 3], [3, 2]]:
+                    return False
+            else:
+                # The input is formatless. It can be considered as NHWC, as this is the way Neutron will look at
+                #  the dimensions. So H and W are the middle dimensions.
+                if dim not in [[1, 2], [2, 1]]:
+                    return False
+
+            return True
 
     @staticmethod
     def _is_supported_in_IR(
@@ -91,15 +138,29 @@ def _normalize_and_to_channel_last_dim(dim: list[int], rank: int) -> list[int]:
         perm = create_channels_last_to_channels_first_permutation(rank, True)
         dim = [perm[d] for d in dim]
 
+        # noinspection PyTypeChecker
         return dim
 
-    # Mean Dim Node format: (Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None)
+    @staticmethod
+    def _get_attrs(node: Node) -> tuple[list[int], bool]:
+        dim = node.args[1]
+        keepdim = node.args[2] if len(node.args) >= 3 else False
+        return dim, keepdim
+
     def convert(self, node: Node):
-        """Convert 'mean.dim' operator to TFLite 'Mean'."""
+        """Convert the 'mean.dim' operator to NeutronIR 'Mean'.
+        The ExecuTorch schema is:
+            mean.dim(
+                Tensor self,
+                int[1]? dim,
+                bool keepdim=False,
+                *,
+                ScalarType? dtype=None
+            ) -> Tensor
+        """
         self.assert_convertible(node)
 
-        dim = node.args[1]
-        keepdim = node.args[2] if len(node.args) >= 3 else False
+        dim, keepdim = self._get_attrs(node)
 
         t_op = self._create_tflite_op_with_io_tensors(node)
         t_op.builtin_options = mean_options.Mean(keepdim)