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

sgemm_tcore: Template-ize kernel code

Browse Source
This commit is contained in:
Hansung Kim
2024-07-29 20:05:58 -07:00
parent 5f342914bd
commit 5d5a6fbad2
3 changed files with 117 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

107
tests/regression/sgemm_tcore/util.hpp
View File

@@ -8,7 +8,7 @@
// Constraints on parameters:
// * Memory:
// (BM + BN) * BK * sizeof(float) <= sharedmem size.
// (BM + BN) * BK * sizeof(T) <= sharedmem size.
// BM * BK == BN * BK >= threadblock size >= NT * CORES_PER_CLUSTER
// When larger, the kernel runs a sequential loop to read into sharedmem;
// but smaller case is not handled.
@@ -147,7 +147,8 @@ inline void vx_wmma(const int dest_reg) {
}
// `local_k` is assumed to be multiple of TCK
inline void vx_wmma_load_a(volatile const float *smem_A, const int local_k,
template <typename T>
inline void vx_wmma_load_a(volatile const T *smem_A, const int local_k,
const int warp_row, const int wm_iter, const int thread_in_warp) {
const int tid = thread_in_warp;
const int tg = tid / 4;
@@ -167,16 +168,16 @@ inline void vx_wmma_load_a(volatile const float *smem_A, const int local_k,
// @perf: bank conflicts
// f8-f15 stores a single row of A
const volatile float *smem_addr;
const volatile T *smem_addr;
smem_addr = &smem_A[(WM * warp_row + TCM * wm_iter + row) * smem_A_cols + local_k];
asm volatile("flw f0, %0(%1)" ::"i"(0 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f1, %0(%1)" ::"i"(1 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f2, %0(%1)" ::"i"(2 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f3, %0(%1)" ::"i"(3 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f4, %0(%1)" ::"i"(4 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f5, %0(%1)" ::"i"(5 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f6, %0(%1)" ::"i"(6 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f7, %0(%1)" ::"i"(7 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f0, %0(%1)" ::"i"(0 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f1, %0(%1)" ::"i"(1 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f2, %0(%1)" ::"i"(2 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f3, %0(%1)" ::"i"(3 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f4, %0(%1)" ::"i"(4 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f5, %0(%1)" ::"i"(5 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f6, %0(%1)" ::"i"(6 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f7, %0(%1)" ::"i"(7 * sizeof(T)), "r"(smem_addr));
// asm volatile("flw f0, %0" ::"m"(smem_A[A_offset + (local_k + 0)]));
// asm volatile("flw f1, %0" ::"m"(smem_A[A_offset + (local_k + 1)]));
// asm volatile("flw f2, %0" ::"m"(smem_A[A_offset + (local_k + 2)]));
@@ -188,16 +189,16 @@ inline void vx_wmma_load_a(volatile const float *smem_A, const int local_k,
} else {
// read smem A tile as-is; bank-conflict-free AS load
// f8-f15 stores a single row of A
const volatile float *smem_addr;
const volatile T *smem_addr;
smem_addr = &smem_A[((local_k + 0) * smem_AS_cols) + (WM * warp_row + TCM * wm_iter) + row];
asm volatile("flw f0, %0(%1)" :: "i"(smem_AS_cols * 0 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f1, %0(%1)" :: "i"(smem_AS_cols * 1 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f2, %0(%1)" :: "i"(smem_AS_cols * 2 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f3, %0(%1)" :: "i"(smem_AS_cols * 3 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f4, %0(%1)" :: "i"(smem_AS_cols * 4 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f5, %0(%1)" :: "i"(smem_AS_cols * 5 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f6, %0(%1)" :: "i"(smem_AS_cols * 6 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f7, %0(%1)" :: "i"(smem_AS_cols * 7 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f0, %0(%1)" :: "i"(smem_AS_cols * 0 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f1, %0(%1)" :: "i"(smem_AS_cols * 1 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f2, %0(%1)" :: "i"(smem_AS_cols * 2 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f3, %0(%1)" :: "i"(smem_AS_cols * 3 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f4, %0(%1)" :: "i"(smem_AS_cols * 4 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f5, %0(%1)" :: "i"(smem_AS_cols * 5 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f6, %0(%1)" :: "i"(smem_AS_cols * 6 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f7, %0(%1)" :: "i"(smem_AS_cols * 7 * sizeof(T)), "r"(smem_addr));
// asm volatile("flw f0, %0" ::"m"(smem_A[((local_k + 0) * smem_AS_cols) + (WM * warp_row + TCM * wm_iter) + row]));
// asm volatile("flw f1, %0" ::"m"(smem_A[((local_k + 1) * smem_AS_cols) + (WM * warp_row + TCM * wm_iter) + row]));
@@ -211,7 +212,8 @@ inline void vx_wmma_load_a(volatile const float *smem_A, const int local_k,
}
// `local_k` is assumed to be multiple of TCK
inline void vx_wmma_load_b(const volatile float *smem_B, const int local_k,
template <typename T>
inline void vx_wmma_load_b(const volatile T *smem_B, const int local_k,
const int warp_col, const int wn_iter,
const int thread_in_warp) {
const int tid = thread_in_warp;
@@ -225,16 +227,16 @@ inline void vx_wmma_load_b(const volatile float *smem_B, const int local_k,
constexpr int smem_B_cols = BN;
// f8-f15 stores a single column of B
const volatile float *smem_addr;
const volatile T *smem_addr;
smem_addr = &smem_B[((local_k + 0) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col];
asm volatile("flw f8, %0(%1)" :: "i"(smem_B_cols * 0 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f9, %0(%1)" :: "i"(smem_B_cols * 1 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f10, %0(%1)" :: "i"(smem_B_cols * 2 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f11, %0(%1)" :: "i"(smem_B_cols * 3 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f12, %0(%1)" :: "i"(smem_B_cols * 4 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f13, %0(%1)" :: "i"(smem_B_cols * 5 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f14, %0(%1)" :: "i"(smem_B_cols * 6 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f15, %0(%1)" :: "i"(smem_B_cols * 7 * sizeof(float)), "r"(smem_addr));
asm volatile("flw f8, %0(%1)" :: "i"(smem_B_cols * 0 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f9, %0(%1)" :: "i"(smem_B_cols * 1 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f10, %0(%1)" :: "i"(smem_B_cols * 2 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f11, %0(%1)" :: "i"(smem_B_cols * 3 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f12, %0(%1)" :: "i"(smem_B_cols * 4 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f13, %0(%1)" :: "i"(smem_B_cols * 5 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f14, %0(%1)" :: "i"(smem_B_cols * 6 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f15, %0(%1)" :: "i"(smem_B_cols * 7 * sizeof(T)), "r"(smem_addr));
// asm volatile("flw f8, %0" ::"m"(smem_B[((local_k + 0) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col]));
// asm volatile("flw f9, %0" ::"m"(smem_B[((local_k + 1) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col]));
@@ -269,10 +271,11 @@ inline void initialize_C(const int dest_reg) {
}
}
template <typename T>
inline void write_results(const int thread_in_warp, const int warp_col,
const int warp_row, const int wn_iter,
const int wm_iter, const int dim_n,
float *C, const int threadblock_id_x,
T *C, const int threadblock_id_x,
const int threadblock_id_y) {
int tid = thread_in_warp;
@@ -284,22 +287,22 @@ inline void write_results(const int thread_in_warp, const int warp_col,
int local_row = (WM * warp_row + TCM * wm_iter) + tid_row;
int local_col = (WN * warp_col + TCN * wn_iter) + tid_col;
float *global_offset_C = C +
T *global_offset_C = C +
(BM * threadblock_id_y) * dim_n +
BN * threadblock_id_x;
// @perf: this likely causes a lot of gmem bank conflicts
if (wm_iter == 0) {
volatile float *gmem_addr = &global_offset_C[dim_n * (local_row + 0) + (local_col + 0)];
volatile float *gmem_addr_tmp = gmem_addr + (2 * dim_n);
asm volatile ("fsw f16, %0(%1)" :: "i"(0 * sizeof(float)), "r"(gmem_addr));
asm volatile ("fsw f17, %0(%1)" :: "i"(1 * sizeof(float)), "r"(gmem_addr));
asm volatile ("fsw f18, %0(%1)" :: "i"(0 * sizeof(float)), "r"(gmem_addr_tmp));
asm volatile ("fsw f19, %0(%1)" :: "i"(1 * sizeof(float)), "r"(gmem_addr_tmp));
asm volatile ("fsw f20, %0(%1)" :: "i"(4 * sizeof(float)), "r"(gmem_addr));
asm volatile ("fsw f21, %0(%1)" :: "i"(5 * sizeof(float)), "r"(gmem_addr));
asm volatile ("fsw f22, %0(%1)" :: "i"(4 * sizeof(float)), "r"(gmem_addr_tmp));
asm volatile ("fsw f23, %0(%1)" :: "i"(5 * sizeof(float)), "r"(gmem_addr_tmp));
volatile T *gmem_addr = &global_offset_C[dim_n * (local_row + 0) + (local_col + 0)];
volatile T *gmem_addr_tmp = gmem_addr + (2 * dim_n);
asm volatile ("fsw f16, %0(%1)" :: "i"(0 * sizeof(T)), "r"(gmem_addr));
asm volatile ("fsw f17, %0(%1)" :: "i"(1 * sizeof(T)), "r"(gmem_addr));
asm volatile ("fsw f18, %0(%1)" :: "i"(0 * sizeof(T)), "r"(gmem_addr_tmp));
asm volatile ("fsw f19, %0(%1)" :: "i"(1 * sizeof(T)), "r"(gmem_addr_tmp));
asm volatile ("fsw f20, %0(%1)" :: "i"(4 * sizeof(T)), "r"(gmem_addr));
asm volatile ("fsw f21, %0(%1)" :: "i"(5 * sizeof(T)), "r"(gmem_addr));
asm volatile ("fsw f22, %0(%1)" :: "i"(4 * sizeof(T)), "r"(gmem_addr_tmp));
asm volatile ("fsw f23, %0(%1)" :: "i"(5 * sizeof(T)), "r"(gmem_addr_tmp));
// asm volatile ("fsw f16, %0" :: "m"(global_offset_C[dim_n * (local_row + 0) + (local_col + 0)]));
// asm volatile ("fsw f17, %0" :: "m"(global_offset_C[dim_n * (local_row + 0) + (local_col + 1)]));
// asm volatile ("fsw f18, %0" :: "m"(global_offset_C[dim_n * (local_row + 2) + (local_col + 0)]));
@@ -309,16 +312,16 @@ inline void write_results(const int thread_in_warp, const int warp_col,
// asm volatile ("fsw f22, %0" :: "m"(global_offset_C[dim_n * (local_row + 2) + (local_col + 4)]));
// asm volatile ("fsw f23, %0" :: "m"(global_offset_C[dim_n * (local_row + 2) + (local_col + 5)]));
} else {
volatile float *gmem_addr = &global_offset_C[dim_n * (local_row + 0) + (local_col + 0)];
volatile float *gmem_addr_tmp = gmem_addr + (2 * dim_n);
asm volatile ("fsw f24, %0(%1)" :: "i"(0 * sizeof(float)), "r"(gmem_addr));
asm volatile ("fsw f25, %0(%1)" :: "i"(1 * sizeof(float)), "r"(gmem_addr));
asm volatile ("fsw f26, %0(%1)" :: "i"(0 * sizeof(float)), "r"(gmem_addr_tmp));
asm volatile ("fsw f27, %0(%1)" :: "i"(1 * sizeof(float)), "r"(gmem_addr_tmp));
asm volatile ("fsw f28, %0(%1)" :: "i"(4 * sizeof(float)), "r"(gmem_addr));
asm volatile ("fsw f29, %0(%1)" :: "i"(5 * sizeof(float)), "r"(gmem_addr));
asm volatile ("fsw f30, %0(%1)" :: "i"(4 * sizeof(float)), "r"(gmem_addr_tmp));
asm volatile ("fsw f31, %0(%1)" :: "i"(5 * sizeof(float)), "r"(gmem_addr_tmp));
volatile T *gmem_addr = &global_offset_C[dim_n * (local_row + 0) + (local_col + 0)];
volatile T *gmem_addr_tmp = gmem_addr + (2 * dim_n);
asm volatile ("fsw f24, %0(%1)" :: "i"(0 * sizeof(T)), "r"(gmem_addr));
asm volatile ("fsw f25, %0(%1)" :: "i"(1 * sizeof(T)), "r"(gmem_addr));
asm volatile ("fsw f26, %0(%1)" :: "i"(0 * sizeof(T)), "r"(gmem_addr_tmp));
asm volatile ("fsw f27, %0(%1)" :: "i"(1 * sizeof(T)), "r"(gmem_addr_tmp));
asm volatile ("fsw f28, %0(%1)" :: "i"(4 * sizeof(T)), "r"(gmem_addr));
asm volatile ("fsw f29, %0(%1)" :: "i"(5 * sizeof(T)), "r"(gmem_addr));
asm volatile ("fsw f30, %0(%1)" :: "i"(4 * sizeof(T)), "r"(gmem_addr_tmp));
asm volatile ("fsw f31, %0(%1)" :: "i"(5 * sizeof(T)), "r"(gmem_addr_tmp));
}
}