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

sgemm_wg: 128x128 config

Browse Source
This commit is contained in:
Hansung Kim
2024-04-24 21:08:31 -07:00
parent 689043b45e
commit 793779aa6c
2 changed files with 10 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/regression/sgemm_wg/kernel.cpp
View File

@@ -16,11 +16,11 @@
// (BM*BN) / (TM*TN) == threadblock size >= NT * CORES_PER_CLUSTER // (BM*BN) / (TM*TN) == threadblock size >= NT * CORES_PER_CLUSTER
// * Combining BM * BK >= (BM*BN) / (TM*TN) == threadblock yields // * Combining BM * BK >= (BM*BN) / (TM*TN) == threadblock yields
// BM <= BK*TM*TN // BM <= BK*TM*TN
#define BM 8 #define BM 32
#define BN BM #define BN BM
#define BK 2 #define BK 8
#define TM 2 #define TM 4
#define TN 2 #define TN 4
void threadblock_barrier(unsigned int tid_in_threadblock, unsigned int barrier_id, unsigned int count) { void threadblock_barrier(unsigned int tid_in_threadblock, unsigned int barrier_id, unsigned int count) {
vx_fence(); vx_fence();
@@ -80,14 +80,14 @@ void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
// //
// Make sure global offset values for A and B are contiguous between // Make sure global offset values for A and B are contiguous between
// neighboring threads to ensure GMEM coalescing. // neighboring threads to ensure GMEM coalescing.
// #pragma GCC unroll 1 #pragma GCC unroll 2
for (uint32_t load_offset = 0; load_offset < BM; load_offset += stride_a) { for (uint32_t load_offset = 0; load_offset < BM; load_offset += stride_a) {
const uint32_t global_a_offset = const uint32_t global_a_offset =
dim_k * (global_a_row + load_offset) + (k + local_a_col); dim_k * (global_a_row + load_offset) + (k + local_a_col);
local_a[BK * (local_a_row + load_offset) + local_a_col] = local_a[BK * (local_a_row + load_offset) + local_a_col] =
A[global_a_offset]; A[global_a_offset];
} }
// #pragma GCC unroll 1 #pragma GCC unroll 2
for (uint32_t load_offset = 0; load_offset < BK; load_offset += stride_b) { for (uint32_t load_offset = 0; load_offset < BK; load_offset += stride_b) {
const uint32_t global_b_offset = const uint32_t global_b_offset =
dim_n * (k + local_b_row + load_offset) + global_b_col; dim_n * (k + local_b_row + load_offset) + global_b_col;
@@ -99,7 +99,7 @@ void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
threadblock_dim_y); threadblock_dim_y);
// Compute single tile*tile matmul // Compute single tile*tile matmul
// #pragma GCC unroll 2 #pragma GCC unroll 4
for (uint32_t local_k = 0; local_k < BK; local_k++) { for (uint32_t local_k = 0; local_k < BK; local_k++) {
// First, pump data from SMEM->RF // First, pump data from SMEM->RF
#pragma GCC unroll TM #pragma GCC unroll TM

6
tests/regression/sgemm_wg/main.cpp
View File

@@ -166,9 +166,9 @@ int main(int argc, char *argv[]) {
RT_CHECK(vx_dev_open(&device)); RT_CHECK(vx_dev_open(&device));
// FIXME: hardcoded // FIXME: hardcoded
uint32_t dim_m = 32; uint32_t dim_m = 128;
uint32_t dim_n = 32; uint32_t dim_n = 128;
uint32_t dim_k = 32; uint32_t dim_k = 128;
generate_source_matrix(dim_m, dim_n, dim_k); generate_source_matrix(dim_m, dim_n, dim_k);
generate_reference_matmul(dim_m, dim_n, dim_k); generate_reference_matmul(dim_m, dim_n, dim_k);