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

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

@@ -7,7 +7,7 @@
#include "include/gemmini.h" #include "include/gemmini.h"
#include "gemmini_mmio.h" #include "gemmini_mmio.h"
#define GEMMINI_DMA 1 #define GEMMINI_DMA 0
#if SMEM_SIZE == 0x4000 #if SMEM_SIZE == 0x4000
#define SMEM_ADDR_Q0 ((float * const) 0xff000000) #define SMEM_ADDR_Q0 ((float * const) 0xff000000)
#define SMEM_ADDR_Q1 ((float * const) 0xff001000) #define SMEM_ADDR_Q1 ((float * const) 0xff001000)
@@ -37,9 +37,10 @@
#error "threadblock size too big for cluster" #error "threadblock size too big for cluster"
#endif #endif
template <typename T>
inline void global_dmem_load(const uint32_t dim_n, const uint32_t dim_k, inline void global_dmem_load(const uint32_t dim_n, const uint32_t dim_k,
const uint32_t k, const float *A, const float *B, const uint32_t k, const T *A, const T *B,
volatile float *local_a, volatile float *local_b, volatile T *local_a, volatile T *local_b,
const uint32_t tid_in_threadblock, const uint32_t tid_in_threadblock,
const uint32_t threadblock_id_x, const uint32_t threadblock_id_x,
const uint32_t threadblock_id_y) { const uint32_t threadblock_id_y) {
@@ -64,8 +65,8 @@ inline void global_dmem_load(const uint32_t dim_n, const uint32_t dim_k,
const uint32_t global_a_row = BM * threadblock_id_y + local_a_row; const uint32_t global_a_row = BM * threadblock_id_y + local_a_row;
// 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 = threads_in_threadblock / BK; constexpr uint32_t row_stride_a = threads_in_threadblock / BK;
const float *global_a = A + dim_k * global_a_row + (k + local_a_col); const T *global_a = A + dim_k * global_a_row + (k + local_a_col);
volatile float *local_a_tmp = local_a + BK * local_a_row + local_a_col; volatile T *local_a_tmp = local_a + BK * local_a_row + local_a_col;
#pragma GCC unroll 1 #pragma GCC unroll 1
for (uint32_t local_row_offset = 0; local_row_offset < BM; for (uint32_t local_row_offset = 0; local_row_offset < BM;
@@ -83,11 +84,11 @@ inline void global_dmem_load(const uint32_t dim_n, const uint32_t dim_k,
if constexpr (!GMEM_COALESCED_A) { if constexpr (!GMEM_COALESCED_A) {
constexpr uint32_t row_stride_as = threads_in_threadblock / BM; constexpr uint32_t row_stride_as = threads_in_threadblock / BM;
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;
const float *global_a = A + dim_k * global_a_row + (k + local_as_row); const T *global_a = A + dim_k * global_a_row + (k + local_as_row);
// FIXME experimenting with global coalescing // FIXME experimenting with global coalescing
// const uint32_t global_a_row = BM * threadblock_id_y + local_as_row; // const uint32_t global_a_row = BM * threadblock_id_y + local_as_row;
// const float *global_a = A + dim_k * global_a_row + (k + local_as_col); // const T *global_a = A + dim_k * global_a_row + (k + local_as_col);
volatile float *local_a_tmp = local_a + BM * local_as_row + local_as_col; volatile T *local_a_tmp = local_a + BM * local_as_row + local_as_col;
static_assert( static_assert(
row_stride_as * 8 <= BK, row_stride_as * 8 <= BK,
@@ -126,22 +127,22 @@ inline void global_dmem_load(const uint32_t dim_n, const uint32_t dim_k,
asm volatile ("flw ft7, (%0)" :: "r"(global_a)); asm volatile ("flw ft7, (%0)" :: "r"(global_a));
global_a += row_stride_as; global_a += row_stride_as;
asm volatile ("fsw ft0, %0(%1)" :: "i"(BM * row_stride_as * 0 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft0, %0(%1)" :: "i"(BM * row_stride_as * 0 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft1, %0(%1)" :: "i"(BM * row_stride_as * 1 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft1, %0(%1)" :: "i"(BM * row_stride_as * 1 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft2, %0(%1)" :: "i"(BM * row_stride_as * 2 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft2, %0(%1)" :: "i"(BM * row_stride_as * 2 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft3, %0(%1)" :: "i"(BM * row_stride_as * 3 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft3, %0(%1)" :: "i"(BM * row_stride_as * 3 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft4, %0(%1)" :: "i"(BM * row_stride_as * 4 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft4, %0(%1)" :: "i"(BM * row_stride_as * 4 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft5, %0(%1)" :: "i"(BM * row_stride_as * 5 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft5, %0(%1)" :: "i"(BM * row_stride_as * 5 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft6, %0(%1)" :: "i"(BM * row_stride_as * 6 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft6, %0(%1)" :: "i"(BM * row_stride_as * 6 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft7, %0(%1)" :: "i"(BM * row_stride_as * 7 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft7, %0(%1)" :: "i"(BM * row_stride_as * 7 * sizeof(T)), "r"(local_a_tmp));
local_a_tmp += BM * row_stride_as * 8; local_a_tmp += BM * row_stride_as * 8;
} }
} else { } else {
constexpr uint32_t row_stride_a = threads_in_threadblock / BK; constexpr uint32_t row_stride_a = threads_in_threadblock / BK;
const uint32_t global_a_row = BM * threadblock_id_y + local_a_row; const uint32_t global_a_row = BM * threadblock_id_y + local_a_row;
const float *global_a = A + dim_k * global_a_row + (k + local_a_col); const T *global_a = A + dim_k * global_a_row + (k + local_a_col);
// NOTE that SMEM writes are transposed // NOTE that SMEM writes are transposed
volatile float *local_a_tmp = local_a + BM * local_a_col + local_a_row; volatile T *local_a_tmp = local_a + BM * local_a_col + local_a_row;
static_assert( static_assert(
row_stride_a * 8 <= BM, row_stride_a * 8 <= BM,
@@ -177,14 +178,14 @@ inline void global_dmem_load(const uint32_t dim_n, const uint32_t dim_k,
global_a += dim_k * row_stride_a; global_a += dim_k * row_stride_a;
// stride along columns // stride along columns
asm volatile ("fsw ft0, %0(%1)" :: "i"(row_stride_a * 0 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft0, %0(%1)" :: "i"(row_stride_a * 0 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft1, %0(%1)" :: "i"(row_stride_a * 1 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft1, %0(%1)" :: "i"(row_stride_a * 1 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft2, %0(%1)" :: "i"(row_stride_a * 2 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft2, %0(%1)" :: "i"(row_stride_a * 2 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft3, %0(%1)" :: "i"(row_stride_a * 3 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft3, %0(%1)" :: "i"(row_stride_a * 3 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft4, %0(%1)" :: "i"(row_stride_a * 4 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft4, %0(%1)" :: "i"(row_stride_a * 4 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft5, %0(%1)" :: "i"(row_stride_a * 5 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft5, %0(%1)" :: "i"(row_stride_a * 5 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft6, %0(%1)" :: "i"(row_stride_a * 6 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft6, %0(%1)" :: "i"(row_stride_a * 6 * sizeof(T)), "r"(local_a_tmp));
asm volatile ("fsw ft7, %0(%1)" :: "i"(row_stride_a * 7 * sizeof(float)), "r"(local_a_tmp)); asm volatile ("fsw ft7, %0(%1)" :: "i"(row_stride_a * 7 * sizeof(T)), "r"(local_a_tmp));
local_a_tmp += row_stride_a * 8; local_a_tmp += row_stride_a * 8;
} }
} }
@@ -192,8 +193,8 @@ inline void global_dmem_load(const uint32_t dim_n, const uint32_t dim_k,
constexpr uint32_t row_stride_b = threads_in_threadblock / BN; constexpr uint32_t row_stride_b = threads_in_threadblock / 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;
const float *global_b = B + dim_n * (k + local_b_row) + global_b_col; const T *global_b = B + dim_n * (k + local_b_row) + global_b_col;
volatile float *local_b_tmp = local_b + BN * local_b_row + local_b_col; volatile T *local_b_tmp = local_b + BN * local_b_row + local_b_col;
static_assert( static_assert(
row_stride_b * 8 <= BK, row_stride_b * 8 <= BK,
@@ -232,21 +233,22 @@ inline void global_dmem_load(const uint32_t dim_n, const uint32_t dim_k,
asm volatile ("flw ft7, (%0)" :: "r"(global_b)); asm volatile ("flw ft7, (%0)" :: "r"(global_b));
global_b += dim_n * row_stride_b; global_b += dim_n * row_stride_b;
asm volatile ("fsw ft0, %0(%1)" :: "i"(BN * row_stride_b * 0 * sizeof(float)), "r"(local_b_tmp)); asm volatile ("fsw ft0, %0(%1)" :: "i"(BN * row_stride_b * 0 * sizeof(T)), "r"(local_b_tmp));
asm volatile ("fsw ft1, %0(%1)" :: "i"(BN * row_stride_b * 1 * sizeof(float)), "r"(local_b_tmp)); asm volatile ("fsw ft1, %0(%1)" :: "i"(BN * row_stride_b * 1 * sizeof(T)), "r"(local_b_tmp));
local_b_tmp += BN * row_stride_b * 2; local_b_tmp += BN * row_stride_b * 2;
asm volatile ("fsw ft2, %0(%1)" :: "i"(BN * row_stride_b * 0 * sizeof(float)), "r"(local_b_tmp)); asm volatile ("fsw ft2, %0(%1)" :: "i"(BN * row_stride_b * 0 * sizeof(T)), "r"(local_b_tmp));
asm volatile ("fsw ft3, %0(%1)" :: "i"(BN * row_stride_b * 1 * sizeof(float)), "r"(local_b_tmp)); asm volatile ("fsw ft3, %0(%1)" :: "i"(BN * row_stride_b * 1 * sizeof(T)), "r"(local_b_tmp));
local_b_tmp += BN * row_stride_b * 2; local_b_tmp += BN * row_stride_b * 2;
asm volatile ("fsw ft4, %0(%1)" :: "i"(BN * row_stride_b * 0 * sizeof(float)), "r"(local_b_tmp)); asm volatile ("fsw ft4, %0(%1)" :: "i"(BN * row_stride_b * 0 * sizeof(T)), "r"(local_b_tmp));
asm volatile ("fsw ft5, %0(%1)" :: "i"(BN * row_stride_b * 1 * sizeof(float)), "r"(local_b_tmp)); asm volatile ("fsw ft5, %0(%1)" :: "i"(BN * row_stride_b * 1 * sizeof(T)), "r"(local_b_tmp));
local_b_tmp += BN * row_stride_b * 2; local_b_tmp += BN * row_stride_b * 2;
asm volatile ("fsw ft6, %0(%1)" :: "i"(BN * row_stride_b * 0 * sizeof(float)), "r"(local_b_tmp)); asm volatile ("fsw ft6, %0(%1)" :: "i"(BN * row_stride_b * 0 * sizeof(T)), "r"(local_b_tmp));
asm volatile ("fsw ft7, %0(%1)" :: "i"(BN * row_stride_b * 1 * sizeof(float)), "r"(local_b_tmp)); asm volatile ("fsw ft7, %0(%1)" :: "i"(BN * row_stride_b * 1 * sizeof(T)), "r"(local_b_tmp));
local_b_tmp += BN * row_stride_b * 2; local_b_tmp += BN * row_stride_b * 2;
} }
} }
template <typename T>
inline void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg, inline void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
const uint32_t tid_in_threadblock, const uint32_t tid_in_threadblock,
const uint32_t threads_per_threadblock, const uint32_t threads_per_threadblock,
@@ -255,10 +257,10 @@ inline void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
const uint32_t threadblock_id_y,*/ const uint32_t threadblock_id_y,*/
const uint32_t threadblocks_per_cluster, const uint32_t threadblocks_per_cluster,
const uint32_t threadblock_id_in_cluster, const uint32_t threadblock_id_in_cluster,
float *sharedmem_per_threadblock) { uint8_t *sharedmem_per_threadblock) {
const float *A = (const float *)arg->addr_a; const T *A = (const T *)arg->addr_a;
const float *B = (const float *)arg->addr_b; const T *B = (const T *)arg->addr_b;
float *C = (float *)arg->addr_c; T *C = (T *)arg->addr_c;
const uint32_t dim_m = arg->dim_m; const uint32_t dim_m = arg->dim_m;
const uint32_t dim_n = arg->dim_n; const uint32_t dim_n = arg->dim_n;
@@ -278,13 +280,13 @@ inline void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
const uint32_t warp_col = warp_id_in_warpgroup % (BN / WN); const uint32_t warp_col = warp_id_in_warpgroup % (BN / WN);
const uint32_t tid_in_warp = tid_in_threadblock % NUM_THREADS; const uint32_t tid_in_warp = tid_in_threadblock % NUM_THREADS;
volatile float *local_a = sharedmem_per_threadblock; volatile T *local_a = reinterpret_cast<T *>(sharedmem_per_threadblock);
constexpr size_t local_a_elems = (BM * BK); constexpr size_t local_a_elems = (BM * BK);
volatile float *local_a_buf = local_a + local_a_elems; volatile T *local_a_buf = local_a + local_a_elems;
volatile float *local_b = local_a_buf + local_a_elems; volatile T *local_b = local_a_buf + local_a_elems;
constexpr size_t local_b_elems = (BK * BN); constexpr size_t local_b_elems = (BK * BN);
volatile float *local_b_buf = local_a_buf + local_b_elems; volatile T *local_b_buf = local_a_buf + local_b_elems;
constexpr uint32_t skips = constexpr uint32_t skips =
loop_matmul_skips(/*skip_lda=*/0, /*skip_ldb=*/0, /*skip_ldd=*/1, loop_matmul_skips(/*skip_lda=*/0, /*skip_ldb=*/0, /*skip_ldd=*/1,
@@ -439,18 +441,18 @@ inline void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
// consumer code: SMEM->RF and compute // consumer code: SMEM->RF and compute
// ---------------------------------------------------------------------- // ----------------------------------------------------------------------
// @perf: this loop spills to stack a lot because of all the flws in // @perf: this loop spills to stack a lot because of all the flws in
const volatile float *local_a_consume; const volatile T *local_a_consume;
const volatile float *local_b_consume; const volatile T *local_b_consume;
if constexpr (GEMMINI_DMA) { if constexpr (GEMMINI_DMA) {
// local_a_consume = (k_index % 2) ? local_a_buf : local_a; // local_a_consume = (k_index % 2) ? local_a_buf : local_a;
// local_b_consume = (k_index % 2) ? local_b_buf : local_b; // local_b_consume = (k_index % 2) ? local_b_buf : local_b;
// FIXME: swap multiply with bitshifts // FIXME: swap multiply with bitshifts
// const uint32_t mask_odd = (block_k & 1) << 31 >> 31; // const uint32_t mask_odd = (block_k & 1) << 31 >> 31;
// const uint32_t mask_even = ((block_k & 1) ^ 1) << 31 >> 31; // const uint32_t mask_even = ((block_k & 1) ^ 1) << 31 >> 31;
// local_a_consume = reinterpret_cast<volatile float *>( // local_a_consume = reinterpret_cast<volatile T *>(
// (mask_odd & reinterpret_cast<uintmax_t>(local_a_buf)) | // (mask_odd & reinterpret_cast<uintmax_t>(local_a_buf)) |
// (mask_even & reinterpret_cast<uintmax_t>(local_a))); // (mask_even & reinterpret_cast<uintmax_t>(local_a)));
// local_b_consume = reinterpret_cast<volatile float *>( // local_b_consume = reinterpret_cast<volatile T *>(
// (mask_odd & reinterpret_cast<uintmax_t>(local_b_buf)) | // (mask_odd & reinterpret_cast<uintmax_t>(local_b_buf)) |
// (mask_even & reinterpret_cast<uintmax_t>(local_b))); // (mask_even & reinterpret_cast<uintmax_t>(local_b)));
local_a_consume = local_a + (block_k & 1) * (local_a_elems); local_a_consume = local_a + (block_k & 1) * (local_a_elems);
@@ -539,19 +541,20 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
const uint32_t problem_size = (dim_m * dim_n) / (ELEM_PER_THREAD); const uint32_t problem_size = (dim_m * dim_n) / (ELEM_PER_THREAD);
const uint32_t num_threadblocks = problem_size / threads_per_threadblock; const uint32_t num_threadblocks = problem_size / threads_per_threadblock;
using float_type = float;
// "static" shared memory allocation. This would determine threadblock // "static" shared memory allocation. This would determine threadblock
// occupancy of a single cluster // occupancy of a single cluster
float *sharedmem_per_threadblock = uint8_t *sharedmem_per_threadblock = reinterpret_cast<uint8_t *>(
(float *)DEV_SMEM_START_ADDR + 2/*overkill for non-dma*/ * (2 * BM * BK) * DEV_SMEM_START_ADDR + sizeof(float_type) * 2 /*overkill for non-dma*/ *
threadblock_id_in_cluster; (2 * BM * BK) * threadblock_id_in_cluster);
thread_block_gemm(arg, tid_in_threadblock, threads_per_threadblock, thread_block_gemm<float_type>(
threadblock_dim_y, arg, tid_in_threadblock, threads_per_threadblock, threadblock_dim_y,
/*threadblock_id_x, threadblock_id_y,*/ /*threadblock_id_x, threadblock_id_y,*/
threadblocks_per_cluster, threadblocks_per_cluster,
// threadblock_id, // threadblock_id,
threadblock_id_in_cluster, threadblock_id_in_cluster, sharedmem_per_threadblock);
sharedmem_per_threadblock);
} }
int main() { int main() {

2
tests/regression/sgemm_tcore/main.cpp
View File

@@ -173,7 +173,7 @@ int main(int argc, char *argv[]) {
uint32_t dim_n = 64; uint32_t dim_n = 64;
uint32_t dim_k = 64; uint32_t dim_k = 64;
using float_type = half; using float_type = float;
generate_source_matrix<float_type>(dim_m, dim_n, dim_k); generate_source_matrix<float_type>(dim_m, dim_n, dim_k);
generate_reference_matmul<float_type>(dim_m, dim_n, dim_k); generate_reference_matmul<float_type>(dim_m, dim_n, dim_k);

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

@@ -8,7 +8,7 @@
// Constraints on parameters: // Constraints on parameters:
// * Memory: // * 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 // BM * BK == BN * BK >= threadblock size >= NT * CORES_PER_CLUSTER
// When larger, the kernel runs a sequential loop to read into sharedmem; // When larger, the kernel runs a sequential loop to read into sharedmem;
// but smaller case is not handled. // 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 // `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 warp_row, const int wm_iter, const int thread_in_warp) {
const int tid = thread_in_warp; const int tid = thread_in_warp;
const int tg = tid / 4; 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 // @perf: bank conflicts
// f8-f15 stores a single row of A // 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]; 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 f0, %0(%1)" ::"i"(0 * sizeof(T)), "r"(smem_addr));
asm volatile("flw f1, %0(%1)" ::"i"(1 * sizeof(float)), "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(float)), "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(float)), "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(float)), "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(float)), "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(float)), "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(float)), "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 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 f1, %0" ::"m"(smem_A[A_offset + (local_k + 1)]));
// asm volatile("flw f2, %0" ::"m"(smem_A[A_offset + (local_k + 2)])); // 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 { } else {
// read smem A tile as-is; bank-conflict-free AS load // read smem A tile as-is; bank-conflict-free AS load
// f8-f15 stores a single row of A // 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]; 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 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(float)), "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(float)), "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(float)), "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(float)), "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(float)), "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(float)), "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(float)), "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 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])); // 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 // `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 warp_col, const int wn_iter,
const int thread_in_warp) { const int thread_in_warp) {
const int tid = 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; constexpr int smem_B_cols = BN;
// f8-f15 stores a single column of B // 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]; 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 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(float)), "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(float)), "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(float)), "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(float)), "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(float)), "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(float)), "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(float)), "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 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]));
@@ -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, inline void write_results(const int thread_in_warp, const int warp_col,
const int warp_row, const int wn_iter, const int warp_row, const int wn_iter,
const int wm_iter, const int dim_n, 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) { const int threadblock_id_y) {
int tid = thread_in_warp; 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_row = (WM * warp_row + TCM * wm_iter) + tid_row;
int local_col = (WN * warp_col + TCN * wn_iter) + tid_col; 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 + (BM * threadblock_id_y) * dim_n +
BN * threadblock_id_x; BN * threadblock_id_x;
// @perf: this likely causes a lot of gmem bank conflicts // @perf: this likely causes a lot of gmem bank conflicts
if (wm_iter == 0) { if (wm_iter == 0) {
volatile float *gmem_addr = &global_offset_C[dim_n * (local_row + 0) + (local_col + 0)]; volatile T *gmem_addr = &global_offset_C[dim_n * (local_row + 0) + (local_col + 0)];
volatile float *gmem_addr_tmp = gmem_addr + (2 * dim_n); volatile T *gmem_addr_tmp = gmem_addr + (2 * dim_n);
asm volatile ("fsw f16, %0(%1)" :: "i"(0 * sizeof(float)), "r"(gmem_addr)); asm volatile ("fsw f16, %0(%1)" :: "i"(0 * sizeof(T)), "r"(gmem_addr));
asm volatile ("fsw f17, %0(%1)" :: "i"(1 * sizeof(float)), "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(float)), "r"(gmem_addr_tmp)); asm volatile ("fsw f18, %0(%1)" :: "i"(0 * sizeof(T)), "r"(gmem_addr_tmp));
asm volatile ("fsw f19, %0(%1)" :: "i"(1 * sizeof(float)), "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(float)), "r"(gmem_addr)); asm volatile ("fsw f20, %0(%1)" :: "i"(4 * sizeof(T)), "r"(gmem_addr));
asm volatile ("fsw f21, %0(%1)" :: "i"(5 * sizeof(float)), "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(float)), "r"(gmem_addr_tmp)); asm volatile ("fsw f22, %0(%1)" :: "i"(4 * sizeof(T)), "r"(gmem_addr_tmp));
asm volatile ("fsw f23, %0(%1)" :: "i"(5 * sizeof(float)), "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 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 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)])); // 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 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)])); // asm volatile ("fsw f23, %0" :: "m"(global_offset_C[dim_n * (local_row + 2) + (local_col + 5)]));
} else { } else {
volatile float *gmem_addr = &global_offset_C[dim_n * (local_row + 0) + (local_col + 0)]; volatile T *gmem_addr = &global_offset_C[dim_n * (local_row + 0) + (local_col + 0)];
volatile float *gmem_addr_tmp = gmem_addr + (2 * dim_n); volatile T *gmem_addr_tmp = gmem_addr + (2 * dim_n);
asm volatile ("fsw f24, %0(%1)" :: "i"(0 * sizeof(float)), "r"(gmem_addr)); asm volatile ("fsw f24, %0(%1)" :: "i"(0 * sizeof(T)), "r"(gmem_addr));
asm volatile ("fsw f25, %0(%1)" :: "i"(1 * sizeof(float)), "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(float)), "r"(gmem_addr_tmp)); asm volatile ("fsw f26, %0(%1)" :: "i"(0 * sizeof(T)), "r"(gmem_addr_tmp));
asm volatile ("fsw f27, %0(%1)" :: "i"(1 * sizeof(float)), "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(float)), "r"(gmem_addr)); asm volatile ("fsw f28, %0(%1)" :: "i"(4 * sizeof(T)), "r"(gmem_addr));
asm volatile ("fsw f29, %0(%1)" :: "i"(5 * sizeof(float)), "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(float)), "r"(gmem_addr_tmp)); asm volatile ("fsw f30, %0(%1)" :: "i"(4 * sizeof(T)), "r"(gmem_addr_tmp));
asm volatile ("fsw f31, %0(%1)" :: "i"(5 * sizeof(float)), "r"(gmem_addr_tmp)); asm volatile ("fsw f31, %0(%1)" :: "i"(5 * sizeof(T)), "r"(gmem_addr_tmp));
} }
} }