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;
#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<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,
(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() {