Enable 16-bit activations and 8 bit weigths in Cadence Quantizer for Matmul (#15929)

Summary:
Pull Request resolved: #15929

# Context
We continue from D84284794 to add support for 16-bit activations. Note that right now, all though they support 16-bit activations already, it's only if the weights are also 16-bits. To do this, we need to change the way we template some functions.

# Current Behavior
Right now, we're composing two macros together, the `ET_FORALL_JARVIS_QUANTIZED_TYPES_WITH_INT16` macro:

https://www.internalfb.com/code/fbsource/[9e8c6d8466107f58aa3de1b9e4ec71c49d670a8f]/fbcode/on_device_ai/Assistant/Jarvis/min_runtime/operators/generic/operators.h?lines=22-25

and the function macro(`quantized_linear` chosen for example):

https://www.internalfb.com/code/fbsource/[9e8c6d8466107f58aa3de1b9e4ec71c49d670a8f]/fbcode/on_device_ai/Assistant/Jarvis/min_runtime/operators/generic/quantized_linear_out.cpp?lines=30-41

so together, it just becomes a switch statement, calling the `quantized_linear` function with the correct template parameter.

However, note that it assumes that both the input activations and weights are the same dtype, which is not the case.

# This Diff
We fix this checking for our datatypes, and calling the functions with the correct data types, as in D86538176.

Differential Revision: D86644079

Author: RahulC7

backends/cadence/aot/quantizer/quantizer.py

@@ -372,3 +372,30 @@ def __init__(self, quantizers: Optional[list[Quantizer]] = None) -> None:
         # Add 16-bit quantizers for LinearPattern
         quantizers.append(CadenceAtenQuantizer(LinearPattern(), qconfig_A16))
         super().__init__(quantizers)
+
+
+class CadenceWith16BitConvActivationsQuantizer(CadenceQuantizer):
+    """
+    Quantizer including A16 conv
+    """
+
+    def __init__(self, quantizers: Optional[list[Quantizer]] = None) -> None:
+        if quantizers is None:
+            quantizers = []
+        # Add 16-bit quantizers for Conv patterns
+        quantizers.append(CadenceAtenQuantizer(Conv1dPattern(), qconfig_A16))
+        quantizers.append(CadenceAtenQuantizer(Conv2dPattern(), qconfig_A16))
+        super().__init__(quantizers)
+
+
+class CadenceWith16BitMatmulActivationsQuantizer(CadenceQuantizer):
+    """
+    Quantizer including A16 matmul
+    """
+
+    def __init__(self, quantizers: Optional[list[Quantizer]] = None) -> None:
+        if quantizers is None:
+            quantizers = []
+        # Add 16-bit quantizers for MatmulPattern
+        quantizers.append(CadenceAtenQuantizer(MatmulPattern(), qconfig_A16))
+        super().__init__(quantizers)

backends/cadence/hifi/operators/op_quantized_matmul_out.cpp

@@ -8,6 +8,7 @@
 
 #include <executorch/backends/cadence/hifi/kernels/kernels.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
+#include <on_device_ai/Assistant/Jarvis/min_runtime/operators/generic/operators.h>
 #include <stdlib.h>
 
 using executorch::aten::ScalarType;
@@ -192,8 +193,20 @@ void quantized_matmul_out(
   size_t leading_dim = X.size(X.dim() - 2);
   size_t out_dim = Y.size(Y.dim() - 1 - transposed);
   size_t in_dim = X.size(X.dim() - 1);
-
-  if (out.scalar_type() == exec_aten::ScalarType::Byte) {
+  if (out.scalar_type() == exec_aten::ScalarType::Short) {
+    ::impl::generic::native::quantized_matmul_out(
+        ctx,
+        X,
+        X_zero_point,
+        Y,
+        Y_zero_point,
+        bias,
+        out_multiplier,
+        out_shift,
+        out_zero_point,
+        transposed,
+        out);
+  } else if (out.scalar_type() == exec_aten::ScalarType::Byte) {
     _typed_quantized_matmul<uint8_t>(
         ctx,
         X,

backends/cadence/hifi/operators/operators.h

@@ -83,6 +83,19 @@ void quantized_linear_per_tensor_out(
     const ::executorch::aten::optional<::executorch::aten::Tensor>& offset,
     ::executorch::aten::Tensor& out);
 
+void quantized_matmul_out(
+    ::executorch::runtime::KernelRuntimeContext& ctx,
+    const ::executorch::aten::Tensor& X,
+    int64_t X_zero_point,
+    const ::executorch::aten::Tensor& Y,
+    int64_t Y_zero_point,
+    const ::executorch::aten::optional<::executorch::aten::Tensor>& bias,
+    int64_t out_multiplier,
+    int64_t out_shift,
+    int64_t out_zero_point,
+    bool transposed,
+    ::executorch::aten::Tensor& out);
+
 void quantized_conv2d_nhwc_out(
     ::executorch::runtime::KernelRuntimeContext& ctx,
     const ::executorch::aten::Tensor& input,

backends/cadence/hifi/operators/targets.bzl

@@ -90,7 +90,6 @@ OPERATORS = [
     "quantized_linear_out",
     "quantized_linear_asym8sxasym8s_asym8s_per_tensor_out",
     "quantized_linear_asym8uxasym8u_asym8u_per_tensor_out",
-    "quantized_matmul_out",
     "quantized_matmul_asym8sxasym8s_asym8s_out",
     "quantized_matmul_asym8uxasym8u_asym8u_out",
     "quantized_relu_out",
@@ -122,3 +121,6 @@ def define_common_targets():
     # Define build targets for all operators registered in the tables above.
     for op in OPERATORS:
         define_operator(op)
+
+    # quantized_matmul_out needs additional dependency for int16 support
+    define_operator("quantized_matmul_out", deps=["fbcode//on_device_ai/Assistant/Jarvis/min_runtime/operators:quantize_matmul_out", "fbcode//on_device_ai/Assistant/Jarvis/min_runtime/operators:headers",])

backends/cadence/hifi/operators/tests/test_op_quantized_matmul_out.cpp

@@ -0,0 +1,145 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gtest/gtest.h>
+#include <sys/times.h>
+
+#include <executorch/kernels/test/TestUtil.h>
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
+#include <executorch/runtime/platform/runtime.h>
+
+#include <executorch/backends/cadence/hifi/operators/operators.h>
+
+namespace impl {
+namespace HiFi {
+namespace native {
+namespace {
+
+using ::executorch::aten::Scalar;
+using ::executorch::aten::ScalarType;
+using ::executorch::aten::Tensor;
+using ::executorch::aten::TensorImpl;
+using ::executorch::runtime::Error;
+using ::executorch::runtime::KernelRuntimeContext;
+using ::executorch::runtime::runtime_init;
+using ::executorch::runtime::testing::TensorFactory;
+
+class HiFiQuantizedMatmulTest : public OperatorTest {
+ public:
+ protected:
+  void quantized_matmul_out(
+      const Tensor& X,
+      int64_t X_zero_point,
+      const Tensor& Y,
+      int64_t Y_zero_point,
+      const std::optional<Tensor>& bias,
+      int64_t out_multiplier,
+      int64_t out_shift,
+      int64_t out_zero_point,
+      bool transposed,
+      Tensor& output) {
+    return ::impl::HiFi::native::quantized_matmul_out(
+        context_,
+        X,
+        X_zero_point,
+        Y,
+        Y_zero_point,
+        bias,
+        out_multiplier,
+        out_shift,
+        out_zero_point,
+        transposed,
+        output);
+  }
+};
+
+// Test quantized_matmul_out with int16 activations and int8 weights
+TEST_F(HiFiQuantizedMatmulTest, QuantizedMatmulInt16Test) {
+  TensorFactory<ScalarType::Short> tf_int16;
+  TensorFactory<ScalarType::Int> tf_int32;
+  TensorFactory<ScalarType::Char> tf_int8;
+
+  // Simple 2D case: X [64, 33] x Y [33, 128] = output [64, 128]
+  // Using simple values for testing
+  Tensor X = tf_int16.ones({64, 33});
+  Tensor Y = tf_int8.ones({33, 128});
+  // Bias not used
+  Tensor bias = tf_int32.full({128}, -30);
+  Tensor output = tf_int16.zeros({64, 128});
+
+  int64_t X_zero_point = 0;
+  int64_t Y_zero_point = 0;
+  int64_t out_multiplier = 1073741824; // 0.5 * 2^31
+  int64_t out_shift = 0;
+  int64_t out_zero_point = 0;
+
+  quantized_matmul_out(
+      X,
+      X_zero_point,
+      Y,
+      Y_zero_point,
+      bias, // pass bias tensor
+      out_multiplier,
+      out_shift,
+      out_zero_point,
+      false, // transposed
+      output);
+
+  // Verify the output is correct
+  // With all ones input and weights, inner dimension is 33
+  // Matmul result: 33, with out_multiplier = 0.5 * 2^31 (scales by 0.5)
+  // Expected value: 33 * 0.5 = 16.5 ≈ 16
+  EXPECT_EQ(output.const_data_ptr<int16_t>()[0], 16);
+}
+
+// Test quantized_matmul_out with transposed Y (int16 activations and int8
+// weights)
+TEST_F(HiFiQuantizedMatmulTest, QuantizedMatmulInt16TransposedTest) {
+  TensorFactory<ScalarType::Short> tf_int16;
+  TensorFactory<ScalarType::Int> tf_int32;
+  TensorFactory<ScalarType::Char> tf_int8;
+
+  // Transposed case: X [64, 33] x Y^T [128, 33] = output [64, 128]
+  Tensor X = tf_int16.ones({64, 33});
+  Tensor Y = tf_int8.ones({128, 33}); // Transposed
+  // Bias not used
+  Tensor bias = tf_int32.full({128}, -30);
+  Tensor output = tf_int16.zeros({64, 128});
+
+  int64_t X_zero_point = 0;
+  int64_t Y_zero_point = 0;
+  int64_t out_multiplier = 1073741824; // 0.5 * 2^31
+  int64_t out_shift = 0;
+  int64_t out_zero_point = 0;
+
+  quantized_matmul_out(
+      X,
+      X_zero_point,
+      Y,
+      Y_zero_point,
+      bias, // pass bias tensor
+      out_multiplier,
+      out_shift,
+      out_zero_point,
+      true, // transposed
+      output);
+
+  // Verify the output is correct
+  // With all ones input and weights, inner dimension is 33
+  // Matmul result: 33, with out_multiplier = 0.5 * 2^31 (scales by 0.5)
+  // Expected value: 33 * 0.5 = 16.5 ≈ 16
+  EXPECT_EQ(output.const_data_ptr<int16_t>()[0], 16);
+}
+
+} // namespace
+} // namespace native
+} // namespace HiFi
+} // namespace impl