Add conv2d support for IntxUnpackedToInt8Tensor (#3371)

**Summary:** This enables `Int8DynamicActivationIntxWeightConfig`
and `IntxWeightOnlyConfig` for conv2d.

**Test Plan:**
```
python test/quantization/quantize_/workflows/intx/test_intx_unpacked_to_int8_tensor.py -k test_conv2d
```

Author: andrewor14

test/quantization/quantize_/workflows/intx/test_intx_unpacked_to_int8_tensor.py

@@ -61,6 +61,18 @@ def test_linear(self):
         error = compute_error(original, quantized)
         self.assertTrue(error > 20)
 
+    def test_conv2d(self):
+        dtype = torch.bfloat16
+        device = "cpu"
+        input = torch.randn(1, 128, 224, 224, dtype=dtype, device=device)
+        conv = torch.nn.Conv2d(128, 64, 3, dtype=dtype, device=device)
+        original = conv(input)
+        is_conv = lambda n, _: isinstance(n, torch.nn.Conv2d)
+        quantize_(conv, self.config, filter_fn=is_conv)
+        quantized = conv(input)
+        error = compute_error(original, quantized)
+        self.assertGreater(error, 15)
+
     def test_hqq_intx_weight_only_config(self):
         dtype = torch.bfloat16
         device = "cpu"

torchao/quantization/quant_api.py

@@ -2297,20 +2297,32 @@ def _intx_weight_only_quantize_tensor(
     intx_packing_format = config.intx_packing_format
     intx_choose_qparams_algorithm = config.intx_choose_qparams_algorithm
 
-    assert weight.dim() == 2, (
-        f"IntxWeightOnlyConfig only works for 2-d Tensor, got: {weight.dim()}"
-    )
+    if weight.dim() == 2:
+        input_dim = -1
+    elif weight.dim() == 4:
+        # conv2d: N, C_in, H, W
+        input_dim = 1
+    else:
+        raise ValueError(
+            f"IntxWeightOnlyConfig only works for 2-d and 4-d Tensors, got: {weight.dim()}"
+        )
+
     if isinstance(granularity, PerGroup):
         group_size = granularity.group_size
     elif isinstance(granularity, PerAxis):
         assert granularity.axis == 0, (
             f"axis must be 0 with PerAxis, but got {granularity.axis}"
         )
-        group_size = weight.shape[-1]
+        group_size = weight.shape[input_dim]
     else:
         raise ValueError(f"granularity must be PerGroup or PerAxis, got {granularity}")
 
-    block_size = (1, group_size)
+    if weight.dim() == 2:
+        block_size = (1, group_size)
+    else:
+        # conv2d: N, C_in, H, W
+        assert weight.dim() == 4
+        block_size = (1, group_size, 1, 1)
 
     if config.version == 2:
         if config.intx_packing_format == IntxPackingFormat.UNPACKED_TO_INT8:

torchao/quantization/quantize_/workflows/intx/intx_unpacked_to_int8_tensor.py

@@ -343,6 +343,36 @@ def _(func, types, args, kwargs):
     return torch.nn.functional.linear(input_tensor, weight_tensor, bias)
 
 
+@implements(aten.conv2d.default)
+@implements_torch_function(torch.nn.functional.conv2d)
+def _(func, types, args, kwargs):
+    (
+        input_tensor,
+        weight_tensor,
+        bias,
+        stride,
+        padding,
+        dilation,
+        groups,
+    ) = fill_defaults(args, 7, [None, [1, 1], [0, 0], [1, 1], 1])
+    assert isinstance(weight_tensor, IntxUnpackedToInt8Tensor)
+
+    # Apply dynamic activation quant
+    if weight_tensor.activation_quantization is not None:
+        if (
+            weight_tensor.activation_quantization
+            == IntxUnpackedToInt8TensorActivationQuantization.INT8_ASYM_PER_TOKEN
+        ):
+            input_tensor = _apply_int8_act_asym_per_token_quant_dequant(input_tensor)
+        else:
+            raise NotImplementedError(
+                f"Unsupported activation quantization: {weight_tensor.activation_quantization}"
+            )
+
+    weight_tensor = weight_tensor.dequantize()
+    return func(input_tensor, weight_tensor, bias, stride, padding, dilation, groups)
+
+
 @implements(aten.embedding.default)
 @implements_torch_function(torch.nn.functional.embedding)
 def _(func, types, args, kwargs):