From c114a7a4abf61202d1fec829b39b23899ba6d2d2 Mon Sep 17 00:00:00 2001 From: Richard Yan Date: Fri, 8 Nov 2024 20:55:27 -0800 Subject: [PATCH] new gemm kernel --- kernel/include/gemmini_mmio.h | 47 +++++- tests/regression/sgemm_gemmini_dma/common.h | 2 + tests/regression/sgemm_gemmini_dma/kernel.cpp | 138 ++++++++---------- 3 files changed, 110 insertions(+), 77 deletions(-) diff --git a/kernel/include/gemmini_mmio.h b/kernel/include/gemmini_mmio.h index ed55236c..894d5fc6 100644 --- a/kernel/include/gemmini_mmio.h +++ b/kernel/include/gemmini_mmio.h @@ -12,7 +12,7 @@ // 128KB // #define SMEM_SIZE 0x20000 // 256KB -#define SMEM_SIZE 0x40000 +#define SMEM_SIZE 0x20000 #define SMEM_MASK (SMEM_SIZE - 1) #define SMEM_ADDR_END (SMEM_BASE + SMEM_SIZE) @@ -85,6 +85,51 @@ static size_t gemmini_tile_idx[NUM_THREADS * NUM_WARPS * NUM_CORES * NUM_CLUSTER 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, acc, act, 0, 0, false, skips) + +#define GEMMINI_CISC_COMPUTE_HEXADECILES 0 +#define GEMMINI_CISC_SET_AB_STRIDE 8 +#define GEMMINI_CISC_STORE_TO_SPAD 9 +#define GEMMINI_CISC_LOAD_TO_HEXADECILES 10 +#define GEMMINI_CISC_SET_DC_STRIDE 11 +#define GEMMINI_CISC_STORE_TO_GMEM 12 + +// cisc instruction wrappers +inline void gemmini_tile_load_ab(const elem_t * const a_addr, const elem_t * const b_addr, + const uint32_t a_hexadecile, const uint32_t b_hexadecile, + const uint32_t tile_idx_i, const uint32_t tile_idx_j, const uint32_t tile_idx_k, + const uint32_t mat_size_m, const uint32_t mat_size_n, const uint32_t mat_size_k, + const uint32_t tile_size_m, const uint32_t tile_size_n, const uint32_t tile_size_k) { + + ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, + (uint64_t) (a_addr + tile_idx_i * tile_size_m * mat_size_k + tile_idx_k * tile_size_k), + (uint64_t) (b_addr + tile_idx_k * tile_size_k * mat_size_n + tile_idx_j * tile_size_n), k_LOOP_WS_CONFIG_ADDRS_AB) + GEMMINI_CISC_CMD_R((mat_size_n << 20) | (mat_size_k << 8) | GEMMINI_CISC_SET_AB_STRIDE); + 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); +} + +inline void gemmini_tile_store_c_gmem(elem_t * const c_addr, + const uint32_t tile_idx_i, const uint32_t tile_idx_j, + const uint32_t mat_size_m, const uint32_t mat_size_n, + const uint32_t tile_size_m, const uint32_t tile_size_n) { + + elem_t * const dram_c_tile_start = c_addr + tile_idx_i * tile_size_m * mat_size_n + tile_idx_j * tile_size_n; + ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, 0, (uint64_t) dram_c_tile_start, k_LOOP_WS_CONFIG_ADDRS_DC) + + GEMMINI_CISC_CMD_R((mat_size_n << 20) | GEMMINI_CISC_SET_DC_STRIDE); + GEMMINI_CISC_CMD_I(GEMMINI_CISC_STORE_TO_GMEM); + + // ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, 0, BOUND_INST, k_LOOP_WS_CONFIG_BOUNDS) + // ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, 0, mat_size_n, k_LOOP_WS_CONFIG_STRIDES_DC) + // ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, 0, loop_matmul_skips(1, 1, 1, 1, 0), k_LOOP_WS) +} + +inline void gemmini_tile_store_c_spad(const uint32_t c_hexadecile) { + GEMMINI_CISC_CMD_R(((uint32_t) (c_hexadecile << 8)) | GEMMINI_CISC_STORE_TO_SPAD); +} /* inline static void sp_tiled_matmul_full_spad_ws(const uint32_t A_sp_addr_start, const uint32_t B_sp_addr_start, const uint32_t D_sp_addr_start, const uint32_t C_dst_sp_addr_start, size_t I, size_t J, size_t K, size_t pad_I, size_t pad_J, size_t pad_K, diff --git a/tests/regression/sgemm_gemmini_dma/common.h b/tests/regression/sgemm_gemmini_dma/common.h index 5c84f3b7..c443474c 100644 --- a/tests/regression/sgemm_gemmini_dma/common.h +++ b/tests/regression/sgemm_gemmini_dma/common.h @@ -5,6 +5,8 @@ #define KERNEL_ARG_DEV_MEM_ADDR 0x9fff0000 #define DEV_SMEM_START_ADDR 0xff000000 +#define MARK_BEG() asm volatile ("slti x0, x1, -1047") +#define MARK_END() asm volatile ("slti x0, x1, -499") typedef struct { uint32_t dim_m; diff --git a/tests/regression/sgemm_gemmini_dma/kernel.cpp b/tests/regression/sgemm_gemmini_dma/kernel.cpp index 2cbd8de5..2e55cb5e 100644 --- a/tests/regression/sgemm_gemmini_dma/kernel.cpp +++ b/tests/regression/sgemm_gemmini_dma/kernel.cpp @@ -10,8 +10,11 @@ #define TILE_M 128 #define TILE_N 64 #define TILE_K 128 -#define BOUND_INST 0x800040008ULL -#define NUM_THREADS_IN_CLUSTER 512 + +// ampere +// #define NUM_THREADS_IN_CLUSTER 512 +// hopper +#define NUM_THREADS_IN_CLUSTER 256 // fp32 8x8 // #define TILE_M 64 @@ -25,7 +28,6 @@ // #define SPAD_ADDR_Q1 0x200 // #define SPAD_ADDR_Q2 0x400 // #define SPAD_ADDR_Q3 0x600 -// #define BOUND_INST 0x800080008ULL // #define NUM_THREADS_IN_CLUSTER 256 // fp32 4x4 @@ -40,7 +42,6 @@ // #define SPAD_ADDR_Q1 0x80 // #define SPAD_ADDR_Q2 0x100 // #define SPAD_ADDR_Q3 0x180 -// #define BOUND_INST 0x400040004ULL // #define NUM_THREADS_IN_CLUSTER 256 #define NUM_CLUSTERS 1 @@ -49,12 +50,10 @@ #define rd_cycles_force(x) asm volatile ("csrr %0, mcycle" : "=r" (x)) #define rd_cycles(x) rd_cycles_force(x) #define HW_TID() ({uint32_t gtid; asm volatile ("csrr %0, mhartid" : "=r" (gtid)); gtid;}) -#define MARK_BEG() asm volatile ("slti x0, x1, -1047") -#define MARK_END() asm volatile ("slti x0, x1, -499") #define PRINTF(...) sprintf(PRINT_BUF, __VA_ARGS__) // #define PRINTF(...) vx_printf(__VA_ARGS__) #define SWISH(beta, x) ((x) / (1 + exp(-(beta) * (x)))) -//#define POWER +// #define POWER typedef uint16_t smem_elem_t; // typedef float smem_elem_t; @@ -81,22 +80,6 @@ void thread_block_matmul_gemmini(kernel_arg_t *__UNIFORM__ arg, } vx_fence(); - // if (HW_TID() < 128) { - // *((volatile uint32_t *) 0xff000000 + HW_TID()) = HW_TID(); - // for (int i = 0; i < 128; i++) { - // if (HW_TID() == i) { - // volatile uint32_t x = *((volatile uint32_t *) 0xff000000 + HW_TID()); - // if (x != i) { - // PRINTF("%d ", x); - // } - // } - // } - // } - // threadblock_barrier(/*barrier_id=*/0, /*count=*/NUM_WARPS); - // if (HW_TID() == 0) { - // PRINTF("\n finished\n"); - // } - // threadblock_barrier(/*barrier_id=*/0, /*count=*/NUM_WARPS); uint32_t marker0, marker1; rd_cycles_force(marker0); @@ -115,83 +98,86 @@ void thread_block_matmul_gemmini(kernel_arg_t *__UNIFORM__ arg, if (HW_TID() == 0) { gemmini_extended3_config_ld(dim_k * sizeof(elem_t), MVIN_SCALE_IDENTITY, false, 0); gemmini_extended3_config_ld(dim_n * sizeof(elem_t), MVIN_SCALE_IDENTITY, false, 1); - // gemmini_extended3_config_ld(repeating_bias ? 0 : (stride_D * sizeof_D), D_scale_factor, low_D, 2); gemmini_extended_config_st(dim_n * sizeof(elem_t), 0, MVIN_SCALE_IDENTITY); // gemmini_extended_config_st(stride_C * sizeof_C, act & 3, scale); for (uint32_t tile_i = num_tile_rows_per_tb * threadblock_id; tile_i < num_tile_rows_per_tb * (threadblock_id + 1); tile_i += 1) { - for (int tile_j = 0; tile_j < num_tiles_n; tile_j += 1) { - for (int tile_k = 0; tile_k < num_tiles_k; tile_k += 1) { - ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, - (uint64_t) (A + tile_i * TILE_M * dim_k + tile_k * TILE_K), - (uint64_t) (B + tile_k * TILE_K * dim_n + tile_j * TILE_N), k_LOOP_WS_CONFIG_ADDRS_AB) - GEMMINI_CISC_CMD_R((dim_n << 20) | (dim_k << 8) | 8); - if (tile_k & 1) { - GEMMINI_CISC_CMD_I(11); - } else { - GEMMINI_CISC_CMD_I(10); - } - - if (tile_k == 0) { - asm volatile("cisc_start_%=:" ::); - gemmini_fence(); - GEMMINI_CISC_CMD_I(0); - asm volatile("cisc_end_%=:" ::); - } else if (tile_k & 1) { - gemmini_fence(); - GEMMINI_CISC_CMD_I(2); - } else { - gemmini_fence(); - GEMMINI_CISC_CMD_I(1); - } + for (uint32_t tile_j = 0; tile_j < num_tiles_n; tile_j += 1) { + for (uint32_t tile_k = 0; tile_k < num_tiles_k; tile_k += 1) { + uint32_t a_hexadecile = (tile_k & 1) << 2; + uint32_t b_hexadecile = a_hexadecile + 8; + gemmini_tile_load_ab(A, B, + a_hexadecile, b_hexadecile, tile_i, tile_j, tile_k, + 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_fence(); + gemmini_tile_store_c_spad(a_hexadecile); // then activate in spad + */ gemmini_fence(); - gemmini_fence(); - gemmini_fence(); - // mvout to scratchpad for activation - // GEMMINI_CISC_CMD_I(9); - // gemmini_fence(); - // } - - // threadblock_barrier(/*barrier_id=*/0, /*count=*/NUM_WARPS); - // // activate - - // // move out to dram - // if (HW_TID() == 0) { - smem_elem_t * const dram_c_tile_start = C + tile_i * TILE_M * dim_n + tile_j * TILE_N; - ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, 0, BOUND_INST, k_LOOP_WS_CONFIG_BOUNDS) - ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, 0, (uint64_t) dram_c_tile_start, k_LOOP_WS_CONFIG_ADDRS_DC) - ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, 0, dim_n, k_LOOP_WS_CONFIG_STRIDES_DC) - ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, 0, loop_matmul_skips(1, 1, 1, 1, 0), k_LOOP_WS) + gemmini_tile_store_c_gmem(C, tile_i, tile_j, dim_m, dim_n, TILE_M, TILE_N); } } - } - // last thread block complete - if (threadblock_id == NUM_CLUSTERS - 1) { - threadblock_barrier(/*barrier_id=*/0, /*count=*/NUM_WARPS); + MARK_END(); rd_cycles_force(marker1); - if (HW_TID() == 0) { - #ifdef POWER - // PRINTF("%d\n", marker1 - marker0); - #else + + #ifndef POWER + if (HW_TID() == 0) { PRINTF("\ncomplete\n"); PRINTF("total cycles: %d\n", marker1 - marker0); for (int i = 0; i < dim_m; i += 8) { for (int j = 0; j < dim_n; j += 8) { - // PRINTF("%d %d ", (int) (C[i * dim_n + j]), (int) (C[i * dim_n + j + 4])); PRINTF("%04x %04x ", (int) (C[i * dim_n + j]), (int) (C[i * dim_n + j + 4])); } PRINTF("\n"); } - #endif + } + #endif + } else { + if (HW_TID() > 8) { + asm volatile("li x1, 0xa0a0a0a0"); + asm volatile("li x2, 0xa0a0a0a0"); + asm volatile("li x3, 0xa0a0a0a0"); + asm volatile("li x4, 0xa0a0a0a0"); + asm volatile("li x5, 0xa0a0a0a0"); + asm volatile("li x6, 0xa0a0a0a0"); + asm volatile("li x7, 0xa0a0a0a0"); + asm volatile("li x8, 0xa0a0a0a0"); + asm volatile("li x9, 0xa0a0a0a0"); + asm volatile("li x10, 0xa0a0a0a0"); + asm volatile("li x11, 0xa0a0a0a0"); + asm volatile("li x12, 0xa0a0a0a0"); + asm volatile("li x13, 0xa0a0a0a0"); + asm volatile("li x14, 0xa0a0a0a0"); + asm volatile("li x15, 0xa0a0a0a0"); + asm volatile("li x16, 0xa0a0a0a0"); + asm volatile("li x17, 0xa0a0a0a0"); + asm volatile("li x18, 0xa0a0a0a0"); + asm volatile("li x19, 0xa0a0a0a0"); + asm volatile("li x20, 0xa0a0a0a0"); + asm volatile("li x21, 0xa0a0a0a0"); + asm volatile("li x22, 0xa0a0a0a0"); + asm volatile("li x23, 0xa0a0a0a0"); + asm volatile("li x24, 0xa0a0a0a0"); + asm volatile("li x25, 0xa0a0a0a0"); + asm volatile("li x26, 0xa0a0a0a0"); + asm volatile("li x27, 0xa0a0a0a0"); + asm volatile("li x28, 0xa0a0a0a0"); + asm volatile("li x29, 0xa0a0a0a0"); + asm volatile("li x30, 0xa0a0a0a0"); + asm volatile("li x31, 0xa0a0a0a0"); + asm volatile("vx_tmc zero"); } } - threadblock_barrier(/*barrier_id=*/0, /*count=*/NUM_WARPS); + vx_fence(); + // threadblock_barrier(/*barrier_id=*/0, /*count=*/NUM_WARPS); vx_tmc(0); }