[Kernel] Bug fix for small_gemm_transb (#318)

src/kernels/gemm_kernel_ext.cpp

@@ -335,7 +335,7 @@ void small_gemm_transb_1xn_dynk(const TA *A, const TB *B, float *C, int N, int K
  // Each loop compute 'BC' elements in C
  int i = 0;
  for (; i + BC - 1 < N; i += BC) {
- const TA *pA = A + i * ldb;
+ const TA *pA = A;
  const TB *pB = B + i * ldb;
 
  __m512 vc[BC];
@@ -356,7 +356,7 @@ void small_gemm_transb_1xn_dynk(const TA *A, const TB *B, float *C, int N, int K
 
  // Remain elements
  for (; i < N; ++i) {
- const TA *pA = A + i * ldb;
+ const TA *pA = A;
  const TB *pB = B + i * ldb;
  __m512 vc = _mm512_set1_ps(0);
 

tests/ut/gemm_kernel_ext_test.cpp

@@ -19,16 +19,17 @@
 
 #include "gtest/gtest.h"
 
+template <typename TA, typename TB, typename TC>
 static void small_gemm_tranb_ref(
- const float *A, const float *B, float *C, int M, int N, int K, int lda, int ldb, int ldc) {
+ const TA *A, const TB *B, TC *C, int M, int N, int K, int lda, int ldb, int ldc) {
  // Loop over the rows of A
  for (int i = 0; i < M; i++) {
  // Loop over the columns of B
  for (int j = 0; j < N; j++) {
  // Compute the dot product of row i of A with column j of B
  float dot_product = 0;
  for (int k = 0; k < K; k++) {
- dot_product += A[i * lda + k] * B[j * ldb + k];
+ dot_product += (float)A[i * lda + k] * (float)B[j * ldb + k];
  }
  // Store the result in C[i][j]
  C[i * ldc + j] = dot_product;
@@ -54,13 +55,14 @@ static void small_gemm_tranb_ref(
 }
 
 // Test function to compare reference and optimized implementations
+template <typename TA = float, typename TB = float, typename TC = float>
 void test_small_gemm_tranb(int M, int N, int K) {
- float *A_ref = new float[M * K];
- float *B_ref = new float[K * N];
- float *C_ref = new float[M * N];
- float *A_opt = new float[M * K];
- float *B_opt = new float[K * N];
- float *C_opt = new float[M * N];
+ TA *A_ref = new TA[M * K];
+ TB *B_ref = new TB[K * N];
+ TC *C_ref = new TC[M * N];
+ TA *A_opt = new TA[M * K];
+ TB *B_opt = new TB[K * N];
+ TC *C_opt = new TC[M * N];
 
  // Generate random matrices A and B
  std::random_device dev;
@@ -262,6 +264,12 @@ TEST(small_gemm_tranb, small_gemm_tranb_f32) {
  test_bigger_kernel();
 }
 
+TEST(small_gemm_tranb, small_gemm_tranb_bf16fp16f32) {
+ test_small_gemm_tranb<bfloat16_t, float16_t, float>(1, 2, 16);
+ test_small_gemm_tranb<bfloat16_t, float16_t, float>(1, 4, 128);
+ test_small_gemm_tranb<bfloat16_t, float16_t, float>(1, 4, 256);
+}
+
 TEST(small_gemm_tranb, small_gemm_tranb_int8) {
  test_small_gemm_tranb_int8(1, 100, 128);
  test_small_gemm_tranb_int8(2, 101, 256);