Cadence: Support quantized_w8a32_conv (pytorch#15137)

Summary:

As titled

Reviewed By: skrtskrtfb

Differential Revision: D84658444

Author: Andrew Grebenisan

backends/cadence/aot/ops_registrations.py

@@ -67,7 +67,6 @@ def _validate_ref_impl_exists() -> None:
         "cadence::dequantize_per_tensor_asym16u",
         "cadence::linalg_vector_norm",
         "cadence::quantized_conv2d_nchw",  # We should only support per_tensor variant, should remove
-        "cadence::quantized_w8a32_conv",
         "cadence::quantize_per_tensor_asym32s",
         "cadence::quantized_relu",  # We should only support per_tensor variant, should remove
         "cadence::linalg_svd",

backends/cadence/aot/ref_implementations.py

@@ -703,6 +703,36 @@ def quantized_conv2d_nchw_per_tensor(
     )
 
 
+@impl_tracked(m, "quantized_w8a32_conv")
+def quantized_w8a32_conv(
+    src: torch.Tensor,
+    weight: torch.Tensor,
+    w_scale: float,
+    bias: torch.Tensor,
+    b_scale: float,
+) -> torch.Tensor:
+    # src comes in shape [batch, in_channel, in_length]
+    # weight comes in shape [out_ch, in_ch, kernel_dim]
+    # output comes in empty with shape [batch, out_ch, in_length - kernel_dim + 1]
+    # Dequantize weight using scale
+    dequant_weight = weight.float() * w_scale  # Assuming zero_point is 0 for weights
+
+    # Dequantize bias using scale
+    dequant_bias = bias.float() * b_scale
+
+    # Perform 1D convolution
+    # src: [batch, in_channel, in_length]
+    # weight: [out_ch, in_ch, kernel_dim]
+    # bias: [out_ch]
+    output = torch.nn.functional.conv1d(
+        src.float(),
+        dequant_weight,
+        dequant_bias,
+    )
+
+    return output
+
+
 @impl_tracked(m, "quantized_conv2d_nhwc.per_tensor")
 def quantized_conv2d_nhwc_per_tensor(
     input_tensor: torch.Tensor,

backends/cadence/aot/tests/test_ref_implementations.py

@@ -1040,6 +1040,116 @@ def test_quantized_conv_per_tensor(
                 f"Output values don't match expected. Got {output}, expected {expected_output}",
             )
 
+    @expand(
+        [
+            # Test case 1: Basic 1D convolution with int8 weights
+            (
+                "basic_int8_weights",
+                torch.tensor(
+                    [[[1.0, 2.0, 3.0, 4.0, 5.0]]], dtype=torch.float32
+                ),  # src: 1x1x5
+                torch.tensor([[[1, -1, 2]]], dtype=torch.int8),  # weight: 1x1x3
+                0.1,  # w_scale
+                torch.tensor([1], dtype=torch.int8),  # bias: 1
+                0.2,  # b_scale
+                torch.tensor(
+                    [[[0.7, 0.9, 1.1]]], dtype=torch.float32
+                ),  # expected: conv1d result
+            ),
+            # Test case 2: Multiple input channels
+            (
+                "multi_input_channels",
+                torch.tensor(
+                    [[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]], dtype=torch.float32
+                ),  # src: 1x2x3
+                torch.tensor([[[2, 1], [1, 2]]], dtype=torch.int8),  # weight: 1x2x2
+                0.5,  # w_scale
+                torch.tensor([1], dtype=torch.int8),  # bias: 1
+                1.0,  # b_scale
+                torch.tensor([[[10.0, 13.0]]], dtype=torch.float32),  # expected
+            ),
+            # Test case 3: Multiple output channels
+            (
+                "multi_output_channels",
+                torch.tensor(
+                    [[[1.0, 2.0, 3.0, 4.0]]], dtype=torch.float32
+                ),  # src: 1x1x4
+                torch.tensor([[[1, -1]], [[2, 0]]], dtype=torch.int8),  # weight: 2x1x2
+                0.25,  # w_scale
+                torch.tensor([0, 1], dtype=torch.int8),  # bias: 2
+                0.5,  # b_scale
+                torch.tensor(
+                    [[[-0.25, -0.25, -0.25], [1.0, 1.5, 2.0]]], dtype=torch.float32
+                ),  # expected
+            ),
+            # Test case 4: Batch size > 1
+            (
+                "batch_size_2",
+                torch.tensor(
+                    [[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]], dtype=torch.float32
+                ),  # src: 2x1x3
+                torch.tensor([[[1, 1]]], dtype=torch.int8),  # weight: 1x1x2
+                1.0,  # w_scale
+                torch.tensor([0], dtype=torch.int8),  # bias: 1
+                1.0,  # b_scale
+                torch.tensor(
+                    [[[3.0, 5.0]], [[9.0, 11.0]]], dtype=torch.float32
+                ),  # expected
+            ),
+            # Test case 5: Zero weights and bias
+            (
+                "zero_weights_bias",
+                torch.tensor([[[1.0, 2.0, 3.0]]], dtype=torch.float32),  # src: 1x1x3
+                torch.tensor([[[0, 0]]], dtype=torch.int8),  # weight: 1x1x2
+                0.1,  # w_scale
+                torch.tensor([0], dtype=torch.int8),  # bias: 1
+                1.0,  # b_scale
+                torch.tensor([[[0.0, 0.0]]], dtype=torch.float32),  # expected
+            ),
+            # Test case 6: Negative weights
+            (
+                "negative_weights",
+                torch.tensor([[[2.0, 4.0, 6.0]]], dtype=torch.float32),  # src: 1x1x3
+                torch.tensor([[[-2, -1]]], dtype=torch.int8),  # weight: 1x1x2
+                0.5,  # w_scale
+                torch.tensor([2], dtype=torch.float32),  # bias: 1
+                1.0,  # b_scale
+                torch.tensor([[[-2.0, -5.0]]], dtype=torch.float32),  # expected
+            ),
+        ]
+    )
+    def test_quantized_w8a32_conv(
+        self,
+        name: str,
+        src: torch.Tensor,
+        weight: torch.Tensor,
+        w_scale: float,
+        bias: torch.Tensor,
+        b_scale: float,
+        expected_output: torch.Tensor,
+    ) -> None:
+        output = torch.ops.cadence.quantized_w8a32_conv(
+            src, weight, w_scale, bias, b_scale
+        )
+
+        # Verify output properties
+        self.assertEqual(
+            output.dtype,
+            torch.float32,
+            f"Output dtype should be float32 in {name}",
+        )
+        self.assertEqual(
+            output.shape,
+            expected_output.shape,
+            f"Output shape should match expected shape in {name}",
+        )
+
+        # Verify output matches expected values
+        self.assertTrue(
+            torch.allclose(output, expected_output, rtol=1e-4, atol=1e-4),
+            f"Output values don't match expected in {name}. Got {output}, expected {expected_output}",
+        )
+
     @expand(
         [
             # Test case 1: Basic int8 case with negative scale