sgemm_wg: Parameterize threadblock dimensions

tests/regression/sgemm_wg/kernel.cpp

@@ -3,7 +3,13 @@
 #include <vx_spawn.h>
 #include "common.h"
 
-void kernel_body(int task_id, kernel_arg_t* __UNIFORM__ arg) {
+inline void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
+                              const uint32_t tid_in_threadblock_x,
+                              const uint32_t tid_in_threadblock_y,
+                              const uint32_t threadblock_dim_x,
+                              const uint32_t threadblock_dim_y,
+                              const uint32_t threadblock_id_x,
+                              const uint32_t threadblock_id_y) {
   const float *global_a = (const float *)arg->addr_a;
   const float *global_b = (const float *)arg->addr_b;
   float *global_c = (float *)arg->addr_c;
@@ -12,42 +18,67 @@ void kernel_body(int task_id, kernel_arg_t* __UNIFORM__ arg) {
   const uint32_t dim_m = arg->dim_m;
   const uint32_t dim_n = arg->dim_n;
   const uint32_t dim_k = arg->dim_k;
-  const uint32_t block_dim = vx_num_warps();
-  const uint32_t local_row = vx_warp_id();
-  const uint32_t local_col = vx_thread_id();
+
+  // FIXME: assumes local block size is square shape
+  const uint32_t local_row = tid_in_threadblock_y;
+  const uint32_t local_col = tid_in_threadblock_x;
+  const uint32_t global_row = threadblock_id_y * threadblock_dim_y + local_row;
+  const uint32_t global_col = threadblock_id_x * threadblock_dim_x + local_col;
 
   // each thread generates one output element
   float reg_c = 0.0f;
 
-  for (uint32_t k = 0; k < dim_k; k += block_dim) {
+  for (uint32_t k = 0; k < dim_k; k += threadblock_dim_x) {
     float *local_a = (float *)DEV_SMEM_START_ADDR;
-    float *local_b = (float *)DEV_SMEM_START_ADDR + (block_dim * block_dim);
+    size_t local_a_elems = threadblock_dim_x * threadblock_dim_y;
+    float *local_b = (float *)DEV_SMEM_START_ADDR + local_a_elems;
 
-    // FIXME: assumes local block size is square shape
-    // TODO: "local_row" should be global_row
-    uint32_t offset_global_a = dim_k * local_row + (k + local_col);
-    uint32_t offset_global_b = dim_n * (local_row + k) + local_col;
-    local_a[block_dim * local_row + local_col] = global_a[offset_global_a];
-    local_b[block_dim * local_row + local_col] = global_b[offset_global_b];
+    uint32_t offset_global_a = dim_k * global_row + (k + local_col);
+    uint32_t offset_global_b = dim_n * (local_row + k) + global_col;
+    // local_a: threadblock_dim_y rows, threadblock_dim_x cols
+    // local_b: threadblock_dim_x rows, threadblock_dim_y cols
+    // threadblock_dim_x == block_k, threadblock_dim_y == block_m == block_n
+    local_a[threadblock_dim_x * local_row + local_col] = global_a[offset_global_a];
+    local_b[threadblock_dim_y * local_col + local_row] = global_b[offset_global_b];
 
-    vx_barrier(0, vx_num_warps());
+    vx_barrier(0, threadblock_dim_y);
     vx_fence();
 
-    for (uint32_t local_k = 0; local_k < block_dim; local_k++) {
-      reg_c += local_a[block_dim * local_row + local_k] *
-               local_b[block_dim * local_k + local_col];
+    for (uint32_t local_k = 0; local_k < threadblock_dim_x; local_k++) {
+      reg_c += local_a[threadblock_dim_x * local_row + local_k] *
+               local_b[threadblock_dim_y * local_col + local_k];
     }
 
-    vx_barrier(0, vx_num_warps());
+    vx_barrier(0, threadblock_dim_y);
     vx_fence();
   }
 
-  global_c[dim_n * local_row + local_col] = reg_c;
+  global_c[dim_n * global_row + global_col] = reg_c;
+}
+
+void kernel_body(int task_id, kernel_arg_t* __UNIFORM__ arg) {
+  const uint32_t dim_n = arg->dim_n;
+  int tid_x = task_id % dim_n;
+  int tid_y = task_id / dim_n;
+
+  const uint32_t threadblock_dim_x = vx_num_threads();
+  const uint32_t threadblock_dim_y = vx_num_warps();
+  const uint32_t threads_per_threadblock = threadblock_dim_x * threadblock_dim_y;
+  const int threadblock_id = task_id / threads_per_threadblock;
+
+  const uint32_t dim_n_in_blocks = dim_n / threadblock_dim_x;
+  const int threadblock_id_x = threadblock_id % dim_n_in_blocks;
+  const int threadblock_id_y = threadblock_id / dim_n_in_blocks;
+
+  const int tid_in_threadblock_x = vx_thread_id();
+  const int tid_in_threadblock_y = vx_warp_id() % threadblock_dim_y;
+  thread_block_gemm(arg, tid_in_threadblock_x, tid_in_threadblock_y, threadblock_dim_x,
+                    threadblock_dim_y, threadblock_id_x, threadblock_id_y);
 }
 
 int main() {
   kernel_arg_t *arg = (kernel_arg_t *)KERNEL_ARG_DEV_MEM_ADDR;
-  int threads_per_core = vx_num_warps() * vx_num_threads();
-  vx_spawn_tasks(threads_per_core, (vx_spawn_tasks_cb)kernel_body, arg);
+  const uint32_t grid_size = arg->dim_m * arg->dim_n;
+  vx_spawn_tasks(grid_size, (vx_spawn_tasks_cb)kernel_body, arg);
   return 0;
 }

tests/regression/sgemm_wg/main.cpp

@@ -116,7 +116,7 @@ int run_test(const kernel_arg_t& kernel_arg,
     for (uint32_t i = 0; i < dim_m * dim_n; ++i) {
       float ref = ref_data.at(i);
       float cur = buf_ptr[i];
-      if (cur != ref) {
+      if (std::abs((cur - ref) / ref) > 1e-6) {
         std::cout << "error at result #" << std::dec << i
                   << std::hex << ": actual=" << cur << ", expected=" << ref << std::endl;
         ++errors;
@@ -146,9 +146,10 @@ int main(int argc, char *argv[]) {
   std::cout << "open device connection" << std::endl;
   RT_CHECK(vx_dev_open(&device));
 
-  uint32_t dim_m = 4; // FIXME: hardcoded
-  uint32_t dim_n = 4; // FIXME: hardcoded
-  uint32_t dim_k = 128; // FIXME: hardcoded
+  // FIXME: hardcoded
+  uint32_t dim_m = 16;
+  uint32_t dim_n = 16;
+  uint32_t dim_k = 32;
 
   generate_source_matrix(dim_m, dim_n, dim_k);
   generate_reference_matmul(dim_m, dim_n, dim_k);