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

sgemm_tcore: Fix AS transpose

Browse Source
This commit is contained in:
Hansung Kim
2024-05-26 13:51:47 -07:00
parent 200fd3e08c
commit 1e48bad4f9
1 changed files with 55 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -8,6 +8,13 @@
#define NUM_LANES 8 #define NUM_LANES 8
#define USE_TENSOR_CORE 1
#define TC_SINGLE_WARP 0
// number of loop around the inner 0..TCK..BK loop to simulate perfect-DRAM
// scenario
#define BK_LOOP 1
#define TRANSPOSE_AS 1
// Constraints on parameters: // Constraints on parameters:
// * Memory: // * Memory:
// (BM + BN) * BK * sizeof(float) <= sharedmem size. // (BM + BN) * BK * sizeof(float) <= sharedmem size.
@@ -20,28 +27,25 @@
// (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 16 #define BM 8
#define BN 16 #define BN 8
#define BK 8 #define BK 8
#define WM 8
#define WN 8
#define TCM 8 #define TCM 8
#define TCN 8 #define TCN 8
#define TCK 8 #define TCK 8
#define WM 8
#define WN 8
#define WMITER (WM / TCM) #define WMITER (WM / TCM)
#define WNITER (WN / TCN) #define WNITER (WN / TCN)
#if USE_TENSOR_CORE == 1
#define TM 1 #define TM 1
#define TN ((TCM * TCN) / NUM_LANES / TM) #define TN ((TCM * TCN) / NUM_LANES / TM)
// #define TN 1 #else
#define TM 1
#define TN 1
#endif
#define ELEM_PER_THREAD (WMITER * WNITER * TM * TN) #define ELEM_PER_THREAD (WMITER * WNITER * TM * TN)
#define USE_TENSOR_CORE 1
#define TC_SINGLE_WARP 0
// number of loop around the inner 0..TCK..BK loop to simulate perfect-DRAM
// scenario
#define BK_LOOP 16
#define TRANSPOSE_AS 1
inline constexpr void map_operand_32lanes(const int tid, int &row, int &col) { inline constexpr void map_operand_32lanes(const int tid, int &row, int &col) {
const int tg = tid / 4; const int tg = tid / 4;
@@ -137,46 +141,51 @@ inline void vx_wmma() {
inline void vx_wmma_load(volatile float *smem_A, volatile float *smem_B, const int local_k, inline void vx_wmma_load(volatile float *smem_A, volatile float *smem_B, const int local_k,
const int warp_col, const int warp_row, const int wn_iter, const int warp_col, const int warp_row, const int wn_iter,
const int wm_iter, const int thread_in_warp) { const int wm_iter, const int thread_in_warp) {
int tid = thread_in_warp; const int tid = thread_in_warp;
int tg = tid / 4; const int tg = tid / 4;
int row = 0; int row = 0;
int col = 0; int col = 0;
map_operand(tid, row, col); map_operand(tid, row, col);
int smem_A_rows = BM; constexpr int smem_A_rows = BM;
int smem_A_cols = BK; constexpr int smem_A_cols = BK;
int smem_B_rows = BK; constexpr int smem_AS_rows = BK;
int smem_B_cols = BN; constexpr int smem_AS_cols = BM;
constexpr int smem_B_rows = BK;
constexpr int smem_B_cols = BN;
if constexpr (!TRANSPOSE_AS) { if constexpr (!TRANSPOSE_AS) {
int A_offset = (row + WM * warp_row + TCM * wm_iter) * smem_A_cols; int A_offset = (WM * warp_row + TCM * wm_iter + row) * smem_A_cols;
// @perf: bank conflicts // @perf: bank conflicts
asm volatile("flw f0, %0" ::"m"(smem_A[A_offset + (local_k + 0)])); // f8-f15 stores a single row of A
asm volatile("flw f1, %0" ::"m"(smem_A[A_offset + (local_k + 1)])); asm volatile("flw f0, %0" ::"m"(smem_A[A_offset + (local_k + 0)]));
asm volatile("flw f2, %0" ::"m"(smem_A[A_offset + (local_k + 2)])); asm volatile("flw f1, %0" ::"m"(smem_A[A_offset + (local_k + 1)]));
asm volatile("flw f3, %0" ::"m"(smem_A[A_offset + (local_k + 3)])); asm volatile("flw f2, %0" ::"m"(smem_A[A_offset + (local_k + 2)]));
asm volatile("flw f4, %0" ::"m"(smem_A[A_offset + (local_k + 4)])); asm volatile("flw f3, %0" ::"m"(smem_A[A_offset + (local_k + 3)]));
asm volatile("flw f5, %0" ::"m"(smem_A[A_offset + (local_k + 5)])); asm volatile("flw f4, %0" ::"m"(smem_A[A_offset + (local_k + 4)]));
asm volatile("flw f6, %0" ::"m"(smem_A[A_offset + (local_k + 6)])); asm volatile("flw f5, %0" ::"m"(smem_A[A_offset + (local_k + 5)]));
asm volatile("flw f7, %0" ::"m"(smem_A[A_offset + (local_k + 7)])); asm volatile("flw f6, %0" ::"m"(smem_A[A_offset + (local_k + 6)]));
asm volatile("flw f7, %0" ::"m"(smem_A[A_offset + (local_k + 7)]));
} else { } else {
// transposed A // transposed A
asm volatile("flw f0, %0" ::"m"(smem_A[((local_k + 0) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row])); // f8-f15 stores a single row of A
asm volatile("flw f1, %0" ::"m"(smem_A[((local_k + 1) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row])); asm volatile("flw f0, %0" ::"m"(smem_A[((local_k + 0) * smem_AS_cols) + (WM * warp_row + TCM * wm_iter) + row]));
asm volatile("flw f2, %0" ::"m"(smem_A[((local_k + 2) * smem_A_rows) + (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]));
asm volatile("flw f3, %0" ::"m"(smem_A[((local_k + 3) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row])); asm volatile("flw f2, %0" ::"m"(smem_A[((local_k + 2) * smem_AS_cols) + (WM * warp_row + TCM * wm_iter) + row]));
asm volatile("flw f4, %0" ::"m"(smem_A[((local_k + 4) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row])); asm volatile("flw f3, %0" ::"m"(smem_A[((local_k + 3) * smem_AS_cols) + (WM * warp_row + TCM * wm_iter) + row]));
asm volatile("flw f5, %0" ::"m"(smem_A[((local_k + 5) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row])); asm volatile("flw f4, %0" ::"m"(smem_A[((local_k + 4) * smem_AS_cols) + (WM * warp_row + TCM * wm_iter) + row]));
asm volatile("flw f6, %0" ::"m"(smem_A[((local_k + 6) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row])); asm volatile("flw f5, %0" ::"m"(smem_A[((local_k + 5) * smem_AS_cols) + (WM * warp_row + TCM * wm_iter) + row]));
asm volatile("flw f7, %0" ::"m"(smem_A[((local_k + 7) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row])); asm volatile("flw f6, %0" ::"m"(smem_A[((local_k + 6) * smem_AS_cols) + (WM * warp_row + TCM * wm_iter) + row]));
asm volatile("flw f7, %0" ::"m"(smem_A[((local_k + 7) * smem_AS_cols) + (WM * warp_row + TCM * wm_iter) + row]));
// #pragma GCC unroll 8 // #pragma GCC unroll 8
// for (int i = 0; i < 8; i++) { // for (int i = 0; i < 8; i++) {
// asm volatile("flw f0, %0" ::"m"(smem_A[((local_k + i) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row])); // asm volatile("flw f0, %0" ::"m"(smem_A[((local_k + i) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row]));
// } // }
} }
// f8-f15 stores a single column of B
asm volatile("flw f8, %0" ::"m"(smem_B[((local_k + 0) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col])); 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])); asm volatile("flw f9, %0" ::"m"(smem_B[((local_k + 1) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col]));
asm volatile("flw f10, %0" ::"m"(smem_B[((local_k + 2) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col])); asm volatile("flw f10, %0" ::"m"(smem_B[((local_k + 2) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col]));
@@ -295,29 +304,31 @@ void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
// number of rows a full TB can read at a time // number of rows a full TB can read at a time
constexpr uint32_t row_stride_a = (BM * BN) / ELEM_PER_THREAD / BK; constexpr uint32_t row_stride_a = (BM * BN) / ELEM_PER_THREAD / BK;
#pragma GCC unroll 1 #pragma GCC unroll 1
for (uint32_t load_offset = 0; load_offset < BM; load_offset += row_stride_a) { for (uint32_t local_row_offset = 0; local_row_offset < BM;
local_row_offset += row_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 + local_row_offset) + (k + local_a_col);
// NOTE: all threads in TB will do this load; make sure this is not // NOTE: all threads in TB will do this load; make sure this is not
// out-of-bounds of BM*BK // out-of-bounds of BM*BK
local_a[BK * (local_a_row + load_offset) + local_a_col] = local_a[BK * (local_a_row + local_row_offset) + local_a_col] =
A[global_a_offset]; A[global_a_offset];
} }
} else { } else {
const uint32_t global_a_row = BM * threadblock_id_y + local_as_col; const uint32_t global_a_row = BM * threadblock_id_y + local_as_col;
constexpr uint32_t row_stride_a = (BM * BN) / ELEM_PER_THREAD / BM; constexpr uint32_t row_stride_as = (BM * BN) / ELEM_PER_THREAD / BM;
#pragma GCC unroll 1 #pragma GCC unroll 1
for (uint32_t load_offset = 0; load_offset < BK; load_offset += row_stride_a) { for (uint32_t local_row_offset = 0; local_row_offset < BK;
local_row_offset += row_stride_as) {
// @perf: bank conflicts here // @perf: bank conflicts here
const uint32_t global_a_offset = const uint32_t global_a_offset =
dim_k * (global_a_row + load_offset) + (k + local_as_row); dim_k * (global_a_row) + (k + local_as_row + local_row_offset);
local_a[BM * (local_as_row + load_offset) + local_as_col] = local_a[BM * (local_as_row + local_row_offset) + local_as_col] =
A[global_a_offset]; A[global_a_offset];
} }
} }
constexpr uint32_t row_stride_b = (BM * BN) / ELEM_PER_THREAD / BN; constexpr uint32_t row_stride_b = (BM * BN) / ELEM_PER_THREAD / BN;
const uint32_t global_b_col = BN * threadblock_id_x + local_b_col; const uint32_t global_b_col = BN * threadblock_id_x + local_b_col;
#pragma GCC unroll 1 #pragma GCC unroll 1
for (uint32_t load_offset = 0; load_offset < BK; load_offset += row_stride_b) { for (uint32_t load_offset = 0; load_offset < BK; load_offset += row_stride_b) {
const uint32_t global_b_offset = const uint32_t global_b_offset =