Using generic implementation for 16-bit activations and 8-bit weights for matmul in backends

Summary:
# 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 now use the generic implementation in the backends and create a unit test as well as e2e tests.

Differential Revision: D87997149

Author: RahulC7

backends/cadence/aot/quantizer/quantizer.py

@@ -386,3 +386,16 @@ def __init__(self, quantizers: Optional[list[Quantizer]] = None) -> None:
         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/op_quantized_matmul.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/op_quantized_matmul_out.h

@@ -0,0 +1,33 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "executorch/runtime/core/exec_aten/exec_aten.h"
+#include "executorch/runtime/kernel/kernel_runtime_context.h"
+
+namespace impl {
+namespace HiFi {
+namespace native {
+
+::executorch::aten::Tensor& 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);
+
+} // namespace native
+} // namespace HiFi
+} // namespace impl

backends/cadence/hifi/operators/targets.bzl

@@ -2,7 +2,7 @@ load("@fbsource//tools/build_defs:platform_defs.bzl", "CXX")
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
 
-def define_operator(name: str, deps: list[str] | None = None) -> None:
+def define_operator(name: str, deps: list[str] | None = None, exported_headers: list[str] | None = None) -> None:
     op_name = "op_{}".format(name)
 
     # Deps used by all operators.
@@ -21,6 +21,8 @@ def define_operator(name: str, deps: list[str] | None = None) -> None:
     ]
     if deps == None:
         deps = []
+    if exported_headers == None:
+        exported_headers = ["operators.h"]
 
     runtime.cxx_library(
         name = op_name,
@@ -32,7 +34,7 @@ def define_operator(name: str, deps: list[str] | None = None) -> None:
         ],
         compatible_with = ["ovr_config//cpu:xtensa"],
         deps = deps + common_deps,
-        exported_headers = ["operators.h"],
+        exported_headers = exported_headers,
     )
 
 OPERATORS = [
@@ -87,7 +89,6 @@ OPERATORS = [
     "quantized_layer_norm",
     "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",
@@ -124,6 +125,9 @@ def define_common_targets():
     define_operator("quantized_linear_out", deps=["fbcode//on_device_ai/Assistant/Jarvis/min_runtime/operators/generic:op_quantized_linear"])
     define_operator("quantized_linear_per_tensor_out", deps=["fbcode//on_device_ai/Assistant/Jarvis/min_runtime/operators/generic:op_quantized_linear"])
 
-      # quantized_conv2d_nchw_out and quantized_conv2d_nhwc_out need additional dependency for int16 support
+    # quantized_conv2d_nchw_out and quantized_conv2d_nhwc_out need additional dependency for int16 support
     define_operator("quantized_conv2d_nchw_out", deps=["fbcode//on_device_ai/Assistant/Jarvis/min_runtime/operators/generic:op_quantized_conv2d"])
     define_operator("quantized_conv2d_nhwc_out", deps=["fbcode//on_device_ai/Assistant/Jarvis/min_runtime/operators/generic:op_quantized_conv2d"])
+
+    # quantized_matmul_out needs additional dependency for int16 support
+    define_operator("quantized_matmul_out", deps=["fbcode//on_device_ai/Assistant/Jarvis/min_runtime/operators/generic:op_quantized_matmul"], exported_headers=["op_quantized_matmul_out.h"])

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 <executorch/backends/cadence/hifi/operators/op_quantized_matmul_out.h>
+#include <gtest/gtest.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>
+
+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:
+  Tensor& 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
+  Tensor expected = tf_int16.full({64, 128}, 16);
+  EXPECT_TENSOR_EQ(output, expected);
+}
+
+// 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
+  Tensor expected = tf_int16.full({64, 128}, 16);
+  EXPECT_TENSOR_EQ(output, expected);
+}
+
+} // namespace
+} // namespace native
+} // namespace HiFi
+} // namespace impl