diff --git a/backends/nxp/backend/edge_program_converter.py b/backends/nxp/backend/edge_program_converter.py index dfdc054fe41..29cbf43cc12 100644 --- a/backends/nxp/backend/edge_program_converter.py +++ b/backends/nxp/backend/edge_program_converter.py @@ -40,6 +40,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 diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/convolution_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/convolution_converter.py index a55d0316f90..581556ca7ce 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/convolution_converter.py +++ b/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 diff --git a/backends/nxp/neutron_partitioner.py b/backends/nxp/neutron_partitioner.py index 46bb8e6d0a7..b3f2d4ad5a3 100644 --- a/backends/nxp/neutron_partitioner.py +++ b/backends/nxp/neutron_partitioner.py @@ -213,6 +213,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 diff --git a/backends/nxp/quantizer/neutron_quantizer.py b/backends/nxp/quantizer/neutron_quantizer.py index 97204b609f4..744552aec1f 100644 --- a/backends/nxp/quantizer/neutron_quantizer.py +++ b/backends/nxp/quantizer/neutron_quantizer.py @@ -24,6 +24,7 @@ CatPattern, ClampPattern, Conv2dPattern, + Conv2dPatternPadding, ConvTranspose2dPattern, DropoutPattern, ExpPattern, @@ -271,6 +272,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 ), diff --git a/backends/nxp/quantizer/patterns.py b/backends/nxp/quantizer/patterns.py index c4ef35b4a5b..d264d82e0fa 100644 --- a/backends/nxp/quantizer/patterns.py +++ b/backends/nxp/quantizer/patterns.py @@ -528,65 +528,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) @@ -595,8 +536,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] @@ -680,6 +622,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) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py index 5dc3b33e6dc..d9ffc6b2128 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py +++ b/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", [