sgemm_wg: write simple C=A*A matmul

tests/regression/sgemm_wg/common.h

@@ -1,13 +1,15 @@
 #ifndef _COMMON_H_
 #define _COMMON_H_
 
+#include <cstdint>
+
 #define KERNEL_ARG_DEV_MEM_ADDR 0x7fff0000
 #define DEV_SMEM_START_ADDR 0xff000000
 
 typedef struct {
-  uint32_t num_points;
+  uint32_t matrix_dim;
-  uint64_t src_addr;
+  uint64_t addr_a;
-  uint64_t dst_addr;  
+  uint64_t addr_c;
 } kernel_arg_t;
 
 #endif

tests/regression/sgemm_wg/kernel.cpp

@@ -4,18 +4,30 @@
 #include "common.h"
 
 void kernel_body(int task_id, kernel_arg_t* __UNIFORM__ arg) {
-  uint32_t num_points = arg->num_points;
+  const float *global_a = (const float *)arg->addr_a;
-  float *src_ptr = (float *)arg->src_addr;
+  float *global_c = (float *)arg->addr_c;
-  float *dst_ptr = (float *)arg->dst_addr;
 
-  float *local_a = (float *)DEV_SMEM_START_ADDR;
+  // assumes NT == NW == matrix_dim
+  const uint32_t dim = arg->matrix_dim;
+  const uint32_t row = vx_warp_id();
+  const uint32_t col = vx_thread_id();
 
-  local_a[num_points - 1 - task_id] = 2 * src_ptr[num_points - 1 - task_id];
+  float *local_c = (float *)DEV_SMEM_START_ADDR;
-  // local_a[task_id] = 2 * src_ptr[task_id];
+  float *local_a = (float *)DEV_SMEM_START_ADDR + (dim * dim);
+  float *local_b = (float *)DEV_SMEM_START_ADDR + 2 * (dim * dim);
+
+  local_a[dim * row + col] = global_a[dim * row + col];
+  local_c[dim * row + col] = 0.0f;
 
   vx_barrier(0, vx_num_warps());
 
-  dst_ptr[task_id] = local_a[task_id];
+  for (uint32_t k = 0; k < dim; k++) {
+    local_c[dim * row + col] += local_a[dim * row + k] * local_a[dim * k + col];
+  }
+
+  vx_barrier(0, vx_num_warps());
+
+  global_c[dim * row + col] = local_c[dim * row + col];
 }
 
 int main() {

tests/regression/sgemm_wg/main.cpp

@@ -57,8 +57,8 @@ static void parse_args(int argc, char **argv) {
 
 void cleanup() {
   if (device) {
-    vx_mem_free(device, kernel_arg.src_addr);
+    vx_mem_free(device, kernel_arg.addr_a);
-    vx_mem_free(device, kernel_arg.dst_addr);
+    vx_mem_free(device, kernel_arg.addr_c);
     vx_dev_close(device);
   }
 }
@@ -76,7 +76,7 @@ void gen_ref_data(uint32_t len) {
   ref_data.resize(len);
 
   for (uint32_t i = 0; i < len; ++i) {
-    float ref_value = 2 * src_data.at(i);
+    float ref_value = src_data.at(i);
     ref_data.at(i) = ref_value;
   }
 }
@@ -94,7 +94,7 @@ int run_test(const kernel_arg_t& kernel_arg,
 
   // download destination buffer
   std::cout << "download destination buffer" << std::endl;
-  RT_CHECK(vx_copy_from_dev(device, staging_buf.data(), kernel_arg.dst_addr, buf_size));
+  RT_CHECK(vx_copy_from_dev(device, staging_buf.data(), kernel_arg.addr_c, buf_size));
 
   // verify result
   std::cout << "verify result" << std::endl;
@@ -134,18 +134,18 @@ int main(int argc, char *argv[]) {
   std::cout << "open device connection" << std::endl;
   RT_CHECK(vx_dev_open(&device));
 
-  uint32_t num_points = count;
+  uint32_t matrix_size = count;
 
   // generate input data
-  gen_input_data(num_points);
+  gen_input_data(matrix_size);
 
   // generate reference data
-  gen_ref_data(num_points);
+  gen_ref_data(matrix_size);
 
   uint32_t src_buf_size = src_data.size() * sizeof(src_data[0]);
   uint32_t dst_buf_size = ref_data.size() * sizeof(src_data[0]);
 
-  std::cout << "number of points: " << num_points << std::endl;
+  std::cout << "number of elements: " << matrix_size << std::endl;
   std::cout << "buffer size: " << dst_buf_size << " bytes" << std::endl;
 
   // upload program
@@ -154,13 +154,13 @@ int main(int argc, char *argv[]) {
 
   // allocate device memory
   std::cout << "allocate device memory" << std::endl;
-  RT_CHECK(vx_mem_alloc(device, src_buf_size, VX_MEM_TYPE_GLOBAL, &kernel_arg.src_addr));
+  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.dst_addr));
+  RT_CHECK(vx_mem_alloc(device, dst_buf_size, VX_MEM_TYPE_GLOBAL, &kernel_arg.addr_c));
 
-  kernel_arg.num_points = num_points;
+  kernel_arg.matrix_dim = 4; // FIXME: hardcoded
 
-  std::cout << "dev_src=0x" << std::hex << kernel_arg.src_addr << std::endl;
+  std::cout << "dev_src=0x" << std::hex << kernel_arg.addr_a << std::endl;
-  std::cout << "dev_dst=0x" << std::hex << kernel_arg.dst_addr << std::endl;
+  std::cout << "dev_dst=0x" << std::hex << kernel_arg.addr_c << std::endl;
 
   // allocate staging buffer
   {
@@ -195,11 +195,11 @@ int main(int argc, char *argv[]) {
   {
     std::cout << "upload source buffer" << std::endl;
     auto buf_ptr = staging_buf.data();
-    memcpy(buf_ptr, src_data.data(), num_points * sizeof(float));      
+    memcpy(buf_ptr, src_data.data(), matrix_size * sizeof(float));
-    RT_CHECK(vx_copy_to_dev(device, kernel_arg.src_addr, staging_buf.data(), src_buf_size));
+    RT_CHECK(vx_copy_to_dev(device, kernel_arg.addr_a, staging_buf.data(), src_buf_size));
 
     std::cout << "uploading source buffer to device, device mem address="
-              << std::hex << kernel_arg.src_addr << ", size=" << std::dec
+              << std::hex << kernel_arg.addr_a << ", size=" << std::dec
               << src_buf_size << " bytes\n";
     std::ofstream file("input.bin", std::ios::binary | std::ios::out);
     if (!file) {
@@ -214,15 +214,15 @@ int main(int argc, char *argv[]) {
   {
     std::cout << "clear destination buffer" << std::endl;
     auto buf_ptr = (int32_t*)staging_buf.data();
-    for (uint32_t i = 0; i < num_points; ++i) {
+    for (uint32_t i = 0; i < matrix_size; ++i) {
       buf_ptr[i] = 0xdeadbeef;
     }
-    RT_CHECK(vx_copy_to_dev(device, kernel_arg.dst_addr, staging_buf.data(), dst_buf_size));  
+    RT_CHECK(vx_copy_to_dev(device, kernel_arg.addr_c, staging_buf.data(), dst_buf_size));
   }
 
   // run tests
   std::cout << "run tests" << std::endl;
-  RT_CHECK(run_test(kernel_arg, dst_buf_size, num_points));
+  RT_CHECK(run_test(kernel_arg, dst_buf_size, matrix_size));
 
   // cleanup
   std::cout << "cleanup" << std::endl;