feat: prepared benchmark for batch-reduce gemm

Author: Integer-Ctrl

CMakeLists.txt

@@ -151,6 +151,7 @@ set(TEST_ARM_INSTRUCTION_FILES
 set(BENCH_FILES
     matmul_16_6_1.bench.cpp
     matmul_16_6_k.bench.cpp
+    matmul.bench.cpp
 )
 
 foreach(file ${SRC_MAIN_FILES})

src/test/kernels/matmul.bench.cpp

@@ -0,0 +1,73 @@
+#include "matmul.bench.h"
+#include "../../main/Brgemm.h"
+#include <benchmark/benchmark.h>
+
+class GemmFixture : public benchmark::Fixture
+{
+public:
+  std::vector<float> matrix_a, matrix_b, matrix_c;
+  double flops;
+
+  void SetUp(::benchmark::State &state) override
+  {
+    flops = 0;
+
+    int M = state.range(0);
+    int N = state.range(1);
+    int K = state.range(2);
+
+    matrix_a.resize(M * K);
+    matrix_b.resize(K * N);
+    matrix_c.resize(M * N);
+
+    fill_random_matrix_args(matrix_a.data(), M * K);
+    fill_random_matrix_args(matrix_b.data(), K * N);
+    fill_random_matrix_args(matrix_c.data(), M * N);
+  }
+
+  void TearDown(::benchmark::State &state) override
+  {
+    state.counters["FLOPS"] = benchmark::Counter(flops, benchmark::Counter::kIsRate);
+  }
+};
+
+BENCHMARK_DEFINE_F(GemmFixture, BM_matmul)(benchmark::State &state)
+{
+  int M = state.range(0);
+  int N = state.range(1);
+  int K = state.range(2);
+
+  mini_jit::Brgemm brgemm;
+  brgemm.generate(M, N, K, 1, 0, 0, 0, mini_jit::Brgemm::dtype_t::fp32);
+  auto kernel = brgemm.get_kernel();
+
+  for (auto _ : state)
+  {
+    kernel(matrix_a.data(), matrix_b.data(), matrix_c.data(), M, 1, M, 1, 1);
+  }
+
+  flops = M * N * K * 2 * state.iterations();
+}
+
+static void CustomArguments(benchmark::internal::Benchmark *b)
+{
+  for (int M = 16; M <= 64; M += 16)
+    for (int N = 16; N <= 64; N += 16)
+      for (int K : {1, 16, 32, 64, 128})
+        b->Args({M, N, K});
+}
+
+// ########## UNCOMMENT WHEN brgemm.generate() supports m, n < 16 ##########
+// static void CustomArguments(benchmark::internal::Benchmark *b)
+// {
+//   for (int M = 1; M <= 64; M += 1)
+//     for (int N = 1; N <= 64; N += 1)
+//       for (int K : {1, 16, 32, 64, 128})
+//         b->Args({M, N, K});
+// }
+
+BENCHMARK_REGISTER_F(GemmFixture, BM_matmul)
+  ->ArgNames({"M", "N", "K"})
+  ->ReportAggregatesOnly(true)
+  ->Apply(CustomArguments)
+  ->MinWarmUpTime(1.0);  // WarmUp in seconds

src/test/kernels/matmul.bench.h

@@ -1,3 +1,4 @@
+#include <cstdint>
 #include <cstdlib>
 #include <ctime>
 
@@ -14,4 +15,25 @@ template <uint32_t TSize> void fill_random_matrix(float (&matrix)[TSize])
   {
     matrix[i] = (static_cast<float>(std::rand())) / (static_cast<float>(std::rand()));
   }
+}
+
+/**
+ * @brief Fill the given matrix with random values.
+ *
+ * @param matrix The matrix to write to.
+ * @param matrix_size The number of elements in the matrix.
+ */
+inline void fill_random_matrix_args(float *matrix, size_t matrix_size)
+{
+  static bool initialized = false;
+  if (!initialized)
+  {
+    std::srand(std::time(0));  // Seed RNG once
+    initialized = true;
+  }
+
+  for (size_t i = 0; i < matrix_size; i++)
+  {
+    matrix[i] = (static_cast<float>(std::rand())) / (static_cast<float>(std::rand()));
+  }
 }