sgemm_wg: Run multiple threadblock per core

tests/regression/sgemm_wg/kernel.cpp

@@ -9,7 +9,9 @@ inline void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
                               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 uint32_t threadblock_id_y,
+                              const uint32_t threadblock_id_in_core,
+                              float *sharedmem_per_threadblock) {
   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;
@@ -29,19 +31,24 @@ inline void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
   float reg_c = 0.0f;
 
   for (uint32_t k = 0; k < dim_k; k += threadblock_dim_x) {
-    float *local_a = (float *)DEV_SMEM_START_ADDR;
+    float *local_a = sharedmem_per_threadblock;
     size_t local_a_elems = threadblock_dim_x * threadblock_dim_y;
-    float *local_b = (float *)DEV_SMEM_START_ADDR + local_a_elems;
+    float *local_b = sharedmem_per_threadblock + local_a_elems;
 
     uint32_t offset_global_a = dim_k * global_row + (k + local_col);
     uint32_t offset_global_b = dim_n * (local_row + k) + global_col;
+    // FIXME: threadblocks size must be BM*BN, not BM*BK or BN*BK. This means
+    // there is a mismatch between the number of elements in the A/B tile and
+    // the C tile.  This is handled by each thread computing multiple result
+    // elements.
+    //
     // 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, threadblock_dim_y);
+    vx_barrier(threadblock_id_in_core, threadblock_dim_y);
     vx_fence();
 
     for (uint32_t local_k = 0; local_k < threadblock_dim_x; local_k++) {
@@ -49,7 +56,7 @@ inline void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
                local_b[threadblock_dim_y * local_col + local_k];
     }
 
-    vx_barrier(0, threadblock_dim_y);
+    vx_barrier(threadblock_id_in_core, threadblock_dim_y);
     vx_fence();
   }
 
@@ -57,14 +64,19 @@ inline void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
 }
 
 void kernel_body(int task_id, kernel_arg_t* __UNIFORM__ arg) {
+  // @perf: All threads are running these compute whose result is mostly same
+  // across the threadblock
+
   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 threadblocks_per_core = 2;
   const uint32_t threadblock_dim_x = vx_num_threads();
-  const uint32_t threadblock_dim_y = vx_num_warps();
+  const uint32_t threadblock_dim_y = vx_num_warps() / threadblocks_per_core;
   const uint32_t threads_per_threadblock = threadblock_dim_x * threadblock_dim_y;
   const int threadblock_id = task_id / threads_per_threadblock;
+  const int threadblock_id_in_core = threadblock_id % threadblocks_per_core;
 
   const uint32_t dim_n_in_blocks = dim_n / threadblock_dim_x;
   const int threadblock_id_x = threadblock_id % dim_n_in_blocks;
@@ -72,8 +84,14 @@ void kernel_body(int task_id, kernel_arg_t* __UNIFORM__ arg) {
 
   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);
+  float *sharedmem_per_threadblock =
+      (float *)DEV_SMEM_START_ADDR +
+      (2 * threadblock_dim_x * threadblock_dim_y) * threadblock_id_in_core;
+  thread_block_gemm(arg, tid_in_threadblock_x, tid_in_threadblock_y,
+                    threadblock_dim_x, threadblock_dim_y, threadblock_id_x,
+                    threadblock_id_y, threadblock_id_in_core,
+                    sharedmem_per_threadblock);
 }
 
 int main() {

tests/regression/sgemm_wg/main.cpp

@@ -147,8 +147,8 @@ int main(int argc, char *argv[]) {
   RT_CHECK(vx_dev_open(&device));
 
   // FIXME: hardcoded
-  uint32_t dim_m = 16;
+  uint32_t dim_m = 32;
-  uint32_t dim_n = 16;
+  uint32_t dim_n = 32;
   uint32_t dim_k = 32;
 
   generate_source_matrix(dim_m, dim_n, dim_k);