From f1e7407d3a689d943316af08c5eff28f345a7e8f Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Tue, 27 Feb 2024 15:44:04 -0800 Subject: [PATCH] sgemm_wg: Run multiple threadblock per core --- tests/regression/sgemm_wg/kernel.cpp | 34 +++++++++++++++++++++------- tests/regression/sgemm_wg/main.cpp | 4 ++-- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/tests/regression/sgemm_wg/kernel.cpp b/tests/regression/sgemm_wg/kernel.cpp index ec207821..4609b9e6 100644 --- a/tests/regression/sgemm_wg/kernel.cpp +++ b/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() { diff --git a/tests/regression/sgemm_wg/main.cpp b/tests/regression/sgemm_wg/main.cpp index c6252991..229463ef 100644 --- a/tests/regression/sgemm_wg/main.cpp +++ b/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_n = 16; + uint32_t dim_m = 32; + uint32_t dim_n = 32; uint32_t dim_k = 32; generate_source_matrix(dim_m, dim_n, dim_k);