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

Merge branch 'new-cisc' of https://github.com/hansungk/vortex-private into new-cisc

Browse Source
This commit is contained in:
Richard Yan
2025-01-28 17:14:49 -08:00
parent b1e6495630 5ef4c8023e
commit 8c45b8b4b7
9 changed files with 1269 additions and 331 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
kernel/include/gemmini_mmio.h
View File

@@ -11,10 +11,11 @@
// #define SMEM_SIZE 0x4000
// 64KB
// #define SMEM_SIZE 0x10000
// 128KB
// #define SMEM_SIZE 0x20000
// 256KB
// 128KB (FP16 GEMM)
#define SMEM_SIZE 0x20000
// 256KB (FlashAttention)
// #define SMEM_SIZE 0x40000
#define SMEM_MASK (SMEM_SIZE - 1)
#define SMEM_ADDR_END (SMEM_BASE + SMEM_SIZE)
@@ -47,6 +48,7 @@ static size_t gemmini_tile_idx[NUM_THREADS * NUM_WARPS * NUM_CORES * NUM_CLUSTER
#define GEMMINI_RS2_ADDR (GEMMINI_CTRL + 0x18)
#define GEMMINI_INST_ADDR (GEMMINI_CTRL + 0x0)
#define GEMMINI_BUSY_ADDR (GEMMINI_CTRL + 0x20)
#define GEMMINI_OCCUPANCY_ADDR (GEMMINI_CTRL + 0x28)
#undef ROCC_INSTRUCTION_RS1_RS2
#define ROCC_INSTRUCTION_RS1_RS2(x, rs1, rs2, funct) { \
*((volatile uint64_t *) GEMMINI_RS1_ADDR) = (rs1); \
@@ -70,6 +72,8 @@ static size_t gemmini_tile_idx[NUM_THREADS * NUM_WARPS * NUM_CORES * NUM_CLUSTER
//#define gemmini_fence() { while (gemmini_status()); }
#define gemmini_fence() { while (*((volatile uint32_t *) GEMMINI_BUSY_ADDR)) asm volatile ("nop"); }
#define virgo_fence(n) { while (*((volatile uint32_t *) GEMMINI_OCCUPANCY_ADDR) > n) asm volatile ("nop"); }
/* cisc instructions */
/* ================= */
@@ -80,6 +84,8 @@ static size_t gemmini_tile_idx[NUM_THREADS * NUM_WARPS * NUM_CORES * NUM_CLUSTER
#define GEMMINI_CISC_CMD_R(x) asm("csrw 0xacc, %0" :: "r" (x))
#define GEMMINI_CISC_COMPUTE_HEXADECILES 0
#define GEMMINI_CISC_COMPUTE_AND_STORE_TO_SPAD 1
#define GEMMINI_CISC_MANUAL 2
#define GEMMINI_CISC_SET_AB_STRIDE 8
#define GEMMINI_CISC_STORE_TO_SPAD 9
#define GEMMINI_CISC_LOAD_TO_HEXADECILES 10
@@ -101,8 +107,19 @@ inline void gemmini_tile_load_ab(const elem_t * const a_addr, const elem_t * con
GEMMINI_CISC_CMD_R((b_hexadecile << 16) | (a_hexadecile << 8) | GEMMINI_CISC_LOAD_TO_HEXADECILES);
}
inline void gemmini_tile_compute(const uint32_t a_hexadecile, const uint32_t b_hexadecile, const bool accumulate) {
GEMMINI_CISC_CMD_R((accumulate << 24) | (b_hexadecile << 16) | (a_hexadecile << 8) | GEMMINI_CISC_COMPUTE_HEXADECILES);
template <bool store_to_spad = false>
inline void gemmini_tile_compute(const uint32_t a_hexadecile,
const uint32_t b_hexadecile,
const uint32_t d_hexadecile,
const bool accumulate) {
if constexpr (!store_to_spad) {
GEMMINI_CISC_CMD_R((static_cast<uint32_t>(accumulate) << 24) |
(b_hexadecile << 16) | (a_hexadecile << 8) |
GEMMINI_CISC_COMPUTE_HEXADECILES);
} else {
GEMMINI_CISC_CMD_R((d_hexadecile << 24) | (b_hexadecile << 16) |
(a_hexadecile << 8) | GEMMINI_CISC_COMPUTE_AND_STORE_TO_SPAD);
}
}
inline void gemmini_tile_store_c_gmem(elem_t * const c_addr,
@@ -125,6 +142,10 @@ inline void gemmini_tile_store_c_spad(const uint32_t c_hexadecile) {
GEMMINI_CISC_CMD_R(((uint32_t) (c_hexadecile << 8)) | GEMMINI_CISC_STORE_TO_SPAD);
}
inline void gemmini_manual_job() {
GEMMINI_CISC_CMD_I(GEMMINI_CISC_MANUAL);
}
/* inline static void sp_tiled_matmul_full_spad_ws(const uint32_t A_sp_addr_start, const uint32_t B_sp_addr_start,
const uint32_t D_sp_addr_start, const uint32_t C_dst_sp_addr_start,
size_t I, size_t J, size_t K, size_t pad_I, size_t pad_J, size_t pad_K,

2
tests/regression/flash_attention/Makefile
View File

@@ -2,6 +2,8 @@ PROJECT = flash_attention
SRCS = main.cpp common.h
# VX_SRCS = kernel.cpp
# VX_SRCS = kernel.gemmini.warpspec.cpp
VX_SRCS = kernel.gemmini.cpp
VX_INCLUDES = flash_impl.hpp ../sgemm_tcore/sgemm_impl.hpp

33
tests/regression/flash_attention/flash_impl.hpp
View File

@@ -4,6 +4,9 @@
#include <vx_spawn.h>
#include <float.h>
#define MARK_BEG() asm volatile ("slti x0, x1, -1047")
#define MARK_END() asm volatile ("slti x0, x1, -499")
#define B_ROW 64
#define B_COL 64
#define HEADDIM 64
@@ -11,8 +14,12 @@
#define ROW_REMAINDER_LOGIC
constexpr uint32_t ROWMAX_SETS = 3;
constexpr bool WARP_SPECIALIZED = true;
constexpr bool TENSOR_CORE = true;
// constexpr bool WARP_SPECIALIZED = true;
// constexpr bool GEMMINI_WARP_SPECIALIZED = false;
// constexpr bool TENSOR_CORE = true;
constexpr bool WARP_SPECIALIZED = false;
constexpr bool GEMMINI_WARP_SPECIALIZED = false;
constexpr bool TENSOR_CORE = false;
// temporary safety stop for wrong configs
static_assert(NUM_CORES == 4);
@@ -96,7 +103,7 @@ inline void thread_block_copy_rowmax(const float *src, float *dest,
dest[offset] = src[offset];
}
if constexpr (!TENSOR_CORE) {
if constexpr (!TENSOR_CORE && GEMMINI_WARP_SPECIALIZED) {
threadblock_barrier(1, 7);
} else {
threadblock_barrier(threadblock_id_in_cluster,
@@ -128,7 +135,7 @@ inline void thread_block_copy_tile(const float *src, float *dest,
if (row >= B_ROW) {
// WARNING: the number of barrier calls have to exactly match that in the
// outside of the branch to prevent stalls!! FIXME better proof this.
if constexpr (!TENSOR_CORE) {
if constexpr (!TENSOR_CORE && GEMMINI_WARP_SPECIALIZED) {
threadblock_barrier(1, 7);
} else {
threadblock_barrier(threadblock_id_in_cluster,
@@ -151,7 +158,7 @@ inline void thread_block_copy_tile(const float *src, float *dest,
dest[gmem_offset] = src[smem_offset];
}
if constexpr (!TENSOR_CORE) {
if constexpr (!TENSOR_CORE && GEMMINI_WARP_SPECIALIZED) {
threadblock_barrier(1, 7);
} else {
threadblock_barrier(threadblock_id_in_cluster,
@@ -208,7 +215,7 @@ __attribute__((always_inline)) inline void thread_block_online_softmax(
if (row >= B_ROW) {
// WARNING: the number of barrier calls have to exactly match that in the
// outside of the branch to prevent stalls!! FIXME better proof this.
if constexpr (!TENSOR_CORE) {
if constexpr (!TENSOR_CORE && GEMMINI_WARP_SPECIALIZED) {
threadblock_barrier(1, 7);
threadblock_barrier(1, 7);
threadblock_barrier(1, 7);
@@ -295,7 +302,7 @@ __attribute__((always_inline)) inline void thread_block_online_softmax(
warp_smem[tid_in_warp] = per_thread_max;
// sync writes to warp_smem
if constexpr (!TENSOR_CORE) {
if constexpr (!TENSOR_CORE && GEMMINI_WARP_SPECIALIZED) {
threadblock_barrier(1, 7);
} else {
threadblock_barrier(threadblock_id_in_cluster,
@@ -350,7 +357,7 @@ __attribute__((always_inline)) inline void thread_block_online_softmax(
#endif // PARALLEL_ROWMAX
#endif // DUMB_ROWMAX
if constexpr (!TENSOR_CORE) {
if constexpr (!TENSOR_CORE && GEMMINI_WARP_SPECIALIZED) {
threadblock_barrier(1, 7);
} else {
threadblock_barrier(threadblock_id_in_cluster,
@@ -398,7 +405,7 @@ __attribute__((always_inline)) inline void thread_block_online_softmax(
asm volatile("flashattn_exp_p_end_%=:" ::);
if constexpr (!TENSOR_CORE) {
if constexpr (!TENSOR_CORE && GEMMINI_WARP_SPECIALIZED) {
threadblock_barrier(1, 7);
} else {
threadblock_barrier(threadblock_id_in_cluster,
@@ -429,7 +436,7 @@ __attribute__((always_inline)) inline void thread_block_online_softmax(
warp_smem[tid_in_warp] = per_thread_sum;
// sync writes to warp_smem
if constexpr (!TENSOR_CORE) {
if constexpr (!TENSOR_CORE && GEMMINI_WARP_SPECIALIZED) {
threadblock_barrier(1, 7);
} else {
threadblock_barrier(threadblock_id_in_cluster,
@@ -462,7 +469,7 @@ __attribute__((always_inline)) inline void thread_block_online_softmax(
asm volatile("flashattn_rowsum_end_%=:" ::);
if constexpr (!TENSOR_CORE) {
if constexpr (!TENSOR_CORE && GEMMINI_WARP_SPECIALIZED) {
threadblock_barrier(1, 7);
} else {
threadblock_barrier(threadblock_id_in_cluster,
@@ -491,7 +498,7 @@ __attribute__((always_inline)) inline void thread_block_online_softmax(
asm volatile("flashattn_rescale_factor_end_%=:" ::);
if constexpr (!TENSOR_CORE) {
if constexpr (!TENSOR_CORE && GEMMINI_WARP_SPECIALIZED) {
threadblock_barrier(1, 7);
} else {
threadblock_barrier(threadblock_id_in_cluster,
@@ -546,7 +553,7 @@ __attribute__((always_inline)) inline void thread_block_O_rescale(
}
// reconverge after warp divergence
if constexpr (!TENSOR_CORE) {
if constexpr (!TENSOR_CORE && GEMMINI_WARP_SPECIALIZED) {
threadblock_barrier(1, 7);
} else {
threadblock_barrier(threadblock_id_in_cluster,

138
tests/regression/flash_attention/kernel.cpp
View File

@@ -8,6 +8,8 @@
#include "gemmini_mmio.h"
#include "flash_impl.hpp"
#define GEMMINI_NEW_CISC 1
constexpr bool DEBUG = false;
constexpr bool Q_IS_K_MAJOR = true;
@@ -94,6 +96,14 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
uint8_t *smem_per_threadblock = reinterpret_cast<uint8_t *>(
DEV_SMEM_START_ADDR);
constexpr uint32_t smem_start = DEV_SMEM_START_ADDR;
constexpr uint32_t smem_hexadecile_size = (SMEM_SIZE / 16);
// currently assumes the Q/K/V tile sizes exactly match the hexadecile size
static_assert(smem_hexadecile_size == smem_Q_size * sizeof(float));
static_assert(smem_hexadecile_size == smem_K_size * sizeof(float));
static_assert(smem_hexadecile_size == smem_QK_size * sizeof(float));
static_assert(smem_hexadecile_size == smem_V_size * sizeof(float));
static_assert(smem_hexadecile_size == smem_O_size * sizeof(float));
constexpr uint32_t smem_octet0 = 0 * (SMEM_SIZE / 8);
constexpr uint32_t smem_octet1 = 1 * (SMEM_SIZE / 8);
constexpr uint32_t smem_octet2 = 2 * (SMEM_SIZE / 8);
@@ -108,31 +118,31 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
// half
// at the same time, make sure Q and K are in different banks so that they
// can be accessed in parallel for GEMM; same for P and V
constexpr uint32_t smem_Q0_offset = smem_octet0;
constexpr uint32_t smem_Q1_offset = smem_octet4;
constexpr uint32_t smem_K0_offset = smem_octet1;
constexpr uint32_t smem_K1_offset = smem_octet5;
constexpr uint32_t smem_V0_offset = smem_K0_offset + smem_K_size * sizeof(float);
constexpr uint32_t smem_V1_offset = smem_K1_offset + smem_K_size * sizeof(float);
constexpr uint32_t smem_S0_offset = smem_octet2;
constexpr uint32_t smem_S1_offset = smem_octet6;
constexpr uint32_t smem_P0_offset = smem_Q0_offset + smem_Q_size * sizeof(float);
constexpr uint32_t smem_P1_offset = smem_Q1_offset + smem_Q_size * sizeof(float);
constexpr uint32_t smem_O0_offset = smem_octet3;
constexpr uint32_t smem_O1_offset = smem_octet7;
constexpr uint32_t smem_Q0_hexadecile = 2 * 0; // octet0
constexpr uint32_t smem_Q1_hexadecile = 2 * 4; // octet4
constexpr uint32_t smem_K0_hexadecile = 2 * 1; // octet1
constexpr uint32_t smem_K1_hexadecile = 2 * 5; // octet5
constexpr uint32_t smem_V0_hexadecile = smem_K0_hexadecile + 1;
constexpr uint32_t smem_V1_hexadecile = smem_K1_hexadecile + 1;
constexpr uint32_t smem_S0_hexadecile = 2 * 2; // octet2
constexpr uint32_t smem_S1_hexadecile = 2 * 6; // octet6
constexpr uint32_t smem_P0_hexadecile = smem_Q0_hexadecile + 1;
constexpr uint32_t smem_P1_hexadecile = smem_Q1_hexadecile + 1;
constexpr uint32_t smem_O0_hexadecile = 2 * 3; // octet3
constexpr uint32_t smem_O1_hexadecile = 2 * 7; // octet7
float *smem_Q0 = reinterpret_cast<float *>(smem_start + smem_Q0_offset);
float *smem_Q1 = reinterpret_cast<float *>(smem_start + smem_Q1_offset);
float *smem_K0 = reinterpret_cast<float *>(smem_start + smem_K0_offset);
float *smem_K1 = reinterpret_cast<float *>(smem_start + smem_K1_offset);
float *smem_V0 = reinterpret_cast<float *>(smem_start + smem_V0_offset);
float *smem_V1 = reinterpret_cast<float *>(smem_start + smem_V1_offset);
float *smem_S0 = reinterpret_cast<float *>(smem_start + smem_S0_offset);
float *smem_S1 = reinterpret_cast<float *>(smem_start + smem_S1_offset);
float *smem_P0 = reinterpret_cast<float *>(smem_start + smem_P0_offset);
float *smem_P1 = reinterpret_cast<float *>(smem_start + smem_P1_offset);
float *smem_O0 = reinterpret_cast<float *>(smem_start + smem_O0_offset);
float *smem_O1 = reinterpret_cast<float *>(smem_start + smem_O1_offset);
float *smem_Q0 = reinterpret_cast<float *>(smem_start + smem_Q0_hexadecile * smem_hexadecile_size);
float *smem_Q1 = reinterpret_cast<float *>(smem_start + smem_Q1_hexadecile * smem_hexadecile_size);
float *smem_K0 = reinterpret_cast<float *>(smem_start + smem_K0_hexadecile * smem_hexadecile_size);
float *smem_K1 = reinterpret_cast<float *>(smem_start + smem_K1_hexadecile * smem_hexadecile_size);
float *smem_V0 = reinterpret_cast<float *>(smem_start + smem_V0_hexadecile * smem_hexadecile_size);
float *smem_V1 = reinterpret_cast<float *>(smem_start + smem_V1_hexadecile * smem_hexadecile_size);
float *smem_S0 = reinterpret_cast<float *>(smem_start + smem_S0_hexadecile * smem_hexadecile_size);
float *smem_S1 = reinterpret_cast<float *>(smem_start + smem_S1_hexadecile * smem_hexadecile_size);
float *smem_P0 = reinterpret_cast<float *>(smem_start + smem_P0_hexadecile * smem_hexadecile_size);
float *smem_P1 = reinterpret_cast<float *>(smem_start + smem_P1_hexadecile * smem_hexadecile_size);
float *smem_O0 = reinterpret_cast<float *>(smem_start + smem_O0_hexadecile * smem_hexadecile_size);
float *smem_O1 = reinterpret_cast<float *>(smem_start + smem_O1_hexadecile * smem_hexadecile_size);
// allocate rowmax/rowsum storage at the end of the sharedmem address space
constexpr uint32_t smem_rowmax_size = B_ROW * ROWMAX_SETS;
@@ -180,25 +190,32 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
float *smem_scratchpad =
(warpgroup_id % 2) ? smem_scratchpad_1 : smem_scratchpad_0;
#ifdef GEMMINI_NEW_CISC
const auto spad_hex_Q = (warpgroup_id % 2) ? smem_Q1_hexadecile : smem_Q0_hexadecile;
const auto spad_hex_K = (warpgroup_id % 2) ? smem_K1_hexadecile : smem_K0_hexadecile;
const auto spad_hex_V = (warpgroup_id % 2) ? smem_V1_hexadecile : smem_V0_hexadecile;
const auto spad_hex_S = (warpgroup_id % 2) ? smem_S1_hexadecile : smem_S0_hexadecile;
#else
static_assert(sizeof(elem_t) == sizeof(float));
constexpr uint32_t spad_addr_factor = DIM * sizeof(elem_t);
constexpr uint32_t spad_addr_Q0 = smem_Q0_offset / spad_addr_factor;
constexpr uint32_t spad_addr_Q1 = smem_Q1_offset / spad_addr_factor;
constexpr uint32_t spad_addr_K0 = smem_K0_offset / spad_addr_factor;
constexpr uint32_t spad_addr_K1 = smem_K1_offset / spad_addr_factor;
constexpr uint32_t spad_addr_V0 = smem_V0_offset / spad_addr_factor;
constexpr uint32_t spad_addr_V1 = smem_V1_offset / spad_addr_factor;
constexpr uint32_t spad_addr_S0 = smem_S0_offset / spad_addr_factor;
constexpr uint32_t spad_addr_S1 = smem_S1_offset / spad_addr_factor;
constexpr uint32_t spad_addr_P0 = smem_P0_offset / spad_addr_factor;
constexpr uint32_t spad_addr_P1 = smem_P1_offset / spad_addr_factor;
constexpr uint32_t spad_addr_O0 = smem_O0_offset / spad_addr_factor;
constexpr uint32_t spad_addr_O1 = smem_O1_offset / spad_addr_factor;
constexpr uint32_t spad_addr_Q0 = smem_Q0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_Q1 = smem_Q1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_K0 = smem_K0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_K1 = smem_K1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_V0 = smem_V0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_V1 = smem_V1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_S0 = smem_S0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_S1 = smem_S1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_P0 = smem_P0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_P1 = smem_P1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_O0 = smem_O0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_O1 = smem_O1_hexadecile * smem_hexadecile_size / spad_addr_factor;
const auto spad_addr_Q = (warpgroup_id % 2) ? spad_addr_Q1 : spad_addr_Q0;
const auto spad_addr_K = (warpgroup_id % 2) ? spad_addr_K1 : spad_addr_K0;
const auto spad_addr_V = (warpgroup_id % 2) ? spad_addr_V1 : spad_addr_V0;
const auto spad_addr_S = (warpgroup_id % 2) ? spad_addr_S1 : spad_addr_S0;
#endif
// initialize rowmax/rowsum values in sharedmem
thread_block_init_sharedmem(tid_in_warpgroup, threads_per_warpgroup, smem_O,
@@ -219,6 +236,8 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
loop_matmul_skips(/*skip_lda=*/0, /*skip_ldb=*/0, /*skip_ldd=*/1,
/*skip_ex=*/1, /*skip_stc=*/1);
MARK_BEG();
if constexpr (GEMMINI_DMA) {
if (tid_in_warpgroup == 0) {
gemmini_extended_config_ex(WEIGHT_STATIONARY, 0, 0, 1, 0, 0);
@@ -244,8 +263,6 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
// other warps behave differently on the branch condition.
// threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core);
// move Q and K into SMEM before the loop starts
//
static_assert(B_ROW == B_COL, "currently only supports square tiles");
if constexpr (GEMMINI_DMA) {
@@ -254,22 +271,28 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
if (tid_in_warpgroup == 0) {
const float *gmem_Q_tile = gmem_Q + HEADDIM * B_ROW * warpgroup_id;
const float *gmem_K_tile = gmem_K;
// do DMA
//
// move Q and K into SMEM before the loop starts. Note this will be done
// separately for the two warpgroups
//
#ifdef GEMMINI_NEW_CISC
// the target addresses of this should match with spad_addr_Q0 and
// spad_addr_K0 set in this kernel
gemmini_tile_load_ab(gmem_Q_tile, gmem_K_tile, spad_hex_Q,
spad_hex_K, 0 /*tile_idx_i*/,
0 /*tile_idx_j*/, 0 /*tile_idx_k*/, dim_seqlen,
dim_seqlen, HEADDIM, B_ROW, B_COL, HEADDIM);
#else
// configure the GMEM addresses for the DMA to read from
ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, (uint64_t)(gmem_Q_tile),
(uint64_t)(gmem_K_tile),
k_LOOP_WS_CONFIG_ADDRS_AB)
// configure address strides for the DMA
GEMMINI_CISC_CMD_R((dim_seqlen << 16) | (HEADDIM << 8) |
GEMMINI_CISC_CMD_R((dim_seqlen << 20) | (HEADDIM << 8) |
8 /*k_LOOP_WS_CONFIG_STRIDES_AB*/);
gemmini_fence();
// #define GEMMINI_DMA_CISC
#ifdef GEMMINI_DMA_CISC
GEMMINI_CISC_CMD_I(9);
gemmini_fence();
#else
// do DMA
//
// among other things, this also configures CONFIG_BOUNDS so that the
// DMA knows the full matrix dimensions
sp_tiled_matmul_full_spad_ws(
@@ -279,9 +302,11 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
/*pad_I=*/0, /*pad_J=*/0, /*pad_K=*/0,
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips);
gemmini_fence();
#endif
// block until DMA complete
gemmini_fence();
// re-configure DMA for K and V load that will later happen in the loop
// GMEM addr stride for K
gemmini_extended3_config_ld(dim_seqlen * sizeof(elem_t), MVIN_SCALE_IDENTITY,
@@ -330,7 +355,7 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
// "inner loop" along the columns of K^T
const uint32_t k_tiles = (dim_seqlen / B_COL);
for (uint32_t tile_k = 0; tile_k < (4 /* for perf measurement */ * k_tiles);
for (uint32_t tile_k = 0; tile_k < (1 /* for perf measurement */ * k_tiles);
tile_k++) {
// float *smem_P_produce = (tile_k % 2) ? smem_P0 : smem_P1;
// float *smem_P_consume = (tile_k % 2) ? smem_P1 : smem_P0;
@@ -547,13 +572,20 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, (uint64_t)(gmem_K_tile),
(uint64_t)(gmem_V_tile),
k_LOOP_WS_CONFIG_ADDRS_AB)
// do DMA
#ifdef GEMMINI_NEW_CISC
gemmini_tile_load_ab(gmem_K_tile, gmem_V_tile, spad_hex_K, spad_hex_V,
0 /*tile_idx_i*/, 0 /*tile_idx_j*/,
0 /*tile_idx_k*/, HEADDIM /*dim_m of KT*/,
HEADDIM /*dim_n of V*/, dim_seqlen /*dim_k of KT*/,
B_ROW, HEADDIM, B_COL);
#else
// configure address strides for the DMA
// FIXME: unnecessary?
GEMMINI_CISC_CMD_R((HEADDIM /*V*/ << 16) | (dim_seqlen /*KT*/ << 8) |
GEMMINI_CISC_CMD_R((HEADDIM /*V*/ << 20) | (dim_seqlen /*KT*/ << 8) |
8 /*k_LOOP_WS_CONFIG_STRIDES_AB*/);
gemmini_fence();
// do DMA
sp_tiled_matmul_full_spad_ws(
spad_addr_K, spad_addr_V,
/*spad_D=*/0, /*spad_C=*/spad_addr_S,
@@ -561,6 +593,8 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
/*pad_I=*/0, /*pad_J=*/0, /*pad_K=*/0,
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips);
#endif
// FIXME: necessary?
gemmini_fence();
}
} else {
@@ -813,8 +847,6 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
warps_per_warpgroup_per_core);
}
}
#if 0
#endif
}
asm volatile ("tile_loop_finish_%=:" :: );
@@ -824,6 +856,8 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
if (warpgroup_id == 0) {
threadblock_barrier(global_barrier_id, warps_per_threadblock_per_core);
}
MARK_END();
}
int main() {

470
tests/regression/flash_attention/kernel.gemmini.cpp
View File

@@ -10,6 +10,10 @@
#define FENCE_GEMM_II
#define GEMMINI_NEW_CISC 1
static_assert(GEMMINI_NEW_CISC, "NOTE: old non-CISC code is untested; look for "
"any misalignment of fields in ciscArgs.");
constexpr bool DEBUG = false;
static_assert(GEMMINI_DMA && !WARP_SPECIALIZED,
@@ -98,40 +102,43 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
"flashattention kernel assumes 1 threadblock occupancy per cluster");
uint8_t *smem_per_threadblock = reinterpret_cast<uint8_t *>(DEV_SMEM_START_ADDR);
constexpr uint32_t smem_start = DEV_SMEM_START_ADDR;
constexpr uint32_t smem_quart0 = 0 * (SMEM_SIZE / 4);
constexpr uint32_t smem_quart1 = 1 * (SMEM_SIZE / 4);
constexpr uint32_t smem_quart2 = 2 * (SMEM_SIZE / 4);
constexpr uint32_t smem_quart3 = 3 * (SMEM_SIZE / 4);
constexpr uint32_t smem_hexadecile_size = (SMEM_SIZE / 16);
// currently assumes the Q/K/V tile sizes exactly match the hexadecile size
static_assert(smem_hexadecile_size == smem_Q_size * sizeof(float));
static_assert(smem_hexadecile_size == smem_K_size * sizeof(float));
static_assert(smem_hexadecile_size == smem_QK_size * sizeof(float));
static_assert(smem_hexadecile_size == smem_V_size * sizeof(float));
static_assert(smem_hexadecile_size == smem_O_size * sizeof(float));
// Q/V/S in quart0/1, K/P/O in quart2/3
constexpr uint32_t smem_Q0_offset = smem_quart0;
constexpr uint32_t smem_Q1_offset = smem_quart1;
constexpr uint32_t smem_K0_offset = smem_quart2;
constexpr uint32_t smem_K1_offset = smem_quart3;
constexpr uint32_t smem_V0_offset = smem_Q0_offset + smem_Q_size * sizeof(float);
constexpr uint32_t smem_V1_offset = smem_Q1_offset + smem_Q_size * sizeof(float);
constexpr uint32_t smem_Q0_hexadecile = 4 * 0;
constexpr uint32_t smem_Q1_hexadecile = 4 * 1;
constexpr uint32_t smem_K0_hexadecile = 4 * 2;
constexpr uint32_t smem_K1_hexadecile = 4 * 3;
constexpr uint32_t smem_V0_hexadecile = smem_Q0_hexadecile + 1;
constexpr uint32_t smem_V1_hexadecile = smem_Q1_hexadecile + 1;
// put S1/S0 with V0/V1 so that softmax and GEMM-II doesn't cause bank
// conflicts
constexpr uint32_t smem_S0_offset = smem_V1_offset + smem_V_size * sizeof(float);
constexpr uint32_t smem_S1_offset = smem_V0_offset + smem_V_size * sizeof(float);
constexpr uint32_t smem_P0_offset = smem_K0_offset + smem_K_size * sizeof(float);
constexpr uint32_t smem_P1_offset = smem_K1_offset + smem_K_size * sizeof(float);
constexpr uint32_t smem_S0_hexadecile = smem_V1_hexadecile + 1;
constexpr uint32_t smem_S1_hexadecile = smem_V0_hexadecile + 1;
constexpr uint32_t smem_P0_hexadecile = smem_K0_hexadecile + 1;
constexpr uint32_t smem_P1_hexadecile = smem_K1_hexadecile + 1;
// reversed!
constexpr uint32_t smem_O0_offset = smem_P1_offset + smem_QK_size * sizeof(float);
constexpr uint32_t smem_O1_offset = smem_P0_offset + smem_QK_size * sizeof(float); // unused
constexpr uint32_t smem_O0_hexadecile = smem_P1_hexadecile + 1;
constexpr uint32_t smem_O1_hexadecile = smem_P0_hexadecile + 1; // unused
float *smem_Q0 = reinterpret_cast<float *>(smem_start + smem_Q0_offset);
float *smem_Q1 = reinterpret_cast<float *>(smem_start + smem_Q1_offset);
float *smem_K0 = reinterpret_cast<float *>(smem_start + smem_K0_offset);
float *smem_K1 = reinterpret_cast<float *>(smem_start + smem_K1_offset);
float *smem_V0 = reinterpret_cast<float *>(smem_start + smem_V0_offset);
float *smem_V1 = reinterpret_cast<float *>(smem_start + smem_V1_offset);
float *smem_S0 = reinterpret_cast<float *>(smem_start + smem_S0_offset);
float *smem_S1 = reinterpret_cast<float *>(smem_start + smem_S1_offset);
float *smem_P0 = reinterpret_cast<float *>(smem_start + smem_P0_offset);
float *smem_P1 = reinterpret_cast<float *>(smem_start + smem_P1_offset);
float *smem_O0 = reinterpret_cast<float *>(smem_start + smem_O0_offset);
float *smem_O1 = reinterpret_cast<float *>(smem_start + smem_O1_offset);
float *smem_Q0 = reinterpret_cast<float *>(smem_start + smem_Q0_hexadecile * smem_hexadecile_size);
float *smem_Q1 = reinterpret_cast<float *>(smem_start + smem_Q1_hexadecile * smem_hexadecile_size);
float *smem_K0 = reinterpret_cast<float *>(smem_start + smem_K0_hexadecile * smem_hexadecile_size);
float *smem_K1 = reinterpret_cast<float *>(smem_start + smem_K1_hexadecile * smem_hexadecile_size);
float *smem_V0 = reinterpret_cast<float *>(smem_start + smem_V0_hexadecile * smem_hexadecile_size);
float *smem_V1 = reinterpret_cast<float *>(smem_start + smem_V1_hexadecile * smem_hexadecile_size);
float *smem_S0 = reinterpret_cast<float *>(smem_start + smem_S0_hexadecile * smem_hexadecile_size);
float *smem_S1 = reinterpret_cast<float *>(smem_start + smem_S1_hexadecile * smem_hexadecile_size);
float *smem_P0 = reinterpret_cast<float *>(smem_start + smem_P0_hexadecile * smem_hexadecile_size);
float *smem_P1 = reinterpret_cast<float *>(smem_start + smem_P1_hexadecile * smem_hexadecile_size);
float *smem_O0 = reinterpret_cast<float *>(smem_start + smem_O0_hexadecile * smem_hexadecile_size);
float *smem_O1 = reinterpret_cast<float *>(smem_start + smem_O1_hexadecile * smem_hexadecile_size);
// allocate rowmax/rowsum storage at the end of the sharedmem address space
constexpr uint32_t smem_rowmax_size = B_ROW * ROWMAX_SETS;
@@ -168,18 +175,18 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
static_assert(sizeof(elem_t) == sizeof(float));
constexpr uint32_t spad_addr_factor = DIM * sizeof(elem_t);
constexpr uint32_t spad_addr_Q0 = smem_Q0_offset / spad_addr_factor;
constexpr uint32_t spad_addr_Q1 = smem_Q1_offset / spad_addr_factor;
constexpr uint32_t spad_addr_K0 = smem_K0_offset / spad_addr_factor;
constexpr uint32_t spad_addr_K1 = smem_K1_offset / spad_addr_factor;
constexpr uint32_t spad_addr_V0 = smem_V0_offset / spad_addr_factor;
constexpr uint32_t spad_addr_V1 = smem_V1_offset / spad_addr_factor;
constexpr uint32_t spad_addr_S0 = smem_S0_offset / spad_addr_factor;
constexpr uint32_t spad_addr_S1 = smem_S1_offset / spad_addr_factor;
constexpr uint32_t spad_addr_P0 = smem_P0_offset / spad_addr_factor;
constexpr uint32_t spad_addr_P1 = smem_P1_offset / spad_addr_factor;
constexpr uint32_t spad_addr_O0 = smem_O0_offset / spad_addr_factor;
constexpr uint32_t spad_addr_O1 = smem_O1_offset / spad_addr_factor;
constexpr uint32_t spad_addr_Q0 = smem_Q0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_Q1 = smem_Q1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_K0 = smem_K0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_K1 = smem_K1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_V0 = smem_V0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_V1 = smem_V1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_S0 = smem_S0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_S1 = smem_S1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_P0 = smem_P0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_P1 = smem_P1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_O0 = smem_O0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_O1 = smem_O1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t global_barrier_id = NUM_WARPS - 1; // arbitrary
static_assert(warps_per_threadblock_per_core == NUM_WARPS);
@@ -212,6 +219,8 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
loop_matmul_skips(/*skip_lda=*/1, /*skip_ldb=*/1, /*skip_ldd=*/0,
/*skip_ex=*/0, /*skip_stc=*/1);
MARK_BEG();
if (tid_in_warpgroup == 0) {
gemmini_extended_config_ex(WEIGHT_STATIONARY, 0, 0, 1, 0, 0);
@@ -244,22 +253,27 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
// make sure to read from the correct row of Q
const float *gmem_Q_tile = gmem_Q + HEADDIM * B_ROW * warpgroup_id;
const float *gmem_K_tile = gmem_K;
// do DMA
//
// move Q to spad_addr_Q0 for the first iteration
//
#ifdef GEMMINI_NEW_CISC
// the target addresses of this should match with spad_addr_Q0 and
// spad_addr_K0 set in this kernel
gemmini_tile_load_ab(gmem_Q_tile, gmem_K_tile, smem_Q0_hexadecile,
smem_K0_hexadecile, 0 /*tile_idx_i*/, 0 /*tile_idx_j*/,
0 /*tile_idx_k*/, dim_seqlen, dim_seqlen, HEADDIM,
B_ROW, B_COL, HEADDIM);
#else
// configure the GMEM addresses for the DMA to read from
ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, (uint64_t)(gmem_Q_tile),
(uint64_t)(gmem_K_tile), k_LOOP_WS_CONFIG_ADDRS_AB)
// configure address strides for the DMA
GEMMINI_CISC_CMD_R((dim_seqlen << 20) | (HEADDIM << 8) |
8 /*k_LOOP_WS_CONFIG_STRIDES_AB*/);
GEMMINI_CISC_SET_AB_STRIDE);
gemmini_fence();
// #define GEMMINI_DMA_CISC
#ifdef GEMMINI_DMA_CISC
// the target addresses of this should match with spad_addr_Q0 and
// spad_addr_K0 set in this kernel
GEMMINI_CISC_CMD_I(10);
#else
// do DMA
//
// among other things, this also configures CONFIG_BOUNDS so that the
// DMA knows the full matrix dimensions
sp_tiled_matmul_full_spad_ws(
@@ -270,9 +284,18 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips);
#endif
// block until DMA complete
gemmini_fence();
// need to also move Q to spad_addr_Q1 for the next iteration
// also move Q to spad_addr_Q1 for the second iteration
//
#ifdef GEMMINI_NEW_CISC
gemmini_tile_load_ab(gmem_Q_tile, gmem_K_tile, smem_Q1_hexadecile,
smem_K1_hexadecile, 0 /*tile_idx_i*/, 0 /*tile_idx_j*/,
0 /*tile_idx_k*/, dim_seqlen, dim_seqlen, HEADDIM,
B_ROW, B_COL, HEADDIM);
#else
// FIXME: re-configure necessary?
gmem_K_tile = gmem_K + (B_COL * 1);
ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, (uint64_t)(gmem_Q_tile),
@@ -280,9 +303,7 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
GEMMINI_CISC_CMD_R((dim_seqlen << 20) | (HEADDIM << 8) |
8 /*k_LOOP_WS_CONFIG_STRIDES_AB*/);
gemmini_fence();
#ifdef GEMMINI_DMA_CISC
// GEMMINI_CISC_CMD_I(11);
#else
sp_tiled_matmul_full_spad_ws(
spad_addr_Q1, spad_addr_K1/*bogus*/,
/*spad_D=*/0, /*spad_C=*/spad_addr_S0/*bogus*/,
@@ -292,12 +313,12 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips);
#endif
gemmini_fence();
gemmini_fence();
gemmini_fence();
// block until DMA complete
gemmini_fence();
// re-configure DMA for K and V load that will later happen in the loop
// FIXME: not sure necessary with new CISC
//
// GMEM addr stride for K
gemmini_extended3_config_ld(dim_seqlen * sizeof(elem_t),
MVIN_SCALE_IDENTITY, false, 0);
@@ -321,22 +342,16 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
// threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core);
// }
constexpr uint32_t threads_per_warpgroup_simt =
threads_per_warpgroup -
CORES_PER_CLUSTER * NUM_THREADS /*warp 0, 4, 8, 12*/;
constexpr uint32_t warpgroup_id_simt = 1;
constexpr uint32_t barrier_id_simt = 1;
constexpr uint32_t barrier_count_simt = NUM_WARPS - 1;
const uint32_t tid_in_warpgroup_simt =
tid_in_warpgroup - (CORES_PER_CLUSTER * NUM_THREADS);
static_assert(barrier_id_simt == 1 && barrier_count_simt == 7);
asm volatile ("tile_loop_start_%=:" :: );
// "inner loop" along the columns of K^T
const uint32_t k_tiles = (dim_seqlen / B_COL);
for (uint32_t tile_k = 0;
tile_k < (4 /*for perf measurement*/ * k_tiles) + 2 /*pipeline latency*/;
tile_k < (1 /*for perf measurement*/ *
// virgo kernel is fully pipelined around (2 GEMMs | softmax);
// requires two loop iterations to finish one tile compute
(2 * k_tiles)) +
2 /*pipeline latency*/;
tile_k++) {
if constexpr (DEBUG || true) {
threadblock_barrier(global_barrier_id, warps_per_threadblock_per_core);
@@ -373,41 +388,53 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
const auto spad_addr_P_consume = (tile_k & 1) ? spad_addr_P1 : spad_addr_P0;
const auto spad_addr_P_produce = (tile_k & 1) ? spad_addr_P0 : spad_addr_P1;
const auto spad_addr_O = spad_addr_O0; // NOTE: there's only single O tile
const auto spad_hex_Q = smem_Q0_hexadecile;
const auto spad_hex_K_consume = (tile_k & 1) ? smem_K1_hexadecile : smem_K0_hexadecile;
const auto spad_hex_K_produce = (tile_k & 1) ? smem_K0_hexadecile : smem_K1_hexadecile;
const auto spad_hex_V_consume = (tile_k & 1) ? smem_V1_hexadecile : smem_V0_hexadecile;
const auto spad_hex_V_produce = (tile_k & 1) ? smem_V0_hexadecile : smem_V1_hexadecile;
const auto spad_hex_S_consume = (tile_k & 1) ? smem_S1_hexadecile : smem_S0_hexadecile;
const auto spad_hex_S_produce = (tile_k & 1) ? smem_S0_hexadecile : smem_S1_hexadecile;
const auto spad_hex_P_consume = (tile_k & 1) ? smem_P1_hexadecile : smem_P0_hexadecile;
const auto spad_hex_P_produce = (tile_k & 1) ? smem_P0_hexadecile : smem_P1_hexadecile;
const auto spad_hex_O = smem_O0_hexadecile; // NOTE: there's only single O tile
asm volatile ("dbuf_sel_end_%=:" :: );
if (vx_warp_id() == 0 /* warp 0 in every core */) {
if (tile_k >= 2) // delay by 2 iters for pipelining
{
const uint32_t tile_k_ = tile_k - 2;
{
// fence completion of the GEMMs in the previous loop iterations. Note
// this is done at the start of the loop to maximize window of
// overlapping.
//
// NOTE: this ideally needs to be put inside tid_in_warpgroup == 0
// branch, but that triggers a TL source ID re-used assertion we haven't
// looked at yet.
gemmini_fence();
// GEMM II: O = O + P*V
// --------------------
// This is done *before* GEMM I in the software pipeline, working on the
// online softmax result tile from the previous iteration
// do all of GEMM kickoffs before the SIMT compute
//
if (tid_in_warpgroup == 0) {
asm volatile("gemm_pv_start_%=:" ::);
if (tile_k >= 2) // delay GEMM II by 2 iters for pipelining
{
const uint32_t tile_k_ = tile_k - 2;
if (tid_in_warpgroup == 0) {
#if 0
if (tile_k_ == 0) {
gemmini_fence();
GEMMINI_CISC_CMD_I(0);
} else if (tile_k_ & 1) {
gemmini_fence();
GEMMINI_CISC_CMD_I(2);
} else {
gemmini_fence();
GEMMINI_CISC_CMD_I(1);
}
#else
// kickoff matmul
// among other things, this also configures CONFIG_BOUNDS so that the
// DMA knows the full matrix dimensions
// GEMM II: O = O + P*V
// --------------------
// This is done *before* GEMM I in the software pipeline, working on
// the online softmax result tile from the previous iteration
asm volatile("gemm_pv_start_%=:" ::);
// kick off GEMM II
//
// FIXME: perf: prevent GMEM->SMEM load for O tile
gemmini_fence();
gemmini_fence();
gemmini_fence();
gemmini_fence();
#ifdef GEMMINI_NEW_CISC
gemmini_tile_compute</*store_to_spad=*/true>(
spad_hex_P_consume, spad_hex_V_consume, spad_hex_O,
0 /*accumulate.
FIXME: Gemmini doens't support accumulation from a spad tile*/);
#else
sp_tiled_matmul_full_spad_ws(
spad_addr_P_consume, spad_addr_V_consume,
/*spad_D=*/spad_addr_O, /*spad_C=*/spad_addr_O,
@@ -416,37 +443,34 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips_matmul);
#endif
asm volatile("gemm_pv_finish_%=:" ::);
}
// // reconverge from mmio divergence
// threadblock_barrier(warpgroup_id_in_cluster,
// warps_per_warpgroup_per_core);
// GEMM I: S = Q*K
//
// kick off asynchronously; fence later
asm volatile("gemm_qk_start_%=:" ::);
asm volatile("gemm_pv_finish_%=:" ::);
}
// FIXME: remove
// // fence to GEMM II completion
// gemmini_fence();
// GEMM I: S = Q*K
//
// kick off asynchronously; fence later
asm volatile("gemm_qk_start_%=:" ::);
// #ifdef FENCE_GEMM_II
// asm volatile("rescale_fence_write_start_%=:" ::);
// // signal that GEMM II is finished to O rescale step
// *smem_O_flag = 1;
// vx_fence();
// asm volatile("rescale_fence_write_end_%=:" ::);
// #endif
if (tid_in_warpgroup == 0) {
// fence to GEMM II completion
gemmini_fence();
gemmini_fence();
gemmini_fence();
gemmini_fence();
#ifdef FENCE_GEMM_II
// signal that GEMM II is finished to O rescale step
*smem_O_flag = 1;
vx_fence();
#endif
// 0,2,.: opcode 0 (quartile 0/2, no accum)
// 1,3,.: opcode 3 (quartile 1/3, no accum)
// const uint32_t opcode = 3 * (tile_k & 1);
//GEMMINI_CISC_CMD_I(opcode);
// Kick off GEMM I
//
#ifdef GEMMINI_NEW_CISC
gemmini_tile_compute</*store_to_spad=*/true>(
spad_hex_Q, spad_hex_K_consume, spad_hex_S_produce,
0 /*accumulate*/);
#else
sp_tiled_matmul_full_spad_ws(
spad_addr_Q, spad_addr_K_consume,
/*spad_D=*/0, /*spad_C=*/spad_addr_S_produce,
@@ -454,11 +478,8 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
/*pad_I=*/0, /*pad_J=*/0, /*pad_K=*/0,
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips_matmul);
#endif
// gemmini_fence();
// gemmini_fence();
// gemmini_fence();
// gemmini_fence();
asm volatile("gemm_qk_finish_%=:" ::);
// data move for K and V
@@ -474,30 +495,40 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
const float *gmem_V_tile =
gmem_V + (HEADDIM * B_COL * (tile_k - 1 /*dragbehind*/));
#if 0
// fence mvout S to SMEM
gemmini_fence();
ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, (uint64_t)(gmem_K_tile),
(uint64_t)(gmem_V_tile),
k_LOOP_WS_CONFIG_ADDRS_AB)
#endif
// configure address strides for the DMA
// FIXME: unnecessary?
GEMMINI_CISC_CMD_R((HEADDIM /*V*/ << 20) | (dim_seqlen /*KT*/ << 8) |
8 /*k_LOOP_WS_CONFIG_STRIDES_AB*/);
// gemmini_fence();
// do DMA
if (tile_k == 0) {
// commented out as we do two move-ins before the loop starts
//
// // configure address strides for the DMA
// // FIXME: unnecessary?
// GEMMINI_CISC_CMD_R((HEADDIM /*V*/ << 20) | (dim_seqlen /*KT*/ <<
// 8) |
// 8 /*k_LOOP_WS_CONFIG_STRIDES_AB*/);
// gemmini_fence();
//
// we load (k-1)th tile for V; skip V for the 1st iteration,
// sp_tiled_matmul_full_spad_ws(
// spad_addr_K_produce, spad_addr_V_produce,
// /*spad_D=*/0, /*spad_C=*/0,
// /*I=*/(B_ROW / DIM), /*J=*/(HEADDIM / DIM), /*K=*/(B_COL / DIM),
// /*I=*/(B_ROW / DIM), /*J=*/(HEADDIM / DIM), /*K=*/(B_COL /
// DIM),
// /*pad_I=*/0, /*pad_J=*/0, /*pad_K=*/0,
// /*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
// /*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0,
// /*low_D=*/0,
// /*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips_only_a);
} else {
#ifdef GEMMINI_NEW_CISC
gemmini_tile_load_ab(
gmem_K_tile, gmem_V_tile, spad_hex_K_produce, spad_hex_V_produce,
0 /*tile_idx_i*/, 0 /*tile_idx_j*/, 0 /*tile_idx_k*/,
HEADDIM /*dim_m of KT*/, HEADDIM /*dim_n of V*/,
dim_seqlen /*dim_k of KT*/, B_ROW, HEADDIM, B_COL);
#else
// configure address strides for the DMA
// FIXME: unnecessary?
GEMMINI_CISC_CMD_R((HEADDIM /*V*/ << 20) | (dim_seqlen /*KT*/ << 8) |
8 /*k_LOOP_WS_CONFIG_STRIDES_AB*/);
gemmini_fence();
sp_tiled_matmul_full_spad_ws(
spad_addr_K_produce, spad_addr_V_produce,
/*spad_D=*/0, /*spad_C=*/0,
@@ -505,29 +536,19 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
/*pad_I=*/0, /*pad_J=*/0, /*pad_K=*/0,
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips);
#endif
}
// fence everything before going to the next tile
gemmini_fence();
gemmini_fence();
gemmini_fence();
gemmini_fence();
asm volatile("move_k_v_finish_%=:" ::);
}
// threadblock_barrier(warpgroup_id_in_cluster,
// warps_per_warpgroup_per_core);
// reconverge from mmio divergence
threadblock_barrier(warpgroup_id_in_cluster,
warps_per_warpgroup_per_core);
}
asm volatile("move_k_v_finish_%=:" ::);
// NOTE: cannot put barrier here; thread 1-7 in warp 0 will skip the
// branch and call this barrier earlier than when thread 0 finishes.
// Since tmask is not considered, that will be a barrier resolve done too
// early
// threadblock_barrier(0, 1);
} else /* warp_id != 0 */ {
if (tile_k >= 1) // delay by 1 iters for pipelining
{
if (tile_k >= 1) // delay online softmax by 1 iters
{
const uint32_t tile_k_ = tile_k - 1;
@@ -536,59 +557,64 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
if (warpgroup_id == 0) {
if (tile_k_ == 0) {
thread_block_copy_tile<B_ROW, B_COL, GEMMINI_DMA>(
smem_S_consume, gmem_tmp_d0, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
smem_S_consume, gmem_tmp_d0, tid_in_warpgroup,
threads_per_warpgroup, warpgroup_id);
} else if (tile_k_ == 1) {
thread_block_copy_tile<B_ROW, B_COL, GEMMINI_DMA>(
smem_S_consume, gmem_tmp_d1, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
smem_S_consume, gmem_tmp_d1, tid_in_warpgroup,
threads_per_warpgroup, warpgroup_id);
}
threadblock_barrier(barrier_id_simt, barrier_count_simt);
threadblock_barrier(warpgroup_id_in_cluster,
warps_per_warpgroup_per_core);
}
}
// Online softmax
//
thread_block_online_softmax</*block_row_major=*/GEMMINI_DMA>(
smem_S_consume, smem_P_produce, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt, smem_scratchpad,
smem_S_consume, smem_P_produce, tid_in_warpgroup,
threads_per_warpgroup, warpgroup_id, smem_scratchpad,
smem_rowmax, smem_rowsum, smem_O_row_scale);
threadblock_barrier(barrier_id_simt, barrier_count_simt);
threadblock_barrier(warpgroup_id_in_cluster,
warps_per_warpgroup_per_core);
if constexpr (DEBUG) {
if (warpgroup_id == 0) {
if (tile_k_ == 0) {
thread_block_copy_rowmax(
smem_rowmax, gmem_tmp_e0, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
smem_rowmax, gmem_tmp_e0, tid_in_warpgroup,
threads_per_warpgroup, warpgroup_id_in_cluster);
thread_block_copy_rowmax(
smem_rowsum, gmem_tmp_e2, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
smem_rowsum, gmem_tmp_e2, tid_in_warpgroup,
threads_per_warpgroup, warpgroup_id_in_cluster);
} else if (tile_k_ == 1) {
thread_block_copy_rowmax(smem_rowmax, gmem_tmp_e1,
tid_in_warpgroup_simt, threads_per_warpgroup_simt,
warpgroup_id_simt);
tid_in_warpgroup, threads_per_warpgroup,
warpgroup_id_in_cluster);
thread_block_copy_rowmax(smem_rowsum, gmem_tmp_e3,
tid_in_warpgroup_simt, threads_per_warpgroup_simt,
warpgroup_id_simt);
tid_in_warpgroup, threads_per_warpgroup,
warpgroup_id_in_cluster);
}
threadblock_barrier(barrier_id_simt, barrier_count_simt);
threadblock_barrier(warpgroup_id_in_cluster,
warps_per_warpgroup_per_core);
}
}
#ifdef FENCE_GEMM_II
// check flag to make sure GEMM II finished and read-after-write
// dependency on O tile is settled for rescale
if (tid_in_warpgroup_simt == 0) {
while ((*smem_O_flag) != 1)
;
// set it back to 0 for the next tile iteration
*smem_O_flag = 0;
vx_fence();
}
// asm volatile("rescale_fence_read_start_%=:" ::);
// // check flag to make sure GEMM II finished and read-after-write
// // dependency on O tile is settled for rescale
// if (tid_in_warpgroup == 0) {
// while ((*smem_O_flag) != 1)
// ;
// // set it back to 0 for the next tile iteration
// *smem_O_flag = 0;
// vx_fence();
// }
// asm volatile("rescale_fence_read_end_%=:" ::);
#endif
#if 0
@@ -605,78 +631,82 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
#endif
if constexpr (DEBUG) {
if (warpgroup_id == 0) {
gemmini_fence();
gemmini_fence();
if (warpgroup_id_in_cluster == 0) {
if (tid_in_warpgroup == 0) {
gemmini_fence();
gemmini_fence();
}
// reconverge from mmio divergence
threadblock_barrier(warpgroup_id_in_cluster,
warps_per_warpgroup_per_core);
// O after PV
if (tile_k_ == 1 /*wait until GEMM II finshes */) {
thread_block_copy_tile<B_ROW, HEADDIM, GEMMINI_DMA>(
smem_O, gmem_tmp_d6, tid_in_warpgroup_simt, threads_per_warpgroup_simt,
warpgroup_id_simt);
smem_O, gmem_tmp_d6, tid_in_warpgroup, threads_per_warpgroup,
warpgroup_id_in_cluster);
} else if (tile_k_ == 2) {
thread_block_copy_tile<B_ROW, HEADDIM, GEMMINI_DMA>(
smem_O, gmem_tmp_d7, tid_in_warpgroup_simt, threads_per_warpgroup_simt,
warpgroup_id_simt);
smem_O, gmem_tmp_d7, tid_in_warpgroup, threads_per_warpgroup,
warpgroup_id_in_cluster);
}
threadblock_barrier(barrier_id_simt, barrier_count_simt);
threadblock_barrier(warpgroup_id_in_cluster,
warps_per_warpgroup_per_core);
}
}
// Oi rescale
thread_block_O_rescale</*block_row_major=*/GEMMINI_DMA>(
smem_O, smem_O /*in-place*/, smem_O_row_scale,
tid_in_warpgroup_simt, threads_per_warpgroup_simt,
warpgroup_id_simt);
smem_O, smem_O /*in-place*/, smem_O_row_scale, tid_in_warpgroup,
threads_per_warpgroup, warpgroup_id_in_cluster);
// rescale-to-PV-GEMM barrier
threadblock_barrier(barrier_id_simt, barrier_count_simt);
threadblock_barrier(warpgroup_id_in_cluster,
warps_per_warpgroup_per_core);
if constexpr (DEBUG) {
if (warpgroup_id == 0) {
if (warpgroup_id_in_cluster == 0) {
// O before PV
if (tile_k_ == 0) {
thread_block_copy_tile<B_ROW, B_COL, GEMMINI_DMA>(
smem_P_produce, gmem_tmp_d2, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
smem_P_produce, gmem_tmp_d2, tid_in_warpgroup,
threads_per_warpgroup, warpgroup_id_in_cluster);
thread_block_copy_tile<B_ROW, HEADDIM, GEMMINI_DMA>(
smem_O, gmem_tmp_d4, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
smem_O, gmem_tmp_d4, tid_in_warpgroup, threads_per_warpgroup,
warpgroup_id_in_cluster);
} else if (tile_k_ == 1) {
thread_block_copy_tile<B_ROW, B_COL, GEMMINI_DMA>(
smem_P_produce, gmem_tmp_d3, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
smem_P_produce, gmem_tmp_d3, tid_in_warpgroup,
threads_per_warpgroup, warpgroup_id_in_cluster);
thread_block_copy_tile<B_ROW, HEADDIM, GEMMINI_DMA>(
smem_O, gmem_tmp_d5, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
smem_O, gmem_tmp_d5, tid_in_warpgroup, threads_per_warpgroup,
warpgroup_id_in_cluster);
}
threadblock_barrier(barrier_id_simt, barrier_count_simt);
threadblock_barrier(warpgroup_id_in_cluster,
warps_per_warpgroup_per_core);
}
}
}
#if 0
// fence GEMM I after Oi rescale
if (tid_in_warpgroup == 0) {
gemmini_fence();
gemmini_fence();
gemmini_fence();
gemmini_fence();
}
// reconverge from mmio divergence
// intra-warpgroup barrier
threadblock_barrier(warpgroup_id_in_cluster,
warps_per_warpgroup_per_core);
#endif
// intra-warpgroup barrier
threadblock_barrier(barrier_id_simt, barrier_count_simt);
// @perf: instead of fencing here, fence at the start of the loop to
// maximize overlapping
// gemmini_fence();
// // reconverge after mmio
// threadblock_barrier(warpgroup_id_in_cluster,
// warps_per_warpgroup_per_core);
}
}
asm volatile ("tile_loop_finish_%=:" :: );
MARK_END();
}
int main() {

738
tests/regression/flash_attention/kernel.gemmini.warpspec.cpp Normal file
View File

@@ -0,0 +1,738 @@
#include <stdint.h>
#include <vx_intrinsics.h>
#include <vx_print.h>
#include <vx_spawn.h>
#include "common.h"
#include "sgemm_impl.hpp"
#include "include/gemmini.h"
#include "gemmini_mmio.h"
#include "flash_impl.hpp"
#define FENCE_GEMM_II
#define GEMMINI_NEW_CISC 1
static_assert(GEMMINI_NEW_CISC, "NOTE: old non-CISC code is untested; look for "
"any misalignment of fields in ciscArgs.");
constexpr bool DEBUG = false;
static_assert(GEMMINI_DMA && !WARP_SPECIALIZED,
"GEMMINI_DMA should be set and WARP_SPECIALIZED unset");
void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
// @perf: All threads are running these compute whose result is mostly same
// across the threadblock
#ifdef RADIANCE
constexpr uint32_t cores_per_cluster = CORES_PER_CLUSTER;
#else
constexpr uint32_t cores_per_cluster = 1;
#endif
// FIXME: headdim not considered
constexpr uint32_t threads_per_threadblock_theoretical =
(B_ROW * B_COL) / (ELEM_PER_THREAD);
constexpr uint32_t hw_threads_per_cluster =
CORES_PER_CLUSTER * NUM_THREADS * NUM_WARPS;
// cap maximum threadblock size to # of HW threads in cluster, to prevent
// multiple "wave" invocations which slows down the kernel
constexpr uint32_t threads_per_threadblock =
(threads_per_threadblock_theoretical > hw_threads_per_cluster)
? hw_threads_per_cluster
: threads_per_threadblock_theoretical;
constexpr uint32_t threadblocks_per_cluster =
hw_threads_per_cluster / threads_per_threadblock;
constexpr uint32_t warps_per_threadblock_per_core =
NUM_WARPS / threadblocks_per_cluster;
const uint32_t threadblock_id = task_id / threads_per_threadblock;
const uint32_t threadblock_id_in_cluster =
threadblock_id % threadblocks_per_cluster;
const uint32_t tid_in_threadblock = task_id % threads_per_threadblock;
const uint32_t warp_id = tid_in_threadblock / NUM_THREADS;
constexpr uint32_t warps_in_threadblock =
threads_per_threadblock / NUM_THREADS;
// warpgroup context
constexpr uint32_t threads_per_warpgroup =
threads_per_threadblock / (WARP_SPECIALIZED ? 2 : 1);
constexpr uint32_t warpgroups_per_cluster =
threadblocks_per_cluster * (WARP_SPECIALIZED ? 2 : 1);
const uint32_t warps_per_warpgroup_per_core =
NUM_WARPS / warpgroups_per_cluster;
const uint32_t warpgroup_id = task_id / threads_per_warpgroup;
const uint32_t warpgroup_id_in_cluster =
warpgroup_id % warpgroups_per_cluster;
const uint32_t tid_in_warpgroup = tid_in_threadblock % threads_per_warpgroup;
// // warpgroup 0: warp 0
// // warpgroup 1: warp 1~7
// const uint32_t warpgroup_id = (warp_id != 0);
const uint32_t dim_seqlen = arg->dim_seqlen;
const uint32_t dim_headdim = arg->dim_headdim;
// get global memory addresses from kernel arguments
const float *gmem_Q = reinterpret_cast<float *>(arg->addr_q);
const float *gmem_K = reinterpret_cast<float *>(arg->addr_k);
const float *gmem_V = reinterpret_cast<float *>(arg->addr_v);
float *gmem_O = reinterpret_cast<float *>(arg->addr_o);
float *gmem_tmp_d0 = reinterpret_cast<float *>(0xd0000000UL);
float *gmem_tmp_d1 = reinterpret_cast<float *>(0xd1000000UL);
float *gmem_tmp_d2 = reinterpret_cast<float *>(0xd2000000UL);
float *gmem_tmp_d3 = reinterpret_cast<float *>(0xd3000000UL);
float *gmem_tmp_d4 = reinterpret_cast<float *>(0xd4000000UL);
float *gmem_tmp_d5 = reinterpret_cast<float *>(0xd5000000UL);
float *gmem_tmp_d6 = reinterpret_cast<float *>(0xd6000000UL);
float *gmem_tmp_d7 = reinterpret_cast<float *>(0xd7000000UL);
float *gmem_tmp_e0 = reinterpret_cast<float *>(0xe0000000UL);
float *gmem_tmp_e1 = reinterpret_cast<float *>(0xe1000000UL);
float *gmem_tmp_e2 = reinterpret_cast<float *>(0xe2000000UL);
float *gmem_tmp_e3 = reinterpret_cast<float *>(0xe3000000UL);
// static shared memory allocation
// these are in float elements, not bytes
constexpr uint32_t smem_Q_size = B_ROW * HEADDIM;
constexpr uint32_t smem_K_size = B_COL * HEADDIM;
constexpr uint32_t smem_QK_size = B_ROW * B_COL;
constexpr uint32_t smem_V_size = B_COL * HEADDIM;
constexpr uint32_t smem_O_size = B_COL * HEADDIM;
static_assert(
threads_per_threadblock == NUM_WARPS * NUM_THREADS * CORES_PER_CLUSTER,
"flashattention kernel assumes 1 threadblock occupancy per cluster");
uint8_t *smem_per_threadblock = reinterpret_cast<uint8_t *>(DEV_SMEM_START_ADDR);
constexpr uint32_t smem_start = DEV_SMEM_START_ADDR;
constexpr uint32_t smem_hexadecile_size = (SMEM_SIZE / 16);
// currently assumes the Q/K/V tile sizes exactly match the hexadecile size
static_assert(smem_hexadecile_size == smem_Q_size * sizeof(float));
static_assert(smem_hexadecile_size == smem_K_size * sizeof(float));
static_assert(smem_hexadecile_size == smem_QK_size * sizeof(float));
static_assert(smem_hexadecile_size == smem_V_size * sizeof(float));
static_assert(smem_hexadecile_size == smem_O_size * sizeof(float));
// Q/V/S in quart0/1, K/P/O in quart2/3
constexpr uint32_t smem_Q0_hexadecile = 4 * 0;
constexpr uint32_t smem_Q1_hexadecile = 4 * 1;
constexpr uint32_t smem_K0_hexadecile = 4 * 2;
constexpr uint32_t smem_K1_hexadecile = 4 * 3;
constexpr uint32_t smem_V0_hexadecile = smem_Q0_hexadecile + 1;
constexpr uint32_t smem_V1_hexadecile = smem_Q1_hexadecile + 1;
// put S1/S0 with V0/V1 so that softmax and GEMM-II doesn't cause bank
// conflicts
constexpr uint32_t smem_S0_hexadecile = smem_V1_hexadecile + 1;
constexpr uint32_t smem_S1_hexadecile = smem_V0_hexadecile + 1;
constexpr uint32_t smem_P0_hexadecile = smem_K0_hexadecile + 1;
constexpr uint32_t smem_P1_hexadecile = smem_K1_hexadecile + 1;
// reversed!
constexpr uint32_t smem_O0_hexadecile = smem_P1_hexadecile + 1;
constexpr uint32_t smem_O1_hexadecile = smem_P0_hexadecile + 1; // unused
float *smem_Q0 = reinterpret_cast<float *>(smem_start + smem_Q0_hexadecile * smem_hexadecile_size);
float *smem_Q1 = reinterpret_cast<float *>(smem_start + smem_Q1_hexadecile * smem_hexadecile_size);
float *smem_K0 = reinterpret_cast<float *>(smem_start + smem_K0_hexadecile * smem_hexadecile_size);
float *smem_K1 = reinterpret_cast<float *>(smem_start + smem_K1_hexadecile * smem_hexadecile_size);
float *smem_V0 = reinterpret_cast<float *>(smem_start + smem_V0_hexadecile * smem_hexadecile_size);
float *smem_V1 = reinterpret_cast<float *>(smem_start + smem_V1_hexadecile * smem_hexadecile_size);
float *smem_S0 = reinterpret_cast<float *>(smem_start + smem_S0_hexadecile * smem_hexadecile_size);
float *smem_S1 = reinterpret_cast<float *>(smem_start + smem_S1_hexadecile * smem_hexadecile_size);
float *smem_P0 = reinterpret_cast<float *>(smem_start + smem_P0_hexadecile * smem_hexadecile_size);
float *smem_P1 = reinterpret_cast<float *>(smem_start + smem_P1_hexadecile * smem_hexadecile_size);
float *smem_O0 = reinterpret_cast<float *>(smem_start + smem_O0_hexadecile * smem_hexadecile_size);
float *smem_O1 = reinterpret_cast<float *>(smem_start + smem_O1_hexadecile * smem_hexadecile_size);
// allocate rowmax/rowsum storage at the end of the sharedmem address space
constexpr uint32_t smem_rowmax_size = B_ROW * ROWMAX_SETS;
constexpr uint32_t smem_rowsum_size = B_ROW;
constexpr uint32_t smem_O_row_scale_size = B_ROW;
float *smem_cursor = smem_O1 + smem_O_size;
// // FIXME: dangerous
// smem_cursor = reinterpret_cast<float *>(0xff038000);
float *smem_rowmax_0 = smem_cursor;
smem_cursor += smem_rowmax_size;
float *smem_rowmax_1 = smem_cursor;
smem_cursor += smem_rowmax_size;
float *smem_rowsum_0 = smem_cursor;
smem_cursor += smem_rowsum_size;
float *smem_rowsum_1 = smem_cursor;
smem_cursor += smem_rowsum_size;
float *smem_O_row_scale_0 = smem_cursor;
smem_cursor += smem_O_row_scale_size;
float *smem_O_row_scale_1 = smem_cursor;
smem_cursor += smem_O_row_scale_size;
// sharedmem "scratchpad" area to put temporary data, e.g. for tree reduction
// in rowsum
// NOTE: out-of bounds is not checked
constexpr uint32_t smem_scratchpad_size =
threads_per_warpgroup * 2 /*arbitrary slack*/;
float *smem_scratchpad_0 = smem_cursor;
smem_cursor += smem_scratchpad_size;
float *smem_scratchpad_1 = smem_cursor;
smem_cursor += smem_scratchpad_size;
uint32_t *smem_O_flag = reinterpret_cast<uint32_t *>(smem_cursor);
smem_cursor += 1 /* 4Byte */;
static_assert(sizeof(elem_t) == sizeof(float));
constexpr uint32_t spad_addr_factor = DIM * sizeof(elem_t);
constexpr uint32_t spad_addr_Q0 = smem_Q0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_Q1 = smem_Q1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_K0 = smem_K0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_K1 = smem_K1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_V0 = smem_V0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_V1 = smem_V1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_S0 = smem_S0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_S1 = smem_S1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_P0 = smem_P0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_P1 = smem_P1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_O0 = smem_O0_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t spad_addr_O1 = smem_O1_hexadecile * smem_hexadecile_size / spad_addr_factor;
constexpr uint32_t global_barrier_id = NUM_WARPS - 1; // arbitrary
static_assert(warps_per_threadblock_per_core == NUM_WARPS);
// initialize rowmax/rowsum values in sharedmem
thread_block_init_sharedmem(tid_in_warpgroup, threads_per_warpgroup, smem_O0,
smem_rowmax_0, smem_rowsum_0, smem_O_row_scale_0);
thread_block_init_sharedmem(tid_in_warpgroup, threads_per_warpgroup, smem_O1,
smem_rowmax_1, smem_rowsum_1, smem_O_row_scale_1);
threadblock_barrier(global_barrier_id, warps_per_threadblock_per_core);
// skip everything except DMA in the loop FSM
constexpr uint32_t skips =
loop_matmul_skips(/*skip_lda=*/0, /*skip_ldb=*/0, /*skip_ldd=*/1,
/*skip_ex=*/1, /*skip_stc=*/1);
constexpr uint32_t skips_only_a =
loop_matmul_skips(/*skip_lda=*/0, /*skip_ldb=*/1, /*skip_ldd=*/1,
/*skip_ex=*/1, /*skip_stc=*/1);
constexpr uint32_t skips_only_b =
loop_matmul_skips(/*skip_lda=*/1, /*skip_ldb=*/0, /*skip_ldd=*/1,
/*skip_ex=*/1, /*skip_stc=*/1);
constexpr uint32_t skips_mvout_spad =
loop_matmul_skips(/*skip_lda=*/1, /*skip_ldb=*/1, /*skip_ldd=*/1,
/*skip_ex=*/1, /*skip_stc=*/0);
constexpr uint32_t skips_matmul =
loop_matmul_skips(/*skip_lda=*/1, /*skip_ldb=*/1, /*skip_ldd=*/1,
/*skip_ex=*/0, /*skip_stc=*/0);
constexpr uint32_t skips_matmul_preload =
loop_matmul_skips(/*skip_lda=*/1, /*skip_ldb=*/1, /*skip_ldd=*/0,
/*skip_ex=*/0, /*skip_stc=*/1);
MARK_BEG();
if (tid_in_warpgroup == 0) {
gemmini_extended_config_ex(WEIGHT_STATIONARY, 0, 0, 1, 0, 0);
// configure DMA with GMEM address strides
// Q matrix
gemmini_extended3_config_ld(HEADDIM * sizeof(elem_t), MVIN_SCALE_IDENTITY,
false, 0);
// K matrix
gemmini_extended3_config_ld(dim_seqlen * sizeof(elem_t),
MVIN_SCALE_IDENTITY, false, 1);
// configure DMA for Q*K store
gemmini_extended_config_st(B_COL * sizeof(elem_t), 0, MVIN_SCALE_IDENTITY);
gemmini_fence();
}
// NOTE about barriers: Placing barriers around thread-divergent branches may
// cause bugs, because the Vortex core doesn't check for tmask for barriers.
// The compiler might decide to duplicate vx_bar into both paths of a
// conditional branch, which will get evaluated twice because of the way
// branches are handled in SIMT; this might result in stalls especially when
// other warps behave differently on the branch condition.
// threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core);
static_assert(B_ROW == B_COL, "currently only supports square tiles");
// move Q and K into SMEM before the loop starts
//
asm volatile("dma_move_start_%=:" ::);
if (tid_in_warpgroup == 0) {
// make sure to read from the correct row of Q
const float *gmem_Q_tile = gmem_Q + HEADDIM * B_ROW * warpgroup_id;
const float *gmem_K_tile = gmem_K;
// do DMA
//
// move Q to spad_addr_Q0 for the first iteration
//
#ifdef GEMMINI_NEW_CISC
// the target addresses of this should match with spad_addr_Q0 and
// spad_addr_K0 set in this kernel
gemmini_tile_load_ab(gmem_Q_tile, gmem_K_tile, smem_Q0_hexadecile,
smem_K0_hexadecile, 0 /*tile_idx_i*/, 0 /*tile_idx_j*/,
0 /*tile_idx_k*/, dim_seqlen, dim_seqlen, HEADDIM,
B_ROW, B_COL, HEADDIM);
#else
// configure the GMEM addresses for the DMA to read from
ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, (uint64_t)(gmem_Q_tile),
(uint64_t)(gmem_K_tile), k_LOOP_WS_CONFIG_ADDRS_AB)
// configure address strides for the DMA
GEMMINI_CISC_CMD_R((dim_seqlen << 20) | (HEADDIM << 8) |
GEMMINI_CISC_SET_AB_STRIDE);
gemmini_fence();
// among other things, this also configures CONFIG_BOUNDS so that the
// DMA knows the full matrix dimensions
sp_tiled_matmul_full_spad_ws(
spad_addr_Q0, spad_addr_K0,
/*spad_D=*/0, /*spad_C=*/spad_addr_S0/*bogus*/,
/*I=*/(B_ROW / DIM), /*J=*/(B_COL / DIM), /*K=*/(HEADDIM / DIM),
/*pad_I=*/0, /*pad_J=*/0, /*pad_K=*/0,
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips);
#endif
// block until DMA complete
gemmini_fence();
// also move Q to spad_addr_Q1 for the second iteration
//
#ifdef GEMMINI_NEW_CISC
gemmini_tile_load_ab(gmem_Q_tile, gmem_K_tile, smem_Q1_hexadecile,
smem_K1_hexadecile, 0 /*tile_idx_i*/, 0 /*tile_idx_j*/,
0 /*tile_idx_k*/, dim_seqlen, dim_seqlen, HEADDIM,
B_ROW, B_COL, HEADDIM);
#else
// FIXME: re-configure necessary?
gmem_K_tile = gmem_K + (B_COL * 1);
ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, (uint64_t)(gmem_Q_tile),
(uint64_t)(gmem_K_tile), k_LOOP_WS_CONFIG_ADDRS_AB)
GEMMINI_CISC_CMD_R((dim_seqlen << 20) | (HEADDIM << 8) |
8 /*k_LOOP_WS_CONFIG_STRIDES_AB*/);
gemmini_fence();
sp_tiled_matmul_full_spad_ws(
spad_addr_Q1, spad_addr_K1/*bogus*/,
/*spad_D=*/0, /*spad_C=*/spad_addr_S0/*bogus*/,
/*I=*/(B_ROW / DIM), /*J=*/(B_COL / DIM), /*K=*/(HEADDIM / DIM),
/*pad_I=*/0, /*pad_J=*/0, /*pad_K=*/0,
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips);
#endif
// block until DMA complete
gemmini_fence();
// re-configure DMA for K and V load that will later happen in the loop
// FIXME: not sure necessary with new CISC
//
// GMEM addr stride for K
gemmini_extended3_config_ld(dim_seqlen * sizeof(elem_t),
MVIN_SCALE_IDENTITY, false, 0);
// GMEM addr stride for V
gemmini_extended3_config_ld(HEADDIM * sizeof(elem_t), MVIN_SCALE_IDENTITY,
false, 1);
gemmini_fence();
}
asm volatile("dma_move_end_%=:" ::);
// protect write to SMEM
// threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core);
threadblock_barrier(global_barrier_id, warps_per_threadblock_per_core);
// if constexpr (DEBUG) {
// thread_block_copy_tile<B_ROW, HEADDIM>(smem_Q0, gmem_tmp_d0, tid_in_warpgroup,
// threads_per_warpgroup, warpgroup_id_in_cluster);
// thread_block_copy_tile<HEADDIM, B_COL>(smem_K0, gmem_tmp_d1, tid_in_warpgroup,
// threads_per_warpgroup, warpgroup_id_in_cluster);
// threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core);
// }
constexpr uint32_t threads_per_warpgroup_simt =
threads_per_warpgroup -
CORES_PER_CLUSTER * NUM_THREADS /*warp 0, 4, 8, 12*/;
constexpr uint32_t warpgroup_id_simt = 1;
constexpr uint32_t barrier_id_simt = 1;
constexpr uint32_t barrier_count_simt = NUM_WARPS - 1;
const uint32_t tid_in_warpgroup_simt =
tid_in_warpgroup - (CORES_PER_CLUSTER * NUM_THREADS);
static_assert(barrier_id_simt == 1 && barrier_count_simt == 7);
asm volatile ("tile_loop_start_%=:" :: );
// "inner loop" along the columns of K^T
const uint32_t k_tiles = (dim_seqlen / B_COL);
for (uint32_t tile_k = 0;
tile_k < (4 /*for perf measurement*/ *
// virgo kernel is fully pipelined around (2 GEMMs | softmax);
// requires two loop iterations to finish one tile compute
(2 * k_tiles)) +
2 /*pipeline latency*/;
tile_k++) {
if constexpr (DEBUG || true) {
threadblock_barrier(global_barrier_id, warps_per_threadblock_per_core);
}
// select the correct double buffer by tile iteration
// all iterations work on the same Q row tile; no ping-pong necessary
asm volatile ("dbuf_sel_start_%=:" :: );
// FIXME speedup by doing arithmetic
float *smem_Q = smem_Q0;
float *smem_K_consume = (tile_k & 1) ? smem_K1 : smem_K0;
float *smem_K_produce = (tile_k & 1) ? smem_K0 : smem_K1;
float *smem_V_consume = (tile_k & 1) ? smem_V1 : smem_V0;
float *smem_V_produce = (tile_k & 1) ? smem_V0 : smem_V1;
float *smem_S_consume = (tile_k & 1) ? smem_S1 : smem_S0;
float *smem_S_produce = (tile_k & 1) ? smem_S0 : smem_S1;
float *smem_P_consume = (tile_k & 1) ? smem_P1 : smem_P0;
float *smem_P_produce = (tile_k & 1) ? smem_P0 : smem_P1;
// O, rowmax/rowsum etc. is sequentially updated at every iteration; no
// ping-pong necessary
float *smem_O = smem_O0;
float *smem_O_row_scale = smem_O_row_scale_0;
float *smem_rowmax = smem_rowmax_0;
float *smem_rowsum = smem_rowsum_0;
float *smem_scratchpad = smem_scratchpad_0;
const auto spad_addr_Q = spad_addr_Q0;
const auto spad_addr_K_consume = (tile_k & 1) ? spad_addr_K1 : spad_addr_K0;
const auto spad_addr_K_produce = (tile_k & 1) ? spad_addr_K0 : spad_addr_K1;
const auto spad_addr_V_consume = (tile_k & 1) ? spad_addr_V1 : spad_addr_V0;
const auto spad_addr_V_produce = (tile_k & 1) ? spad_addr_V0 : spad_addr_V1;
const auto spad_addr_S_consume = (tile_k & 1) ? spad_addr_S1 : spad_addr_S0;
const auto spad_addr_S_produce = (tile_k & 1) ? spad_addr_S0 : spad_addr_S1;
const auto spad_addr_P_consume = (tile_k & 1) ? spad_addr_P1 : spad_addr_P0;
const auto spad_addr_P_produce = (tile_k & 1) ? spad_addr_P0 : spad_addr_P1;
const auto spad_addr_O = spad_addr_O0; // NOTE: there's only single O tile
const auto spad_hex_Q = smem_Q0_hexadecile;
const auto spad_hex_K_consume = (tile_k & 1) ? smem_K1_hexadecile : smem_K0_hexadecile;
const auto spad_hex_K_produce = (tile_k & 1) ? smem_K0_hexadecile : smem_K1_hexadecile;
const auto spad_hex_V_consume = (tile_k & 1) ? smem_V1_hexadecile : smem_V0_hexadecile;
const auto spad_hex_V_produce = (tile_k & 1) ? smem_V0_hexadecile : smem_V1_hexadecile;
const auto spad_hex_S_consume = (tile_k & 1) ? smem_S1_hexadecile : smem_S0_hexadecile;
const auto spad_hex_S_produce = (tile_k & 1) ? smem_S0_hexadecile : smem_S1_hexadecile;
const auto spad_hex_P_consume = (tile_k & 1) ? smem_P1_hexadecile : smem_P0_hexadecile;
const auto spad_hex_P_produce = (tile_k & 1) ? smem_P0_hexadecile : smem_P1_hexadecile;
const auto spad_hex_O = smem_O0_hexadecile; // NOTE: there's only single O tile
asm volatile ("dbuf_sel_end_%=:" :: );
if (vx_warp_id() == 0 /* warp 0 in every core */) {
if (tile_k >= 2) // delay by 2 iters for pipelining
{
const uint32_t tile_k_ = tile_k - 2;
// GEMM II: O = O + P*V
// --------------------
// This is done *before* GEMM I in the software pipeline, working on the
// online softmax result tile from the previous iteration
asm volatile("gemm_pv_start_%=:" ::);
if (tid_in_warpgroup == 0) {
// kickoff matmul
// FIXME: perf: prevent GMEM->SMEM load for O tile
gemmini_fence();
#ifdef GEMMINI_NEW_CISC
gemmini_tile_compute</*store_to_spad=*/true>(
spad_hex_P_consume, spad_hex_V_consume, spad_hex_O,
0 /*accumulate.
FIXME: Gemmini doens't support accumulation from a spad tile*/);
#else
sp_tiled_matmul_full_spad_ws(
spad_addr_P_consume, spad_addr_V_consume,
/*spad_D=*/spad_addr_O, /*spad_C=*/spad_addr_O,
/*I=*/(B_ROW / DIM), /*J=*/(HEADDIM / DIM), /*K=*/(B_COL / DIM),
/*pad_I=*/0, /*pad_J=*/0, /*pad_K=*/0,
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips_matmul);
#endif
}
// // reconverge from mmio divergence
// threadblock_barrier(warpgroup_id_in_cluster,
// warps_per_warpgroup_per_core);
asm volatile("gemm_pv_finish_%=:" ::);
}
// GEMM I: S = Q*K
//
// kick off asynchronously; fence later
asm volatile("gemm_qk_start_%=:" ::);
if (tid_in_warpgroup == 0) {
// fence to GEMM II completion
gemmini_fence();
#ifdef FENCE_GEMM_II
asm volatile("rescale_fence_write_start_%=:" ::);
// signal that GEMM II is finished to O rescale step
*smem_O_flag = 1;
vx_fence();
asm volatile("rescale_fence_write_end_%=:" ::);
#endif
// Kick off GEMM I
//
#ifdef GEMMINI_NEW_CISC
gemmini_tile_compute</*store_to_spad=*/true>(
spad_hex_Q, spad_hex_K_consume, spad_hex_S_produce,
0 /*accumulate*/);
#else
sp_tiled_matmul_full_spad_ws(
spad_addr_Q, spad_addr_K_consume,
/*spad_D=*/0, /*spad_C=*/spad_addr_S_produce,
/*I=*/(B_ROW / DIM), /*J=*/(B_COL / DIM), /*K=*/(HEADDIM / DIM),
/*pad_I=*/0, /*pad_J=*/0, /*pad_K=*/0,
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips_matmul);
#endif
asm volatile("gemm_qk_finish_%=:" ::);
// data move for K and V
//
// Q stays in SMEM for the entire loop
asm volatile("move_k_v_start_%=:" ::);
// configure GMEM addresses for K and V tiles
// load K for the next iteration
const float *gmem_K_tile = gmem_K + (B_COL * (tile_k + 1 /*runahead*/));
// load V for the *previous* iteration; this will be consumed 2
// iterations later
const float *gmem_V_tile =
gmem_V + (HEADDIM * B_COL * (tile_k - 1 /*dragbehind*/));
#if 0
// fence mvout S to SMEM
gemmini_fence();
ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, (uint64_t)(gmem_K_tile),
(uint64_t)(gmem_V_tile),
k_LOOP_WS_CONFIG_ADDRS_AB)
#endif
// do DMA
if (tile_k == 0) {
// // configure address strides for the DMA
// // FIXME: unnecessary?
// GEMMINI_CISC_CMD_R((HEADDIM /*V*/ << 20) | (dim_seqlen /*KT*/ << 8) |
// 8 /*k_LOOP_WS_CONFIG_STRIDES_AB*/);
// gemmini_fence();
//
// we load (k-1)th tile for V; skip V for the 1st iteration,
// sp_tiled_matmul_full_spad_ws(
// spad_addr_K_produce, spad_addr_V_produce,
// /*spad_D=*/0, /*spad_C=*/0,
// /*I=*/(B_ROW / DIM), /*J=*/(HEADDIM / DIM), /*K=*/(B_COL / DIM),
// /*pad_I=*/0, /*pad_J=*/0, /*pad_K=*/0,
// /*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
// /*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips_only_a);
} else {
#ifdef GEMMINI_NEW_CISC
gemmini_tile_load_ab(
gmem_K_tile, gmem_V_tile, spad_hex_K_produce, spad_hex_V_produce,
0 /*tile_idx_i*/, 0 /*tile_idx_j*/, 0 /*tile_idx_k*/,
HEADDIM /*dim_m of KT*/, HEADDIM /*dim_n of V*/,
dim_seqlen /*dim_k of KT*/, B_ROW, HEADDIM, B_COL);
#else
// configure address strides for the DMA
// FIXME: unnecessary?
GEMMINI_CISC_CMD_R((HEADDIM /*V*/ << 20) | (dim_seqlen /*KT*/ << 8) |
8 /*k_LOOP_WS_CONFIG_STRIDES_AB*/);
gemmini_fence();
sp_tiled_matmul_full_spad_ws(
spad_addr_K_produce, spad_addr_V_produce,
/*spad_D=*/0, /*spad_C=*/0,
/*I=*/(B_ROW / DIM), /*J=*/(HEADDIM / DIM), /*K=*/(B_COL / DIM),
/*pad_I=*/0, /*pad_J=*/0, /*pad_K=*/0,
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips);
#endif
}
// fence everything before going to the next tile
gemmini_fence();
}
// threadblock_barrier(warpgroup_id_in_cluster,
// warps_per_warpgroup_per_core);
asm volatile("move_k_v_finish_%=:" ::);
// NOTE: cannot put barrier here; thread 1-7 in warp 0 will skip the
// branch and call this barrier earlier than when thread 0 finishes.
// Since tmask is not considered, that will be a barrier resolve done too
// early
// threadblock_barrier(0, 1);
} else /* warp_id != 0 */ {
if (tile_k >= 1) // delay online softmax by 1 iters
{
const uint32_t tile_k_ = tile_k - 1;
if constexpr (DEBUG) {
// verify S = Q*K before softmax
if (warpgroup_id == 0) {
if (tile_k_ == 0) {
thread_block_copy_tile<B_ROW, B_COL, GEMMINI_DMA>(
smem_S_consume, gmem_tmp_d0, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
} else if (tile_k_ == 1) {
thread_block_copy_tile<B_ROW, B_COL, GEMMINI_DMA>(
smem_S_consume, gmem_tmp_d1, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
}
threadblock_barrier(barrier_id_simt, barrier_count_simt);
}
}
// Online softmax
//
thread_block_online_softmax</*block_row_major=*/GEMMINI_DMA>(
smem_S_consume, smem_P_produce, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt, smem_scratchpad,
smem_rowmax, smem_rowsum, smem_O_row_scale);
threadblock_barrier(barrier_id_simt, barrier_count_simt);
if constexpr (DEBUG) {
if (warpgroup_id == 0) {
if (tile_k_ == 0) {
thread_block_copy_rowmax(
smem_rowmax, gmem_tmp_e0, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
thread_block_copy_rowmax(
smem_rowsum, gmem_tmp_e2, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
} else if (tile_k_ == 1) {
thread_block_copy_rowmax(smem_rowmax, gmem_tmp_e1,
tid_in_warpgroup_simt, threads_per_warpgroup_simt,
warpgroup_id_simt);
thread_block_copy_rowmax(smem_rowsum, gmem_tmp_e3,
tid_in_warpgroup_simt, threads_per_warpgroup_simt,
warpgroup_id_simt);
}
threadblock_barrier(barrier_id_simt, barrier_count_simt);
}
}
#ifdef FENCE_GEMM_II
asm volatile("rescale_fence_read_start_%=:" ::);
// check flag to make sure GEMM II finished and read-after-write
// dependency on O tile is settled for rescale
if (tid_in_warpgroup_simt == 0) {
while ((*smem_O_flag) != 1)
;
// set it back to 0 for the next tile iteration
*smem_O_flag = 0;
vx_fence();
}
asm volatile("rescale_fence_read_end_%=:" ::);
#endif
#if 0
if (tid_in_warpgroup == 0) {
gemmini_fence();
gemmini_fence();
gemmini_fence();
gemmini_fence();
}
// reconverge from mmio divergence
threadblock_barrier(warpgroup_id_in_cluster,
warps_per_warpgroup_per_core);
#endif
if constexpr (DEBUG) {
if (warpgroup_id == 0) {
gemmini_fence();
gemmini_fence();
// O after PV
if (tile_k_ == 1 /*wait until GEMM II finshes */) {
thread_block_copy_tile<B_ROW, HEADDIM, GEMMINI_DMA>(
smem_O, gmem_tmp_d6, tid_in_warpgroup_simt, threads_per_warpgroup_simt,
warpgroup_id_simt);
} else if (tile_k_ == 2) {
thread_block_copy_tile<B_ROW, HEADDIM, GEMMINI_DMA>(
smem_O, gmem_tmp_d7, tid_in_warpgroup_simt, threads_per_warpgroup_simt,
warpgroup_id_simt);
}
threadblock_barrier(barrier_id_simt, barrier_count_simt);
}
}
// Oi rescale
thread_block_O_rescale</*block_row_major=*/GEMMINI_DMA>(
smem_O, smem_O /*in-place*/, smem_O_row_scale,
tid_in_warpgroup_simt, threads_per_warpgroup_simt,
warpgroup_id_simt);
// rescale-to-PV-GEMM barrier
threadblock_barrier(barrier_id_simt, barrier_count_simt);
if constexpr (DEBUG) {
if (warpgroup_id == 0) {
// O before PV
if (tile_k_ == 0) {
thread_block_copy_tile<B_ROW, B_COL, GEMMINI_DMA>(
smem_P_produce, gmem_tmp_d2, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
thread_block_copy_tile<B_ROW, HEADDIM, GEMMINI_DMA>(
smem_O, gmem_tmp_d4, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
} else if (tile_k_ == 1) {
thread_block_copy_tile<B_ROW, B_COL, GEMMINI_DMA>(
smem_P_produce, gmem_tmp_d3, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
thread_block_copy_tile<B_ROW, HEADDIM, GEMMINI_DMA>(
smem_O, gmem_tmp_d5, tid_in_warpgroup_simt,
threads_per_warpgroup_simt, warpgroup_id_simt);
}
threadblock_barrier(barrier_id_simt, barrier_count_simt);
}
}
}
#if 0
// fence GEMM I after Oi rescale
if (tid_in_warpgroup == 0) {
gemmini_fence();
gemmini_fence();
gemmini_fence();
gemmini_fence();
}
// reconverge from mmio divergence
threadblock_barrier(warpgroup_id_in_cluster,
warps_per_warpgroup_per_core);
#endif
// intra-warpgroup barrier
threadblock_barrier(barrier_id_simt, barrier_count_simt);
}
}
asm volatile ("tile_loop_finish_%=:" :: );
MARK_END();
}
int main() {
kernel_arg_t *arg = (kernel_arg_t *)KERNEL_ARG_DEV_MEM_ADDR;
const uint32_t hw_threads_per_cluster =
CORES_PER_CLUSTER * vx_num_threads() * vx_num_warps();
// fix to 1 threadblock per cluster
const uint32_t grid_size = hw_threads_per_cluster;
#ifdef RADIANCE
vx_spawn_tasks_cluster(grid_size, (vx_spawn_tasks_cb)kernel_body, arg);
#else
// NOTE: This kernel assumes contiguous thread scheduling for efficient shared
// memory allocation, and therefore does not work with original vx_spawn_tasks
vx_spawn_tasks_contiguous(grid_size, (vx_spawn_tasks_cb)kernel_body, arg);
#endif
return 0;
}

3
tests/regression/sgemm_gemmini_dma/kernel.cpp
View File

@@ -113,7 +113,8 @@ void thread_block_matmul_gemmini(kernel_arg_t *__UNIFORM__ arg,
dim_m, dim_n, dim_k, TILE_M, TILE_N, TILE_K);
/* DO STUFF */
gemmini_fence();
gemmini_tile_compute(a_hexadecile, b_hexadecile, tile_k > 0);
gemmini_tile_compute</*store_to_spad=*/false>(
a_hexadecile, b_hexadecile, 0 /*d_hexadecile*/, tile_k > 0);
}
/*

134
tests/regression/sgemm_tcore/sgemm_impl.hpp
View File

@@ -6,7 +6,7 @@
#include "include/gemmini.h"
#include "gemmini_mmio.h"
#define FP_SIZE 16
#define FP_SIZE 32
// "fake" fp16 type that only has the correct data width.
using float16_t = uint16_t;
@@ -19,7 +19,7 @@ using float_type = float16_t;
// Generate kernel for the Hopper-style SMEM-decoupled tensor core. This uses
// asynchronous HGMMA and HGMMA_WAIT instructions.
#define TENSOR_HOPPER 1
#define TENSOR_HOPPER 0
// Constraints on parameters:
// * Memory:
@@ -104,6 +104,12 @@ static_assert(WMITER * WNITER * TCM * TCN * NUM_WARPS * CORES_PER_CLUSTER ==
#define TRANSPOSE_AT_PRODUCE 0
#define TRANSPOSE_AT_CONSUME 0
// if 1, wmma_store() will not respect the register <-> matrix fragment mapping
// scheme and instead do a fast coalesced GMEM writes for move out. This
// doesn't necessarily mean breaking correctness; it means that the final
// result matrix will be stored in a swizzled form in the global memory.
#define WMMA_STORE_FAST 0
#define GEMMINI_DMA 1
#define GEMMINI_DMA_FAST 1
#if SMEM_SIZE == 0x4000
@@ -213,10 +219,9 @@ inline constexpr void map_c_8lanes(const int tid, int &row, int &col) {
col += ((tid % 4) / 2) * 2;
}
inline constexpr void map_c_8lanes_hopper(const int tid, int &row, int &col) {
inline constexpr void map_c_8lanes_coalesced(const int tid, int &row, int &col) {
const int tg = tid / 2;
// FIXME wrong!!!
row = 0;
col = tid;
}
@@ -225,8 +230,8 @@ inline constexpr void map_c(const int tid, int &row, int &col) {
if constexpr (NUM_THREADS == 32) {
map_c_32lanes(tid, row, col);
} else if constexpr (NUM_THREADS == 8) {
if constexpr (TENSOR_HOPPER) {
map_c_8lanes_hopper(tid, row, col);
if constexpr (TENSOR_HOPPER || WMMA_STORE_FAST) {
map_c_8lanes_coalesced(tid, row, col);
} else {
map_c_8lanes(tid, row, col);
}
@@ -664,26 +669,48 @@ wmma_store(const int thread_in_warp, const int warp_col, const int warp_row,
volatile uint8_t *addr = reinterpret_cast<volatile uint8_t *>(
&write_addr[dim_n * (local_row + 0) + (local_col + 0)]);
volatile uint8_t *addr_tworow = addr + (2 * dim_n) * sizeof(float);
asm volatile("fsw f16, %0(%1)" ::"i"(0 * sizeof(float)), "r"(addr));
asm volatile("fsw f17, %0(%1)" ::"i"(1 * sizeof(float)), "r"(addr));
asm volatile("fsw f18, %0(%1)" ::"i"(0 * sizeof(float)), "r"(addr_tworow));
asm volatile("fsw f19, %0(%1)" ::"i"(1 * sizeof(float)), "r"(addr_tworow));
asm volatile("fsw f20, %0(%1)" ::"i"(4 * sizeof(float)), "r"(addr));
asm volatile("fsw f21, %0(%1)" ::"i"(5 * sizeof(float)), "r"(addr));
asm volatile("fsw f22, %0(%1)" ::"i"(4 * sizeof(float)), "r"(addr_tworow));
asm volatile("fsw f23, %0(%1)" ::"i"(5 * sizeof(float)), "r"(addr_tworow));
if constexpr (!WMMA_STORE_FAST) {
asm volatile("fsw f16, %0(%1)" ::"i"(0 * sizeof(float)), "r"(addr));
asm volatile("fsw f17, %0(%1)" ::"i"(1 * sizeof(float)), "r"(addr));
asm volatile("fsw f18, %0(%1)" ::"i"(0 * sizeof(float)), "r"(addr_tworow));
asm volatile("fsw f19, %0(%1)" ::"i"(1 * sizeof(float)), "r"(addr_tworow));
asm volatile("fsw f20, %0(%1)" ::"i"(4 * sizeof(float)), "r"(addr));
asm volatile("fsw f21, %0(%1)" ::"i"(5 * sizeof(float)), "r"(addr));
asm volatile("fsw f22, %0(%1)" ::"i"(4 * sizeof(float)), "r"(addr_tworow));
asm volatile("fsw f23, %0(%1)" ::"i"(5 * sizeof(float)), "r"(addr_tworow));
} else {
asm volatile("fsw f16, %0(%1)" ::"i"(0 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f17, %0(%1)" ::"i"(1 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f18, %0(%1)" ::"i"(2 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f19, %0(%1)" ::"i"(3 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f20, %0(%1)" ::"i"(4 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f21, %0(%1)" ::"i"(5 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f22, %0(%1)" ::"i"(6 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f23, %0(%1)" ::"i"(7 * WN * sizeof(float)), "r"(addr));
}
} else {
volatile uint8_t *addr = reinterpret_cast<volatile uint8_t *>(
&write_addr[dim_n * (local_row + 0) + (local_col + 0)]);
volatile uint8_t *addr_tworow = addr + (2 * dim_n) * sizeof(float);
asm volatile("fsw f24, %0(%1)" ::"i"(0 * sizeof(float)), "r"(addr));
asm volatile("fsw f25, %0(%1)" ::"i"(1 * sizeof(float)), "r"(addr));
asm volatile("fsw f26, %0(%1)" ::"i"(0 * sizeof(float)), "r"(addr_tworow));
asm volatile("fsw f27, %0(%1)" ::"i"(1 * sizeof(float)), "r"(addr_tworow));
asm volatile("fsw f28, %0(%1)" ::"i"(4 * sizeof(float)), "r"(addr));
asm volatile("fsw f29, %0(%1)" ::"i"(5 * sizeof(float)), "r"(addr));
asm volatile("fsw f30, %0(%1)" ::"i"(4 * sizeof(float)), "r"(addr_tworow));
asm volatile("fsw f31, %0(%1)" ::"i"(5 * sizeof(float)), "r"(addr_tworow));
if constexpr (!WMMA_STORE_FAST) {
asm volatile("fsw f24, %0(%1)" ::"i"(0 * sizeof(float)), "r"(addr));
asm volatile("fsw f25, %0(%1)" ::"i"(1 * sizeof(float)), "r"(addr));
asm volatile("fsw f26, %0(%1)" ::"i"(2 * sizeof(float)), "r"(addr_tworow));
asm volatile("fsw f27, %0(%1)" ::"i"(3 * sizeof(float)), "r"(addr_tworow));
asm volatile("fsw f28, %0(%1)" ::"i"(4 * sizeof(float)), "r"(addr));
asm volatile("fsw f29, %0(%1)" ::"i"(5 * sizeof(float)), "r"(addr));
asm volatile("fsw f30, %0(%1)" ::"i"(6 * sizeof(float)), "r"(addr_tworow));
asm volatile("fsw f31, %0(%1)" ::"i"(7 * sizeof(float)), "r"(addr_tworow));
} else {
asm volatile("fsw f24, %0(%1)" ::"i"(0 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f25, %0(%1)" ::"i"(1 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f26, %0(%1)" ::"i"(2 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f27, %0(%1)" ::"i"(3 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f28, %0(%1)" ::"i"(4 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f29, %0(%1)" ::"i"(5 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f30, %0(%1)" ::"i"(6 * WN * sizeof(float)), "r"(addr));
asm volatile("fsw f31, %0(%1)" ::"i"(7 * WN * sizeof(float)), "r"(addr));
}
}
asm volatile ("wmma_store_finish_%=:" :: );
@@ -1150,19 +1177,20 @@ inline void thread_block_gemm(const T *A, const T *B, float *C,
if constexpr (GEMMINI_DMA) {
// pipeline initiation
if (tid_in_threadblock == 0) {
// configure dma gmem address to load from
ROCC_INSTRUCTION_RS1_RS2(
XCUSTOM_ACC,
(uint64_t)(A + block_m * BM * dim_k + /*block_k:*/0 * BK),
(uint64_t)(B + /*block_k:*/0 * BK * dim_n + block_n * BN),
k_LOOP_WS_CONFIG_ADDRS_AB)
// GEMMINI_CISC(8) does k_LOOP_WS_CONFIG_STRIDES_AB
GEMMINI_CISC_CMD_R((dim_n << 20) | (dim_k << 8) | 8);
gemmini_fence();
if (block_m == 0 && block_n == 0) {
if (tid_in_threadblock == 0) {
// configure dma gmem address to load from
ROCC_INSTRUCTION_RS1_RS2(
XCUSTOM_ACC,
(uint64_t)(A + block_m * BM * dim_k + /*block_k:*/ 0 * BK),
(uint64_t)(B + /*block_k:*/ 0 * BK * dim_n + block_n * BN),
k_LOOP_WS_CONFIG_ADDRS_AB)
// GEMMINI_CISC(8) does k_LOOP_WS_CONFIG_STRIDES_AB
GEMMINI_CISC_CMD_R((dim_n << 20) | (dim_k << 8) | 8);
gemmini_fence();
GEMMINI_CISC_CMD_I(10);
gemmini_fence();
GEMMINI_CISC_CMD_I(10);
gemmini_fence();
#if 0
// sp_tiled_matmul_full_spad_ws includes CONFIG_BOUNDS
@@ -1181,10 +1209,11 @@ inline void thread_block_gemm(const T *A, const T *B, float *C,
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips)
gemmini_fence();
#endif
}
}
threadblock_barrier(threadblock_id_in_cluster,
warps_per_threadblock_per_core);
threadblock_barrier(threadblock_id_in_cluster,
warps_per_threadblock_per_core);
}
}
#pragma GCC unroll 1
@@ -1197,12 +1226,28 @@ inline void thread_block_gemm(const T *A, const T *B, float *C,
// this is either done using DMA or SIMT cores depending on GEMMINI_DMA
#if (GEMMINI_DMA == 1)
if ((tid_in_threadblock == 0) && ((block_k * BK) != (dim_k - BK))) {
if (tid_in_threadblock == 0) {
asm volatile("next_index_start_%=:" ::);
const uint32_t next_block_k =
((block_k + 1) * BK == dim_k) ? 0 : block_k + 1;
const uint32_t next_block_n =
(next_block_k == 0)
? (((block_n + 1) * BN == dim_n) ? 0 : block_n + 1)
: block_n;
const uint32_t next_block_m =
(next_block_n == 0)
? (((block_m + 1) == block_m_end) ? block_m_start /*unused*/
: block_m + 1)
: block_m;
asm volatile("next_index_end_%=:" ::);
// configure dma gmem address to load from
ROCC_INSTRUCTION_RS1_RS2(
XCUSTOM_ACC,
(uint64_t)(A + block_m * BM * dim_k + (block_k + 1/*runahead*/) * BK),
(uint64_t)(B + (block_k + 1/*runahead*/) * BK * dim_n + block_n * BN),
(uint64_t)(A + next_block_m * BM * dim_k + next_block_k * BK),
(uint64_t)(B + next_block_k * BK * dim_n + next_block_n * BN),
k_LOOP_WS_CONFIG_ADDRS_AB)
// GEMMINI_CISC(8) does k_LOOP_WS_CONFIG_STRIDES_AB
GEMMINI_CISC_CMD_R((dim_n << 20) | (dim_k << 8) | 8);
@@ -1210,6 +1255,11 @@ inline void thread_block_gemm(const T *A, const T *B, float *C,
// block_k is even: opcode 11 (write to local_a_buf)
// block_k is odd: opcode 10 (write to local_a)
//
// FIXME: This depends on (dim_k / BK) being an even number, since
// the last iteration of the k-loop is prefetching for the first
// iteration of the n-loop. The ping-poing indexing has to match for
// the two loop end to connect.
const uint32_t opcode = 11 - (block_k & 1);
GEMMINI_CISC_CMD_I(opcode);
// // TODO: branch is probably slow
@@ -1349,6 +1399,8 @@ inline void thread_block_gemm(const T *A, const T *B, float *C,
}
if constexpr (write_to_gmem) {
asm volatile("move_out_start_%=:" ::);
if constexpr (TENSOR_HOPPER) {
// wait until all results are accumulated into the RF
vx_wgmma_wait();
@@ -1367,6 +1419,8 @@ inline void thread_block_gemm(const T *A, const T *B, float *C,
}
}
}
asm volatile("move_out_end_%=:" ::);
}
}
asm volatile("loop_mn_end_%=:" ::);

51
tests/regression/sgemm_tcore/switch_args_input.sh Executable file
View File

@@ -0,0 +1,51 @@
#!/bin/sh
#
# Updates symlink to args.bin, input.a.bin, input.b.bin to point to the right
# binary according to the dimension size given as the argument.
if [ "$#" != "2" ]; then
echo "usage: $0 DIMENSION 1|0"
echo "second argument indicates using DMA or not."
exit 1
fi
dim="$1"
dma="$2"
if [ "$2" == "1" ]; then
layout_a="row.swizzle_fp16"
layout_b="row"
else
layout_a="col.swizzle_fp16"
layout_b="row.swizzle_fp16"
fi
check_exists() {
if ! [ -f "$1" ]; then
echo "error: looked for file $1 that does not exist."
exit 1
fi
}
args="args.m$1n$1k$1.bin"
input_a="input.a.rand01.fp16.m$1n$1k$1.$layout_a.bin"
input_b="input.b.rand01.fp16.m$1n$1k$1.$layout_b.bin"
check_exists "$args"
check_exists "$input_a"
check_exists "$input_b"
echo "will symlink:"
echo "args.bin -> $args"
echo "input.a.bin -> $input_a"
echo "input.b.bin -> $input_b"
echo "continue? (Y/N)"
read -r -s -n 1 answer
if [ "$answer" != "Y" ] && [ "$answer" != "y" ]; then
echo "exiting..."
exit 1
fi
ln -sf -v "$args" "args.bin"
ln -sf -v "$input_a" "input.a.bin"
ln -sf -v "$input_b" "input.b.bin"
echo "done."