From 8dadbdd42d6180eeef14fbefcec48145e6cfe9a8 Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Tue, 29 Oct 2024 19:43:22 -0700 Subject: [PATCH 01/22] tensor: Do DMA mvin for next m/n loop at the last k iter This increases util by pulling the DMA wait time out of the K-loop wraparound (next N) and overlapping it with the last K iter. --- tests/regression/sgemm_tcore/sgemm_impl.hpp | 62 +++++++++++++++------ 1 file changed, 44 insertions(+), 18 deletions(-) diff --git a/tests/regression/sgemm_tcore/sgemm_impl.hpp b/tests/regression/sgemm_tcore/sgemm_impl.hpp index afc37542..13226a76 100644 --- a/tests/regression/sgemm_tcore/sgemm_impl.hpp +++ b/tests/regression/sgemm_tcore/sgemm_impl.hpp @@ -1150,19 +1150,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 +1182,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 +1199,27 @@ 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 == block_m_end) ? 0 : block_n + 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 +1227,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 +1371,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 +1391,8 @@ inline void thread_block_gemm(const T *A, const T *B, float *C, } } } + + asm volatile("move_out_end_%=:" ::); } } asm volatile("loop_mn_end_%=:" ::); From 6b39a6fe703ff9ba3759d608000c8abb204e6571 Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Tue, 29 Oct 2024 20:14:33 -0700 Subject: [PATCH 02/22] Add convenience script for switching input/args binaries --- .../sgemm_tcore/switch_args_input.sh | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100755 tests/regression/sgemm_tcore/switch_args_input.sh 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." From 21b6655c101e18c4fdb20f10331be4946b6530ef Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Tue, 29 Oct 2024 22:34:22 -0700 Subject: [PATCH 03/22] sgemm_impl: Implement fast coalesced wmma_store Enables a fairer comparison between core-coupled tensor core to Hopper tensor core, where the latter benefits from coalesced full-throughput moveout to GMEM because it does not use the 1x2 interleaved register mapping. This means the result matrix will be stored swizzled in the GMEM, without breaking correctness. --- tests/regression/sgemm_tcore/sgemm_impl.hpp | 69 ++++++++++++++------- 1 file changed, 48 insertions(+), 21 deletions(-) diff --git a/tests/regression/sgemm_tcore/sgemm_impl.hpp b/tests/regression/sgemm_tcore/sgemm_impl.hpp index 13226a76..5bd694dd 100644 --- a/tests/regression/sgemm_tcore/sgemm_impl.hpp +++ b/tests/regression/sgemm_tcore/sgemm_impl.hpp @@ -104,7 +104,13 @@ static_assert(WMITER * WNITER * TCM * TCN * NUM_WARPS * CORES_PER_CLUSTER == #define TRANSPOSE_AT_PRODUCE 0 #define TRANSPOSE_AT_CONSUME 0 -#define GEMMINI_DMA 1 +// 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 1 + +#define GEMMINI_DMA 0 #define GEMMINI_DMA_FAST 1 #if SMEM_SIZE == 0x4000 #define SMEM_ADDR_Q0 ((float * const) 0xff000000) @@ -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_%=:" :: ); From c001618fb91adf6debcab52dd1cee984a723541b Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Tue, 29 Oct 2024 22:35:56 -0700 Subject: [PATCH 04/22] sgemm_impl: Fix wrong next block_m logic for DMA --- tests/regression/sgemm_tcore/sgemm_impl.hpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/regression/sgemm_tcore/sgemm_impl.hpp b/tests/regression/sgemm_tcore/sgemm_impl.hpp index 5bd694dd..aaf66492 100644 --- a/tests/regression/sgemm_tcore/sgemm_impl.hpp +++ b/tests/regression/sgemm_tcore/sgemm_impl.hpp @@ -1237,7 +1237,8 @@ inline void thread_block_gemm(const T *A, const T *B, float *C, : block_n; const uint32_t next_block_m = (next_block_n == 0) - ? ((block_m == block_m_end) ? 0 : block_n + 1) + ? (((block_m + 1) == block_m_end) ? block_m_start /*unused*/ + : block_m + 1) : block_m; asm volatile("next_index_end_%=:" ::); From 405525501822c069bcc56155789c3b95fab853bd Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Fri, 8 Nov 2024 16:40:16 -0800 Subject: [PATCH 05/22] flash: Fix tcore kernel for CISC arg field changes --- tests/regression/flash_attention/kernel.cpp | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/regression/flash_attention/kernel.cpp b/tests/regression/flash_attention/kernel.cpp index 3c2d463c..17fcf91f 100644 --- a/tests/regression/flash_attention/kernel.cpp +++ b/tests/regression/flash_attention/kernel.cpp @@ -259,7 +259,7 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { (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(); @@ -549,7 +549,7 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { k_LOOP_WS_CONFIG_ADDRS_AB) // 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(); @@ -813,8 +813,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_%=:" :: ); From 4e087a8aab3c3ba8fd51f359c9a24f462352f097 Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Fri, 8 Nov 2024 16:43:08 -0800 Subject: [PATCH 06/22] flash: Fix loop iteration for gemmini Kernel is software-pipelined around 2 GEMMs and softmax; it requires two iterations to fully complete a tile. --- tests/regression/flash_attention/kernel.gemmini.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/regression/flash_attention/kernel.gemmini.cpp b/tests/regression/flash_attention/kernel.gemmini.cpp index ac3788d4..089f0a3f 100644 --- a/tests/regression/flash_attention/kernel.gemmini.cpp +++ b/tests/regression/flash_attention/kernel.gemmini.cpp @@ -336,7 +336,11 @@ 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) + 2 /*pipeline latency*/; + 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); From 1e3d476e70d8e521efe7082fcd73fdc52a4cfaa8 Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Fri, 8 Nov 2024 21:56:26 -0800 Subject: [PATCH 07/22] Switch header configs to flash --- kernel/include/gemmini_mmio.h | 4 ++-- tests/regression/flash_attention/flash_impl.hpp | 6 ++++-- tests/regression/sgemm_tcore/sgemm_impl.hpp | 6 +++--- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/kernel/include/gemmini_mmio.h b/kernel/include/gemmini_mmio.h index ed55236c..0dec66b0 100644 --- a/kernel/include/gemmini_mmio.h +++ b/kernel/include/gemmini_mmio.h @@ -9,9 +9,9 @@ // #define SMEM_SIZE 0x4000 // 64KB // #define SMEM_SIZE 0x10000 -// 128KB +// 128KB (FP16 GEMM) // #define SMEM_SIZE 0x20000 -// 256KB +// 256KB (FlashAttention) #define SMEM_SIZE 0x40000 #define SMEM_MASK (SMEM_SIZE - 1) diff --git a/tests/regression/flash_attention/flash_impl.hpp b/tests/regression/flash_attention/flash_impl.hpp index 47e21c70..2e2bd693 100644 --- a/tests/regression/flash_attention/flash_impl.hpp +++ b/tests/regression/flash_attention/flash_impl.hpp @@ -11,8 +11,10 @@ #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 TENSOR_CORE = true; +constexpr bool WARP_SPECIALIZED = false; +constexpr bool TENSOR_CORE = false; // temporary safety stop for wrong configs static_assert(NUM_CORES == 4); diff --git a/tests/regression/sgemm_tcore/sgemm_impl.hpp b/tests/regression/sgemm_tcore/sgemm_impl.hpp index aaf66492..989c5df9 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: @@ -110,7 +110,7 @@ static_assert(WMITER * WNITER * TCM * TCN * NUM_WARPS * CORES_PER_CLUSTER == // result matrix will be stored in a swizzled form in the global memory. #define WMMA_STORE_FAST 1 -#define GEMMINI_DMA 0 +#define GEMMINI_DMA 1 #define GEMMINI_DMA_FAST 1 #if SMEM_SIZE == 0x4000 #define SMEM_ADDR_Q0 ((float * const) 0xff000000) From 365b1d8e67628f70e9cc764802446508b940187b Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 10:16:40 -0800 Subject: [PATCH 08/22] flash: Add begin end markers --- tests/regression/flash_attention/flash_impl.hpp | 3 +++ tests/regression/flash_attention/kernel.cpp | 4 ++++ tests/regression/flash_attention/kernel.gemmini.cpp | 4 ++++ 3 files changed, 11 insertions(+) diff --git a/tests/regression/flash_attention/flash_impl.hpp b/tests/regression/flash_attention/flash_impl.hpp index 2e2bd693..ec5e6f9c 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 diff --git a/tests/regression/flash_attention/kernel.cpp b/tests/regression/flash_attention/kernel.cpp index 17fcf91f..c3298a61 100644 --- a/tests/regression/flash_attention/kernel.cpp +++ b/tests/regression/flash_attention/kernel.cpp @@ -219,6 +219,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); @@ -822,6 +824,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 089f0a3f..090233cb 100644 --- a/tests/regression/flash_attention/kernel.gemmini.cpp +++ b/tests/regression/flash_attention/kernel.gemmini.cpp @@ -212,6 +212,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); @@ -681,6 +683,8 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { } asm volatile ("tile_loop_finish_%=:" :: ); + + MARK_END(); } int main() { From 952b8debbbfb9e4723c1823faa06643260f92363 Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 16:21:34 -0800 Subject: [PATCH 09/22] flash: Update to use new CISC interface --- .../flash_attention/kernel.gemmini.cpp | 183 +++++++++++------- 1 file changed, 111 insertions(+), 72 deletions(-) diff --git a/tests/regression/flash_attention/kernel.gemmini.cpp b/tests/regression/flash_attention/kernel.gemmini.cpp index 090233cb..c46b0bd1 100644 --- a/tests/regression/flash_attention/kernel.gemmini.cpp +++ b/tests/regression/flash_attention/kernel.gemmini.cpp @@ -10,6 +10,8 @@ #define FENCE_GEMM_II +#define GEMMINI_NEW_CISC + constexpr bool DEBUG = false; static_assert(GEMMINI_DMA && !WARP_SPECIALIZED, @@ -98,40 +100,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 +173,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); @@ -246,22 +251,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( @@ -274,7 +284,14 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { #endif 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), @@ -282,9 +299,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*/, @@ -379,6 +394,17 @@ 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 */) { @@ -394,26 +420,19 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { asm volatile("gemm_pv_start_%=:" ::); 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 + // 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, @@ -452,7 +471,11 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { // 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); +#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, @@ -460,6 +483,7 @@ 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(); @@ -487,14 +511,16 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { (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) { + // // 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, @@ -504,6 +530,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_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, @@ -511,6 +549,7 @@ 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 From 6990fcc1e66820935f4d1244191c9cf505bee4d7 Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 16:43:45 -0800 Subject: [PATCH 10/22] Add compute-and-mvout-to-spad API --- kernel/include/gemmini_mmio.h | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/kernel/include/gemmini_mmio.h b/kernel/include/gemmini_mmio.h index 6043e0be..f1ca0e77 100644 --- a/kernel/include/gemmini_mmio.h +++ b/kernel/include/gemmini_mmio.h @@ -87,6 +87,7 @@ static size_t gemmini_tile_idx[NUM_THREADS * NUM_WARPS * NUM_CORES * NUM_CLUSTER #define GEMMINI_CISC_COMPUTE_HEXADECILES 0 +#define GEMMINI_CISC_COMPUTE_AND_STORE_TO_SPAD 1 #define GEMMINI_CISC_SET_AB_STRIDE 8 #define GEMMINI_CISC_STORE_TO_SPAD 9 #define GEMMINI_CISC_LOAD_TO_HEXADECILES 10 @@ -107,8 +108,18 @@ 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((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, From ad75561efe1c759b2e1435ce95b32954c03d77d3 Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 16:44:17 -0800 Subject: [PATCH 11/22] flash: Reduce fence calls to improve util --- .../flash_attention/kernel.gemmini.cpp | 32 ++++++++----------- 1 file changed, 14 insertions(+), 18 deletions(-) diff --git a/tests/regression/flash_attention/kernel.gemmini.cpp b/tests/regression/flash_attention/kernel.gemmini.cpp index c46b0bd1..e97ea635 100644 --- a/tests/regression/flash_attention/kernel.gemmini.cpp +++ b/tests/regression/flash_attention/kernel.gemmini.cpp @@ -10,7 +10,9 @@ #define FENCE_GEMM_II -#define GEMMINI_NEW_CISC +#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; @@ -282,6 +284,8 @@ 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(); // also move Q to spad_addr_Q1 for the second iteration @@ -309,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); @@ -424,9 +428,6 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { // 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, @@ -458,16 +459,17 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { if (tid_in_warpgroup == 0) { // fence to GEMM II completion gemmini_fence(); - gemmini_fence(); - gemmini_fence(); - 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 + // // 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); @@ -485,10 +487,6 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { /*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 @@ -511,7 +509,6 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { (uint64_t)(gmem_V_tile), k_LOOP_WS_CONFIG_ADDRS_AB) #endif - // gemmini_fence(); // do DMA if (tile_k == 0) { @@ -554,9 +551,6 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { // fence everything before going to the next tile gemmini_fence(); - gemmini_fence(); - gemmini_fence(); - gemmini_fence(); } // threadblock_barrier(warpgroup_id_in_cluster, @@ -625,6 +619,7 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { } #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) { @@ -634,6 +629,7 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { *smem_O_flag = 0; vx_fence(); } + asm volatile("rescale_fence_read_end_%=:" ::); #endif #if 0 From ac42f2dbba187df00de6267f9257efdba27db5bf Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 16:49:39 -0800 Subject: [PATCH 12/22] sgemm_gemmini_dma: Update with new compute API --- tests/regression/sgemm_gemmini_dma/kernel.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/regression/sgemm_gemmini_dma/kernel.cpp b/tests/regression/sgemm_gemmini_dma/kernel.cpp index 2e55cb5e..8a1e603f 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); } /* From 673e07ed43abcd95899e5ed45e6fe68b170497fb Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 19:08:39 -0800 Subject: [PATCH 13/22] flash: Add non-warp-specialized gemmini flash kernel --- tests/regression/flash_attention/Makefile | 4 +- .../regression/flash_attention/flash_impl.hpp | 23 +- .../kernel.gemmini.nowarpspec.cpp | 716 ++++++++++++++++++ 3 files changed, 731 insertions(+), 12 deletions(-) create mode 100644 tests/regression/flash_attention/kernel.gemmini.nowarpspec.cpp diff --git a/tests/regression/flash_attention/Makefile b/tests/regression/flash_attention/Makefile index 0456e983..2f2c1c6a 100644 --- a/tests/regression/flash_attention/Makefile +++ b/tests/regression/flash_attention/Makefile @@ -2,7 +2,9 @@ PROJECT = flash_attention SRCS = main.cpp common.h -VX_SRCS = kernel.gemmini.cpp +# VX_SRCS = kernel.cpp +# VX_SRCS = kernel.gemmini.cpp +VX_SRCS = kernel.gemmini.nowarpspec.cpp VX_INCLUDES = flash_impl.hpp ../sgemm_tcore/sgemm_impl.hpp OPTS ?= -n16 diff --git a/tests/regression/flash_attention/flash_impl.hpp b/tests/regression/flash_attention/flash_impl.hpp index ec5e6f9c..096333e5 100644 --- a/tests/regression/flash_attention/flash_impl.hpp +++ b/tests/regression/flash_attention/flash_impl.hpp @@ -17,6 +17,7 @@ constexpr uint32_t ROWMAX_SETS = 3; // constexpr bool WARP_SPECIALIZED = true; // 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 @@ -101,7 +102,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, @@ -133,7 +134,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, @@ -156,7 +157,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, @@ -213,7 +214,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); @@ -300,7 +301,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, @@ -355,7 +356,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, @@ -403,7 +404,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, @@ -434,7 +435,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, @@ -467,7 +468,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, @@ -496,7 +497,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, @@ -551,7 +552,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.gemmini.nowarpspec.cpp b/tests/regression/flash_attention/kernel.gemmini.nowarpspec.cpp new file mode 100644 index 00000000..79079811 --- /dev/null +++ b/tests/regression/flash_attention/kernel.gemmini.nowarpspec.cpp @@ -0,0 +1,716 @@ +#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); + // } + + 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 (tile_k >= 2) // delay GEMM II 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) { + // FIXME: remove + // // 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*/)); + + // 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 + } + } + + // reconverge from mmio divergence + threadblock_barrier(warpgroup_id_in_cluster, + warps_per_warpgroup_per_core); + + asm volatile("move_k_v_finish_%=:" ::); + + // FIXME: remove for nowarpspec + // + // 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); + } + + { + 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, + threads_per_warpgroup, warpgroup_id); + } else if (tile_k_ == 1) { + thread_block_copy_tile( + smem_S_consume, gmem_tmp_d1, tid_in_warpgroup, + threads_per_warpgroup, warpgroup_id); + } + + 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, + threads_per_warpgroup, warpgroup_id, smem_scratchpad, + smem_rowmax, smem_rowsum, smem_O_row_scale); + + 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, + threads_per_warpgroup, warpgroup_id_in_cluster); + thread_block_copy_rowmax( + 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, threads_per_warpgroup, + warpgroup_id_in_cluster); + thread_block_copy_rowmax(smem_rowsum, gmem_tmp_e3, + tid_in_warpgroup, threads_per_warpgroup, + warpgroup_id_in_cluster); + } + + threadblock_barrier(warpgroup_id_in_cluster, + warps_per_warpgroup_per_core); + } + } + +#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 == 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_in_cluster == 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, threads_per_warpgroup, + warpgroup_id_in_cluster); + } else if (tile_k_ == 2) { + thread_block_copy_tile( + smem_O, gmem_tmp_d7, tid_in_warpgroup, threads_per_warpgroup, + warpgroup_id_in_cluster); + } + + 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, + threads_per_warpgroup, warpgroup_id_in_cluster); + + // rescale-to-PV-GEMM barrier + threadblock_barrier(warpgroup_id_in_cluster, + warps_per_warpgroup_per_core); + + if constexpr (DEBUG) { + 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, + threads_per_warpgroup, warpgroup_id_in_cluster); + thread_block_copy_tile( + 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, + threads_per_warpgroup, warpgroup_id_in_cluster); + thread_block_copy_tile( + smem_O, gmem_tmp_d5, tid_in_warpgroup, threads_per_warpgroup, + warpgroup_id_in_cluster); + } + + threadblock_barrier(warpgroup_id_in_cluster, + warps_per_warpgroup_per_core); + } + } + } + + // intra-warpgroup barrier + threadblock_barrier(warpgroup_id_in_cluster, + warps_per_warpgroup_per_core); + + // fence everything before going to the next tile + gemmini_fence(); + } + } + + 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; +} From 76a6aaf085c74c96a240c037801d58b250712b5e Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 19:09:09 -0800 Subject: [PATCH 14/22] flash: doc update --- tests/regression/flash_attention/kernel.gemmini.cpp | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/regression/flash_attention/kernel.gemmini.cpp b/tests/regression/flash_attention/kernel.gemmini.cpp index e97ea635..f2a2e471 100644 --- a/tests/regression/flash_attention/kernel.gemmini.cpp +++ b/tests/regression/flash_attention/kernel.gemmini.cpp @@ -470,9 +470,6 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { // Kick off GEMM I // - // 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); #ifdef GEMMINI_NEW_CISC gemmini_tile_compute( spad_hex_Q, spad_hex_K_consume, spad_hex_S_produce, @@ -566,7 +563,7 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { } 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; From 8fe6d918f2cd7458a06129e222cb26611f3e375b Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 19:49:20 -0800 Subject: [PATCH 15/22] flash: Update tcore kernel to use new CISC --- tests/regression/flash_attention/kernel.cpp | 126 ++++++++++++-------- 1 file changed, 79 insertions(+), 47 deletions(-) diff --git a/tests/regression/flash_attention/kernel.cpp b/tests/regression/flash_attention/kernel.cpp index c3298a61..c5efbf3b 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, @@ -246,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) { @@ -256,6 +271,19 @@ 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), @@ -265,13 +293,6 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { 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( @@ -281,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, @@ -549,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*/ << 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, @@ -563,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 { From 1c9b02215652a9d4c9c59556a9b28c822867afdf Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 19:58:45 -0800 Subject: [PATCH 16/22] flash: Rename nowarpspec to default --- tests/regression/flash_attention/Makefile | 4 +- .../regression/flash_attention/flash_impl.hpp | 1 + .../flash_attention/kernel.gemmini.cpp | 152 ++++++++---------- ...rpspec.cpp => kernel.gemmini.warpspec.cpp} | 152 ++++++++++-------- 4 files changed, 155 insertions(+), 154 deletions(-) rename tests/regression/flash_attention/{kernel.gemmini.nowarpspec.cpp => kernel.gemmini.warpspec.cpp} (87%) diff --git a/tests/regression/flash_attention/Makefile b/tests/regression/flash_attention/Makefile index 2f2c1c6a..3a25e4f3 100644 --- a/tests/regression/flash_attention/Makefile +++ b/tests/regression/flash_attention/Makefile @@ -3,8 +3,8 @@ PROJECT = flash_attention SRCS = main.cpp common.h # VX_SRCS = kernel.cpp -# VX_SRCS = kernel.gemmini.cpp -VX_SRCS = kernel.gemmini.nowarpspec.cpp +# VX_SRCS = kernel.gemmini.warpspec.cpp +VX_SRCS = kernel.gemmini.cpp VX_INCLUDES = flash_impl.hpp ../sgemm_tcore/sgemm_impl.hpp OPTS ?= -n16 diff --git a/tests/regression/flash_attention/flash_impl.hpp b/tests/regression/flash_attention/flash_impl.hpp index 096333e5..6f210b75 100644 --- a/tests/regression/flash_attention/flash_impl.hpp +++ b/tests/regression/flash_attention/flash_impl.hpp @@ -15,6 +15,7 @@ constexpr uint32_t ROWMAX_SETS = 3; // 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; diff --git a/tests/regression/flash_attention/kernel.gemmini.cpp b/tests/regression/flash_attention/kernel.gemmini.cpp index f2a2e471..79079811 100644 --- a/tests/regression/flash_attention/kernel.gemmini.cpp +++ b/tests/regression/flash_attention/kernel.gemmini.cpp @@ -342,16 +342,6 @@ 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 @@ -411,8 +401,8 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { 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 + { + if (tile_k >= 2) // delay GEMM II by 2 iters for pipelining { const uint32_t tile_k_ = tile_k - 2; @@ -457,16 +447,17 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { asm volatile("gemm_qk_start_%=:" ::); if (tid_in_warpgroup == 0) { - // fence to GEMM II completion - gemmini_fence(); + // FIXME: remove + // // 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 +// #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 // @@ -499,14 +490,6 @@ 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 - // do DMA if (tile_k == 0) { // // configure address strides for the DMA @@ -545,24 +528,24 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { /*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); + // reconverge from mmio divergence + threadblock_barrier(warpgroup_id_in_cluster, + warps_per_warpgroup_per_core); asm volatile("move_k_v_finish_%=:" ::); + // FIXME: remove for nowarpspec + // // 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; @@ -572,46 +555,49 @@ 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); } } @@ -619,7 +605,7 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { 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) { + if (tid_in_warpgroup == 0) { while ((*smem_O_flag) != 1) ; // set it back to 0 for the next tile iteration @@ -643,74 +629,66 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { #endif if constexpr (DEBUG) { - if (warpgroup_id == 0) { + if (warpgroup_id_in_cluster == 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); + 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); + // fence everything before going to the next tile + gemmini_fence(); } } diff --git a/tests/regression/flash_attention/kernel.gemmini.nowarpspec.cpp b/tests/regression/flash_attention/kernel.gemmini.warpspec.cpp similarity index 87% rename from tests/regression/flash_attention/kernel.gemmini.nowarpspec.cpp rename to tests/regression/flash_attention/kernel.gemmini.warpspec.cpp index 79079811..f2a2e471 100644 --- a/tests/regression/flash_attention/kernel.gemmini.nowarpspec.cpp +++ b/tests/regression/flash_attention/kernel.gemmini.warpspec.cpp @@ -342,6 +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 @@ -401,8 +411,8 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { const auto spad_hex_O = smem_O0_hexadecile; // NOTE: there's only single O tile asm volatile ("dbuf_sel_end_%=:" :: ); - { - if (tile_k >= 2) // delay GEMM II by 2 iters for pipelining + 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; @@ -447,17 +457,16 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { asm volatile("gemm_qk_start_%=:" ::); if (tid_in_warpgroup == 0) { - // FIXME: remove - // // fence to GEMM II completion - // gemmini_fence(); + // 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 +#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 // @@ -490,6 +499,14 @@ 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 + // do DMA if (tile_k == 0) { // // configure address strides for the DMA @@ -528,24 +545,24 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { /*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips); #endif } + + // fence everything before going to the next tile + gemmini_fence(); } - // reconverge from mmio divergence - threadblock_barrier(warpgroup_id_in_cluster, - warps_per_warpgroup_per_core); + // threadblock_barrier(warpgroup_id_in_cluster, + // warps_per_warpgroup_per_core); asm volatile("move_k_v_finish_%=:" ::); - // FIXME: remove for nowarpspec - // // 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; @@ -555,49 +572,46 @@ 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, - threads_per_warpgroup, warpgroup_id); + 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, - threads_per_warpgroup, warpgroup_id); + smem_S_consume, gmem_tmp_d1, tid_in_warpgroup_simt, + threads_per_warpgroup_simt, warpgroup_id_simt); } - threadblock_barrier(warpgroup_id_in_cluster, - warps_per_warpgroup_per_core); + threadblock_barrier(barrier_id_simt, barrier_count_simt); } } // Online softmax // thread_block_online_softmax( - smem_S_consume, smem_P_produce, tid_in_warpgroup, - threads_per_warpgroup, warpgroup_id, smem_scratchpad, + 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(warpgroup_id_in_cluster, - warps_per_warpgroup_per_core); + 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, - threads_per_warpgroup, warpgroup_id_in_cluster); + 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, - threads_per_warpgroup, warpgroup_id_in_cluster); + 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, threads_per_warpgroup, - warpgroup_id_in_cluster); + tid_in_warpgroup_simt, threads_per_warpgroup_simt, + warpgroup_id_simt); thread_block_copy_rowmax(smem_rowsum, gmem_tmp_e3, - tid_in_warpgroup, threads_per_warpgroup, - warpgroup_id_in_cluster); + tid_in_warpgroup_simt, threads_per_warpgroup_simt, + warpgroup_id_simt); } - threadblock_barrier(warpgroup_id_in_cluster, - warps_per_warpgroup_per_core); + threadblock_barrier(barrier_id_simt, barrier_count_simt); } } @@ -605,7 +619,7 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { 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) { + if (tid_in_warpgroup_simt == 0) { while ((*smem_O_flag) != 1) ; // set it back to 0 for the next tile iteration @@ -629,66 +643,74 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { #endif if constexpr (DEBUG) { - if (warpgroup_id_in_cluster == 0) { + 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, threads_per_warpgroup, - warpgroup_id_in_cluster); + 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, threads_per_warpgroup, - warpgroup_id_in_cluster); + smem_O, gmem_tmp_d7, tid_in_warpgroup_simt, threads_per_warpgroup_simt, + warpgroup_id_simt); } - threadblock_barrier(warpgroup_id_in_cluster, - warps_per_warpgroup_per_core); + 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, - threads_per_warpgroup, warpgroup_id_in_cluster); + 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(warpgroup_id_in_cluster, - warps_per_warpgroup_per_core); + threadblock_barrier(barrier_id_simt, barrier_count_simt); if constexpr (DEBUG) { - if (warpgroup_id_in_cluster == 0) { + if (warpgroup_id == 0) { // O before PV if (tile_k_ == 0) { thread_block_copy_tile( - smem_P_produce, gmem_tmp_d2, tid_in_warpgroup, - threads_per_warpgroup, warpgroup_id_in_cluster); + 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, threads_per_warpgroup, - warpgroup_id_in_cluster); + 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, - threads_per_warpgroup, warpgroup_id_in_cluster); + 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, threads_per_warpgroup, - warpgroup_id_in_cluster); + smem_O, gmem_tmp_d5, tid_in_warpgroup_simt, + threads_per_warpgroup_simt, warpgroup_id_simt); } - threadblock_barrier(warpgroup_id_in_cluster, - warps_per_warpgroup_per_core); + threadblock_barrier(barrier_id_simt, barrier_count_simt); } } } - // intra-warpgroup barrier +#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 - // fence everything before going to the next tile - gemmini_fence(); + // intra-warpgroup barrier + threadblock_barrier(barrier_id_simt, barrier_count_simt); } } From fcd8b0b892e8432f6a6aa057c9360e690e48f9ee Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 20:37:58 -0800 Subject: [PATCH 17/22] flash: Disable rescale flag check GEMM-II finishes much earlier than softmax for this to be a problem. --- .../flash_attention/kernel.gemmini.cpp | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/regression/flash_attention/kernel.gemmini.cpp b/tests/regression/flash_attention/kernel.gemmini.cpp index 79079811..882d8f96 100644 --- a/tests/regression/flash_attention/kernel.gemmini.cpp +++ b/tests/regression/flash_attention/kernel.gemmini.cpp @@ -602,17 +602,17 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { } #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 == 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_%=:" ::); + // 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 From 68054689c96d0405e0c62bd296cadadd97f0a886 Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 20:59:26 -0800 Subject: [PATCH 18/22] flash: Move fence to start of loop; wrap all MMIO in one tid=0 branch --- .../flash_attention/kernel.gemmini.cpp | 100 +++++++++--------- 1 file changed, 52 insertions(+), 48 deletions(-) diff --git a/tests/regression/flash_attention/kernel.gemmini.cpp b/tests/regression/flash_attention/kernel.gemmini.cpp index 882d8f96..e3335861 100644 --- a/tests/regression/flash_attention/kernel.gemmini.cpp +++ b/tests/regression/flash_attention/kernel.gemmini.cpp @@ -402,22 +402,28 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { asm volatile ("dbuf_sel_end_%=:" :: ); { - if (tile_k >= 2) // delay GEMM II by 2 iters for pipelining - { - const uint32_t tile_k_ = tile_k - 2; + // do all of GEMM kickoffs before the SIMT compute + // + if (tid_in_warpgroup == 0) { + // 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. + 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 + if (tile_k >= 2) // delay GEMM II by 2 iters for pipelining + { + const uint32_t tile_k_ = tile_k - 2; - asm volatile("gemm_pv_start_%=:" ::); + // 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 - if (tid_in_warpgroup == 0) { - // kickoff matmul + asm volatile("gemm_pv_start_%=:" ::); + // kick off GEMM II + // // 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, @@ -432,32 +438,26 @@ 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_%=:" ::); - } - - // GEMM I: S = Q*K - // - // kick off asynchronously; fence later - asm volatile("gemm_qk_start_%=:" ::); - - if (tid_in_warpgroup == 0) { // FIXME: remove // // 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 + // #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 // @@ -492,9 +492,12 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { // 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) | + // GEMMINI_CISC_CMD_R((HEADDIM /*V*/ << 20) | (dim_seqlen /*KT*/ << + // 8) | // 8 /*k_LOOP_WS_CONFIG_STRIDES_AB*/); // gemmini_fence(); // @@ -502,9 +505,11 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { // 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 @@ -528,22 +533,14 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { /*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips); #endif } + + asm volatile("move_k_v_finish_%=:" ::); } // reconverge from mmio divergence threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core); - - asm volatile("move_k_v_finish_%=:" ::); - - // FIXME: remove for nowarpspec - // - // 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); - } + } { if (tile_k >= 1) // delay online softmax by 1 iters @@ -630,8 +627,13 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { if constexpr (DEBUG) { if (warpgroup_id_in_cluster == 0) { - gemmini_fence(); - gemmini_fence(); + 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 */) { @@ -687,8 +689,10 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core); - // fence everything before going to the next tile - gemmini_fence(); + + // instead of fencing here, we fence at the start of the loop to maximize + // overlapping + // gemmini_fence(); } } From cb916ead39f598c619e9218b9ec4d23edda38f1a Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 20:59:58 -0800 Subject: [PATCH 19/22] Fix potential bitwidth bug in compute API --- kernel/include/gemmini_mmio.h | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/kernel/include/gemmini_mmio.h b/kernel/include/gemmini_mmio.h index f1ca0e77..45030b94 100644 --- a/kernel/include/gemmini_mmio.h +++ b/kernel/include/gemmini_mmio.h @@ -114,8 +114,9 @@ inline void gemmini_tile_compute(const uint32_t a_hexadecile, const uint32_t d_hexadecile, const bool accumulate) { if constexpr (!store_to_spad) { - GEMMINI_CISC_CMD_R((accumulate << 24) | (b_hexadecile << 16) | - (a_hexadecile << 8) | GEMMINI_CISC_COMPUTE_HEXADECILES); + 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); From 4448f31fdc0251824bda134e50de265fa6abdf46 Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sat, 9 Nov 2024 22:04:45 -0800 Subject: [PATCH 20/22] fence: Fix moving fence to start of loop For unknown reasons, guarding the fence with a tid == 0 branch causes a TL source ID re-used assertion. Just call the fence from all thread/warps as a workaround. At least, all threads in a warp will coalesce into one request. --- .../flash_attention/kernel.gemmini.cpp | 22 +++++++++++++------ 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/regression/flash_attention/kernel.gemmini.cpp b/tests/regression/flash_attention/kernel.gemmini.cpp index e3335861..b1ec5b29 100644 --- a/tests/regression/flash_attention/kernel.gemmini.cpp +++ b/tests/regression/flash_attention/kernel.gemmini.cpp @@ -402,13 +402,18 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { asm volatile ("dbuf_sel_end_%=:" :: ); { + // 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(); + // do all of GEMM kickoffs before the SIMT compute // if (tid_in_warpgroup == 0) { - // 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. - gemmini_fence(); if (tile_k >= 2) // delay GEMM II by 2 iters for pipelining { @@ -689,10 +694,13 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core); - - // instead of fencing here, we fence at the start of the loop to maximize - // overlapping + // @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); } } From 7d7cb5f60aa6d7d5e214b1ab6b6775b1bd8ab5bc Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Sun, 10 Nov 2024 22:44:02 -0800 Subject: [PATCH 21/22] flash: Disable perf loop multiplier --- tests/regression/flash_attention/kernel.cpp | 2 +- tests/regression/flash_attention/kernel.gemmini.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/regression/flash_attention/kernel.cpp b/tests/regression/flash_attention/kernel.cpp index c5efbf3b..9120d683 100644 --- a/tests/regression/flash_attention/kernel.cpp +++ b/tests/regression/flash_attention/kernel.cpp @@ -355,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; diff --git a/tests/regression/flash_attention/kernel.gemmini.cpp b/tests/regression/flash_attention/kernel.gemmini.cpp index b1ec5b29..9208e863 100644 --- a/tests/regression/flash_attention/kernel.gemmini.cpp +++ b/tests/regression/flash_attention/kernel.gemmini.cpp @@ -347,7 +347,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*/ * + 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)) + From 5ef4c8023e8d2a86c46878ddcedee5620bd11ae1 Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Mon, 11 Nov 2024 14:06:15 -0800 Subject: [PATCH 22/22] sgemm_impl: Disable wmma fast store Doesn't seem to have a big impact on tcore util. --- tests/regression/sgemm_tcore/sgemm_impl.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/regression/sgemm_tcore/sgemm_impl.hpp b/tests/regression/sgemm_tcore/sgemm_impl.hpp index 989c5df9..4d121a0f 100644 --- a/tests/regression/sgemm_tcore/sgemm_impl.hpp +++ b/tests/regression/sgemm_tcore/sgemm_impl.hpp @@ -108,7 +108,7 @@ static_assert(WMITER * WNITER * TCM * TCN * NUM_WARPS * CORES_PER_CLUSTER == // 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 1 +#define WMMA_STORE_FAST 0 #define GEMMINI_DMA 1 #define GEMMINI_DMA_FAST 1