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,70 +48,166 @@
 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_new_flow(
+        node: Node,
+        parameters_mapping: dict[str, Parameter],
+    ) -> bool:
+        (
+            inp_node,
+            w_node,
+            b_node,
+            stride,
+            padding,
+            dilation,
+            transposed,
+            _,
+            groups,
+        ) = 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 if hasattr(inp_node, "meta") else inp_node.shape
+        )
+        inp_channels = inp_node_shape[1]
+
+        if kernel_h * kernel_w * inp_channels > 65535:
+            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:
+        if custom_delegation_options.use_new_flow_neutron_c:
+            return ConvolutionConverter._is_supported_on_target_new_flow(
+                node, parameters_mapping
+            )
+
         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)
+        _, w_node, _, stride, padding, dilation, transposed, _, groups = (
+            ConvolutionConverter._get_convolution_arguments(node)
         )
 
         if node_t_params["batch_size"] != 1:
             # Only batch size 1 is supported on neutron.
             return False
 
-        if conv_params.transposed:
+        if transposed:
             # TransposeConv2d with groups > 1 is not supported
             # TODO: split into multiple convs with groups = 1
-            if conv_params.groups > 1:
+            if groups > 1:
                 return False
-            if not node_is_effectively_static_tensor(weights, parameters_mapping):
+            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.
                 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"]
+                dilation != [1, 1]
+                or padding[0] != 0
+                or padding[1] >= node_t_params["kernel_width"]
                 or (
-                    conv_params.padding[1] != 0 and node_t_params["inp_height"] != 1
+                    padding[1] != 0 and node_t_params["inp_height"] != 1
                 )  # Slice added by explicit padding
-                or conv_params.stride[0] != 1
+                or stride[0] != 1
                 or (
                     (
-                        conv_params.stride[1] != node_t_params["kernel_width"] / 2
+                        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"]
+                    and stride[1] != node_t_params["kernel_width"]
                 )
-                or conv_params.stride[1] % 2 != 0
+                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
-        elif conv_params.groups == 1:  # Regular convolution.
+        elif groups == 1:  # Regular convolution.
             pass
         elif conv_utils.group_conv_convertible_as_depthwise(
-            node, conv_params.groups
+            node, groups
         ):  # Depthwise convolution.
             # Only supported if the weights are static, because TFLite `DepthwiseConv2D` 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):
+            if not node_is_effectively_static_tensor(w_node, parameters_mapping):
                 return False
         elif conv_utils.group_conv_convertible_into_multiple_convolutions(
-            node, conv_params.groups
+            node, groups
         ):  # Separable conv.
             # Requires addition of `Split` and `Concatenation` operators, which are not supported on Neutron.
             return False
@@ -149,10 +245,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 +262,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 +472,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 +503,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)
         )