Cadence ops: Support quantized_w8a32_linear (pytorch#15171)

Summary:
Pull Request resolved: pytorch#15171

As titled.

Differential Revision: D84745967

Author: Andrew Grebenisan

backends/cadence/aot/ops_registrations.py

@@ -53,7 +53,6 @@ def _validate_ref_impl_exists() -> None:
     # 1. be removed
     # 2. have a reference implementation added to ref_implementations.py
     _WARN_ONLY = {
-        "cadence::quantized_w8a32_linear",
         "cadence::quantized_add",  # We should only support per_tensor variant, should remove
         "cadence::_softmax_f32_f32",
         "cadence::requantize",  # We should only support per_tensor variant, should remove
@@ -2702,10 +2701,13 @@ def quantized_w8a32_linear_meta(
     b_scale: float,
 ) -> torch.Tensor:
     # src comes in shape [leading_dims, in_dim]
-    # weight comes in shape [in_dim, out_dim]
+    # weight comes in shape [out_dim, in_dim]
     # output comes in empty with shape [leading_dims, out_dim]
     src_shape = list(src.shape)
     weight_shape = weight.shape
+    assert (src_shape[-1] % 4) == 0
+    if len(src_shape) >= 2:
+        assert src_shape[-2] == 1
     assert len(weight_shape) == 2
     assert src_shape[-1] == weight_shape[-1]
     src_shape[-1] = weight_shape[0]

backends/cadence/aot/ref_implementations.py

@@ -866,6 +866,9 @@ def quantized_w8a32_conv(
     # 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
+    assert weight.dtype == torch.int8
+    assert bias.dtype == torch.int8
+
     dequant_weight = weight.float() * w_scale
 
     # Dequantize bias using scale
@@ -884,6 +887,33 @@ def quantized_w8a32_conv(
     return output
 
 
+@impl_tracked(m, "quantized_w8a32_linear")
+def quantized_w8a32_linear(
+    src: torch.Tensor,
+    weight: torch.Tensor,
+    w_scale: float,
+    bias: torch.Tensor,
+    b_scale: float,
+) -> torch.Tensor:
+    # src comes in shape [leading_dims, in_dim]
+    # weight comes in shape [out_dim, in_dim]
+    # output comes in empty with shape [leading_dims, out_dim]
+    assert weight.dtype == torch.int8
+    assert bias.dtype == torch.int8
+    if len(src.shape) >= 2:
+        assert src.shape[-2] == 1
+    dequant_weight = weight.float() * w_scale
+    dequant_bias = bias.float() * b_scale
+
+    output = torch.nn.functional.linear(
+        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

@@ -1188,7 +1188,7 @@ def test_quantized_conv_per_tensor(
                     dtype=torch.int8,
                 ),  # weight: 4x4x3
                 0.5,  # w_scale
-                torch.tensor([2, 2, 2, 2], dtype=torch.float32),  # bias: 4
+                torch.tensor([2, 2, 2, 2], dtype=torch.int8),  # bias: 4
                 1.0,  # b_scale
                 torch.tensor(
                     [
@@ -1236,6 +1236,82 @@ def test_quantized_w8a32_conv(
             f"Output values don't match expected in {name}. Got {output}, expected {expected_output}",
         )
 
+    @expand(
+        [
+            (
+                "multi_input_features",
+                torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32),  # src: 1x3
+                torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.int8),  # weight: 2x3
+                0.5,  # w_scale
+                torch.tensor([0, 1], dtype=torch.int8),  # bias: 2
+                1.0,  # b_scale
+                torch.tensor([[3.5, 5.0]], dtype=torch.float32),  # expected
+            ),
+            (
+                "batch_size_2",
+                torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32),  # src: 2x2
+                torch.tensor([[1, 1], [2, -1]], dtype=torch.int8),  # weight: 2x2
+                1.0,  # w_scale
+                torch.tensor([0, 0], dtype=torch.int8),  # bias: 2
+                1.0,  # b_scale
+                torch.tensor([[3.0, 0.0], [7.0, 2.0]], dtype=torch.float32),  # expected
+            ),
+            (
+                "3d_input",
+                torch.tensor(
+                    [[[1.0, 2.0], [3.0, 4.0]]], dtype=torch.float32
+                ),  # src: 1x2x2
+                torch.tensor([[1, 1], [2, -1]], dtype=torch.int8),  # weight: 2x2
+                1.0,  # w_scale
+                torch.tensor([0, 1], dtype=torch.int8),  # bias: 2
+                1.0,  # b_scale
+                torch.tensor(
+                    [[[3.0, 1.0], [7.0, 3.0]]], dtype=torch.float32
+                ),  # expected
+            ),
+            (
+                "negative_weights",
+                torch.tensor([[2.0, 4.0]], dtype=torch.float32),  # src: 1x2
+                torch.tensor([[-2, -1], [-3, -2]], dtype=torch.int8),  # weight: 2x2
+                0.5,  # w_scale
+                torch.tensor([2, 1], dtype=torch.int8),  # bias: 2
+                1.0,  # b_scale
+                torch.tensor([[-2.0, -6.0]], dtype=torch.float32),  # expected
+            ),
+        ]
+    )
+    def test_quantized_w8a32_linear(
+        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_linear(
+            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