Dynamic Dispatch for Kernels + Support MKL-based kernels w/ Fallback

kernels/avx/matmul_avx.h

@@ -0,0 +1,37 @@
+#ifndef MATMUL_OPERATOR_AVX_H
+#define MATMUL_OPERATOR_AVX_H
+
+#include "matmul.h"
+#include <iostream>
+
+namespace matmul {
+
+class MatmulOperatorAVX : public MatmulOperator {
+   public:
+    void mat_mul_accelerator_transposed_fastover_column(const struct matmul_params* params) override;
+    void mat_mul_accelerator_transposed_fastover_column_bias(const struct matmul_params* params) override;
+
+    // int8 operations
+    void mat_mul_accelerator_int8_fast_32unroll_over_column(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll_nobias(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_batch(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_ofp32(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_ofp32_batch(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32_over_column(const struct matmul_params* params) override;
+
+    void mat_mul_accelerator_int8_int4_fast_no_offset(struct matmul_params* params) override;
+
+    void mat_mul_accelerator_int4_fast(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int4_fast_no_offset(const struct matmul_params* params) override;
+};
+
+inline MatmulOperator& CreateMatmulOperatorAVX() {
+    static MatmulOperatorAVX instance;
+    return instance;
+}
+
+}  // namespace matmul
+
+#endif 

kernels/avx/matmul_avx_fp32.cc

@@ -4,8 +4,9 @@
 #include <pthread.h>
 #include <stdio.h>
 #include <xmmintrin.h>  // intel SSE intrinsic
+#include <iostream>
 
-#include "../matmul.h"
+#include "matmul_avx.h"
 
 namespace matmul {
 
@@ -60,7 +61,7 @@ void *mat_mul_transposed_fastover_column_func(void *args) {
     return NULL;
 }
 
-void MatmulOperator::mat_mul_accelerator_transposed_fastover_column(const struct matmul_params *params) {
+void MatmulOperatorAVX::mat_mul_accelerator_transposed_fastover_column(const struct matmul_params *params) {
     int i, j, k;
 
     int num_thread = params->opt_params.num_thread;
@@ -112,7 +113,7 @@ void fp32_ref_matmul_bias(const struct matmul_params *params) {
     }
 }
 
-void MatmulOperator::mat_mul_accelerator_transposed_fastover_column_bias(const struct matmul_params *params) {
+void MatmulOperatorAVX::mat_mul_accelerator_transposed_fastover_column_bias(const struct matmul_params *params) {
     fp32_ref_matmul_bias(params);
 }
 

kernels/avx/matmul_avx_int4.cc

@@ -3,7 +3,7 @@
 
 #include <cassert>
 
-#include "../matmul.h"
+#include "matmul_avx.h"
 
 static inline __m256i bytes_from_nibbles_32(const uint8_t *rsi) {
     // Load 16 bytes from memory
@@ -675,7 +675,7 @@ static void *fast_zp_no_offset_over_column_func_v5(void *args) {
 }
 
 namespace matmul {
-void MatmulOperator::mat_mul_accelerator_int4_fast(const struct matmul_params *params) {
+void MatmulOperatorAVX::mat_mul_accelerator_int4_fast(const struct matmul_params *params) {
     const int num_thread = params->opt_params.num_thread;
     int i, j, k;
     pthread_t thread_pool[num_thread];
@@ -693,7 +693,7 @@ void MatmulOperator::mat_mul_accelerator_int4_fast(const struct matmul_params *p
     for (j = 0; j < num_thread; j++) pthread_join(thread_pool[j], NULL);
 };
 
-void MatmulOperator::mat_mul_accelerator_int4_fast_no_offset(const struct matmul_params *params) {
+void MatmulOperatorAVX::mat_mul_accelerator_int4_fast_no_offset(const struct matmul_params *params) {
     const int num_thread = params->opt_params.num_thread;
     int i, j, k;
     pthread_t thread_pool[num_thread];

kernels/avx/matmul_avx_int8.cc

@@ -5,7 +5,7 @@
 #include <cstdlib>
 #include <iostream>
 
-#include "../matmul.h"
+#include "matmul_avx.h"
 
 inline void assign_8int32(int *ptr, int &acc) {
     acc = (ptr[0] + ptr[1] + ptr[2] + ptr[3] + ptr[4] + ptr[5] + ptr[6] + ptr[7]);
@@ -381,7 +381,7 @@ void *mat_mul_accelerator_int8_thread_func_2x2_32unroll(void *args) {
     return NULL;
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll(const struct matmul_params *params) {
+void MatmulOperatorAVX::mat_mul_accelerator_int8_fast_2x2_32unroll(const struct matmul_params *params) {
     int j, num_thread = params->opt_params.num_thread;
 
     assert(params->A.column % 64 == 0);
@@ -478,7 +478,7 @@ void *mat_mul_accelerator_int8_fast_32unroll_over_column_thread_func(void *args)
     return NULL;
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_32unroll_over_column(const struct matmul_params *params) {
+void MatmulOperatorAVX::mat_mul_accelerator_int8_fast_32unroll_over_column(const struct matmul_params *params) {
     int j, num_thread = params->opt_params.num_thread;
 
     if (num_thread > params->C.column) num_thread = params->C.column;
@@ -610,7 +610,7 @@ void *mat_mul_accelerator_int8_thread_func_2x2_32unroll_nobias(void *args) {
     return NULL;
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias(const struct matmul_params *params) {
+void MatmulOperatorAVX::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias(const struct matmul_params *params) {
     int j, num_thread = params->opt_params.num_thread;
 
     assert((params->C.column) % 2 == 0);
@@ -681,7 +681,7 @@ void *mat_mul_accelerator_int8_thread_func_2x2_32unroll_nobias_batch(void *args)
     return NULL;
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_batch(const struct matmul_params *params) {
+void MatmulOperatorAVX::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_batch(const struct matmul_params *params) {
     int j, num_thread = params->opt_params.num_thread;
 
     assert((params->C.column) % 2 == 0);
@@ -791,7 +791,7 @@ void *mat_mul_accelerator_int8_thread_func_2x2_32unroll_nobias_ofp32(void *args)
     return NULL;
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_ofp32(const struct matmul_params *params) {
+void MatmulOperatorAVX::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_ofp32(const struct matmul_params *params) {
     int j, num_thread = params->opt_params.num_thread;
 
     assert(params->A.column % 32 == 0);
@@ -851,7 +851,7 @@ void *mat_mul_accelerator_int8_thread_func_2x2_32unroll_nobias_ofp32_batch(void
     return NULL;
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_ofp32_batch(const struct matmul_params *params) {
+void MatmulOperatorAVX::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_ofp32_batch(const struct matmul_params *params) {
     int j, num_thread = params->opt_params.num_thread;
 
     assert(params->A.column % 32 == 0);
@@ -940,7 +940,7 @@ void *mat_mul_accelerator_int8_thread_func_2x2_32unroll_bfp32_ofp32(void *args)
     return NULL;
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32(const struct matmul_params *params) {
+void MatmulOperatorAVX::mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32(const struct matmul_params *params) {
     int j, num_thread = params->opt_params.num_thread;
 
     assert(params->A.column % 64 == 0);
@@ -1211,8 +1211,9 @@ void *mat_mul_accelerator_int8_thread_func_2x2_32unroll_bfp32_ofp32_over_column(
     return NULL;
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32_over_column(
+void MatmulOperatorAVX::mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32_over_column(
     const struct matmul_params *params) {
+
     int j, num_thread = params->opt_params.num_thread;
 
     if (num_thread > params->C.column) num_thread = params->C.column;
@@ -1241,4 +1242,5 @@ void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32_over
     }
 }
 
-}  // namespace matmul
+
+}

kernels/avx/matmul_avx_int8_int4.cc

@@ -4,7 +4,7 @@
 #include <cassert>
 #include <cmath>
 
-#include "../matmul.h"
+#include "matmul_avx.h"
 
 #include "pthread_pool.h"
 
@@ -322,7 +322,7 @@ static void quantize_fp_to_int8_block_size32(float *x, int size, int8_t *qx, flo
 
 namespace matmul {
 
-void MatmulOperator::mat_mul_accelerator_int8_int4_fast_no_offset(struct matmul_params *params) {
+void MatmulOperatorAVX::mat_mul_accelerator_int8_int4_fast_no_offset(struct matmul_params *params) {
     // const int num_thread = 4;
     const int num_thread = params->opt_params.num_thread;
     int i, j, k;

kernels/cuda/gemv_cuda.cu

@@ -17,7 +17,7 @@
 #include <iostream>
 #include <stdio.h>
 
-#include "../matmul.h"
+#include "matmul_cuda.h"
 #include "ops/linear.h"
 
 // #include <cuda_runtime.h>
@@ -210,7 +210,7 @@ namespace matmul{
   Returns:
     out_feats: tensor of shape [B, OC];
   */
-  void MatmulOperator::gemv_forward_cuda(const struct matmul_params *params)
+  void MatmulOperatorCUDA::gemv_forward_cuda(const struct matmul_params *params)
   {
     const struct matrix *A = &params->A, *B = &params->B, *C = &params->C;
 
@@ -259,11 +259,11 @@ namespace matmul{
     PROFILE_END("gemv_forward_cuda");
   }
 
-  void MatmulOperator::mat_mul_accelerator_int4_fast(const struct matmul_params *params) {
+  void MatmulOperatorCUDA::mat_mul_accelerator_int4_fast(const struct matmul_params *params) {
     // TODO: remove this
   };
 
-  void MatmulOperator::mat_mul_accelerator_int4_fast_no_offset(const struct matmul_params *params) {
+  void MatmulOperatorCUDA::mat_mul_accelerator_int4_fast_no_offset(const struct matmul_params *params) {
     // TODO: remove this
   };
 

kernels/cuda/matmul_cuda.h

@@ -0,0 +1,38 @@
+#ifndef MATMUL_OPERATOR_CUDA_H
+#define MATMUL_OPERATOR_CUDA_H
+
+#include "matmul.h"
+#include <iostream>
+
+namespace matmul {
+
+class MatmulOperatorCUDA : public MatmulOperator {
+   public:
+    void mat_mul_accelerator_transposed_fastover_column(const struct matmul_params* params) override;
+
+    // int8 operations
+    void mat_mul_accelerator_int8_fast_32unroll_over_column(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll_nobias(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_batch(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_ofp32(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_ofp32_batch(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32_over_column(const struct matmul_params* params) override;
+
+    void mat_mul_accelerator_int4_fast(const struct matmul_params* params) override;
+    void mat_mul_accelerator_int4_fast_no_offset(const struct matmul_params* params) override;
+
+    void gemv_forward_cuda(const struct matmul_params* params) override;
+    void naive_mat_mul_fp16_int4(const struct matmul_params* params) override;
+};
+
+// Declaring as static to prevent linker errors due to both cc and cu files
+static inline MatmulOperator& CreateMatmulOperatorCUDA() {
+    static MatmulOperatorCUDA instance;
+    return instance;
+}
+
+}  // namespace matmul
+
+#endif 

kernels/cuda/matmul_int4.cu

@@ -1,11 +1,11 @@
 #include <cstdlib>
 #include <iostream>
 
-#include "../matmul.h"
+#include "matmul_cuda.h"
 
 namespace matmul {
 
-void MatmulOperator::naive_mat_mul_fp16_int4(const struct matmul_params *params) {
+void MatmulOperatorCUDA::naive_mat_mul_fp16_int4(const struct matmul_params *params) {
     const struct matrix *A = &params->A, *B = &params->B, *C = &params->C;
     const int block_size = params->block_size;
     // CHECK_MATRICES_int4weight(A, B, C);

kernels/cuda/matmul_ref_fp32.cc

@@ -4,7 +4,7 @@
 #include <cmath>
 #include <cstdlib>
 
-#include "../matmul.h"
+#include "matmul_cuda.h"
 
 namespace matmul {
 void fp32_ref_matmul(const struct matmul_params *params) {
@@ -28,7 +28,8 @@ void fp32_ref_matmul(const struct matmul_params *params) {
     }
 }
 
-void MatmulOperator::mat_mul_accelerator_transposed_fastover_column(const struct matmul_params *params) {
+void MatmulOperatorCUDA::mat_mul_accelerator_transposed_fastover_column(const struct matmul_params *params) {
+    std::cout<<"mat_mul_accelerator_transposed_fastover_column, fp32"<<std::endl;
     fp32_ref_matmul(params);
 }
 

kernels/cuda/matmul_ref_int8.cc

@@ -4,7 +4,7 @@
 #include <cmath>
 #include <cstdlib>
 
-#include "../matmul.h"
+#include "matmul_cuda.h"
 
 namespace matmul {
 void int8_ref_matmul(const struct matmul_params *params) {
@@ -157,35 +157,35 @@ void int8_ref_matmul_nobias_ofp32_batch(const struct matmul_params *params) {
     }
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll(const struct matmul_params *params) {
+void MatmulOperatorCUDA::mat_mul_accelerator_int8_fast_2x2_32unroll(const struct matmul_params *params) {
     int8_ref_matmul(params);
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias(const struct matmul_params *params) {
+void MatmulOperatorCUDA::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias(const struct matmul_params *params) {
     int8_ref_matmul_nobias(params);
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_batch(const struct matmul_params *params) {
+void MatmulOperatorCUDA::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_batch(const struct matmul_params *params) {
     int8_ref_matmul_nobias_batch(params);
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_32unroll_over_column(const struct matmul_params *params) {
+void MatmulOperatorCUDA::mat_mul_accelerator_int8_fast_32unroll_over_column(const struct matmul_params *params) {
     int8_ref_matmul(params);
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32(const struct matmul_params *params) {
+void MatmulOperatorCUDA::mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32(const struct matmul_params *params) {
     int8_ref_matmul_bfp32_ofp32(params);
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_ofp32(const struct matmul_params *params) {
+void MatmulOperatorCUDA::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_ofp32(const struct matmul_params *params) {
     int8_ref_matmul_nobias_ofp32(params);
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_ofp32_batch(const struct matmul_params *params) {
+void MatmulOperatorCUDA::mat_mul_accelerator_int8_fast_2x2_32unroll_nobias_ofp32_batch(const struct matmul_params *params) {
     int8_ref_matmul_nobias_ofp32_batch(params);
 }
 
-void MatmulOperator::mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32_over_column(
+void MatmulOperatorCUDA::mat_mul_accelerator_int8_fast_2x2_32unroll_bfp32_ofp32_over_column(
     const struct matmul_params *params) {
     int8_ref_matmul_bfp32_ofp32(params);
 }