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

Merge branch 'kernels' into tensor_core

Browse Source
This commit is contained in:
Hansung Kim
2024-05-08 13:25:31 -07:00
parent 5821bfd10d a606a9ef42
commit 6ba6a1e2e5
55 changed files with 3384 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
tests/kernel/gemmini_mmio/Makefile Normal file
View File

@@ -0,0 +1,54 @@
XLEN ?= 32
ifeq ($(XLEN),64)
RISCV_TOOLCHAIN_PATH ?= /opt/riscv64-gnu-toolchain
CFLAGS += -march=rv64imafd -mabi=lp64d
else
RISCV_TOOLCHAIN_PATH ?= /opt/riscv-gnu-toolchain
CFLAGS += -march=rv32imaf -mabi=ilp32f
endif
RISCV_PREFIX ?= riscv$(XLEN)-unknown-elf
VORTEX_KN_PATH ?= $(realpath ../../../kernel)
GEMMINI_SW_PATH ?= $(realpath /scratch/yrh/chipyard/generators/gemmini/software/gemmini-rocc-tests)
CC = $(RISCV_TOOLCHAIN_PATH)/bin/$(RISCV_PREFIX)-gcc
AR = $(RISCV_TOOLCHAIN_PATH)/bin/$(RISCV_PREFIX)-gcc-ar
DP = $(RISCV_TOOLCHAIN_PATH)/bin/$(RISCV_PREFIX)-objdump
CP = $(RISCV_TOOLCHAIN_PATH)/bin/$(RISCV_PREFIX)-objcopy
SIM_DIR = ../../../sim
CFLAGS += -O3 -funroll-loops -v -mcmodel=medany -fno-exceptions -nostartfiles -fdata-sections -ffunction-sections
CFLAGS += -I$(VORTEX_KN_PATH)/include -I$(VORTEX_KN_PATH)/../hw -I$(GEMMINI_SW_PATH)
LDFLAGS += -lm -Wl,-Bstatic,--gc-sections,-T,$(VORTEX_KN_PATH)/linker/vx_link$(XLEN).ld,--defsym=STARTUP_ADDR=0x80000000 $(VORTEX_KN_PATH)/libvortexrt.a
PROJECT = gemmini_mmio
SRCS = main.cpp
all: $(PROJECT).elf $(PROJECT).bin $(PROJECT).dump
$(PROJECT).dump: $(PROJECT).elf
$(DP) -D $(PROJECT).elf > $(PROJECT).dump
$(PROJECT).bin: $(PROJECT).elf
$(CP) -O binary $(PROJECT).elf $(PROJECT).bin
$(PROJECT).elf: $(SRCS)
$(CC) $(CFLAGS) $(SRCS) $(LDFLAGS) -o $(PROJECT).elf
run-rtlsim: $(PROJECT).bin
$(SIM_DIR)/rtlsim/rtlsim $(PROJECT).bin
run-simx: $(PROJECT).bin
$(SIM_DIR)/simx/simx $(PROJECT).bin
.depend: $(SRCS)
$(CC) $(CFLAGS) -MM $^ > .depend;
clean:
rm -rf *.elf *.bin *.dump .depend

162
tests/kernel/gemmini_mmio/gemmini_mmio.h Normal file
View File

@@ -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

144
tests/kernel/gemmini_mmio/main.cpp Normal file
View File

@@ -0,0 +1,144 @@
#include <stdio.h>
#include <stdint.h>
#include <vx_intrinsics.h>
#include <vx_print.h>
#include <vx_spawn.h>
#include "include/gemmini.h"
#include "gemmini_mmio.h"
#define rd_cycles(x) asm volatile ("csrr %0, mcycle" : "=r" (x))
int main() {
int cid;
asm volatile ("csrr %0, 0xcc2" : "=r" (cid));
if (cid > 0) vx_tmc(0);
vx_tmc(0xff);
// load up A and B and C
const uint32_t spad_A = 0x00000000;
const uint32_t spad_B = 0x00000080; // 16B word addressed
const uint32_t acc_C = 0x80000000; // accmem + accumulate
const uint32_t spad_C = 0x00000100;
volatile float *smem_A = (float *) SPAD_TO_SMEM(spad_A); // 0xff000000; // byte addressed
float *smem_B = (float *) SPAD_TO_SMEM(spad_B); // 0xff000200;
float *smem_C = (float *) SPAD_TO_SMEM(spad_C); // 0xff000400;
int I = 32 / DIM;
int J = 32 / DIM;
int K = 32 / DIM;
char *print_buf = (char *) PRINT_BUF;
// int cid = vx_core_id();
int nc = vx_num_cores();
int nt = vx_num_threads();
int tid = vx_thread_id();
vx_tmc_one();
gemmini_config_ld(0);
gemmini_extended_config_ex(WEIGHT_STATIONARY, 0, 0, 1, 0, 0);
gemmini_config_st(0);
/* sprintf(print_buf, "A spad: 0x%x-0x%x, smem: 0x%x-%x\n", spad_A, spad_A + I * K * DIM, (uint32_t) smem_A, (uint32_t) smem_A + sizeof(float) * I * K * DIM * DIM);
sprintf(print_buf, "B spad: 0x%x-0x%x, smem: 0x%x-%x\n", spad_B, spad_B + K * J * DIM, (uint32_t) smem_B, (uint32_t) smem_B + sizeof(float) * K * J * DIM * DIM);
sprintf(print_buf, "C spad: 0x%x-0x%x, smem: 0x%x-%x\n", spad_C, spad_C + I * J * DIM, (uint32_t) smem_C, (uint32_t) smem_C + sizeof(float) * I * J * DIM * DIM); */
sprintf(print_buf, "DIM %d\n", DIM);
sprintf(print_buf, "num cores %d\n", nc);
sprintf(print_buf, "num threads %d\n", nt);
sprintf(print_buf, "thread ids ");
vx_tmc(-1);
sprintf(print_buf, "%d", tid);
uint32_t start_cycles, end_cycles;
rd_cycles(start_cycles);
// load A with 128->1 in row-major order
for (int t = 0; t < DIM * DIM / nt; t++) {
int n = tid + t * nt;
int x = n / DIM;
int y = n % DIM;
for (int k = 0; k < K; k++) {
for (int i = 0; i < I; i++) {
int tile_byte_offset = (i * K + k) * DIM * DIM;
smem_A[tile_byte_offset + n] = (float) ((I * K * DIM * DIM - ((i * DIM + x) * DIM * K + (k * DIM + y))) % 64);
// smem_A[tile_byte_offset + x * DIM + y] = (float) ((I * K * DIM * DIM - ((i * DIM + x) * DIM * K + (k * DIM + y))) % 64);
}
}
}
// load B with 0->191 in row-major order
for (int t = 0; t < DIM * DIM / nt; t++) {
int n = tid + t * nt;
int x = n / DIM;
int y = n % DIM;
for (int k = 0; k < K; k++) {
for (int j = 0; j < J; j++) {
int tile_byte_offset = (k * J + j) * DIM * DIM;
smem_B[tile_byte_offset + n] = (float) (((k * DIM + x) * DIM * J + (j * DIM + y)) % 64);
}
// smem_B[tile_byte_offset + x * DIM + y] = (float) (((k * DIM + x) * DIM * J + (j * DIM + y)) % 64);
}
}
rd_cycles(end_cycles);
// for (int i = 0; i < I * J * DIM * DIM; i++) smem_C[i] = 1.f;
vx_tmc_one();
sprintf(print_buf, "\ndata loading took %d cycles for %d floats\n", end_cycles - start_cycles, DIM * DIM * (I * K + J * K));
gemmini_fence();
// sprintf(print_buf, "\nA in\n");
// for (int i = 0; i < I * DIM; i++) {
// for (int j = 0; j < K * DIM; j++) {
// sprintf(print_buf, "%d ", (int) (smem_A[SMEM_MAT_OFFSET(i, j, K * DIM)]));
// }
// sprintf(print_buf, "\n");
// }
// sprintf(print_buf, "\nB in\n");
// for (int i = 0; i < K * DIM; i++) {
// for (int j = 0; j < J * DIM; j++) {
// sprintf(print_buf, "%d ", (int) (smem_B[SMEM_MAT_OFFSET(i, j, J * DIM)]));
// }
// sprintf(print_buf, "\n");
// if (i == 2) i = K * DIM - 3;
// }
uint32_t fence_cycles;
rd_cycles(start_cycles);
sp_tiled_matmul_full_spad_ws(spad_A, spad_B, /*spad_D=*/0, spad_C,
/*I=*/I, /*J=*/J, /*K=*/K, /*pad_I=*/0, /*pad_J=*/0, /*pad_K=*/0,
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
/*no_bias=*/1, /*repeating_bias=*/0, /*act=*/NO_ACTIVATION);
rd_cycles(fence_cycles);
gemmini_fence();
rd_cycles(end_cycles);
sprintf(print_buf, "gemmini cycles taken: %d, fence cycles: %d\n", end_cycles - start_cycles, end_cycles - fence_cycles);
// check results
for (int i = 0; i < I * DIM; i++) {
for (int j = 0; j < J * DIM; j++) {
int sum = 0;
for (int k = 0; k < K * DIM; k++) sum += ((I * K * DIM * DIM - i * K * DIM - k) % 64) * ((k * J * DIM + j) % 64);
if ((int) (smem_C[SMEM_MAT_OFFSET(i, j, J * DIM)] * 10) != (int) (sum * 10)) {
sprintf(print_buf, "TEST FAILED (actual/reference)\n");
for (int ii = 0; ii < I * DIM; ii++) {
for (int jj = 0; jj < J * DIM; jj++) {
sum = 0;
for (int k = 0; k < K * DIM; k++) sum += ((I * K * DIM * DIM - ii * K * DIM - k) % 64) * ((k * J * DIM + jj) % 64);
sprintf(print_buf, "%d/%d ", (int) (smem_C[SMEM_MAT_OFFSET(ii, jj, J * DIM)]), (int) sum);
}
sprintf(print_buf, "\n");
}
return 1;
}
}
}
sprintf(print_buf, "TEST PASSED\n");
vx_tmc(0);
return 0;
}