diff --git a/kernel/include/gemmini_mmio.h b/kernel/include/gemmini_mmio.h index 08e1a792..ef783c8a 100644 --- a/kernel/include/gemmini_mmio.h +++ b/kernel/include/gemmini_mmio.h @@ -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 +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(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, diff --git a/tests/regression/flash_attention/Makefile b/tests/regression/flash_attention/Makefile index 0456e983..3a25e4f3 100644 --- a/tests/regression/flash_attention/Makefile +++ b/tests/regression/flash_attention/Makefile @@ -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 diff --git a/tests/regression/flash_attention/flash_impl.hpp b/tests/regression/flash_attention/flash_impl.hpp index 47e21c70..6f210b75 100644 --- a/tests/regression/flash_attention/flash_impl.hpp +++ b/tests/regression/flash_attention/flash_impl.hpp @@ -4,6 +4,9 @@ #include #include +#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, diff --git a/tests/regression/flash_attention/kernel.cpp b/tests/regression/flash_attention/kernel.cpp index 3c2d463c..9120d683 100644 --- a/tests/regression/flash_attention/kernel.cpp +++ b/tests/regression/flash_attention/kernel.cpp @@ -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( 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(smem_start + smem_Q0_offset); - float *smem_Q1 = reinterpret_cast(smem_start + smem_Q1_offset); - float *smem_K0 = reinterpret_cast(smem_start + smem_K0_offset); - float *smem_K1 = reinterpret_cast(smem_start + smem_K1_offset); - float *smem_V0 = reinterpret_cast(smem_start + smem_V0_offset); - float *smem_V1 = reinterpret_cast(smem_start + smem_V1_offset); - float *smem_S0 = reinterpret_cast(smem_start + smem_S0_offset); - float *smem_S1 = reinterpret_cast(smem_start + smem_S1_offset); - float *smem_P0 = reinterpret_cast(smem_start + smem_P0_offset); - float *smem_P1 = reinterpret_cast(smem_start + smem_P1_offset); - float *smem_O0 = reinterpret_cast(smem_start + smem_O0_offset); - float *smem_O1 = reinterpret_cast(smem_start + smem_O1_offset); + float *smem_Q0 = reinterpret_cast(smem_start + smem_Q0_hexadecile * smem_hexadecile_size); + float *smem_Q1 = reinterpret_cast(smem_start + smem_Q1_hexadecile * smem_hexadecile_size); + float *smem_K0 = reinterpret_cast(smem_start + smem_K0_hexadecile * smem_hexadecile_size); + float *smem_K1 = reinterpret_cast(smem_start + smem_K1_hexadecile * smem_hexadecile_size); + float *smem_V0 = reinterpret_cast(smem_start + smem_V0_hexadecile * smem_hexadecile_size); + float *smem_V1 = reinterpret_cast(smem_start + smem_V1_hexadecile * smem_hexadecile_size); + float *smem_S0 = reinterpret_cast(smem_start + smem_S0_hexadecile * smem_hexadecile_size); + float *smem_S1 = reinterpret_cast(smem_start + smem_S1_hexadecile * smem_hexadecile_size); + float *smem_P0 = reinterpret_cast(smem_start + smem_P0_hexadecile * smem_hexadecile_size); + float *smem_P1 = reinterpret_cast(smem_start + smem_P1_hexadecile * smem_hexadecile_size); + float *smem_O0 = reinterpret_cast(smem_start + smem_O0_hexadecile * smem_hexadecile_size); + float *smem_O1 = reinterpret_cast(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() { diff --git a/tests/regression/flash_attention/kernel.gemmini.cpp b/tests/regression/flash_attention/kernel.gemmini.cpp index ac3788d4..9208e863 100644 --- a/tests/regression/flash_attention/kernel.gemmini.cpp +++ b/tests/regression/flash_attention/kernel.gemmini.cpp @@ -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(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(smem_start + smem_Q0_offset); - float *smem_Q1 = reinterpret_cast(smem_start + smem_Q1_offset); - float *smem_K0 = reinterpret_cast(smem_start + smem_K0_offset); - float *smem_K1 = reinterpret_cast(smem_start + smem_K1_offset); - float *smem_V0 = reinterpret_cast(smem_start + smem_V0_offset); - float *smem_V1 = reinterpret_cast(smem_start + smem_V1_offset); - float *smem_S0 = reinterpret_cast(smem_start + smem_S0_offset); - float *smem_S1 = reinterpret_cast(smem_start + smem_S1_offset); - float *smem_P0 = reinterpret_cast(smem_start + smem_P0_offset); - float *smem_P1 = reinterpret_cast(smem_start + smem_P1_offset); - float *smem_O0 = reinterpret_cast(smem_start + smem_O0_offset); - float *smem_O1 = reinterpret_cast(smem_start + smem_O1_offset); + float *smem_Q0 = reinterpret_cast(smem_start + smem_Q0_hexadecile * smem_hexadecile_size); + float *smem_Q1 = reinterpret_cast(smem_start + smem_Q1_hexadecile * smem_hexadecile_size); + float *smem_K0 = reinterpret_cast(smem_start + smem_K0_hexadecile * smem_hexadecile_size); + float *smem_K1 = reinterpret_cast(smem_start + smem_K1_hexadecile * smem_hexadecile_size); + float *smem_V0 = reinterpret_cast(smem_start + smem_V0_hexadecile * smem_hexadecile_size); + float *smem_V1 = reinterpret_cast(smem_start + smem_V1_hexadecile * smem_hexadecile_size); + float *smem_S0 = reinterpret_cast(smem_start + smem_S0_hexadecile * smem_hexadecile_size); + float *smem_S1 = reinterpret_cast(smem_start + smem_S1_hexadecile * smem_hexadecile_size); + float *smem_P0 = reinterpret_cast(smem_start + smem_P0_hexadecile * smem_hexadecile_size); + float *smem_P1 = reinterpret_cast(smem_start + smem_P1_hexadecile * smem_hexadecile_size); + float *smem_O0 = reinterpret_cast(smem_start + smem_O0_hexadecile * smem_hexadecile_size); + float *smem_O1 = reinterpret_cast(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( + 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( + 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( - 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( - 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( - 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( - 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( - 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( - 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( - 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( - 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( - 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( - 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() { diff --git a/tests/regression/flash_attention/kernel.gemmini.warpspec.cpp b/tests/regression/flash_attention/kernel.gemmini.warpspec.cpp new file mode 100644 index 00000000..f2a2e471 --- /dev/null +++ b/tests/regression/flash_attention/kernel.gemmini.warpspec.cpp @@ -0,0 +1,738 @@ +#include +#include +#include +#include +#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(arg->addr_q); + const float *gmem_K = reinterpret_cast(arg->addr_k); + const float *gmem_V = reinterpret_cast(arg->addr_v); + float *gmem_O = reinterpret_cast(arg->addr_o); + + float *gmem_tmp_d0 = reinterpret_cast(0xd0000000UL); + float *gmem_tmp_d1 = reinterpret_cast(0xd1000000UL); + float *gmem_tmp_d2 = reinterpret_cast(0xd2000000UL); + float *gmem_tmp_d3 = reinterpret_cast(0xd3000000UL); + float *gmem_tmp_d4 = reinterpret_cast(0xd4000000UL); + float *gmem_tmp_d5 = reinterpret_cast(0xd5000000UL); + float *gmem_tmp_d6 = reinterpret_cast(0xd6000000UL); + float *gmem_tmp_d7 = reinterpret_cast(0xd7000000UL); + float *gmem_tmp_e0 = reinterpret_cast(0xe0000000UL); + float *gmem_tmp_e1 = reinterpret_cast(0xe1000000UL); + float *gmem_tmp_e2 = reinterpret_cast(0xe2000000UL); + float *gmem_tmp_e3 = reinterpret_cast(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(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(smem_start + smem_Q0_hexadecile * smem_hexadecile_size); + float *smem_Q1 = reinterpret_cast(smem_start + smem_Q1_hexadecile * smem_hexadecile_size); + float *smem_K0 = reinterpret_cast(smem_start + smem_K0_hexadecile * smem_hexadecile_size); + float *smem_K1 = reinterpret_cast(smem_start + smem_K1_hexadecile * smem_hexadecile_size); + float *smem_V0 = reinterpret_cast(smem_start + smem_V0_hexadecile * smem_hexadecile_size); + float *smem_V1 = reinterpret_cast(smem_start + smem_V1_hexadecile * smem_hexadecile_size); + float *smem_S0 = reinterpret_cast(smem_start + smem_S0_hexadecile * smem_hexadecile_size); + float *smem_S1 = reinterpret_cast(smem_start + smem_S1_hexadecile * smem_hexadecile_size); + float *smem_P0 = reinterpret_cast(smem_start + smem_P0_hexadecile * smem_hexadecile_size); + float *smem_P1 = reinterpret_cast(smem_start + smem_P1_hexadecile * smem_hexadecile_size); + float *smem_O0 = reinterpret_cast(smem_start + smem_O0_hexadecile * smem_hexadecile_size); + float *smem_O1 = reinterpret_cast(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(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(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(smem_Q0, gmem_tmp_d0, tid_in_warpgroup, + // threads_per_warpgroup, warpgroup_id_in_cluster); + // thread_block_copy_tile(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( + 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( + 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( + 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( + 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( + 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( + 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( + 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( + 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( + smem_P_produce, gmem_tmp_d2, tid_in_warpgroup_simt, + threads_per_warpgroup_simt, warpgroup_id_simt); + thread_block_copy_tile( + 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( + smem_P_produce, gmem_tmp_d3, tid_in_warpgroup_simt, + threads_per_warpgroup_simt, warpgroup_id_simt); + thread_block_copy_tile( + 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; +} diff --git a/tests/regression/sgemm_gemmini_dma/kernel.cpp b/tests/regression/sgemm_gemmini_dma/kernel.cpp index efcb5b52..8ae36bd8 100644 --- a/tests/regression/sgemm_gemmini_dma/kernel.cpp +++ b/tests/regression/sgemm_gemmini_dma/kernel.cpp @@ -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( + a_hexadecile, b_hexadecile, 0 /*d_hexadecile*/, tile_k > 0); } /* diff --git a/tests/regression/sgemm_tcore/sgemm_impl.hpp b/tests/regression/sgemm_tcore/sgemm_impl.hpp index afc37542..4d121a0f 100644 --- a/tests/regression/sgemm_tcore/sgemm_impl.hpp +++ b/tests/regression/sgemm_tcore/sgemm_impl.hpp @@ -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( &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( &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_%=:" ::); diff --git a/tests/regression/sgemm_tcore/switch_args_input.sh b/tests/regression/sgemm_tcore/switch_args_input.sh new file mode 100755 index 00000000..392794e4 --- /dev/null +++ b/tests/regression/sgemm_tcore/switch_args_input.sh @@ -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."