diff --git a/tests/regression/flash_attention/kernel.cpp b/tests/regression/flash_attention/kernel.cpp index 64d0c302..72d25a61 100644 --- a/tests/regression/flash_attention/kernel.cpp +++ b/tests/regression/flash_attention/kernel.cpp @@ -2,6 +2,7 @@ #include #include #include +#include #include "common.h" #include "sgemm_impl.hpp" #include "include/gemmini.h" @@ -13,7 +14,7 @@ using float_type = float16_t; #define B_ROW BM #define B_COL BN -inline void thread_block_flashattn(float *S, float *gmem, +inline void thread_block_flashattn(float *S, const uint32_t tid_in_threadblock, const uint32_t threads_per_threadblock, const uint32_t threadblock_id_in_cluster, @@ -37,29 +38,121 @@ inline void thread_block_flashattn(float *S, float *gmem, // asm volatile("fmv.s %0, f22" : "=f"(ft[6])); // asm volatile("fmv.s %0, f23" : "=f"(ft[7])); - // row-max - // - // one warp handles one row in tile; iterate enough times to cover all the - // rows + volatile float *gmem_tmp0 = reinterpret_cast(0xd0000000UL); + volatile float *gmem_tmp1 = reinterpret_cast(0xe0000000UL); + volatile float *gmem_tmp2 = reinterpret_cast(0xf0000000UL); + for (int warp_offset = 0; warp_offset < B_ROW; warp_offset += warps_in_threadblock) { const uint32_t row = warp_offset + warp_id; const uint32_t first_thread_offset = B_COL * row; - uint32_t thread_offset = first_thread_offset + tid_in_warp; - float curr_max = S[first_thread_offset]; - constexpr uint32_t load_iter = B_COL / NUM_THREADS; + // rowmax + // + // two-level tree reduction: reduce each row into NUM_THREADS intermediate + // maxes, then reduce it to one global max + // one warp handles one row in tile + +// #define DUMB_ROWMAX +#ifdef DUMB_ROWMAX + if (tid_in_warp == 0) { + float max = S[first_thread_offset]; #pragma GCC unroll - for (int iter = 0; iter < load_iter; iter++) { + for (int i = 0; i < B_COL; i++) { + asm volatile("fmax.s %0, %1, %2" + : "=f"(max) + : "f"(max), "f"(S[first_thread_offset + i])); + } + sharedmem_row_max_sum[row] = max; + gmem_tmp0[row] = max; + } + +#else + static_assert((B_ROW % NUM_THREADS) == 0, + "B_ROW must be a multiple of NUM_THREADS"); + constexpr uint32_t per_row_iter = B_COL / NUM_THREADS; + uint32_t thread_offset = first_thread_offset + tid_in_warp; + float per_thread_max = FLT_MIN; +#pragma GCC unroll + for (int i = 0; i < per_row_iter; i++) { + const float next = S[thread_offset]; asm volatile("fmax.s %0, %1, %2" - : "=f"(curr_max) - : "f"(curr_max), "f"(S[thread_offset])); + : "=f"(per_thread_max) + : "f"(per_thread_max), "f"(next)); thread_offset += NUM_THREADS; } - // get max value across the same-warp threads using smem - // NOTE: be careful with out-of-bounds + // stage per-thread max value in smem + // FIXME: we could warp_id instead of row here, but we need another barrier + // at the end of the loop iteration to prevent write-after-read hazard + // FIXME: threadblock_id needs to be in here too float *warp_smem = sharedmem_scratchpad + (row * NUM_THREADS); - warp_smem[tid_in_warp] = curr_max; + warp_smem[tid_in_warp] = per_thread_max; + + // sync writes to warp_smem + threadblock_barrier(threadblock_id_in_cluster, + warps_per_threadblock_per_core); + + // elect 0-th thread to reduce all other thread's values in the warp + if (tid_in_warp == 0) { + for (int iter = 1; iter < NUM_THREADS; iter++) { + float other = warp_smem[iter]; + asm volatile("fmax.s %0, %1, %2" + : "=f"(per_thread_max) + : "f"(per_thread_max), "f"(other)); + } + sharedmem_row_max_sum[row] = per_thread_max; + gmem_tmp0[row] = per_thread_max; + } +#endif + + // FIXME: unnecessary? + threadblock_barrier(threadblock_id_in_cluster, + warps_per_threadblock_per_core); + + // exponential + // + // B_ROW / (B_ROW * B_COL / (exp_elem * threads_per_threadblock)) + // const uint32_t row_stride = + // (exp_elem_per_thread * threads_per_threadblock) / B_COL; + + thread_offset = first_thread_offset + tid_in_warp; + + // broadcast rowmax to all threads in the warp + const float row_max = sharedmem_row_max_sum[row]; + +#pragma GCC unroll + for (int i = 0; i < per_row_iter; i++) { + float val = S[thread_offset]; + + // FIXME: placeholder for proper exp + val = val; + + // update S in-place to P + // S[thread_offset] = val; + gmem_tmp1[thread_offset] = val; + gmem_tmp2[thread_offset] = val - row_max; + + thread_offset += NUM_THREADS; + } + + threadblock_barrier(threadblock_id_in_cluster, + warps_per_threadblock_per_core); + + // rowsum + // + // two-level tree reduction, similar to rowmax + +#if 0 + float per_thread_sum = 0.0f; +#pragma GCC unroll + for (int i = 0; i < per_row_iter; i++) { + per_thread_sum += S[thread_offset]; + thread_offset += NUM_THREADS; + } + // stage per-thread sum value in smem + // FIXME: threadblock_id needs to be in here too + warp_smem = sharedmem_scratchpad + (row * NUM_THREADS); + warp_smem[tid_in_warp] = per_thread_sum; // sync writes to warp_smem threadblock_barrier(threadblock_id_in_cluster, @@ -69,21 +162,16 @@ inline void thread_block_flashattn(float *S, float *gmem, if (tid_in_warp == 0) { for (int iter = 1; iter < NUM_THREADS; iter++) { float other = warp_smem[iter]; - asm volatile("fmax.s %0, %1, %2" - : "=f"(curr_max) - : "f"(curr_max), "f"(other)); + per_thread_sum += other; } - sharedmem_row_max_sum[row] = curr_max; + sharedmem_row_max_sum[row] = per_thread_sum; + gmem_tmp2[row] = per_thread_sum; } - } +#endif - // exponential - // - // FIXME: placeholder for proper exp - constexpr uint32_t exp_elem_per_thread = 1; - // B_ROW / (B_ROW * B_COL / (exp_elem * threads_per_threadblock)) - const uint32_t row_stride = - (exp_elem_per_thread * threads_per_threadblock) / B_COL; + threadblock_barrier(threadblock_id_in_cluster, + warps_per_threadblock_per_core); + } asm volatile("thread_block_flashattn_finish_%=:" ::); } @@ -135,17 +223,18 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { reinterpret_cast(SMEM_ADDR_END) - sharedmem_row_max_sum_size; // sharedmem "scratchpad" area to put temporary data, e.g. for tree reduction // in rowsum - // FIXME: size is arbitrary, and out-of bounds is not checked - constexpr uint32_t sharedmem_scratchpad_size = 0x1000; + // NOTE: out-of bounds is not checked + constexpr uint32_t sharedmem_scratchpad_size = + sizeof(float) * B_ROW * NUM_THREADS * 2 /*arbitrary slack*/; uint8_t *sharedmem_scratchpad = sharedmem_row_max_sum - sharedmem_scratchpad_size; - thread_block_gemm( - (const float_type *)arg->addr_a, (const float_type *)arg->addr_b, - (float *)smem_S /*write result to SMEM */, arg->dim_m, arg->dim_n, - arg->dim_k, tid_in_threadblock, threads_per_threadblock, - threadblocks_per_cluster, threadblock_id_in_cluster, - sharedmem_per_threadblock); + // thread_block_gemm( + // (const float_type *)arg->addr_a, (const float_type *)arg->addr_b, + // (float *)smem_S /*write result to SMEM */, arg->dim_m, arg->dim_n, + // arg->dim_k, tid_in_threadblock, threads_per_threadblock, + // threadblocks_per_cluster, threadblock_id_in_cluster, + // sharedmem_per_threadblock); // protect writes of GEMM results before softmax const uint32_t warps_per_threadblock_per_core = @@ -153,10 +242,10 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { threadblock_barrier(threadblock_id_in_cluster, warps_per_threadblock_per_core); - thread_block_flashattn( - (float *)smem_S, (float *)arg->addr_c, tid_in_threadblock, - threads_per_threadblock, threadblock_id_in_cluster, - (float *)sharedmem_scratchpad_size, (float *)sharedmem_row_max_sum); + thread_block_flashattn((float *)arg->addr_a /* smem_S, */, tid_in_threadblock, + threads_per_threadblock, threadblock_id_in_cluster, + (float *)sharedmem_scratchpad, + (float *)sharedmem_row_max_sum); } int main() {