NXP backend: added support for same/valid padding for conv2d (pytorch#20957)

### Summary

Added support for `same`/`valid/` padding for `conv2d`.

### Test plan

tests can be manually run using `pytest -c /dev/null
backends/nxp/tests/`

cc @robert-kalmar @JakeStevens @digantdesai @rascani @MartinPavella

Author: novak-vaclav

backends/nxp/backend/edge_program_converter.py

@@ -41,6 +41,7 @@
     exir_ops.edge.dim_order_ops._clone_dim_order.default: CloneConverter,  # noqa F405
     exir_ops.edge.aten.constant_pad_nd.default: ConstantPadNDConverter,  # noqa F405
     exir_ops.edge.aten.convolution.default: ConvolutionConverter,  # noqa F405
+    exir_ops.edge.aten.convolution.out: ConvolutionConverter,  # noqa F405
     exir_ops.edge.aten.exp.default: ExpConverter,  # noqa F405
     exir_ops.edge.aten.hardtanh.default: HardTanhConverter,  # noqa F405
     exir_ops.edge.aten.leaky_relu.default: LeakyReluConverter,  # noqa F405

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

@@ -241,17 +241,15 @@ def _is_supported_in_IR(
     ) -> bool:
         input_tensor_rank = len(node.meta["val"].shape)
         dimensions = input_tensor_rank - 2
-        is_transposed = node.args[6]
-        output_padding = node.args[7]
-        groups = node.args[8]
+        conv_params = ConvolutionConverter._get_conv_params(node)
 
-        if is_transposed and conv_utils.group_conv_convertible_as_depthwise(
-            node, groups
+        if conv_params.transposed and conv_utils.group_conv_convertible_as_depthwise(
+            node, conv_params.groups
         ):
             # TFLite does not support transposed depthwise convolution
             return False
 
-        if not is_transposed and output_padding != [0] * dimensions:
+        if not conv_params.transposed and conv_params.out_padding != [0] * dimensions:
             return False
 
         if input_tensor_safe(node, 2) is None:
@@ -284,14 +282,21 @@ def _compute_slicing_params(
     def _get_conv_params(
         conv_node: Node,
     ) -> ConvParameters:
+        def _normalize_ls_arg(ls):
+            # sometimes, `conv2d` args can be a list of one element. In such case, convert it to 2d arg
+            # example: padding = [0] => [0, 0]
+            return [ls[0], ls[0]] if len(ls) == 1 else ls
+
         x, w, b, stride, padding, dilation, transposed, out_padding, groups = (
             conv_node.args
         )
 
-        stride = list(stride)
-        padding = list(padding)
-        dilation = list(dilation)
-        out_padding = None if out_padding is None else list(out_padding)
+        stride = _normalize_ls_arg(list(stride))
+        padding = _normalize_ls_arg(list(padding))
+        dilation = _normalize_ls_arg(list(dilation))
+        out_padding = (
+            None if out_padding is None else _normalize_ls_arg(list(out_padding))
+        )
 
         return ConvParameters(
             x, w, b, stride, padding, dilation, transposed, out_padding, groups

backends/nxp/neutron_partitioner.py

@@ -214,6 +214,7 @@ def tag_qdq_clusters(self, nodes: list[torch.fx.Node]):
     exir_ops.edge.dim_order_ops._clone_dim_order.default: CloneConverter,  # noqa F405
     exir_ops.edge.aten.constant_pad_nd.default: ConstantPadNDConverter,  # noqa F405
     exir_ops.edge.aten.convolution.default: ConvolutionConverter,  # noqa F405
+    exir_ops.edge.aten.convolution.out: ConvolutionConverter,  # noqa F405
     exir_ops.edge.aten.exp.default: ExpConverter,  # noqa F405
     exir_ops.edge.aten.hardtanh.default: HardTanhConverter,  # noqa F405
     exir_ops.edge.aten.leaky_relu.default: LeakyReluConverter,  # noqa F405

backends/nxp/quantizer/neutron_quantizer.py

@@ -25,6 +25,7 @@
     CatPattern,
     ClampPattern,
     Conv2dPattern,
+    Conv2dPatternPadding,
     ConvTranspose2dPattern,
     DropoutPattern,
     ExpPattern,
@@ -275,6 +276,7 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec, is_qat: bool = False)
                 OpQuantizer(CatPattern(is_qat=is_qat), static_qconfig),
                 OpQuantizer(ClampPattern(self, is_qat=is_qat), static_qconfig),
                 OpQuantizer(Conv2dPattern(self, is_qat=is_qat), static_qconfig),
+                OpQuantizer(Conv2dPatternPadding(self, is_qat=is_qat), static_qconfig),
                 OpQuantizer(
                     ConvTranspose2dPattern(self, is_qat=is_qat), static_qconfig
                 ),

backends/nxp/quantizer/patterns.py

@@ -537,65 +537,6 @@ def _is_batch_norm(node_: Node) -> bool:
 
 
 class ConvPattern(QuantizationPattern):
-    @abstractmethod
-    def partition_types(self) -> list[OpOverload]:
-        pass
-
-    def get_anchors(
-        self, gm: fx.GraphModule, fused_partition: list[fx.GraphModule]
-    ) -> PartitionAnchors:
-        conv_node = fused_partition[0].nodes[-1]
-
-        bias_quantization_qspec = DerivedQuantizationSpec(
-            derived_from=[
-                (conv_node.args[0], conv_node),
-                (conv_node.args[1], conv_node),
-            ],
-            derive_qparams_fn=get_bias_qparams,
-            dtype=torch.int32,
-            quant_min=-(2**31) + 1,
-            quant_max=2**31 - 1,
-            qscheme=torch.per_channel_symmetric,
-            ch_axis=0,
-        )
-
-        weight_observer_or_fake_quant_ctr = (
-            FakeQuantize.with_args(observer=MovingAveragePerChannelMinMaxObserver)
-            if self.is_qat
-            else PerChannelMinMaxObserver
-        )
-        weight_quantization_spec = QuantizationSpec(
-            dtype=torch.int8,
-            observer_or_fake_quant_ctr=weight_observer_or_fake_quant_ctr,
-            quant_min=-127,
-            quant_max=127,
-            qscheme=torch.per_channel_symmetric,
-            ch_axis=0,
-        )
-
-        # Keep bias empty if not supplied
-        bias = []
-        if len(conv_node.args) > 2 and conv_node.args[2] is not None:
-            bias = [(conv_node, NodeArgsIdx(2), bias_quantization_qspec)]
-
-        output_specs = [(conv_node,)]
-        # In order for QAT to be numerically correct, there should be no quantization between
-        # convolution node and batch norm node.
-        if self.is_qat:
-            conv_users = conv_node.users
-            possibly_bn = list(conv_users.keys())[0] if len(conv_users) == 1 else None
-            if possibly_bn and _is_batch_norm(possibly_bn):
-                output_specs = []
-
-        return PartitionAnchors(
-            inputs=[(conv_node, NodeArgsIdx(0))],
-            weights=[(conv_node, NodeArgsIdx(1), weight_quantization_spec)],
-            biases=bias,
-            output=output_specs,
-        )
-
-
-class Conv2dPattern(ConvPattern):
     def __init__(self, neutron_quantizer, is_qat: bool = False):
         super().__init__(is_qat=is_qat)
 
@@ -604,8 +545,9 @@ def __init__(self, neutron_quantizer, is_qat: bool = False):
             self.neutron_quantizer.neutron_target_spec.neutron_target_info
         )
 
+    @abstractmethod
     def partition_types(self) -> list[OpOverload]:
-        return [torch.ops.aten.conv2d.default]
+        pass
 
     def get_anchors(
         self, gm: fx.GraphModule, fused_partition: list[fx.GraphModule]
@@ -689,6 +631,16 @@ def get_anchors(
         )
 
 
+class Conv2dPattern(ConvPattern):
+    def partition_types(self) -> list[OpOverload]:
+        return [torch.ops.aten.conv2d.default]
+
+
+class Conv2dPatternPadding(ConvPattern):
+    def partition_types(self) -> list[OpOverload]:
+        return [torch.ops.aten.conv2d.padding]
+
+
 class ConvTranspose2dPattern(QuantizationPattern):
     def __init__(self, neutron_quantizer, is_qat: bool = False):
         super().__init__(is_qat=is_qat)

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

@@ -35,7 +35,7 @@ def assert_delegated_and_correct(
     input_shape,
     mocker,
     request,
-    use_qat,
+    use_qat=False,
     exp_delegated_ops=None,
     exp_non_delegated_ops=None,
     et_ref_model=ReferenceModel.QUANTIZED_EXECUTORCH_CPP,
@@ -558,19 +558,6 @@ def test__tr_no_deleg(
 
 
 class TestConv:
-    @staticmethod
-    def _conv_id(ins, oc, ks=3, s=2, d=1, p=0, b=True, g=1):
-        return (
-            f"ic={ins}, "
-            f"oc={oc}, "
-            f"ks={ks}, "
-            f"s={s}, "
-            f"d={d}, "
-            f"p={p}, "
-            f"b={b}, "
-            f"g={g}"
-        )
-
     @pytest.mark.parametrize(
         "input_shape, out_channels",
         [
@@ -1105,6 +1092,104 @@ def test__d_fwd_misc_arg(
 
         assert_delegated_and_correct(model, input_shape, mocker, request, use_qat)
 
+    @pytest.mark.parametrize(
+        "input_shape, kernel_size, stride, dilation, padding, bias",
+        [
+            pytest.param(
+                ins := (2, 5, 13, 11),
+                ks := 3,
+                s := 2,
+                d := 2,
+                p := "valid",
+                b := True,
+                id=f"string padding arg: {_conv_id(ins, ins[1], ks=ks, s=s, d=d, p=p, b=b)}",
+            ),
+            pytest.param(
+                ins := (2, 5, 13, 11),
+                ks := 3,
+                s := 1,
+                d := 2,
+                p := "same",
+                b := True,
+                id=f"string padding arg: {_conv_id(ins, ins[1], ks=ks, s=s, d=d, p=p, b=b)}",
+            ),
+        ],
+    )
+    def test__fwd_str_arg(
+        self,
+        input_shape,
+        kernel_size,
+        stride,
+        dilation,
+        padding,
+        bias,
+        request,
+        mocker,
+    ):
+        out_channels = input_shape[1]
+
+        model = Conv2dModule(
+            in_channels=input_shape[1],
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            padding=padding,
+            bias=bias,
+        )
+
+        assert_delegated_and_correct(model, input_shape, mocker, request)
+
+    @pytest.mark.parametrize(
+        "input_shape, kernel_size, stride, dilation, padding, bias",
+        [
+            pytest.param(
+                ins := (2, 5, 13, 11),
+                ks := 3,
+                s := 2,
+                d := 2,
+                p := "valid",
+                b := True,
+                id=f"string padding arg, depthwise: {_conv_id(ins, ins[1], ks=ks, s=s, d=d, p=p, b=b, g=ins[1])}",
+            ),
+            pytest.param(
+                ins := (2, 5, 13, 11),
+                ks := 3,
+                s := 1,
+                d := 2,
+                p := "same",
+                b := True,
+                id=f"string padding arg, depthwise: {_conv_id(ins, ins[1], ks=ks, s=s, d=d, p=p, b=b, g=ins[1])}",
+            ),
+        ],
+    )
+    def test__d_fwd_str_arg(
+        self,
+        input_shape,
+        kernel_size,
+        stride,
+        dilation,
+        padding,
+        bias,
+        request,
+        mocker,
+    ):
+        out_channels = input_shape[1]
+        group = input_shape[1]
+
+        model = Conv2dModule(
+            in_channels=input_shape[1],
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            padding=padding,
+            bias=bias,
+            group=group,
+        )
+
+        assert_delegated_and_correct(model, input_shape, mocker, request)
+
     @pytest.mark.parametrize(
         "input_shape, out_channels, kernel_size, stride, dilation, padding",
         [