Add gemm kernel to interface

Differential Revision: D71833068

Pull Request resolved: #2040

Author: kimishpatel

torchao/experimental/kernels/cpu/aarch64/matmul/matmul.h

@@ -10,6 +10,7 @@
 
 #pragma once
 
+#include <cassert>
 #if defined(__aarch64__) && defined(__ARM_NEON)
 
 #include <arm_neon.h>
@@ -106,6 +107,83 @@ void kernel(
     const int rhs_qparams_stride);
 
 } // namespace fp32_a_input_channelwise_8bit_b_4x16x4_f32
+
+namespace fp32_a_input_channelwise_8bit_b_f32 {
+
+template <bool b_has_zeros, bool a_transposed, bool b_tranposed>
+void kernel(
+    int m,
+    int n,
+    int k,
+    const float* lhs,
+    int lhs_stride_m,
+    const int8_t* rhs,
+    int rhs_stride_n,
+    float32_t* output,
+    int out_stride_m,
+    const int8_t* rhs_zero_points,
+    const float* rhs_scales,
+    const float beta,
+    const int rhs_qparams_stride);
+
+template <bool b_has_zeros, bool a_transposed, bool b_tranposed>
+void kernel(
+    int m,
+    int n,
+    int k,
+    const float* lhs,
+    int lhs_stride_m,
+    const int8_t* rhs,
+    int rhs_stride_n,
+    float32_t* output,
+    int out_stride_m,
+    const int8_t* rhs_zero_points,
+    const float* rhs_scales,
+    const float beta,
+    const int rhs_qparams_stride) {
+  assert(n >= 16);
+  if (m > 16) {
+    auto remaining_m = m % 16;
+    auto m_for_gemm_kernel = m - remaining_m;
+    fp32_a_input_channelwise_8bit_b_4x16x4_f32::
+        kernel<b_has_zeros, a_transposed, b_tranposed>(
+            m_for_gemm_kernel,
+            n,
+            k,
+            lhs,
+            lhs_stride_m,
+            rhs,
+            rhs_stride_n,
+            output,
+            out_stride_m,
+            rhs_zero_points,
+            rhs_scales,
+            beta,
+            rhs_qparams_stride);
+    output += m_for_gemm_kernel * out_stride_m;
+    lhs += m_for_gemm_kernel * lhs_stride_m;
+    m = remaining_m;
+  }
+  if (m > 0) {
+    fp32_a_input_channelwise_8bit_b_1x16x4_f32::
+        kernel<b_has_zeros, a_transposed, b_tranposed>(
+            m,
+            n,
+            k,
+            lhs,
+            lhs_stride_m,
+            rhs,
+            rhs_stride_n,
+            output,
+            out_stride_m,
+            rhs_zero_points,
+            rhs_scales,
+            beta,
+            rhs_qparams_stride);
+  }
+}
+
+} // namespace fp32_a_input_channelwise_8bit_b_f32
 } // namespace torchao::kernels::cpu::aarch64::quantized_matmul
 
 #include <torchao/experimental/kernels/cpu/aarch64/matmul/channelwise_8bit_a_channelwise_8bit_b_1x16x16_f32_smlal-impl.h>

torchao/experimental/kernels/cpu/interface/quantized_matmul.h

@@ -12,9 +12,7 @@
 #include <torchao/experimental/kernels/cpu/fallback/matmul/fp32_a_channelwise_8bit_b_fp32_c.h>
 
 #if defined(__aarch64__) && defined(__ARM_NEON)
-#include <torchao/experimental/kernels/cpu/aarch64/matmul/channelwise_8bit_a_channelwise_8bit_b_1x16x16_f32_smlal-impl.h>
-#include <torchao/experimental/kernels/cpu/aarch64/matmul/channelwise_8bit_a_channelwise_8bit_b_1x8x16_f32_neondot-impl.h>
-#include <torchao/experimental/kernels/cpu/aarch64/matmul/fp32_a_input_channelwise_8bit_b_1x16x4_f32_impl.h>
+#include <torchao/experimental/kernels/cpu/aarch64/matmul/matmul.h>
 #endif // defined(__aarch64__) && defined(__ARM_NEON)
 
 namespace torchao::kernels::cpu::quantized_matmul {
@@ -138,8 +136,8 @@ get_fp32_a_input_channelwise_8bit_b_f32_c_matmul(
   if (!a_transposed && !b_transposed && n >= 16) {
     a_stride_m = k;
     b_stride_n = n;
-    return aarch64::quantized_matmul::
-        fp32_a_input_channelwise_8bit_b_1x16x4_f32::kernel<true, false, false>;
+    return aarch64::quantized_matmul::fp32_a_input_channelwise_8bit_b_f32::
+        kernel<true, false, false>;
   }
 #endif // defined(__aarch64__) && defined(__ARM_NEON)
   assert(!a_transposed);

torchao/experimental/kernels/cpu/interface/test_qmatmul_interface.cpp

@@ -624,6 +624,22 @@ TEST_P(
       /*m=*/4, /*k=*/5, /*n=*/3, beta(), *this, 32);
 }
 
+TEST_P(
+    FP32A_QuantizedB_FP32C_Interface_Test,
+    BTranposedWithZeroPointsOddSizes2) {
+  generate(19, 37, 35, true, false, false);
+  test_fp32_a_input_channelwise_8bit_b(
+      /*m=*/19, /*k=*/37, /*n=*/35, beta(), *this);
+}
+
+TEST_P(
+    FP32A_QuantizedB_FP32C_Interface_Test,
+    BTranposedWithZeroPointsOddSizesStrided2) {
+  generate(23, 37, 50, true, false, false, 32);
+  test_fp32_a_input_channelwise_8bit_b(
+      /*m=*/23, /*k=*/37, /*n=*/50, beta(), *this, 32);
+}
+
 INSTANTIATE_TEST_SUITE_P(
     F32AInt8BFP32CTest,
     FP32A_QuantizedB_FP32C_Interface_Test,