From 9d71fa44a7579a7756eecd3fbe017dc0611b89cf Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Mon, 2 Sep 2024 17:03:46 -0700 Subject: [PATCH] sgemm_tcore: Fix invocation with compile time threadblock size --- tests/regression/sgemm_tcore/kernel.cpp | 25 ++++++++++------- tests/regression/sgemm_tcore/sgemm_impl.hpp | 30 ++++++++++----------- 2 files changed, 29 insertions(+), 26 deletions(-) diff --git a/tests/regression/sgemm_tcore/kernel.cpp b/tests/regression/sgemm_tcore/kernel.cpp index 10857722..4f115aa6 100644 --- a/tests/regression/sgemm_tcore/kernel.cpp +++ b/tests/regression/sgemm_tcore/kernel.cpp @@ -17,16 +17,20 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { constexpr uint32_t cores_per_cluster = 1; #endif - uint32_t threads_per_threadblock = (BM * BN) / (ELEM_PER_THREAD); - const uint32_t hw_threads_per_cluster = - cores_per_cluster * vx_num_threads() * vx_num_warps(); + constexpr uint32_t threads_per_threadblock_theoretical = + (BM * BN) / (ELEM_PER_THREAD); + constexpr uint32_t hw_threads_per_cluster = + CORES_PER_CLUSTER * NUM_THREADS * NUM_WARPS; // cap maximum threadblock size to # of HW threads in cluster, to prevent // multiple "wave" invocations which slows down the kernel - if (threads_per_threadblock > hw_threads_per_cluster) { - threads_per_threadblock = hw_threads_per_cluster; - } - const uint32_t threadblocks_per_cluster = + constexpr uint32_t threads_per_threadblock = + (threads_per_threadblock_theoretical > hw_threads_per_cluster) + ? hw_threads_per_cluster + : threads_per_threadblock_theoretical; + constexpr uint32_t threadblocks_per_cluster = hw_threads_per_cluster / threads_per_threadblock; + constexpr uint32_t warps_per_threadblock_per_core = + NUM_WARPS / threadblocks_per_cluster; const int threadblock_id = task_id / threads_per_threadblock; const int threadblock_id_in_cluster = @@ -47,11 +51,12 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { DEV_SMEM_START_ADDR + sizeof(float_type) * 2 /*overkill for non-dma*/ * (2 * BM * BK) * threadblock_id_in_cluster); - thread_block_gemm( + thread_block_gemm( (const float_type *)arg->addr_a, (const float_type *)arg->addr_b, (float *)arg->addr_c, arg->dim_m, arg->dim_n, arg->dim_k, - tid_in_threadblock, threads_per_threadblock, threadblocks_per_cluster, - threadblock_id_in_cluster, sharedmem_per_threadblock); + tid_in_threadblock, threadblocks_per_cluster, threadblock_id_in_cluster, + sharedmem_per_threadblock); } int main() { diff --git a/tests/regression/sgemm_tcore/sgemm_impl.hpp b/tests/regression/sgemm_tcore/sgemm_impl.hpp index 13818d43..785a538c 100644 --- a/tests/regression/sgemm_tcore/sgemm_impl.hpp +++ b/tests/regression/sgemm_tcore/sgemm_impl.hpp @@ -736,29 +736,27 @@ __attribute__((always_inline)) inline void thread_block_gemm_single_tile( } } -template +template < + typename T, uint32_t threads_per_threadblock, bool write_to_gmem = true, + // by default, A/B tiles are placed at the start of the smem + uint32_t smem_a_offset = 0, // byte offset of A tile in shared + // memory + uint32_t smem_a_dbuf_offset = 0, // byte offset of A + // double-buffer tile in shared + // memory + uint32_t smem_b_offset = sizeof(float) * BM * BK, // byte offset of B tile + // in shared memory + uint32_t smem_b_dbuf_offset = sizeof(float) * BM * + BK // byte offset of B double-buffer + // tile in shared memory + > inline void thread_block_gemm(const T *A, const T *B, float *C, const uint32_t dim_m, const uint32_t dim_n, const uint32_t dim_k, const uint32_t tid_in_threadblock, - const uint32_t threads_per_threadblock, const uint32_t threadblocks_per_cluster, const uint32_t threadblock_id_in_cluster, uint8_t *sharedmem_per_threadblock) { - // no double-buffering const uint32_t threads_per_warpgroup = threads_per_threadblock; const uint32_t warp_id_in_warpgroup = tid_in_threadblock / NUM_THREADS; const uint32_t warp_row = warp_id_in_warpgroup / (BN / WN);