From 5f79e8a3f1fdbf6c78248094bbce3bb836c9d4c4 Mon Sep 17 00:00:00 2001
From: Hansung Kim <hansung_kim@berkeley.edu>
Date: Mon, 12 Feb 2024 22:29:38 -0800
Subject: [PATCH] sgemm_wg: reference matmul in cpu

---
 tests/regression/sgemm_wg/main.cpp | 37 ++++++++++++++++--------------
 1 file changed, 20 insertions(+), 17 deletions(-)

diff --git a/tests/regression/sgemm_wg/main.cpp b/tests/regression/sgemm_wg/main.cpp
index 9252b7b0..d12216e4 100644
--- a/tests/regression/sgemm_wg/main.cpp
+++ b/tests/regression/sgemm_wg/main.cpp
@@ -63,27 +63,32 @@ void cleanup() {
   }
 }
 
-void gen_input_data(uint32_t len) {
-  src_data.resize(len);
+void generate_source_matrix(uint32_t dim) {
+  src_data.resize(dim * dim);
 
-  for (uint32_t i = 0; i < len; ++i) {
+  for (uint32_t i = 0; i < dim * dim; ++i) {
     src_data[i] = static_cast<float>(i);
     std::cout << i << ": value=" << src_data[i] << std::endl;
   }
 }
 
-void gen_ref_data(uint32_t len) {
-  ref_data.resize(len);
+void generate_reference_matmul(uint32_t dim) {
+  ref_data.resize(dim * dim);
 
-  for (uint32_t i = 0; i < len; ++i) {
-    float ref_value = src_data.at(i);
-    ref_data.at(i) = ref_value;
+  for (uint32_t i = 0; i < dim; ++i) {
+    for (uint32_t j = 0; j < dim; ++j) {
+      float ref = 0.0f;
+      for (uint32_t k = 0; k < dim; ++k) {
+        ref += src_data[dim * i + k] * src_data[dim * k + j];
+      }
+      ref_data.at(dim * i + j) = ref;
+    }
   }
 }
 
 int run_test(const kernel_arg_t& kernel_arg,
              uint32_t buf_size,
-             uint32_t num_points) {
+             uint32_t dim) {
   // start device
   std::cout << "start device" << std::endl;
   RT_CHECK(vx_start(device));
@@ -101,7 +106,7 @@ int run_test(const kernel_arg_t& kernel_arg,
   {
     int errors = 0;
     auto buf_ptr = (float*)staging_buf.data();
-    for (uint32_t i = 0; i < num_points; ++i) {
+    for (uint32_t i = 0; i < dim * dim; ++i) {
       float ref = ref_data.at(i);
       float cur = buf_ptr[i];
       if (cur != ref) {
@@ -135,12 +140,10 @@ int main(int argc, char *argv[]) {
   RT_CHECK(vx_dev_open(&device));
 
   uint32_t matrix_size = count;
+  uint32_t matrix_dim = 4; // FIXME: hardcoded
 
-  // generate input data
-  gen_input_data(matrix_size);
-
-  // generate reference data
-  gen_ref_data(matrix_size);
+  generate_source_matrix(matrix_dim);
+  generate_reference_matmul(matrix_dim);
 
   uint32_t src_buf_size = src_data.size() * sizeof(src_data[0]);
   uint32_t dst_buf_size = ref_data.size() * sizeof(src_data[0]);
@@ -157,7 +160,7 @@ int main(int argc, char *argv[]) {
   RT_CHECK(vx_mem_alloc(device, src_buf_size, VX_MEM_TYPE_GLOBAL, &kernel_arg.addr_a));
   RT_CHECK(vx_mem_alloc(device, dst_buf_size, VX_MEM_TYPE_GLOBAL, &kernel_arg.addr_c));
 
-  kernel_arg.matrix_dim = 4; // FIXME: hardcoded
+  kernel_arg.matrix_dim = matrix_dim;
 
   std::cout << "dev_src=0x" << std::hex << kernel_arg.addr_a << std::endl;
   std::cout << "dev_dst=0x" << std::hex << kernel_arg.addr_c << std::endl;
@@ -222,7 +225,7 @@ int main(int argc, char *argv[]) {
 
   // run tests
   std::cout << "run tests" << std::endl;
-  RT_CHECK(run_test(kernel_arg, dst_buf_size, matrix_size));
+  RT_CHECK(run_test(kernel_arg, dst_buf_size, kernel_arg.matrix_dim));
 
   // cleanup
   std::cout << "cleanup" << std::endl;