Use GEMM kernel for KleidiAI to accelerate FP16Benchmark

Makefile.FP16Benchmark.aarch64

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright 2024 Arm Limited and/or its affiliate <[email protected]>
+# SPDX-License-Identifier: BSD-3-Clause
+
+CC           := g++
+SRCEXT       := cc
+OBJEXT       := o
+BUILD_DIR    := obj
+KLEIDIAI_DIR := external/kleidiai/kai/ukernels/matmul/matmul_f32_f32_f16p
+
+CFLAGS    := -O3 -mcpu=native -std=c++20 -fvisibility=hidden -fopenmp
+CDEFINES  := -DFBGEMM_ENABLE_KLEIDIAI=1 -DCPUINFO_SUPPORTED_PLATFORM=1 -DFBGEMM_FP16_FALLBACK_TO_REF_KERNEL=1
+INCLUDES  := -I./include -I./ -I./external/cpuinfo/include -I./external/asmjit/src -I./external/googletest/googletest/include -I./src
+
+SRC_DIR  := src
+BENCH_DIR := bench
+TEST_DIR := test
+SOURCES := src/FbgemmFP16.cc src/FbgemmFPCommon.cc src/Utils.cc src/RefImplementations.cc src/TransposeUtils.cc src/FbgemmFP16UKernelsSve128.cc
+SRC_OBJECTS := $(patsubst $(SRC_DIR)/%,$(BUILD_DIR)/%,$(SOURCES:.$(SRCEXT)=.$(OBJEXT)))
+KLEIDIAI_OBJECTS := $(BUILD_DIR)/KleidiAIFP16UKernelsNeon.$(OBJEXT)
+
+LIB := -lcpuinfo -fopenmp
+LIBDIR := -L/usr/lib/aarch64-linux-gnu
+
+BENCH_TARGET := FP16Benchmark
+TEST_TARGET := FP16Test
+
+all: $(BENCH_TARGET) $(TEST_TARGET)
+
+$(TEST_TARGET): $(SRC_OBJECTS) $(KLEIDIAI_OBJECTS) $(BUILD_DIR)/BenchUtils.$(OBJEXT) $(BUILD_DIR)/FP16Test.$(OBJEXT)
+	   $(CC) -o $(TEST_TARGET) $^ $(LIBDIR) $(LIB) -lgtest
+
+$(BENCH_TARGET): $(SRC_OBJECTS) $(KLEIDIAI_OBJECTS) $(BUILD_DIR)/BenchUtils.$(OBJEXT) $(BUILD_DIR)/FP16Benchmark.$(OBJEXT)
+	   $(CC) -o $(BENCH_TARGET) $^ $(LIBDIR) $(LIB)
+
+$(BUILD_DIR)/%.$(OBJEXT): $(SRC_DIR)/%.$(SRCEXT)
+			  @mkdir -p $(dir $@)
+			  $(CC) $(CDEFINES) $(CFLAGS) $(INCLUDES) -c -o $@ $<
+
+$(BUILD_DIR)/%.$(OBJEXT): $(BENCH_DIR)/%.$(SRCEXT)
+			  $(CC) $(CDEFINES) $(CFLAGS) $(INCLUDES) -c -o $@ $<
+
+$(BUILD_DIR)/%.$(OBJEXT): $(KLEIDIAI_DIR)/%.$(SRCEXT)
+			  $(CC) $(CDEFINES) $(CFLAGS) $(INCLUDES) -c -o $@ $<
+
+$(BUILD_DIR)/%.$(OBJEXT): $(TEST_DIR)/%.$(SRCEXT)
+			  $(CC) $(CDEFINES) $(CFLAGS) $(INCLUDES) -c -o $@ $<

bench/BenchUtils.h

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright 2024 Arm Limited and/or its affiliates <[email protected]>
  * All rights reserved.
  *
  * This source code is licensed under the BSD-style license found in the
@@ -70,7 +71,9 @@ NOINLINE float cache_evict(const T& vec) {
   float dummy = 0.0f;
   for (std::size_t i = 0; i < dataSize; i += CACHE_LINE_SIZE) {
     dummy += data[i] * 1.0f;
+#ifndef __aarch64__
     _mm_mfence();
+#endif
 #ifndef _MSC_VER
     asm volatile("" ::: "memory");
 #endif
@@ -513,6 +516,7 @@ void performance_test(
         k,
         gflops * repetitions,
         gbs * repetitions);
+
   }
 }
 

include/fbgemm/FbgemmFPCommon.h

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright 2024 Arm Limited and/or its affiliates <[email protected]>
  * All rights reserved.
  *
  * This source code is licensed under the BSD-style license found in the
@@ -25,6 +26,9 @@ using partition_array_t = std::array<std::array<std::array<int, 2>, 2>, 121>;
 extern partition_array_t partition_avx2;
 extern partition_array_t partition_avx512;
 extern partition_array_t partition_sve128;
+#ifdef FBGEMM_ENABLE_KLEIDIAI
+extern partition_array_t partition_neon;
+#endif
 
 template <typename T>
 struct GemmParams {
@@ -35,7 +39,11 @@ struct GemmParams {
   float* C;
   uint64_t ldc;
   uint64_t b_block_cols;
+#ifdef FBGEMM_ENABLE_KLEIDIAI
+  uint64_t lda;
+#else
   uint64_t b_block_size;
+#endif
 };
 
 template <typename T>
@@ -155,8 +163,12 @@ void cblas_gemm_compute(
         for (auto m2 = m_start; m2 < m_end; m2 += kernel_nrows) {
           assert(kernel_nrows * kb < static_cast<int64_t>(scratchpad->size()));
           if (m != 1) {
+#ifdef FBGEMM_ENABLE_KLEIDIAI
+            gp.A = const_cast<float*>(&A[m2 * k + k_ind]);
+#else
             PackA(kernel_nrows, kb, &A[m2 * k + k_ind], k, scratchpad->data());
             gp.A = scratchpad->data();
+#endif
           } else {
             // When m == 1, it is actually vector matrix multiplication. We
             // don't need to do the transposition for packA here. Instead, we
@@ -172,7 +184,11 @@ void cblas_gemm_compute(
           gp.C = &C[m2 * ldc];
           gp.ldc = ldc * sizeof(C[0]);
           gp.b_block_cols = nbcol;
+#ifdef FBGEMM_ENABLE_KLEIDIAI
+          gp.lda = k * sizeof(A[0]);
+#else
           gp.b_block_size = gp.k * Bp.blockColSize() * sizeof(gp.B[0]);
+#endif
 
           if ((n % Bp.blockColSize()) == 0) {
             int64_t jb_begin, jb_end;

include/fbgemm/FbgemmPackMatrixB.h

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright 2024 Arm Limited and/or its affiliates <[email protected]>
  * All rights reserved.
  *
  * This source code is licensed under the BSD-style license found in the
@@ -60,7 +61,13 @@ class PackedGemmMatrixB {
       const float alpha,
       const float* smat,
       const int brow = 512)
-      : nrow_(nrow), ncol_(ncol), brow_(brow), kernel_ncol_blocks_(2) {
+      : nrow_(nrow), ncol_(ncol), brow_(brow),
+#ifdef FBGEMM_ENABLE_KLEIDIAI
+      kernel_ncol_blocks_(1)
+#else
+      kernel_ncol_blocks_(2)
+#endif
+  {
     initializeParam();
     initializeMemory();
     // copy source matrix into packed matrix

src/FbgemmFP16.cc

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright 2024 Arm Limited and/or its affiliates <[email protected]>
  * All rights reserved.
  *
  * This source code is licensed under the BSD-style license found in the
@@ -15,6 +16,9 @@
 #include "./FbgemmFP16UKernelsAvx512.h"
 #include "./FbgemmFP16UKernelsAvx512_256.h"
 #include "./FbgemmFP16UKernelsSve128.h"
+#ifdef FBGEMM_ENABLE_KLEIDIAI
+#include "./KleidiAIFP16UKernelsNeon.h"
+#endif
 #include "fbgemm/Fbgemm.h"
 #include "fbgemm/FbgemmFPCommon.h"
 
@@ -27,12 +31,15 @@ namespace {
 // the restrictions of ymm register numbers (16).
 constexpr kernel_array_t<float16> kernel_fp16_avx2 = {
     nullptr,
+#ifndef __aarch64__
     gemmkernel_1x2_Avx2_fp16_fA0fB0fC0,
     gemmkernel_2x2_Avx2_fp16_fA0fB0fC0,
     gemmkernel_3x2_Avx2_fp16_fA0fB0fC0,
     gemmkernel_4x2_Avx2_fp16_fA0fB0fC0,
     gemmkernel_5x2_Avx2_fp16_fA0fB0fC0,
-    gemmkernel_6x2_Avx2_fp16_fA0fB0fC0};
+    gemmkernel_6x2_Avx2_fp16_fA0fB0fC0
+#endif
+};
 
 constexpr kernel_array_t<float16> kernel_fp16_sve128 = {
     nullptr,
@@ -53,8 +60,23 @@ constexpr kernel_array_t<float16> kernel_fp16_sve128 = {
 #endif
 };
 
+constexpr kernel_array_t<float16> kernel_fp16_neon = {
+  nullptr,
+#if defined(__aarch64__) && defined(FBGEMM_ENABLE_KLEIDIAI)
+    kleidiai::gemmkernel_1x1_Neon_fp16_fA0fB0fC0,
+    kleidiai::gemmkernel_2x1_Neon_fp16_fA0fB0fC0,
+    kleidiai::gemmkernel_3x1_Neon_fp16_fA0fB0fC0,
+    kleidiai::gemmkernel_4x1_Neon_fp16_fA0fB0fC0,
+    kleidiai::gemmkernel_5x1_Neon_fp16_fA0fB0fC0,
+    kleidiai::gemmkernel_6x1_Neon_fp16_fA0fB0fC0,
+    kleidiai::gemmkernel_7x1_Neon_fp16_fA0fB0fC0,
+    kleidiai::gemmkernel_8x1_Neon_fp16_fA0fB0fC0,
+#endif
+};
+
 constexpr kernel_array_t<float16> kernel_fp16_avx512_256 = {
     nullptr,
+#ifndef __aarch64__
     gemmkernel_1x2_Avx2_fp16_fA0fB0fC0,
     gemmkernel_2x2_Avx2_fp16_fA0fB0fC0,
     gemmkernel_3x2_Avx2_fp16_fA0fB0fC0,
@@ -68,7 +90,9 @@ constexpr kernel_array_t<float16> kernel_fp16_avx512_256 = {
     gemmkernel_11x2_Avx512_256_fp16_fA0fB0fC0,
     gemmkernel_12x2_Avx512_256_fp16_fA0fB0fC0,
     gemmkernel_13x2_Avx512_256_fp16_fA0fB0fC0,
-    gemmkernel_14x2_Avx512_256_fp16_fA0fB0fC0};
+    gemmkernel_14x2_Avx512_256_fp16_fA0fB0fC0
+#endif
+};
 
 constexpr kernel_array_t<float16> kernel_fp16_avx512 = {
 #ifndef __aarch64__
@@ -104,10 +128,18 @@ const isa_descriptor<float16>& getIsaHandlers(inst_set_t isa, float16) {
       std::make_tuple(kernel_fp16_avx512_256, partition_avx512);
   static isa_descriptor<float16> sve128_descriptor =
       std::make_tuple(kernel_fp16_sve128, partition_sve128);
+#ifdef FBGEMM_ENABLE_KLEIDIAI
+  static isa_descriptor<float16> neon_descriptor =
+      std::make_tuple(kernel_fp16_neon, partition_neon);
+#endif
 
   switch (isa) {
     case inst_set_t::sve:
+#ifdef FBGEMM_ENABLE_KLEIDIAI
+      return neon_descriptor;
+#else
       return sve128_descriptor;
+#endif
     case inst_set_t::anyarch:
     case inst_set_t::avx2:
       return avx2_descriptor;

src/FbgemmFPCommon.cc

@@ -1,5 +1,6 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright 2024 Arm Limited and/or its affiliates <[email protected]>
  * All rights reserved.
  *
  * This source code is licensed under the BSD-style license found in the
@@ -301,6 +302,135 @@ partition_array_t partition_sve128 = {
   }
 };
 
+partition_array_t partition_neon = {
+  // NOTE: clang-format wants to use a different formatting but the current
+  // formatting should be easier to read.
+  {
+    {{ { 0, 0 }, { 0, 0 } } }, // 0
+    {{ { 1, 1 }, { 0, 0 } } }, // 1
+    {{ { 2, 1 }, { 0, 0 } } }, // 2
+    {{ { 3, 1 }, { 0, 0 } } }, // 3
+    {{ { 4, 1 }, { 0, 0 } } }, // 4
+    {{ { 5, 1 }, { 0, 0 } } }, // 5
+    {{ { 6, 1 }, { 0, 0 } } }, // 6
+    {{ { 7, 1 }, { 0, 0 } } }, // 7
+    {{ { 8, 1 }, { 0, 0 } } }, // 8
+    {{ { 5, 1 }, { 4, 1 } } }, // 9
+    {{ { 5, 2 }, { 0, 0 } } }, // 10
+    {{ { 6, 1 }, { 5, 1 } } }, // 11
+    {{ { 6, 2 }, { 0, 0 } } }, // 12
+    {{ { 7, 1 }, { 6, 1 } } }, // 13
+    {{ { 8, 1 }, { 6, 1 } } }, // 14
+    {{ { 8, 1 }, { 7, 1 } } }, // 15
+    {{ { 8, 2 }, { 0, 0 } } }, // 16
+    {{ { 8, 2 }, { 1, 1 } } }, // 17
+    {{ { 6, 3 }, { 0, 0 } } }, // 18
+    {{ { 8, 2 }, { 3, 1 } } }, // 19
+    {{ { 5, 4 }, { 0, 0 } } }, // 20
+    {{ { 5, 3 }, { 6, 1 } } }, // 21
+    {{ { 8, 2 }, { 6, 1 } } }, // 22
+    {{ { 8, 2 }, { 7, 1 } } }, // 23
+    {{ { 8, 3 }, { 0, 0 } } }, // 24
+    {{ { 8, 3 }, { 1, 1 } } }, // 25
+    {{ { 8, 3 }, { 2, 1 } } }, // 26
+    {{ { 8, 3 }, { 3, 1 } } }, // 27
+    {{ { 8, 3 }, { 4, 1 } } }, // 28
+    {{ { 8, 3 }, { 5, 1 } } }, // 29
+    {{ { 8, 3 }, { 6, 1 } } }, // 30
+    {{ { 8, 3 }, { 7, 1 } } }, // 31
+    {{ { 8, 4 }, { 0, 0 } } }, // 32
+    {{ { 8, 4 }, { 1, 1 } } }, // 33
+    {{ { 8, 4 }, { 2, 1 } } }, // 34
+    {{ { 8, 4 }, { 3, 1 } } }, // 35
+    {{ { 8, 4 }, { 4, 1 } } }, // 36
+    {{ { 8, 4 }, { 5, 1 } } }, // 37
+    {{ { 8, 4 }, { 6, 1 } } }, // 38
+    {{ { 8, 4 }, { 7, 1 } } }, // 39
+    {{ { 8, 5 }, { 0, 0 } } }, // 40
+    {{ { 8, 5 }, { 1, 1 } } }, // 41
+    {{ { 8, 5 }, { 2, 1 } } }, // 42
+    {{ { 8, 5 }, { 3, 1 } } }, // 43
+    {{ { 8, 5 }, { 4, 1 } } }, // 44
+    {{ { 8, 5 }, { 5, 1 } } }, // 45
+    {{ { 8, 5 }, { 6, 1 } } }, // 46
+    {{ { 8, 5 }, { 7, 1 } } }, // 47
+    {{ { 8, 6 }, { 0, 0 } } }, // 48
+    {{ { 8, 6 }, { 1, 1 } } }, // 49
+    {{ { 8, 6 }, { 2, 1 } } }, // 50
+    {{ { 8, 6 }, { 3, 1 } } }, // 51
+    {{ { 8, 6 }, { 4, 1 } } }, // 52
+    {{ { 8, 6 }, { 5, 1 } } }, // 53
+    {{ { 8, 6 }, { 6, 1 } } }, // 54
+    {{ { 8, 6 }, { 7, 1 } } }, // 55
+    {{ { 8, 7 }, { 0, 0 } } }, // 56
+    {{ { 8, 7 }, { 1, 1 } } }, // 57
+    {{ { 8, 7 }, { 2, 1 } } }, // 58
+    {{ { 8, 7 }, { 3, 1 } } }, // 59
+    {{ { 8, 7 }, { 4, 1 } } }, // 60
+    {{ { 8, 7 }, { 5, 1 } } }, // 61
+    {{ { 8, 7 }, { 6, 1 } } }, // 62
+    {{ { 8, 7 }, { 7, 1 } } }, // 63
+    {{ { 8, 8 }, { 0, 0 } } }, // 64
+    {{ { 8, 8 }, { 1, 1 } } }, // 65
+    {{ { 8, 8 }, { 2, 1 } } }, // 66
+    {{ { 8, 8 }, { 3, 1 } } }, // 67
+    {{ { 8, 8 }, { 4, 1 } } }, // 68
+    {{ { 8, 8 }, { 5, 1 } } }, // 69
+    {{ { 8, 8 }, { 6, 1 } } }, // 70
+    {{ { 8, 8 }, { 7, 1 } } }, // 71
+    {{ { 8, 9 }, { 0, 0 } } }, // 72
+    {{ { 8, 9 }, { 1, 1 } } }, // 73
+    {{ { 8, 9 }, { 2, 1 } } }, // 74
+    {{ { 8, 9 }, { 3, 1 } } }, // 75
+    {{ { 8, 9 }, { 4, 1 } } }, // 76
+    {{ { 8, 9 }, { 5, 1 } } }, // 77
+    {{ { 8, 9 }, { 6, 1 } } }, // 78
+    {{ { 8, 9 }, { 7, 1 } } }, // 79
+    {{ { 8, 10 }, { 0, 0 } } }, // 80
+    {{ { 8, 10 }, { 1, 1 } } }, // 81
+    {{ { 8, 10 }, { 2, 1 } } }, // 82
+    {{ { 8, 10 }, { 3, 1 } } }, // 83
+    {{ { 8, 10 }, { 4, 1 } } }, // 84
+    {{ { 8, 10 }, { 5, 1 } } }, // 85
+    {{ { 8, 10 }, { 6, 1 } } }, // 86
+    {{ { 8, 10 }, { 7, 1 } } }, // 87
+    {{ { 8, 11 }, { 0, 0 } } }, // 88
+    {{ { 8, 11 }, { 1, 1 } } }, // 89
+    {{ { 8, 11 }, { 2, 1 } } }, // 90
+    {{ { 8, 11 }, { 3, 1 } } }, // 91
+    {{ { 8, 11 }, { 4, 1 } } }, // 92
+    {{ { 8, 11 }, { 5, 1 } } }, // 93
+    {{ { 8, 11 }, { 6, 1 } } }, // 94
+    {{ { 8, 11 }, { 7, 1 } } }, // 95
+    {{ { 8, 12 }, { 0, 0 } } }, // 96
+    {{ { 8, 12 }, { 1, 1 } } }, // 97
+    {{ { 8, 12 }, { 2, 1 } } }, // 98
+    {{ { 8, 12 }, { 3, 1 } } }, // 99
+    {{ { 8, 12 }, { 4, 1 } } }, // 100
+    {{ { 8, 12 }, { 5, 1 } } }, // 101
+    {{ { 8, 12 }, { 6, 1 } } }, // 102
+    {{ { 8, 12 }, { 7, 1 } } }, // 103
+    {{ { 8, 13 }, { 0, 0 } } }, // 104
+    {{ { 8, 13 }, { 1, 1 } } }, // 105
+    {{ { 8, 13 }, { 2, 1 } } }, // 106
+    {{ { 8, 13 }, { 3, 1 } } }, // 107
+    {{ { 8, 13 }, { 4, 1 } } }, // 108
+    {{ { 8, 13 }, { 5, 1 } } }, // 109
+    {{ { 8, 13 }, { 6, 1 } } }, // 110
+    {{ { 8, 13 }, { 7, 1 } } }, // 111
+    {{ { 8, 14 }, { 0, 0 } } }, // 112
+    {{ { 8, 14 }, { 1, 1 } } }, // 113
+    {{ { 8, 14 }, { 2, 1 } } }, // 114
+    {{ { 8, 14 }, { 3, 1 } } }, // 115
+    {{ { 8, 14 }, { 4, 1 } } }, // 116
+    {{ { 8, 14 }, { 5, 1 } } }, // 117
+    {{ { 8, 14 }, { 6, 1 } } }, // 118
+    {{ { 8, 14 }, { 7, 1 } } }, // 119
+    {{ { 8, 15 }, { 0, 0 } } }, // 120
+  }
+};
+
+
 partition_array_t partition_avx512 = {
   // NOTE: clang-format wants to use a different formatting but the current
   // formatting should be easier to read.