wu-arch/kernels
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

sgemm_tcore: Fix invocation with compile time threadblock size

Browse Source
This commit is contained in:
Hansung Kim
2024-09-02 17:03:46 -07:00
parent 70273fd00d
commit 9d71fa44a7
2 changed files with 29 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
tests/regression/sgemm_tcore/kernel.cpp
View File

@@ -17,16 +17,20 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
constexpr uint32_t cores_per_cluster = 1; constexpr uint32_t cores_per_cluster = 1;
#endif #endif
uint32_t threads_per_threadblock = (BM * BN) / (ELEM_PER_THREAD); constexpr uint32_t threads_per_threadblock_theoretical =
const uint32_t hw_threads_per_cluster = (BM * BN) / (ELEM_PER_THREAD);
cores_per_cluster * vx_num_threads() * vx_num_warps(); 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 // cap maximum threadblock size to # of HW threads in cluster, to prevent
// multiple "wave" invocations which slows down the kernel // multiple "wave" invocations which slows down the kernel
if (threads_per_threadblock > hw_threads_per_cluster) { constexpr uint32_t threads_per_threadblock =
threads_per_threadblock = hw_threads_per_cluster; (threads_per_threadblock_theoretical > hw_threads_per_cluster)
} ? hw_threads_per_cluster
const uint32_t threadblocks_per_cluster = : threads_per_threadblock_theoretical;
constexpr uint32_t threadblocks_per_cluster =
hw_threads_per_cluster / threads_per_threadblock; 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 = task_id / threads_per_threadblock;
const int threadblock_id_in_cluster = 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*/ * DEV_SMEM_START_ADDR + sizeof(float_type) * 2 /*overkill for non-dma*/ *
(2 * BM * BK) * threadblock_id_in_cluster); (2 * BM * BK) * threadblock_id_in_cluster);
thread_block_gemm<float_type, /*write_to_gmem=*/true>( thread_block_gemm<float_type, threads_per_threadblock,
/*write_to_gmem=*/true>(
(const float_type *)arg->addr_a, (const float_type *)arg->addr_b, (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, (float *)arg->addr_c, arg->dim_m, arg->dim_n, arg->dim_k,
tid_in_threadblock, threads_per_threadblock, threadblocks_per_cluster, tid_in_threadblock, threadblocks_per_cluster, threadblock_id_in_cluster,
threadblock_id_in_cluster, sharedmem_per_threadblock); sharedmem_per_threadblock);
} }
int main() { int main() {

30
tests/regression/sgemm_tcore/sgemm_impl.hpp
View File

@@ -736,29 +736,27 @@ __attribute__((always_inline)) inline void thread_block_gemm_single_tile(
} }
} }
template <typename T, bool write_to_gmem = true, template <
// by default, A/B tiles are placed at the start of the smem typename T, uint32_t threads_per_threadblock, bool write_to_gmem = true,
uint32_t smem_a_offset = 0, // byte offset of A tile in shared // by default, A/B tiles are placed at the start of the smem
// memory uint32_t smem_a_offset = 0, // byte offset of A tile in shared
uint32_t smem_a_dbuf_offset = 0, // byte offset of A // memory
// double-buffer tile in shared uint32_t smem_a_dbuf_offset = 0, // byte offset of A
// memory // double-buffer tile in shared
uint32_t smem_b_offset = sizeof(float) * BM * // memory
BK, // byte offset of B tile uint32_t smem_b_offset = sizeof(float) * BM * BK, // byte offset of B tile
// in shared memory // in shared memory
uint32_t smem_b_dbuf_offset = sizeof(float) * BM * uint32_t smem_b_dbuf_offset = sizeof(float) * BM *
BK // byte offset of B double-buffer BK // byte offset of B double-buffer
// tile in shared memory // tile in shared memory
> >
inline void thread_block_gemm(const T *A, const T *B, float *C, 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_m, const uint32_t dim_n,
const uint32_t dim_k, const uint32_t dim_k,
const uint32_t tid_in_threadblock, const uint32_t tid_in_threadblock,
const uint32_t threads_per_threadblock,
const uint32_t threadblocks_per_cluster, const uint32_t threadblocks_per_cluster,
const uint32_t threadblock_id_in_cluster, const uint32_t threadblock_id_in_cluster,
uint8_t *sharedmem_per_threadblock) { uint8_t *sharedmem_per_threadblock) {
// no double-buffering
const uint32_t threads_per_warpgroup = threads_per_threadblock; 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_id_in_warpgroup = tid_in_threadblock / NUM_THREADS;
const uint32_t warp_row = warp_id_in_warpgroup / (BN / WN); const uint32_t warp_row = warp_id_in_warpgroup / (BN / WN);