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 "gemmini_mmio.h"
#define GEMMINI_DMA 1
#define GEMMINI_DMA 0
#if SMEM_SIZE == 0x4000
#define SMEM_ADDR_Q0 ((float * const) 0xff000000)
#define SMEM_ADDR_Q1 ((float * const) 0xff001000)
@@ -37,9 +37,10 @@
#error "threadblock size too big for cluster"
#endif
template <typename T>
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,
volatile float *local_a, volatile float *local_b,
const uint32_t k, const T *A, const T *B,
volatile T *local_a, volatile T *local_b,
const uint32_t tid_in_threadblock,
const uint32_t threadblock_id_x,
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;
// number of rows a full TB can read at a time
constexpr uint32_t row_stride_a = threads_in_threadblock / BK;
const float *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;
const T *global_a = A + dim_k * global_a_row + (k + local_a_col);
volatile T *local_a_tmp = local_a + BK * local_a_row + local_a_col;
#pragma GCC unroll 1
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) {
constexpr uint32_t row_stride_as = threads_in_threadblock / BM;
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
// 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);
volatile float *local_a_tmp = local_a + BM * local_as_row + local_as_col;
// const T *global_a = A + dim_k * global_a_row + (k + local_as_col);
volatile T *local_a_tmp = local_a + BM * local_as_row + local_as_col;
static_assert(
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));
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 ft1, %0(%1)" :: "i"(BM * row_stride_as * 1 * sizeof(float)), "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 ft3, %0(%1)" :: "i"(BM * row_stride_as * 3 * sizeof(float)), "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 ft5, %0(%1)" :: "i"(BM * row_stride_as * 5 * sizeof(float)), "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 ft7, %0(%1)" :: "i"(BM * row_stride_as * 7 * 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(T)), "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(T)), "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(T)), "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(T)), "r"(local_a_tmp));
local_a_tmp += BM * row_stride_as * 8;
}
} else {
constexpr uint32_t row_stride_a = threads_in_threadblock / BK;
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
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(
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;
// stride along columns
asm volatile ("fsw ft0, %0(%1)" :: "i"(row_stride_a * 0 * sizeof(float)), "r"(local_a_tmp));
asm volatile ("fsw ft1, %0(%1)" :: "i"(row_stride_a * 1 * sizeof(float)), "r"(local_a_tmp));
asm volatile ("fsw ft2, %0(%1)" :: "i"(row_stride_a * 2 * sizeof(float)), "r"(local_a_tmp));
asm volatile ("fsw ft3, %0(%1)" :: "i"(row_stride_a * 3 * sizeof(float)), "r"(local_a_tmp));
asm volatile ("fsw ft4, %0(%1)" :: "i"(row_stride_a * 4 * sizeof(float)), "r"(local_a_tmp));
asm volatile ("fsw ft5, %0(%1)" :: "i"(row_stride_a * 5 * sizeof(float)), "r"(local_a_tmp));
asm volatile ("fsw ft6, %0(%1)" :: "i"(row_stride_a * 6 * sizeof(float)), "r"(local_a_tmp));
asm volatile ("fsw ft7, %0(%1)" :: "i"(row_stride_a * 7 * 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(T)), "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(T)), "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(T)), "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(T)), "r"(local_a_tmp));
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;
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;
volatile float *local_b_tmp = local_b + BN * local_b_row + local_b_col;
const T *global_b = B + dim_n * (k + local_b_row) + global_b_col;
volatile T *local_b_tmp = local_b + BN * local_b_row + local_b_col;
static_assert(
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));
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 ft1, %0(%1)" :: "i"(BN * row_stride_b * 1 * 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(T)), "r"(local_b_tmp));
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 ft3, %0(%1)" :: "i"(BN * row_stride_b * 1 * 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(T)), "r"(local_b_tmp));
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 ft5, %0(%1)" :: "i"(BN * row_stride_b * 1 * 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(T)), "r"(local_b_tmp));
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 ft7, %0(%1)" :: "i"(BN * row_stride_b * 1 * 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(T)), "r"(local_b_tmp));
local_b_tmp += BN * row_stride_b * 2;
}
}
template <typename T>
inline void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
const uint32_t tid_in_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 threadblocks_per_cluster,
const uint32_t threadblock_id_in_cluster,
float *sharedmem_per_threadblock) {
const float *A = (const float *)arg->addr_a;
const float *B = (const float *)arg->addr_b;
float *C = (float *)arg->addr_c;
uint8_t *sharedmem_per_threadblock) {
const T *A = (const T *)arg->addr_a;
const T *B = (const T *)arg->addr_b;
T *C = (T *)arg->addr_c;
const uint32_t dim_m = arg->dim_m;
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 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);
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);
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 =
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
// ----------------------------------------------------------------------
// @perf: this loop spills to stack a lot because of all the flws in
const volatile float *local_a_consume;
const volatile float *local_b_consume;
const volatile T *local_a_consume;
const volatile T *local_b_consume;
if constexpr (GEMMINI_DMA) {
// local_a_consume = (k_index % 2) ? local_a_buf : local_a;
// local_b_consume = (k_index % 2) ? local_b_buf : local_b;
// FIXME: swap multiply with bitshifts
// const uint32_t mask_odd = (block_k & 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_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_even & reinterpret_cast<uintmax_t>(local_b)));
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 num_threadblocks = problem_size / threads_per_threadblock;
using float_type = float;
// "static" shared memory allocation. This would determine threadblock
// occupancy of a single cluster
float *sharedmem_per_threadblock =
(float *)DEV_SMEM_START_ADDR + 2/*overkill for non-dma*/ * (2 * BM * BK) *
threadblock_id_in_cluster;
uint8_t *sharedmem_per_threadblock = reinterpret_cast<uint8_t *>(
DEV_SMEM_START_ADDR + sizeof(float_type) * 2 /*overkill for non-dma*/ *
(2 * BM * BK) * threadblock_id_in_cluster);
thread_block_gemm(arg, tid_in_threadblock, threads_per_threadblock,
threadblock_dim_y,
/*threadblock_id_x, threadblock_id_y,*/
threadblocks_per_cluster,
// threadblock_id,
threadblock_id_in_cluster,
sharedmem_per_threadblock);
thread_block_gemm<float_type>(
arg, tid_in_threadblock, threads_per_threadblock, threadblock_dim_y,
/*threadblock_id_x, threadblock_id_y,*/
threadblocks_per_cluster,
// threadblock_id,
threadblock_id_in_cluster, sharedmem_per_threadblock);
}
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_k = 64;
using float_type = half;
using float_type = float;
generate_source_matrix<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:
// * 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));
}
}