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

flash: Add begin end markers

Browse Source
This commit is contained in:
Hansung Kim
2024-11-09 10:16:40 -08:00
parent 1e3d476e70
commit 365b1d8e67
3 changed files with 11 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -4,6 +4,9 @@
#include <vx_spawn.h> #include <vx_spawn.h>
#include <float.h> #include <float.h>
#define MARK_BEG() asm volatile ("slti x0, x1, -1047")
#define MARK_END() asm volatile ("slti x0, x1, -499")
#define B_ROW 64 #define B_ROW 64
#define B_COL 64 #define B_COL 64
#define HEADDIM 64 #define HEADDIM 64

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

@@ -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, loop_matmul_skips(/*skip_lda=*/0, /*skip_ldb=*/0, /*skip_ldd=*/1,
/*skip_ex=*/1, /*skip_stc=*/1); /*skip_ex=*/1, /*skip_stc=*/1);
MARK_BEG();
if constexpr (GEMMINI_DMA) { if constexpr (GEMMINI_DMA) {
if (tid_in_warpgroup == 0) { if (tid_in_warpgroup == 0) {
gemmini_extended_config_ex(WEIGHT_STATIONARY, 0, 0, 1, 0, 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) { if (warpgroup_id == 0) {
threadblock_barrier(global_barrier_id, warps_per_threadblock_per_core); threadblock_barrier(global_barrier_id, warps_per_threadblock_per_core);
} }
MARK_END();
} }
int main() { int main() {

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

@@ -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, loop_matmul_skips(/*skip_lda=*/1, /*skip_ldb=*/1, /*skip_ldd=*/0,
/*skip_ex=*/0, /*skip_stc=*/1); /*skip_ex=*/0, /*skip_stc=*/1);
MARK_BEG();
if (tid_in_warpgroup == 0) { if (tid_in_warpgroup == 0) {
gemmini_extended_config_ex(WEIGHT_STATIONARY, 0, 0, 1, 0, 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_%=:" :: ); asm volatile ("tile_loop_finish_%=:" :: );
MARK_END();
} }
int main() { int main() {