feat: added support for conv2d using new NeutronC flow

Author: novak-vaclav

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

@@ -48,78 +48,171 @@
 from torch.nn import Parameter
 
 
+# The arguments of the conv are:
+# convolution(
+#   Tensor input, Tensor weight, Tensor? bias,
+#   SymInt[] stride, SymInt[] padding, SymInt[] dilation,
+#   bool transposed, SymInt[] output_padding, SymInt groups
+# ) -> Tensor
+Stride = Padding = Dilation = OutPadding = list[int]
+Transposed = bool
+Groups = int
+ConvolutionArgs = tuple[
+    Node, Node, Node | None, Stride, Padding, Dilation, Transposed, OutPadding, Groups
+]
+
+
 class ConvolutionConverter(NodeConverter):
     @staticmethod
-    def _is_supported_on_target(
+    def _is_supported_on_target_regular_conv(
+        node: Node,
+        parameters_mapping: dict[str, Parameter],
+    ) -> bool:
+        (
+            inp_node,
+            w_node,
+            b_node,
+            stride,
+            _,
+            dilation,
+            _,
+            _,
+            _,
+        ) = ConvolutionConverter._get_convolution_arguments(node)
+
+        # Input must be INT8/UINT8
+        # Output must be INT8/UINT8
+        inp_out_supported_types = [torch.int8, torch.uint8]
+        if not NodeConverter.uses_quantization_type_for_io(
+            node, inp_out_supported_types, [0], [0]
+        ):
+            return False
+
+        # Weights must be INT8
+        w_supported_types = [torch.int8]
+        if not NodeConverter.uses_quantization_type_for_io(
+            node, w_supported_types, [1], []
+        ):
+            return False
+
+        # Bias must be INT32
+        if b_node is not None:
+            b_supported_types = [torch.int32]
+            if not NodeConverter.uses_quantization_type_for_io(
+                node, b_supported_types, [2], []
+            ):
+                return False
+
+        # Weights must be constant
+        if not node_is_effectively_static_tensor(w_node, parameters_mapping):
+            return False
+
+        # Bias must be constant (if present)
+        if b_node is not None and not node_is_effectively_static_tensor(
+            b_node, parameters_mapping
+        ):
+            return False
+
+        # kernelH <= 4096, kernelW <= 4096
+        # strideH <= 4096, strideW <= 4096
+        # dilationH <= 4096, dilationW <= 4096
+        w_node_shape = w_node.meta["val"].shape
+
+        kernel_h = w_node_shape[2]
+        kernel_w = w_node_shape[3]
+        stride_h = stride[0]
+        stride_w = stride[1]
+        dilation_h = dilation[0]
+        dilation_w = dilation[1]
+
+        dim_sizes = [kernel_h, kernel_w, stride_h, stride_w, dilation_h, dilation_w]
+
+        if any(dim > 4096 for dim in dim_sizes):
+            return False
+
+        # kernelH * kernelW * inpC <= 65535
+        inp_node_shape = inp_node.meta["val"].shape
+        inp_channels = (
+            inp_node_shape[1] if len(inp_node_shape) == 4 else inp_node_shape[0]
+        )
+
+        if kernel_h * kernel_w * inp_channels > 65535:
+            return False
+
+        return True
+
+    @staticmethod
+    def _is_supported_on_target_transp_conv(
         node: Node,
         neutron_target_spec: NeutronTargetSpec,
         parameters_mapping: dict[str, Parameter],
-        custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
+        # TODO: EIEX-894 update the requirements of delegation for new Neutron flow
+        _, w_node, _, stride, padding, dilation, transposed, _, groups = (
+            ConvolutionConverter._get_convolution_arguments(node)
+        )
+
         num_macs = neutron_target_spec.get_num_macs()
         node_t_params = get_node_tensor_params(node)
-        weights = node.args[1]
-        conv_params = ConvParameters(
-            *ConvolutionConverter._get_convolution_arguments(node)
-        )
 
         if node_t_params["batch_size"] != 1:
-            # Only batch size 1 is supported on neutron.
+            # Only TransposeConv2d with batch size = 1 is supported on neutron.
             return False
 
-        if conv_params.transposed:
-            # TransposeConv2d with groups > 1 is not supported
-            # TODO: split into multiple convs with groups = 1
-            if conv_params.groups > 1:
-                return False
-            if not node_is_effectively_static_tensor(weights, parameters_mapping):
-                # Only supported if the weights are static, because TFLite `TransposeConv` uses permuted
-                #  weights. In case the weights are dynamic, a Transpose operator would have to be added, which
-                #  is not supported on Neutron.
-                return False
-            # neutron-library/src/utils/NeutronLibraryInterrogation.cpp#876 TransposeConv2DKernelKind
-            if (
-                conv_params.dilation != [1, 1]
-                or conv_params.padding[0] != 0
-                or conv_params.padding[1] >= node_t_params["kernel_width"]
-                or (
-                    conv_params.padding[1] != 0 and node_t_params["inp_height"] != 1
-                )  # Slice added by explicit padding
-                or conv_params.stride[0] != 1
-                or (
-                    (
-                        conv_params.stride[1] != node_t_params["kernel_width"] / 2
-                        or node_t_params["out_height"] != 1
-                    )
-                    and conv_params.stride[1] != node_t_params["kernel_width"]
-                )
-                or conv_params.stride[1] % 2 != 0
-                or node_t_params["inp_channels"] % num_macs != 0
-                or node_t_params["out_channels"] % num_macs != 0
-                or node_t_params["kernel_width"] % 2 != 0
-                or node_t_params["kernel_height"] != 1
-            ):
-                return False
-        elif conv_params.groups == 1:  # Regular convolution.
-            pass
-        elif conv_utils.group_conv_convertible_as_depthwise(
-            node, conv_params.groups
-        ):  # Depthwise convolution.
-            # Only supported if the weights are static, because TFLite `DepthwiseConv2D` uses permuted
+        # TransposeConv2d with groups > 1 is not supported
+        # TODO: split into multiple convs with groups = 1
+        if groups > 1:
+            return False
+        if not node_is_effectively_static_tensor(w_node, parameters_mapping):
+            # Only supported if the weights are static, because TFLite `TransposeConv` uses permuted
             #  weights. In case the weights are dynamic, a Transpose operator would have to be added, which
             #  is not supported on Neutron.
-            if not node_is_effectively_static_tensor(weights, parameters_mapping):
-                return False
-        elif conv_utils.group_conv_convertible_into_multiple_convolutions(
-            node, conv_params.groups
-        ):  # Separable conv.
-            # Requires addition of `Split` and `Concatenation` operators, which are not supported on Neutron.
             return False
-        else:  # Unexpected case (should never happen).
+        # neutron-library/src/utils/NeutronLibraryInterrogation.cpp#876 TransposeConv2DKernelKind
+        if (
+            dilation != [1, 1]
+            or padding[0] != 0
+            or padding[1] >= node_t_params["kernel_width"]
+            or (
+                padding[1] != 0 and node_t_params["inp_height"] != 1
+            )  # Slice added by explicit padding
+            or stride[0] != 1
+            or (
+                (
+                    stride[1] != node_t_params["kernel_width"] / 2
+                    or node_t_params["out_height"] != 1
+                )
+                and stride[1] != node_t_params["kernel_width"]
+            )
+            or stride[1] % 2 != 0
+            or node_t_params["inp_channels"] % num_macs != 0
+            or node_t_params["out_channels"] % num_macs != 0
+            or node_t_params["kernel_width"] % 2 != 0
+            or node_t_params["kernel_height"] != 1
+        ):
             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:
+        is_transposed = (ConvolutionConverter._get_convolution_arguments(node))[6]
+
+        if is_transposed:
+            return ConvolutionConverter._is_supported_on_target_transp_conv(
+                node, neutron_target_spec, parameters_mapping
+            )
+
+        else:
+            return ConvolutionConverter._is_supported_on_target_regular_conv(
+                node, parameters_mapping
+            )
+
     @staticmethod
     def _is_supported_in_IR(
         node: Node,
@@ -149,10 +242,6 @@ def _is_supported_in_IR(
 
         return True
 
-    Stride = Padding = Dilation = OutPadding = list[int]
-    Transposed = bool
-    Groups = int
-
     def _compute_slicing_params(
         self, output_shape, explicit_padding
     ) -> tuple[list[int], list[int]]:
@@ -170,14 +259,14 @@ def _compute_slicing_params(
     @staticmethod
     def _get_convolution_arguments(
         conv_node: Node,
-    ) -> (Stride, Padding, Dilation, Transposed, OutPadding, Groups):
-        # The arguments of the conv are:
-        # [x, w, b, stride, padding, dilation, transposed, output padding, groups]
-        # https://github.com/pytorch/pytorch/blob/v2.6.0/aten/src/ATen/native/Convolution.cpp#L286-L291
-        _, _, _, stride, padding, dilation, transposed, out_padding, groups = (
+    ) -> ConvolutionArgs:
+        x, w, b, stride, padding, dilation, transposed, out_padding, groups = (
             conv_node.args
         )
         return (
+            x,
+            w,
+            b,
             list(stride),
             list(padding),
             list(dilation),
@@ -380,16 +469,8 @@ def _convert_2d_conv(
 
             elif conv_utils.group_conv_convertible_into_multiple_convolutions(
                 t_op, conv_params.groups
-            ):  # Convert to separated `Conv2D`.
-                t_op.builtin_options = conv_2d_options.Conv2D()
-
-                return conv_utils.create_separated_convolutions_based_on_group(
-                    t_op,
-                    conv_params,
-                    self.builder,
-                    self._convert_unpadded_2D,
-                    conv_utils.conv_op_factory,
-                )
+            ):
+                raise RuntimeError("NXP backend: Group convolution was not decomposed.")
 
             else:
                 # Convert to regular `Conv2D`.
@@ -419,7 +500,7 @@ def _convert_2d_conv(
     def convert(self, node: Node):
         self.assert_convertible(node)
 
-        stride, padding, dilation, transposed, out_padding, groups = (
+        _, _, _, stride, padding, dilation, transposed, out_padding, groups = (
             self._get_convolution_arguments(node)
         )
 

backends/nxp/tests/generic_tests/test_batch_norm_fusion.py

@@ -112,7 +112,7 @@ def test_batch_norm_conv_fusing__full_pipeline__1d(bias: bool):
         module, tuple(input_shape)
     ).exported_program()
 
-    assert len(edge_program.graph.nodes) == 21
+    assert len(edge_program.graph.nodes) == 7
     assert not graph_contains_any_of_ops(edge_program.graph, batch_norm_target_ops)
 
 

backends/nxp/tests/generic_tests/test_neutron_converter_manager.py

@@ -6,38 +6,11 @@
 import multiprocessing
 import pickle
 
-import torch
-from executorch import exir
-from executorch.backends.nxp.backend.edge_program_converter import (
-    EdgeProgramToIRConverter,
-)
 from executorch.backends.nxp.backend.neutron_converter_manager import (
     NeutronConverterManager,
 )
-from executorch.backends.nxp.backend.node_format_inference import NodeFormatInference
 from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program
-from executorch.backends.nxp.tests.models import Conv2dModule, LinearModule
-
-
-def test_conv2d_neutron_conversion():
-    model = Conv2dModule()
-
-    example_input = (torch.ones(1, 4, 32, 32),)
-    exir_program = torch.export.export(model, example_input)
-    edge_program_manager = exir.to_edge(exir_program)
-
-    NodeFormatInference(edge_program_manager.exported_program()).identify_node_formats()
-    edge_program_converter = EdgeProgramToIRConverter()
-    tflite_model, *_ = edge_program_converter.convert_program(
-        edge_program_manager.exported_program()
-    )
-
-    neutron_converter_manager = NeutronConverterManager()
-    neutron_model = neutron_converter_manager.convert(tflite_model, "imxrt700", False)
-
-    assert len(
-        neutron_model
-    ), "Produced NeutronGraph-based TFLite model has zero length!"
+from executorch.backends.nxp.tests.models import LinearModule
 
 
 def test_conv2d_neutron_conversion__prefetching(mocker):