diff --git a/userspace/npu-probe/.gitignore b/userspace/npu-probe/.gitignore index 65c2a17..bff7fc4 100644 --- a/userspace/npu-probe/.gitignore +++ b/userspace/npu-probe/.gitignore @@ -7,3 +7,5 @@ verify/golden.diff hwtest openprobe + +gemmtest diff --git a/userspace/npu-probe/gemmtest.c b/userspace/npu-probe/gemmtest.c new file mode 100644 index 0000000..aed4d98 --- /dev/null +++ b/userspace/npu-probe/gemmtest.c @@ -0,0 +1,122 @@ +/* SPDX-License-Identifier: MIT + * + * gemmtest.c — arbitrary-size INT8 matmul on the rocket NPU by TILING into + * single-op tiles (the U6 core primitive). Y[M][N] = X[M][K]*W[K][N]. + * + * Tiles M via rkt_gemm_plan; weights packed once; each M-chunk is one verified + * single-op matmul (fresh BOs per tile); outputs assembled + checked vs CPU. + */ +#define _GNU_SOURCE +#include +#include +#include +#include +#include + +#include "librocket.h" +#include "rocket_accel.h" +#include "rkt_matmul.h" +#include "rkt_operands.h" +#include "rkt_gemm.h" +#include "coredpu_defs.h" + +#define DEV "/dev/accel/accel0" +enum { M = 4096, K = 16, N = 8, IZP = 128, WZP = 128, OZP = 128 }; +#define TILE_M 1024 + +struct bo { uint32_t h; uint64_t dma, off; void *m; uint32_t sz; }; +static int mk(int fd, struct bo *b, uint32_t s) +{ + b->sz = (s + 0xfff) & ~0xfffu; + if (rocket_create_bo(fd, b->sz, &b->h, &b->dma, &b->off)) return -1; + b->m = rocket_mmap_bo(fd, b->off, b->sz); + return b->m == MAP_FAILED ? -1 : 0; +} +static void wr(int fd, struct bo *b, const void *s, unsigned n) +{ + rocket_prep_bo(fd, b->h, INT64_MAX); + memset(b->m, 0, b->sz); + if (s) memcpy(b->m, s, n); + rocket_fini_bo(fd, b->h); +} + +int main(void) +{ + uint8_t *X = malloc((size_t)M * K), Wt[N * K]; + int32_t Bin[N]; + memset(X, IZP + 1, (size_t)M * K); + for (int i = 0; i < N * K; i++) Wt[i] = WZP + 1; + for (int n = 0; n < N; n++) Bin[n] = 0; + + int fd = rocket_open(DEV); + if (fd < 0) { printf("open %d\n", fd); return 1; } + + unsigned wsz = rkt_packed_weights_size(1, 1, K, N); + uint8_t *wpk = calloc(1, wsz); + int32_t bpk[N]; + rkt_pack_weights(Wt, 1, 1, K, N, WZP, wpk); + rkt_compute_biases(Wt, Bin, 1, 1, K, N, WZP, IZP, bpk); + struct bo guard = {0}, wtb = {0}, biasb = {0}; + mk(fd, &guard, 0x1000); + mk(fd, &wtb, wsz); mk(fd, &biasb, sizeof bpk); + wr(fd, &wtb, wpk, wsz); wr(fd, &biasb, bpk, sizeof bpk); + + uint8_t *ipk = calloc(1, rkt_raw_input_size(1, TILE_M, K)); + uint8_t *Y = calloc((size_t)M, N); + + struct rkt_gemm_tile tiles[64]; + int nt = rkt_gemm_plan(M, N, TILE_M, N, tiles, 64); + if (nt <= 0) { printf("plan %d\n", nt); return 1; } + printf("M=%d N=%d K=%d -> %d M-tiles of <=%d rows\n", M, N, K, nt, TILE_M); + + for (int t = 0; t < nt; t++) { + unsigned r = tiles[t].row, m = tiles[t].m; + /* fresh BOs per tile (avoid cross-job reuse state) */ + struct bo inb = {0}, outb = {0}, regb = {0}; + mk(fd, &inb, rkt_raw_input_size(1, m, K)); + mk(fd, &outb, rkt_raw_output_size(1, m, N)); + mk(fd, ®b, 0x1000); + + rkt_pack_input(X + (size_t)r * K, 1, m, K, IZP, ipk); + wr(fd, &inb, ipk, rkt_raw_input_size(1, m, K)); + wr(fd, &outb, NULL, 0); + + uint64_t rc[512]; + int n = rkt_build_matmul_regcmd(rc, 512, m, N, K, inb.dma, wtb.dma, + outb.dma, IZP, WZP, OZP); + if (n < 0) { printf("tile %d build -1 (m=%u)\n", t, m); return 1; } + for (int i = 0; i < n; i++) + if ((uint32_t)(rc[i] & 0xffff) == REG_DPU_RDMA_RDMA_BS_BASE_ADDR) + rc[i] = (rc[i] & ~(0xffffffffULL << 16)) | + ((uint64_t)(uint32_t)biasb.dma << 16); + wr(fd, ®b, rc, (unsigned)n * sizeof(uint64_t)); + + struct drm_rocket_task task = { .regcmd = (uint32_t)regb.dma, .regcmd_count = (uint32_t)n }; + uint32_t ih[] = { inb.h, wtb.h, biasb.h, regb.h }, oh[] = { outb.h }; + struct drm_rocket_job job; + memset(&job, 0, sizeof job); + job.tasks = (uintptr_t)&task; job.task_count = 1; + job.task_struct_size = sizeof(struct drm_rocket_task); + job.in_bo_handles = (uintptr_t)ih; job.in_bo_handle_count = 4; + job.out_bo_handles = (uintptr_t)oh; job.out_bo_handle_count = 1; + if (rocket_submit(fd, &job, 1)) { printf("submit fail t%d\n", t); return 1; } + if (rocket_prep_bo(fd, outb.h, INT64_MAX)) { printf("wait fail t%d\n", t); return 2; } + + uint8_t *got = malloc((size_t)m * N); + rkt_unpack_output(outb.m, 1, m, N, got); + memcpy(Y + (size_t)r * N, got, (size_t)m * N); + free(got); + printf(" tile %d: rows[%u..%u] out[0]=%u\n", t, r, r + m - 1, Y[(size_t)r * N]); + rocket_close_bo(fd, inb.h); rocket_close_bo(fd, outb.h); rocket_close_bo(fd, regb.h); + } + + long bad = 0; + for (long i = 0; i < (long)M * N; i++) + if (Y[i] > 145 || Y[i] < 143) bad++; + printf("Y[0]=%u Y[last]=%u | %s (%ld/%ld bad vs 144)\n", + Y[0], Y[(long)(M - 1) * N], + bad ? "MISMATCH" : "TILED NPU MATMUL CORRECT", + bad, (long)M * N); + rocket_close_bo(fd, guard.h); rocket_close_bo(fd, wtb.h); rocket_close_bo(fd, biasb.h); + return bad ? 3 : 0; +}