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() {

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,
// 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
>
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);