Add support for quantized transposed conv

Summary: As title. Updates the tests accordingly.

Differential Revision: D68939306

Author: digantdesai

backends/xnnpack/quantizer/xnnpack_quantizer.py

@@ -249,8 +249,9 @@ class XNNPACKQuantizer(Quantizer):
     STATIC_OPS = [
         "linear_relu",
         "linear",
-        "conv_relu",
         "conv",
+        "conv_transpose",
+        "conv_relu",
         "conv_transpose_relu",
         "adaptive_avg_pool2d",
         # TODO: move this to BoltNNQuantizer?

backends/xnnpack/quantizer/xnnpack_quantizer_utils.py

@@ -285,18 +285,17 @@ def _annotate_linear_relu(
     return annotated_partitions
 
 
-@register_annotator("conv")
-def _annotate_conv(
+def _do_annotate_conv(
     gm: torch.fx.GraphModule,
     quantization_config: Optional[QuantizationConfig],
     filter_fn: Optional[Callable[[Node], bool]] = None,
+    is_conv_transpose: bool = False,
 ) -> Optional[list[list[Node]]]:
     annotated_partitions = []
+    is_conv_node = _is_conv_transpose_node if is_conv_transpose else _is_conv_node
+
     for n in gm.graph.nodes:
-        if n.op != "call_function" or n.target not in [
-            torch.ops.aten.conv1d.default,
-            torch.ops.aten.conv2d.default,
-        ]:
+        if not is_conv_node(n):
             continue
         conv_node = n
 
@@ -392,6 +391,27 @@ def _do_annotate_conv_relu(
         annotated_partitions.append(partition)
     return annotated_partitions
 
+@register_annotator("conv")
+def _annotate_conv(
+    gm: torch.fx.GraphModule,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[Node], bool]] = None,
+) -> Optional[list[list[Node]]]:
+    return _do_annotate_conv(
+        gm, quantization_config, filter_fn, is_conv_transpose=False
+    )
+
+
+@register_annotator("conv_transpose")
+def _annotate_transpose_conv(
+    gm: torch.fx.GraphModule,
+    quantization_config: Optional[QuantizationConfig],
+    filter_fn: Optional[Callable[[Node], bool]] = None,
+) -> Optional[list[list[Node]]]:
+    return _do_annotate_conv(
+        gm, quantization_config, filter_fn, is_conv_transpose=True
+    )
+
 
 @register_annotator("conv_relu")
 def _annotate_conv_relu(

backends/xnnpack/test/ops/test_conv2d.py

@@ -243,15 +243,14 @@ def test_qs8_conv2d_test(self) -> None:
                 self._test(
                     Conv2d(bias=has_bias, transpose=transpose),
                     quant_config=get_symmetric_quantization_config(),
-                    check_quantized=not transpose,  # XNNPackQuantizer does not quantize this pattern yet
                 )
 
     def test_qs8_conv2d_per_channel(self) -> None:
         for transpose in (True, False):
             self._test(
                 Conv2d(transpose=transpose),
                 quant_config=get_symmetric_quantization_config(is_per_channel=True),
-                check_quantized=not transpose,  # XNNPackQuantizer does not quantize this pattern yet
+                delegated=not transpose,  # XNNPACK does not support per input channel quantization for transpose convolutions with groups > 1
             )
 
     def test_fp32_conv2d_seq(self) -> None:
@@ -264,7 +263,6 @@ def test_qs8_conv2d_seq(self) -> None:
                 Conv2dSeq(transpose=transpose),
                 conv_count=2,
                 quant_config=get_symmetric_quantization_config(),
-                check_quantized=not transpose,  # XNNPackQuantizer does not quantize this pattern yet
             )
 
     def test_fp32_conv2d_single_int_params(self):
@@ -282,7 +280,6 @@ def test_fp32_conv2d_depthwise(self):
         # - Groups must equal In Channels
         # - Out Channels must be a positive multiple of In Channels
         for transpose in (True, False):
-
             self._test(
                 Conv2d(groups=2, in_channels=2, out_channels=6, transpose=transpose)
             )
@@ -292,7 +289,6 @@ def test_qs8_conv2d_depthwise(self):
             self._test(
                 Conv2d(groups=2, in_channels=2, out_channels=6, transpose=transpose),
                 quant_config=get_symmetric_quantization_config(),
-                check_quantized=not transpose,  # XNNPackQuantizer does not quantize this pattern yet
             )
 
     def test_fp32_conv2d_bn(self):
@@ -384,7 +380,6 @@ def test_qs8_conv2d_bn(self):
                 Conv2dBatchNorm(transpose=transpose),
                 quant_config=get_symmetric_quantization_config(),
                 conv_count=2,
-                check_quantized=not transpose,  # XNNPackQuantizer does not quantize this pattern yet
             )
 
     def test_qs8_conv2d_relu(self):
@@ -415,7 +410,7 @@ def get_inputs(self):
             self._test(
                 ConvReLU(transpose=transpose),
                 quant_config=get_symmetric_quantization_config(is_per_channel=True),
-                delegated=not transpose,
+                delegated=not transpose,  # XNNPACK does not support per input channel quantization for transpose convolutions with groups > 1
             )
 
     def test_qs8_conv2d_dw_relu(self):
@@ -467,9 +462,8 @@ def get_inputs(self):
                     quant_config=get_symmetric_quantization_config(
                         is_per_channel=per_channel_quant
                     ),
-                    # xnnpack only supports per output channel quantization for transposed convolutions
-                    # XNNPackQuantizer quantizes per input channel currently
-                    delegated=not transpose or not per_channel_quant,
+                    # XNNPACK does not support per input channel quantization for transpose convolutions with groups > 1
+                    delegated=not (transpose and per_channel_quant),
                 )
 
     def test_qs8_conv2d_relu_seq(self):