From 449d99f0bb4b3e978e5bae107e9f59f61e8c58b6 Mon Sep 17 00:00:00 2001 From: Richard Yan Date: Tue, 16 Apr 2024 17:15:22 -0700 Subject: [PATCH] dram gemm kernel --- kernel/include/gemmini_mmio.h | 162 ++++++++++++ tests/regression/sgemm_gemmini/kernel.cpp | 286 +++++++++++++++++----- 2 files changed, 387 insertions(+), 61 deletions(-) create mode 100644 kernel/include/gemmini_mmio.h diff --git a/kernel/include/gemmini_mmio.h b/kernel/include/gemmini_mmio.h new file mode 100644 index 00000000..b9dde44d --- /dev/null +++ b/kernel/include/gemmini_mmio.h @@ -0,0 +1,162 @@ +#ifndef GEMMINI_MMIO_H +#define GEMMINI_MMIO_H +#ifndef GEMMINI_PARAMS_H +#error INCLUDE GEMMINI.H FIRST +#endif + +#define SMEM_BASE 0xff000000 +#define SMEM_SIZE 0x4000 +#define SMEM_MASK (SMEM_SIZE - 1) +#define SMEM_ADDR_END 0xff008000 + +#define SPAD_BASE 0x0 +#define SPAD_ROW_SIZE (DIM * sizeof(elem_t)) +#define SPAD_NUM_ROWS (SMEM_SIZE / SPAD_ROW_SIZE) +#define SPAD_MASK (SPAD_NUM_ROWS - 1) + +#define PRINT_BUF ((char *) (SMEM_ADDR_END)) +#define GEMMINI_RS1_ADDR 0xff007010 +#define GEMMINI_RS2_ADDR 0xff007018 +#define GEMMINI_INST_ADDR 0xff007000 +#define GEMMINI_BUSY_ADDR 0xff007020 + +#define SMEM_TO_SPAD(smem_addr) (SPAD_BASE + ((smem_addr) & SMEM_MASK) / SPAD_ROW_SIZE) +#define SPAD_TO_SMEM(spad_addr) (SMEM_BASE + ((spad_addr) & SPAD_MASK) * SPAD_ROW_SIZE) + +// convert normal matrix i,j into tiled smem offset +// top_in_tiles = i / DIM +// left_in_tiles = j / DIM +// num_tiles_before_current = top_in_tiles * (J / DIM) + left_in_tiles +// smem_addr = num_tiles_before_current * DIM * DIM + (i % DIM) * DIM + (j % DIM) +#define SMEM_MAT_OFFSET(i, j, J) \ + (((i) / DIM * (J) / DIM + (j) / DIM) * DIM * DIM + ((i) % DIM) * DIM + ((j) % DIM)) + +// #define fence() { for (int i = 0; i < 10; i++) *((volatile uint32_t *) (0xFFFF0000)) = 0xdeadbeef; } +#undef gemmini_fence +#define gemmini_fence() { while (*((volatile uint32_t *) GEMMINI_BUSY_ADDR)) asm volatile ("nop"); } + +#undef ROCC_INSTRUCTION_RS1_RS2 +#define ROCC_INSTRUCTION_RS1_RS2(x, rs1, rs2, funct) { \ + /* printf("function %d\n", funct); */ \ + uint32_t instruction = (0x7B) | (0 << 7) | (3 << 12) | (1 << 15) | (2 << 20) | ((uint32_t) (funct) << 25); \ + *((volatile uint64_t *) GEMMINI_RS1_ADDR) = (volatile uint64_t) (rs1); \ + *((volatile uint64_t *) GEMMINI_RS2_ADDR) = (volatile uint64_t) (rs2); \ + /* *((volatile uint32_t*) GEMMINI_RS2_ADDR) = (uint32_t) ((uint64_t) (rs2) & 0xFFFFFFFFULL); */ \ + /* *((volatile uint32_t*) (GEMMINI_RS2_ADDR + 4)) = (uint32_t) ((uint64_t) (rs2) >> 32); */ \ + /* gemmini_fence(); */ \ + *((volatile uint32_t*) GEMMINI_INST_ADDR) = instruction; \ + /* sprintf((char *) PRINT_BUF, "%llx %llx %d\n", rs1, rs2, funct); */ \ +} + +static void sp_tiled_matmul_full_spad_ws(const uint32_t A_sp_addr_start, const uint32_t B_sp_addr_start, + const uint32_t D_sp_addr_start, const uint32_t C_dst_sp_addr_start, + size_t I, size_t J, size_t K, size_t pad_I, size_t pad_J, size_t pad_K, + bool a_transpose, bool b_transpose, + bool full_C, bool low_D, + bool no_bias, bool repeating_bias, + int act) { + + gemmini_loop_ws_spad(I, J, K, pad_I, pad_J, pad_K, + A_sp_addr_start, B_sp_addr_start + K * J * DIM, NULL, C_dst_sp_addr_start, + a_transpose, b_transpose, + full_C, low_D, false, + act, 0, 0, false); + /* + return; + + + // const uint32_t A_sp_addr_start = 0; + // const uint32_t B_sp_addr_start = BANK_NUM * BANK_ROWS - K * J * DIM; + // const uint32_t D_sp_addr_start = 1 << (ADDR_LEN-1); + const uint32_t C_sp_addr_start = 2 << (ADDR_LEN-2) | (full_C << (ADDR_LEN-3)); + // const int D_blocks = low_D ? (J <= MAX_BLOCK_LEN ? J : MAX_BLOCK_LEN) : + // (J <= MAX_BLOCK_LEN_ACC ? J : MAX_BLOCK_LEN_ACC); + const int C_blocks = 1; //full_C ? 1 : (J <= MAX_BLOCK_LEN ? J : MAX_BLOCK_LEN); + // const size_t sizeof_D = low_D ? sizeof(elem_t) : sizeof(acc_t); + const size_t sizeof_C = full_C ? sizeof(acc_t) : sizeof(elem_t); + gemmini_fence(); + + if (a_transpose || b_transpose || (I < 4)) { + for (size_t k = 0; k < K; k++) { + for (size_t j = 0; j < J; j++) { + for (size_t i = 0; i < I; i++) { + const uint32_t A_sp_addr = a_transpose ? (A_sp_addr_start + (k*I + i)*DIM) : + (A_sp_addr_start + (i*K + k)*DIM); + const uint32_t B_sp_addr = b_transpose ? (B_sp_addr_start + (j*K + k)*DIM) : + (B_sp_addr_start + (k*J + j)*DIM); + const uint32_t C_sp_addr = C_sp_addr_start + (i*J + j)*DIM; + // Compute + uint32_t pre_sp_addr = i == 0 ? B_sp_addr : GARBAGE_ADDR; + uint32_t out_sp_addr = C_sp_addr | ((k == 0 ? 0 : 1) << (ADDR_LEN-2)); + gemmini_extended_preload(pre_sp_addr, out_sp_addr, DIM, DIM, DIM, DIM); + if (i == 0) { // First iteration + gemmini_extended_compute_preloaded(A_sp_addr, GARBAGE_ADDR, DIM, DIM, DIM, DIM); + } else { // All other iterations + gemmini_extended_compute_accumulated(A_sp_addr, GARBAGE_ADDR, DIM, DIM, DIM, DIM); + } + if (k == K - 1) { + // Move-out C (if not normalizing) + // if (((act != LAYERNORM) && (act != SOFTMAX)) && (j == J-1 || j % C_blocks == C_blocks-1)) { + const size_t rounded_j = j; // (j / C_blocks) * C_blocks; + const uint32_t rounded_C_sp_addr = C_sp_addr; // C_sp_addr_start + (i*J + rounded_j)*DIM; + + const uint32_t C_dst_sp_addr = ((uint32_t) C_dst_sp_addr_start) + (i * J + rounded_j) * DIM; // * DIM * sizeof_C; + + // const size_t blocks = rounded_j + C_blocks <= J ? C_blocks : J-rounded_j; + constexpr size_t cols = DIM; // blocks * DIM - (rounded_j + blocks >= J ? pad_J : 0); + constexpr size_t rows = DIM; // DIM - (i == I - 1 ? pad_I : 0); + + gemmini_extended_mvout_spad(C_dst_sp_addr, 1, rounded_C_sp_addr, cols, rows); + // } + } + } + } + } + } else { + for (size_t k = 0; k < K; k++) { + for (size_t j = 0; j < J; j++) { + uint32_t A_sp_addr = A_sp_addr_start + k * DIM; // (i*K + k)*DIM; + const uint32_t B_sp_addr = B_sp_addr_start + (k*J + j)*DIM; + uint32_t C_sp_addr = C_sp_addr_start + j * DIM; // (i*J + j)*DIM; + for (size_t i = 0; i < I; i += 4) { + // Compute + // constexpr uint32_t pre_sp_addr = i == 0 ? B_sp_addr : GARBAGE_ADDR; + const uint32_t out_sp_addr = C_sp_addr | ((k == 0 ? 0 : 1) << (ADDR_LEN-2)); + if (i == 0) { // First iteration + gemmini_extended_preload(B_sp_addr, out_sp_addr, DIM, DIM, DIM, DIM); + gemmini_extended_compute_preloaded(A_sp_addr, GARBAGE_ADDR, DIM, DIM, DIM, DIM); + gemmini_extended_preload(GARBAGE_ADDR, out_sp_addr + J * DIM, DIM, DIM, DIM, DIM); + gemmini_extended_compute_accumulated(A_sp_addr + K * DIM, GARBAGE_ADDR, DIM, DIM, DIM, DIM); + gemmini_extended_preload(GARBAGE_ADDR, out_sp_addr + 2 * J * DIM, DIM, DIM, DIM, DIM); + gemmini_extended_compute_accumulated(A_sp_addr + 2 * K * DIM, GARBAGE_ADDR, DIM, DIM, DIM, DIM); + gemmini_extended_preload(GARBAGE_ADDR, out_sp_addr + 3 * J * DIM, DIM, DIM, DIM, DIM); + gemmini_extended_compute_accumulated(A_sp_addr + 3 * K * DIM, GARBAGE_ADDR, DIM, DIM, DIM, DIM); + } else { // All other iterations + gemmini_extended_preload(GARBAGE_ADDR, out_sp_addr, DIM, DIM, DIM, DIM); + gemmini_extended_compute_accumulated(A_sp_addr, GARBAGE_ADDR, DIM, DIM, DIM, DIM); + gemmini_extended_preload(GARBAGE_ADDR, out_sp_addr + J * DIM, DIM, DIM, DIM, DIM); + gemmini_extended_compute_accumulated(A_sp_addr + K * DIM, GARBAGE_ADDR, DIM, DIM, DIM, DIM); + gemmini_extended_preload(GARBAGE_ADDR, out_sp_addr + 2 * J * DIM, DIM, DIM, DIM, DIM); + gemmini_extended_compute_accumulated(A_sp_addr + 2 * K * DIM, GARBAGE_ADDR, DIM, DIM, DIM, DIM); + gemmini_extended_preload(GARBAGE_ADDR, out_sp_addr + 3 * J * DIM, DIM, DIM, DIM, DIM); + gemmini_extended_compute_accumulated(A_sp_addr + 3 * K * DIM, GARBAGE_ADDR, DIM, DIM, DIM, DIM); + } + if (k == K - 1) { + for (int x = 0; x < 3; x++) gemmini_fence(); + gemmini_extended_mvout_spad((uint32_t) C_dst_sp_addr_start + (i * J + j) * DIM, 1, C_sp_addr, DIM, DIM); + gemmini_extended_mvout_spad((uint32_t) C_dst_sp_addr_start + ((i + 1) * J + j) * DIM, 1, C_sp_addr + J * DIM, DIM, DIM); + gemmini_extended_mvout_spad((uint32_t) C_dst_sp_addr_start + ((i + 2) * J + j) * DIM, 1, C_sp_addr + 2 * J * DIM, DIM, DIM); + gemmini_extended_mvout_spad((uint32_t) C_dst_sp_addr_start + ((i + 3) * J + j) * DIM, 1, C_sp_addr + 3 * J * DIM, DIM, DIM); + } + A_sp_addr += 4 * K * DIM; + C_sp_addr += 4 * J * DIM; + } + } + } + } + gemmini_fence(); + */ +} + + +#endif diff --git a/tests/regression/sgemm_gemmini/kernel.cpp b/tests/regression/sgemm_gemmini/kernel.cpp index 34c72d00..dfe15327 100644 --- a/tests/regression/sgemm_gemmini/kernel.cpp +++ b/tests/regression/sgemm_gemmini/kernel.cpp @@ -24,16 +24,20 @@ #define THREAD_ELEMS 8 // elements per thread in a tile #define THREAD_STRIDE 8 // threads per core -#define SMEM_ADDR_0K ((float *) 0xff000000) -#define SMEM_ADDR_4K ((float *) 0xff001000) -#define SMEM_ADDR_8K ((float *) 0xff002000) -#define SMEM_ADDR_12K ((float *) 0xff003000) +#define SMEM_ADDR_0K ((float * const) 0xff000000) +#define SMEM_ADDR_4K ((float * const) 0xff001000) +#define SMEM_ADDR_8K ((float * const) 0xff002000) +#define SMEM_ADDR_12K ((float * const) 0xff003000) #define SPAD_ADDR_0K 0x0 #define SPAD_ADDR_4K 0x80 #define SPAD_ADDR_8K 0x100 #define SPAD_ADDR_12K 0x180 +#define HARDCODE +#define PRINTF(...) sprintf(PRINT_BUF, __VA_ARGS__) +//#define PRINTF(...) vx_printf(__VA_ARGS__) + // #define DEBUG_PRINT #define rd_cycles(x) asm volatile ("csrr %0, mcycle" : "=r" (x)) @@ -55,13 +59,27 @@ void thread_block_matmul_gemmini(kernel_arg_t *__UNIFORM__ arg, const uint32_t num_tiles_n = dim_n / TILE_N; const uint32_t num_tiles_k = dim_k / TILE_K; // TODO: make this into constexpr by subbing architectural params with macros - const uint32_t num_threads_in_cluster = vx_num_threads() * vx_num_warps() * CORES_PER_CLUSTER; + // const uint32_t num_threads_in_cluster = vx_num_threads() * vx_num_warps() * CORES_PER_CLUSTER; + constexpr uint32_t num_threads_in_cluster = 128; + constexpr uint32_t a_elems_per_thread = TILE_MK / num_threads_in_cluster; + constexpr uint32_t b_elems_per_thread = TILE_NK / num_threads_in_cluster; + constexpr uint32_t c_elems_per_thread = TILE_MN / num_threads_in_cluster; const uint32_t hw_tid = tid_in_threadblock % num_threads_in_cluster; - const uint32_t a_elems_per_thread = TILE_MK / num_threads_in_cluster; - const uint32_t b_elems_per_thread = TILE_NK / num_threads_in_cluster; - const uint32_t c_elems_per_thread = TILE_MN / num_threads_in_cluster; const uint32_t thread_load_offset = hw_tid; - const uint32_t thread_load_stride = num_threads_in_cluster; + constexpr uint32_t thread_load_stride = num_threads_in_cluster; + + // the dram coordinates are (i1 + i0, j1 + j0). i0 and j0 are both spatially mapped only. + const uint32_t j0 = hw_tid % DIM; + const uint32_t i0 = (hw_tid / DIM) % DIM; + + // j1 is both spatially and temporally mapped. j1 increases every iteration. + const uint32_t j1_idx = (hw_tid / DIM / DIM) * DIM; // A: % TILE_K, B: % TILE_N, C: % TILE_N + // every iteratioon, j1 increases by j1_stride + constexpr uint32_t j1_stride = (num_threads_in_cluster / DIM / DIM) * DIM; // mod TILE_W after stride + + // i1 is only temporally mapped. i1 increments every one or more iterations + constexpr uint32_t i1_stride = DIM; // step per increment (increment doesnt happen every iteration) + constexpr uint32_t i1_iters = (DIM * DIM * (TILE_K / DIM)) / num_threads_in_cluster; // num of iters before striding uint32_t marker0, marker1, marker2, marker3, marker4; uint32_t marker5, marker6, marker7, marker8, marker9; @@ -70,10 +88,9 @@ void thread_block_matmul_gemmini(kernel_arg_t *__UNIFORM__ arg, gemmini_config_ld(0); gemmini_extended_config_ex(WEIGHT_STATIONARY, 0, 0, 1, 0, 0); gemmini_config_st(0); - sprintf(PRINT_BUF, "start\n"); + PRINTF("start\n"); } - // TODO: check for tb id rd_cycles(marker0); @@ -82,6 +99,7 @@ void thread_block_matmul_gemmini(kernel_arg_t *__UNIFORM__ arg, tile_i += 1) { for (int tile_j = 0; tile_j < num_tiles_n; tile_j += 1) { float * const smem_c_tile_start = SMEM_ADDR_4K; + float * const smem_acc_tile_start = SMEM_ADDR_8K; float * const dram_c_tile_start = C + tile_i * TILE_M * dim_n + tile_j * TILE_N; for (int tile_k = 0; tile_k < num_tiles_k; tile_k += 1) { @@ -93,57 +111,153 @@ void thread_block_matmul_gemmini(kernel_arg_t *__UNIFORM__ arg, rd_cycles(marker1); + #ifdef HARDCODE + #if (TILE_MK / NUM_THREADS / NUM_WARPS / CORES_PER_CLUSTER) != 8 + #error CANNOT UNROLL + #endif // preload A matrix -#pragma GCC unroll 8 // TODO: macro computed - for (int thread_i = 0; thread_i < a_elems_per_thread; thread_i++) { - uint32_t elem_offset = thread_load_offset + thread_load_stride * thread_i; - smem_a_tile_start[SMEM_MAT_OFFSET(elem_offset / TILE_K, elem_offset % TILE_K, TILE_K)] = \ - dram_a_tile_start[elem_offset / TILE_K * dim_k + elem_offset % TILE_K]; - } + { + constexpr uint32_t every_iter = j1_stride; + const uint32_t every_2iters = i1_stride * dim_k; + const uint32_t runtime_const = i0 * dim_k + j1_idx + j0; + smem_a_tile_start[0 * num_threads_in_cluster + hw_tid] = \ + dram_a_tile_start[runtime_const + every_iter * 0 + every_2iters * 0]; + smem_a_tile_start[1 * num_threads_in_cluster + hw_tid] = \ + dram_a_tile_start[runtime_const + every_iter * 1 + every_2iters * 0]; + smem_a_tile_start[2 * num_threads_in_cluster + hw_tid] = \ + dram_a_tile_start[runtime_const + every_iter * 0 + every_2iters * 1]; + smem_a_tile_start[3 * num_threads_in_cluster + hw_tid] = \ + dram_a_tile_start[runtime_const + every_iter * 1 + every_2iters * 1]; + smem_a_tile_start[4 * num_threads_in_cluster + hw_tid] = \ + dram_a_tile_start[runtime_const + every_iter * 0 + every_2iters * 2]; + smem_a_tile_start[5 * num_threads_in_cluster + hw_tid] = \ + dram_a_tile_start[runtime_const + every_iter * 1 + every_2iters * 2]; + smem_a_tile_start[6 * num_threads_in_cluster + hw_tid] = \ + dram_a_tile_start[runtime_const + every_iter * 0 + every_2iters * 3]; + smem_a_tile_start[7 * num_threads_in_cluster + hw_tid] = \ + dram_a_tile_start[runtime_const + every_iter * 1 + every_2iters * 3]; + /* const float v0 = dram_a_tile_start[runtime_const + every_iter * 0 + every_2iters * 0]; + const float v1 = dram_a_tile_start[runtime_const + every_iter * 1 + every_2iters * 0]; + const float v2 = dram_a_tile_start[runtime_const + every_iter * 0 + every_2iters * 1]; + const float v3 = dram_a_tile_start[runtime_const + every_iter * 1 + every_2iters * 1]; + const float v4 = dram_a_tile_start[runtime_const + every_iter * 0 + every_2iters * 2]; + const float v5 = dram_a_tile_start[runtime_const + every_iter * 1 + every_2iters * 2]; + const float v6 = dram_a_tile_start[runtime_const + every_iter * 0 + every_2iters * 3]; + const float v7 = dram_a_tile_start[runtime_const + every_iter * 1 + every_2iters * 3]; -#ifdef DEBUG_PRINT + smem_a_tile_start[0 * num_threads_in_cluster + hw_tid] = v0; + smem_a_tile_start[1 * num_threads_in_cluster + hw_tid] = v1; + smem_a_tile_start[2 * num_threads_in_cluster + hw_tid] = v2; + smem_a_tile_start[3 * num_threads_in_cluster + hw_tid] = v3; + smem_a_tile_start[4 * num_threads_in_cluster + hw_tid] = v4; + smem_a_tile_start[5 * num_threads_in_cluster + hw_tid] = v5; + smem_a_tile_start[6 * num_threads_in_cluster + hw_tid] = v6; + smem_a_tile_start[7 * num_threads_in_cluster + hw_tid] = v7; */ + } + #else + #pragma GCC unroll 8 // TODO: macro computed + for (uint32_t thread_i = 0, j1 = 0, i1 = 0; + thread_i < a_elems_per_thread; + thread_i += 1, + j1 = (j1 + j1_stride) % TILE_K, + i1 = (thread_i % i1_iters == 0) ? i1 + i1_stride : i1) { + smem_a_tile_start[thread_i * num_threads_in_cluster + hw_tid] = \ + dram_a_tile_start[(0 + i0) * dim_k + j1 + j1_idx + j0]; + } + // for (int thread_i = 0; thread_i < a_elems_per_thread; thread_i++) { + // uint32_t elem_offset = thread_load_offset + thread_load_stride * thread_i; + // smem_a_tile_start[SMEM_MAT_OFFSET(elem_offset / TILE_K, elem_offset % TILE_K, TILE_K)] = \ + // dram_a_tile_start[elem_offset / TILE_K * dim_k + elem_offset % TILE_K]; + // } + #endif + + #ifdef DEBUG_PRINT if (hw_tid == 0) { - sprintf(PRINT_BUF, "\nA %d %d\n", tile_i, tile_k); + PRINTF("\nA %d %d\n", tile_i, tile_k); for (int i = 0; i < TILE_M; i += 8) { for (int j = 0; j < TILE_K; j += 8) { uint32_t mat_offset = SMEM_MAT_OFFSET(i, j, TILE_K); - sprintf(PRINT_BUF, "%x %x ", + PRINTF("%x %x ", (int) (smem_a_tile_start[mat_offset]), (int) (smem_a_tile_start[mat_offset + 4]) ); } - sprintf(PRINT_BUF, "\n"); + PRINTF("\n"); } } -#endif + #endif + + threadblock_barrier(0, /*barrier_id=*/threadblock_id, /*count=*/NUM_WARPS); // preload B matrix -#pragma GCC unroll 8 + #ifdef HARDCODE + #if (TILE_NK / NUM_THREADS / NUM_WARPS / CORES_PER_CLUSTER) != 8 + #error CANNOT UNROLL + #endif + constexpr uint32_t every_iter = j1_stride; + const uint32_t every_2iters = i1_stride * dim_n; + const uint32_t runtime_const = i0 * dim_n + j1_idx + j0; + smem_b_tile_start[0 * num_threads_in_cluster + hw_tid] = \ + dram_b_tile_start[runtime_const + every_iter * 0 + every_2iters * 0]; + smem_b_tile_start[1 * num_threads_in_cluster + hw_tid] = \ + dram_b_tile_start[runtime_const + every_iter * 1 + every_2iters * 0]; + smem_b_tile_start[2 * num_threads_in_cluster + hw_tid] = \ + dram_b_tile_start[runtime_const + every_iter * 0 + every_2iters * 1]; + smem_b_tile_start[3 * num_threads_in_cluster + hw_tid] = \ + dram_b_tile_start[runtime_const + every_iter * 1 + every_2iters * 1]; + smem_b_tile_start[4 * num_threads_in_cluster + hw_tid] = \ + dram_b_tile_start[runtime_const + every_iter * 0 + every_2iters * 2]; + smem_b_tile_start[5 * num_threads_in_cluster + hw_tid] = \ + dram_b_tile_start[runtime_const + every_iter * 1 + every_2iters * 2]; + smem_b_tile_start[6 * num_threads_in_cluster + hw_tid] = \ + dram_b_tile_start[runtime_const + every_iter * 0 + every_2iters * 3]; + smem_b_tile_start[7 * num_threads_in_cluster + hw_tid] = \ + dram_b_tile_start[runtime_const + every_iter * 1 + every_2iters * 3]; + /* const float v0 = dram_a_tile_start[runtime_const + every_iter * 0 + every_2iters * 0]; + const float v1 = dram_a_tile_start[runtime_const + every_iter * 1 + every_2iters * 0]; + const float v2 = dram_a_tile_start[runtime_const + every_iter * 0 + every_2iters * 1]; + const float v3 = dram_a_tile_start[runtime_const + every_iter * 1 + every_2iters * 1]; + const float v4 = dram_a_tile_start[runtime_const + every_iter * 0 + every_2iters * 2]; + const float v5 = dram_a_tile_start[runtime_const + every_iter * 1 + every_2iters * 2]; + const float v6 = dram_a_tile_start[runtime_const + every_iter * 0 + every_2iters * 3]; + const float v7 = dram_a_tile_start[runtime_const + every_iter * 1 + every_2iters * 3]; + + smem_a_tile_start[0 * num_threads_in_cluster + hw_tid] = v0; + smem_a_tile_start[1 * num_threads_in_cluster + hw_tid] = v1; + smem_a_tile_start[2 * num_threads_in_cluster + hw_tid] = v2; + smem_a_tile_start[3 * num_threads_in_cluster + hw_tid] = v3; + smem_a_tile_start[4 * num_threads_in_cluster + hw_tid] = v4; + smem_a_tile_start[5 * num_threads_in_cluster + hw_tid] = v5; + smem_a_tile_start[6 * num_threads_in_cluster + hw_tid] = v6; + smem_a_tile_start[7 * num_threads_in_cluster + hw_tid] = v7; */ + #else + #pragma GCC unroll 8 for (int thread_i = 0; thread_i < b_elems_per_thread; thread_i++) { uint32_t elem_offset = thread_load_offset + thread_load_stride * thread_i; smem_b_tile_start[SMEM_MAT_OFFSET(elem_offset / TILE_N, elem_offset % TILE_N, TILE_N)] = \ dram_b_tile_start[elem_offset / TILE_N * dim_n + elem_offset % TILE_N]; } + #endif -#ifdef DEBUG_PRINT + #ifdef DEBUG_PRINT if (hw_tid == 0) { - sprintf(PRINT_BUF, "\nB %d %d\n", tile_k, tile_j); + PRINTF("\nB %d %d\n", tile_k, tile_j); for (int i = 0; i < TILE_K; i += 8) { for (int j = 0; j < TILE_N; j += 8) { uint32_t mat_offset = SMEM_MAT_OFFSET(i, j, TILE_N); - sprintf(PRINT_BUF, "%x %x ", + PRINTF("%x %x ", (int) (smem_b_tile_start[mat_offset]), (int) (smem_b_tile_start[mat_offset + 4]) ); } - sprintf(PRINT_BUF, "\n"); + PRINTF("\n"); } } -#endif - rd_cycles(marker2); + #endif + rd_cycles(marker2); // cluster wide barrier to wait for A and B loads to complete - threadblock_barrier(0, /*barrier_id=*/threadblock_id, /*count=*/num_threads_in_cluster); + threadblock_barrier(0, /*barrier_id=*/threadblock_id, /*count=*/NUM_WARPS); rd_cycles(marker3); if (hw_tid == 0) { sp_tiled_matmul_full_spad_ws(SPAD_ADDR_0K, SPAD_ADDR_12K, /*spad_D=*/0, SPAD_ADDR_4K, @@ -153,57 +267,92 @@ void thread_block_matmul_gemmini(kernel_arg_t *__UNIFORM__ arg, gemmini_fence(); } rd_cycles(marker4); - threadblock_barrier(0, /*barrier_id=*/threadblock_id, /*count=*/num_threads_in_cluster); + threadblock_barrier(0, /*barrier_id=*/threadblock_id, /*count=*/NUM_WARPS); rd_cycles(marker5); // accumulate C matrix if (tile_k == 0) { -#pragma GCC unroll 8 + #pragma GCC ivdep + #pragma GCC unroll 8 for (int thread_i = 0; thread_i < c_elems_per_thread; thread_i++) { uint32_t elem_offset = thread_load_offset + thread_load_stride * thread_i; - *(SMEM_ADDR_8K + elem_offset) = smem_c_tile_start[elem_offset]; + smem_acc_tile_start[elem_offset] = smem_c_tile_start[elem_offset]; } } else { -#pragma GCC unroll 8 - for (int thread_i = 0; thread_i < c_elems_per_thread; thread_i++) { - uint32_t elem_offset = thread_load_offset + thread_load_stride * thread_i; - *(SMEM_ADDR_8K + elem_offset) += smem_c_tile_start[elem_offset]; + #if (TILE_NK / NUM_THREADS / NUM_WARPS / CORES_PER_CLUSTER) != 8 + #error CANNOT UNROLL + #endif + for (int thread_i = 0; thread_i < c_elems_per_thread; thread_i += 8) { + constexpr uint32_t s = num_threads_in_cluster; + smem_acc_tile_start[hw_tid + s * 0] += smem_c_tile_start[hw_tid + s * 0]; + smem_acc_tile_start[hw_tid + s * 1] += smem_c_tile_start[hw_tid + s * 1]; + smem_acc_tile_start[hw_tid + s * 2] += smem_c_tile_start[hw_tid + s * 2]; + smem_acc_tile_start[hw_tid + s * 3] += smem_c_tile_start[hw_tid + s * 3]; + smem_acc_tile_start[hw_tid + s * 4] += smem_c_tile_start[hw_tid + s * 4]; + smem_acc_tile_start[hw_tid + s * 5] += smem_c_tile_start[hw_tid + s * 5]; + smem_acc_tile_start[hw_tid + s * 6] += smem_c_tile_start[hw_tid + s * 6]; + smem_acc_tile_start[hw_tid + s * 7] += smem_c_tile_start[hw_tid + s * 7]; } } rd_cycles(marker6); -#ifdef DEBUG_PRINT + #ifdef DEBUG_PRINT if (hw_tid == 0) { - sprintf(PRINT_BUF, "\nC %d %d %d\n", tile_i, tile_j, tile_k); + PRINTF("\nC %d %d %d\n", tile_i, tile_j, tile_k); for (int i = 0; i < TILE_M; i += 8) { for (int j = 0; j < TILE_N; j += 8) { uint32_t mat_offset = SMEM_MAT_OFFSET(i, j, TILE_N); - sprintf(PRINT_BUF, "%d %d ", + PRINTF("%d %d ", (int) (smem_c_tile_start[mat_offset]), (int) (smem_c_tile_start[mat_offset + 4]) ); } - sprintf(PRINT_BUF, "\n"); + PRINTF("\n"); } } -#endif + #endif } rd_cycles(marker7); // move out to dram - #pragma GCC unroll 8 // TODO: macro computed + + #ifdef HARDCODE + #if (TILE_MN / NUM_THREADS / NUM_WARPS / CORES_PER_CLUSTER) != 8 + #error CANNOT UNROLL + #endif + constexpr uint32_t every_iter = j1_stride; + const uint32_t every_2iters = i1_stride * dim_n; + const uint32_t runtime_const = i0 * dim_n + j1_idx + j0; + dram_c_tile_start[runtime_const + every_iter * 0 + every_2iters * 0] = \ + smem_acc_tile_start[0 * num_threads_in_cluster + hw_tid]; + dram_c_tile_start[runtime_const + every_iter * 1 + every_2iters * 0] = \ + smem_acc_tile_start[1 * num_threads_in_cluster + hw_tid]; + dram_c_tile_start[runtime_const + every_iter * 0 + every_2iters * 1] = \ + smem_acc_tile_start[2 * num_threads_in_cluster + hw_tid]; + dram_c_tile_start[runtime_const + every_iter * 1 + every_2iters * 1] = \ + smem_acc_tile_start[3 * num_threads_in_cluster + hw_tid]; + dram_c_tile_start[runtime_const + every_iter * 0 + every_2iters * 2] = \ + smem_acc_tile_start[4 * num_threads_in_cluster + hw_tid]; + dram_c_tile_start[runtime_const + every_iter * 1 + every_2iters * 2] = \ + smem_acc_tile_start[5 * num_threads_in_cluster + hw_tid]; + dram_c_tile_start[runtime_const + every_iter * 0 + every_2iters * 3] = \ + smem_acc_tile_start[6 * num_threads_in_cluster + hw_tid]; + dram_c_tile_start[runtime_const + every_iter * 1 + every_2iters * 3] = \ + smem_acc_tile_start[7 * num_threads_in_cluster + hw_tid]; + #else + #pragma GCC unroll 8 for (int thread_i = 0; thread_i < c_elems_per_thread; thread_i++) { uint32_t elem_offset = thread_load_offset + thread_load_stride * thread_i; dram_c_tile_start[elem_offset / TILE_N * dim_n + elem_offset % TILE_N] = \ *(SMEM_ADDR_8K + SMEM_MAT_OFFSET(elem_offset / TILE_N, elem_offset % TILE_N, TILE_N)); } + #endif rd_cycles(marker8); /* if (hw_tid == 0) { sprintf(PRINT_BUF, "\nC %d %d\n", tile_i, tile_j); for (int i = 0; i < TILE_M; i += 8) { for (int j = 0; j < TILE_N; j += 8) { - uint32_t mat_offset = SMEM_MAT_OFFSET(i, j, TILE_N); sprintf(PRINT_BUF, "%d %d ", (int) (C[(tile_i * TILE_M + i) * dim_n + tile_j * TILE_N + j]), (int) (C[(tile_i * TILE_M + i) * dim_n + tile_j * TILE_N + j + 4]) @@ -216,26 +365,42 @@ void thread_block_matmul_gemmini(kernel_arg_t *__UNIFORM__ arg, } // last thread block complete if (threadblock_id == NUM_CLUSTERS - 1) { - threadblock_barrier(0, /*barrier_id=*/0, /*count=*/num_threads_in_cluster); + threadblock_barrier(0, /*barrier_id=*/0, /*count=*/NUM_WARPS); rd_cycles(marker9); if (hw_tid == 0) { - sprintf(PRINT_BUF, "complete\n"); - sprintf(PRINT_BUF, "total cycles: %d\n", marker9 - marker0); - sprintf(PRINT_BUF, "single tile cycles: %d\n", marker6 - marker1); - sprintf(PRINT_BUF, "A/B tile load cycles: %d\n", marker2 - marker1); - sprintf(PRINT_BUF, "gemmini cycles: %d\n", marker4 - marker3); - sprintf(PRINT_BUF, "first barrier: %d\n", marker3 - marker2); - sprintf(PRINT_BUF, "second barrier: %d\n", marker5 - marker4); - sprintf(PRINT_BUF, "accumulation cycles: %d\n", marker6 - marker5); - sprintf(PRINT_BUF, "dram mvout cycles: %d\n", marker8 - marker7); + PRINTF("\ncomplete\n"); + PRINTF("total cycles: %d\n", marker9 - marker0); + PRINTF("tile start: %d\n", marker1); + PRINTF("single tile cycles: %d\n", marker6 - marker1); + PRINTF("A/B tile load cycles: %d\n", marker2 - marker1); + PRINTF("first barrier: %d\n", marker3 - marker2); + PRINTF("gemmini cycles: %d\n", marker4 - marker3); + PRINTF("second barrier: %d\n", marker5 - marker4); + PRINTF("accumulation cycles: %d\n", marker6 - marker5); + PRINTF("dram mvout cycles: %d\n", marker8 - marker7); } - threadblock_barrier(0, /*barrier_id=*/0, /*count=*/num_threads_in_cluster); + threadblock_barrier(0, /*barrier_id=*/1, /*count=*/NUM_WARPS); if (hw_tid == num_threads_in_cluster - 1) { - sprintf(PRINT_BUF, "single tile cycles: %d\n", marker6 - marker1); - sprintf(PRINT_BUF, "A/B tile load cycles: %d\n", marker2 - marker1); - sprintf(PRINT_BUF, "gemmini cycles: %d\n", marker4 - marker3); - sprintf(PRINT_BUF, "first barrier: %d\n", marker3 - marker2); - sprintf(PRINT_BUF, "second barrier: %d\n", marker5 - marker4); + PRINTF("\ntile start: %d\n", marker1); + PRINTF("single tile cycles: %d\n", marker6 - marker1); + PRINTF("A/B tile load cycles: %d\n", marker2 - marker1); + PRINTF("gemmini cycles: %d\n", marker4 - marker3); + PRINTF("first barrier: %d\n", marker3 - marker2); + PRINTF("second barrier: %d\n", marker5 - marker4); + PRINTF("accumulation cycles: %d\n", marker6 - marker5); + PRINTF("dram mvout cycles: %d\n", marker8 - marker7); + } + threadblock_barrier(0, /*barrier_id=*/2, /*count=*/NUM_WARPS); + if (hw_tid == 0) { + for (int i = 0; i < dim_m; i += 8) { + for (int j = 0; j < dim_n; j += 8) { + sprintf(PRINT_BUF, "%d %d ", + (int) (C[i * dim_n + j]), + (int) (C[i * dim_n + j + 4]) + ); + } + PRINTF("\n"); + } } vx_tmc_one(); } @@ -254,7 +419,6 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) { int main() { kernel_arg_t *arg = (kernel_arg_t *)KERNEL_ARG_DEV_MEM_ADDR; - sprintf(PRINT_BUF, "m=%d, n=%d\n", arg->dim_m, arg->dim_n); const uint32_t num_threads_in_cluster = vx_num_threads() * vx_num_warps() * CORES_PER_CLUSTER; const uint32_t grid_size = num_threads_in_cluster * NUM_CLUSTERS;