NXP backend: Add prelu support using new neutron flow

.github/workflows/pull.yml

@@ -1427,7 +1427,7 @@ jobs:
       docker-image: ci-image:executorch-ubuntu-22.04-clang12
       submodules: 'recursive'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-      timeout: 120
+      timeout: 150
       script: |
         set -eux
 

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

@@ -6,6 +6,7 @@
 import logging
 import operator
 from abc import ABC, abstractmethod
+from math import prod
 from typing import Callable
 
 import torch
@@ -465,29 +466,30 @@ def uses_shape_broadcasting(node: Node) -> bool:
         )
 
     @staticmethod
-    def at_least_one_input_shape_matches_the_output_shape(node: Node) -> bool:
-        """Determine if given PyTorch fx Node uses at least one input shape broadcasting for it's input nodes or not.
+    def inputs_satisfy_broadcast_condition(node: Node) -> bool:
+        """Determine if given PyTorch fx Node has inputs that satisfy broadcasting conditions for Neutron or not.
 
         :param node: PyTorch fx Node with 'all_input_nodes' initialized.
-        :return: True, if at least one input has the same shape as the output node.
+        :return: True, if at least one input has the same number of elements as the output node.
                  False otherwise.
         """
 
         if node.all_input_nodes is None:
             logger.e(
                 logger.Code.INTERNAL_ERROR,
-                "node_converter.at_least_one_input_shape_matches_the_output_shape(): 'all_input_nodes' are None!",
+                "node_converter.inputs_satisfy_broadcast_condition(): 'all_input_nodes' are None!",
             )
 
         if len(node.all_input_nodes) == 0:
             logger.e(
                 logger.Code.INTERNAL_ERROR,
-                "node_converter.at_least_one_input_shape_matches_the_output_shape(): Operator has no inputs!",
+                "node_converter.inputs_satisfy_broadcast_condition(): Operator has no inputs!",
             )
 
         output_shape = node.meta["val"].shape
+        num_elements = prod(output_shape)
 
         return any(
-            input_tensor.meta["val"].shape == output_shape
+            prod(input_tensor.meta["val"].shape) == num_elements
             for input_tensor in node.all_input_nodes
         )

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

@@ -26,7 +26,7 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        if not NodeConverter.at_least_one_input_shape_matches_the_output_shape(node):
+        if not NodeConverter.inputs_satisfy_broadcast_condition(node):
             return False
 
         supported_types = [torch.int8, torch.uint8]

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

@@ -26,7 +26,7 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        if not NodeConverter.at_least_one_input_shape_matches_the_output_shape(node):
+        if not NodeConverter.inputs_satisfy_broadcast_condition(node):
             return False
 
         supported_types = [torch.int8, torch.uint8]

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

@@ -3,7 +3,10 @@
 # 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.data_format import NXP_NODE_FORMAT
+from executorch.backends.nxp.backend.edge_helper import input_rank
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
     NodeConverter,
@@ -24,38 +27,17 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        node_shape = node.meta["val"].shape
-        rank = len(node_shape)
-
-        # According to Neutron spec., PReLU can be done only on 4D tensors
-        if rank != 4:
-            return False
-
-        ch_idx, h_idx, w_idx = PReLUConverter._get_channel_spatial_indices(node)
-        # According to Neutron spec., size of channels must be divisible by num_macs.
-        num_macs = neutron_target_spec.get_num_macs()
-        if node_shape[ch_idx] % num_macs != 0:
+        if not NodeConverter.inputs_satisfy_broadcast_condition(node):
             return False
 
-        # According to Neutron spec., height * width cannot be greater than a given constant.
-        if node_shape[w_idx] * node_shape[h_idx] > 4096:
+        supported_types = [torch.int8, torch.uint8]
+        if not NodeConverter.uses_quantization_type_for_io(
+            node, supported_types, [0, 1], [0]
+        ):
             return False
 
         return True
 
-    @staticmethod
-    def _get_channel_spatial_indices(node: Node):
-        if node.meta[NXP_NODE_FORMAT].is_channels_first():
-            ch_idx = 1
-            h_idx = 2
-            w_idx = 3
-        else:
-            ch_idx = 3
-            h_idx = 1
-            w_idx = 2
-
-        return ch_idx, h_idx, w_idx
-
     @staticmethod
     def _is_supported_in_IR(
         node: Node,
@@ -65,6 +47,12 @@ def _is_supported_in_IR(
         if len(node.args) != 2:
             return False
 
+        if input_rank(node, 0) > 2 and tuple(
+            node.all_input_nodes[1].meta["val"].shape
+        ) != (1,):
+            if not node.meta[NXP_NODE_FORMAT].is_channels_first():
+                # Should never happen.
+                return False
         return True
 
     def convert(self, node: Node):

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

@@ -26,7 +26,7 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        if not NodeConverter.at_least_one_input_shape_matches_the_output_shape(node):
+        if not NodeConverter.inputs_satisfy_broadcast_condition(node):
             return False
 
         supported_types = [torch.int8, torch.uint8]

backends/nxp/backend/node_format_inference.py

@@ -10,6 +10,7 @@
 
 from executorch.backends.nxp.backend.data_format import DataFormat, NXP_NODE_FORMAT
 from executorch.backends.nxp.backend.edge_helper import (
+    input_rank,
     is_channels_last_dim_order,
     try_get_arg,
 )
@@ -26,6 +27,7 @@
     MaxPool2DWithIndices,
     MeanDim,
     PermuteCopy,
+    Prelu,
     QuantizePerTensor,
     SumDimIntList,
     UpsampleBilinear2D,
@@ -54,6 +56,11 @@ class NodeFormatInference:
         UpsampleNearest2D: {"inputs": [0]},
     }
 
+    ops_conditionally_with_channels_first_nodes = {
+        Prelu: {"inputs": [0]},
+        torch.ops.aten.prelu.default: {"inputs": [0]},
+    }
+
     # A set of Edge Aten ops, which have the ability to change the format (for example - input nodes
     # are channels first but output is formatless).
     ops_that_can_change_tensor_format = {
@@ -64,6 +71,8 @@ class NodeFormatInference:
         SumDimIntList,
     }
 
+    prelu_targets = [Prelu, torch.ops.aten.prelu.default]
+
     _type_changed_during_last_run: bool
 
     # Mapping between Node and its ancestors (inputs)
@@ -160,6 +169,13 @@ def _infer_format_of_nodes(self, node: Node):
                     f"Node format inference for node type: {op_type} not found!"
                 )
 
+        elif node.target in self.prelu_targets:
+            num_parameters_shape = tuple(node.args[1].meta["val"].shape)
+            if input_rank(node, 0) > 2 and num_parameters_shape != (1,):
+                self._handle_node_which_uses_channels_first_format(node)
+            else:
+                self._handle_node_which_can_use_any_node_format(node)
+
         elif node.op != "call_function" or (
             hasattr(node, "target") and node.target in self._known_targets
         ):
@@ -248,10 +264,14 @@ def _handle_node_which_uses_channels_first_format(self, node: Node):
         Function for assigning format to nodes that require channels first input (Conv, MaxPool etc.)
         """
         op_type = self._get_node_op_type(node)
+        ops_using_channels_first_format = (
+            self.ops_with_channels_first_nodes
+            | self.ops_conditionally_with_channels_first_nodes
+        )
 
         for index, ancestor_node in enumerate(self._node_inputs[node]):
             # Go through input nodes and assign them correct format
-            if index in self.ops_with_channels_first_nodes[op_type]["inputs"]:
+            if index in ops_using_channels_first_format[op_type]["inputs"]:
                 self._assign_format_to_node(ancestor_node, DataFormat.CHANNELS_FIRST)
 
                 # We need to propagate channels first format up to already visited nodes

backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py

@@ -115,6 +115,10 @@ def test__basic_nsys_inference_qat(self, mocker, request):
                 [ModelInputSpec((1, 4, 8, 8)), ModelInputSpec((8, 8))],
                 id="2 inputs 4D + 2D.",
             ),
+            pytest.param(
+                [ModelInputSpec((10,)), ModelInputSpec((1, 1))],
+                id="2 inputs 1D + 2D, num_elems of input == num_elems of output",
+            ),
         ],
     )
     def test__broadcast(self, mocker, request, input_spec):

backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py

@@ -96,6 +96,10 @@ def test__basic_nsys_inference_qat(self, mocker, request, x_input_shape):
             pytest.param(
                 [ModelInputSpec((4,)), ModelInputSpec((4, 4))], id="2 inputs 1D+2D."
             ),
+            pytest.param(
+                [ModelInputSpec((10,)), ModelInputSpec((1, 1))],
+                id="2 inputs 1D + 2D, num_elems of input == num_elems of output",
+            ),
         ],
     )
     def test__broadcast(self, input_spec, mocker, request):