marfrit/rosenblatt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rosenblatt: full 2D-tiled arbitrary matmul on NPU (M and N), gemma-scale verified

Browse Source 
gemmtest generalized to tile BOTH M and N via rkt_gemm_plan: each (m x n,
full-K) tile extracts its own rows/cols, packs standalone, runs one verified
single-op matmul, assembles into Y. Verified: M=2048 N=256 K=16 (4 tiles) and
gemma-scale M=12 K=2048 N=256 (6x128 tiles) -> output matches CPU exactly.

So an arbitrary (M,N,K) INT8 matmul now runs correctly on the mainline rocket
NPU. Per-op sizing: m limited by CBUF input banks (K=16 -> m<=1024; large K
shrinks m since weights eat banks), K*n <= ~327k. Next: weight-reuse across
M-tiles (Mesa reuse_weights_cbuf) for perf (large-K forces tiny m today), then
U6 ggml-backend wiring (MUL_MAT dispatch -> this primitive).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Markus Fritsche
2026-07-14 09:01:10 +02:00
parent 2270f749bf
commit 562373ab83
1 changed files with 59 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
userspace/npu-probe/gemmtest.c
+59 -48
View File
@@ -1,10 +1,14 @@
/* 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].
* gemmtest.c — arbitrary (M,N,K) INT8 matmul on the rocket NPU by 2D TILING
* into single-op tiles (the U6 core primitive). Y[M][N] = X[M][K]*W[K][N],
* W given in conv order Wc[N][K] (Wc[oc][ic]).
*
* 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.
* Tiles M and N via rkt_gemm_plan; each (m x n, full-K) tile is one verified
* single-op matmul on fresh BOs, with the tile's own rows/cols packed
* standalone; outputs assembled into Y and checked vs a CPU reference.
*
* HW limits used: per-op m <= ~1024 rows (K=16 CBUF), K*n <= ~327k weights.
*/
#define _GNU_SOURCE
#include <stdint.h>
@@ -21,8 +25,9 @@
#include "coredpu_defs.h"
#define DEV "/dev/accel/accel0"
enum { M = 4096, K = 16, N = 8, IZP = 128, WZP = 128, OZP = 128 };
enum { M = 2048, N = 256, K = 16, IZP = 128, WZP = 128, OZP = 128 };
#define TILE_M 1024
#define TILE_N 128
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)
@@ -42,57 +47,63 @@ static void wr(int fd, struct bo *b, const void *s, unsigned n)
int main(void)
{
uint8_t *X = malloc((size_t)M * K), Wt[N * K];
int32_t Bin[N];
/* full operands: X = signed +1, Wc = signed +1 => each Y = K = 16 -> 144 */
uint8_t *X = malloc((size_t)M * K), *Wc = malloc((size_t)N * K);
int32_t *Bin = calloc(N, sizeof(int32_t));
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;
memset(Wc, WZP + 1, (size_t)N * K);
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);
uint8_t *Xtile = malloc((size_t)TILE_M * K);
uint8_t *Wtile = malloc((size_t)TILE_N * K);
uint8_t *ipk = calloc(1, rkt_raw_input_size(1, TILE_M, K));
uint8_t *wpk = calloc(1, rkt_packed_weights_size(1, 1, K, TILE_N));
int32_t bpk[TILE_N];
struct bo guard = {0};
mk(fd, &guard, 0x1000);
struct rkt_gemm_tile tiles[64];
int nt = rkt_gemm_plan(M, N, TILE_M, N, tiles, 64);
struct rkt_gemm_tile tiles[256];
int nt = rkt_gemm_plan(M, N, TILE_M, TILE_N, tiles, 256);
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);
printf("M=%d N=%d K=%d -> %d tiles (%dx%d)\n", M, N, K, nt, TILE_M, TILE_N);
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, &regb, 0x1000);
unsigned r = tiles[t].row, cc = tiles[t].col, m = tiles[t].m, n = tiles[t].n;
/* extract this tile's rows/cols and pack standalone */
for (unsigned y = 0; y < m; y++)
memcpy(Xtile + (size_t)y * K, X + (size_t)(r + y) * K, K);
for (unsigned j = 0; j < n; j++)
memcpy(Wtile + (size_t)j * K, Wc + (size_t)(cc + j) * K, K);
rkt_pack_input(Xtile, 1, m, K, IZP, ipk);
rkt_pack_weights(Wtile, 1, 1, K, n, WZP, wpk);
rkt_compute_biases(Wtile, Bin + cc, 1, 1, K, n, WZP, IZP, bpk);
rkt_pack_input(X + (size_t)r * K, 1, m, K, IZP, ipk);
struct bo inb = {0}, wtb = {0}, bib = {0}, outb = {0}, regb = {0};
mk(fd, &inb, rkt_raw_input_size(1, m, K));
mk(fd, &wtb, rkt_packed_weights_size(1, 1, K, n));
mk(fd, &bib, n * sizeof(int32_t));
mk(fd, &outb, rkt_raw_output_size(1, m, n));
mk(fd, &regb, 0x1000);
wr(fd, &inb, ipk, rkt_raw_input_size(1, m, K));
wr(fd, &wtb, wpk, rkt_packed_weights_size(1, 1, K, n));
wr(fd, &bib, bpk, n * sizeof(int32_t));
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++)
int nw = rkt_build_matmul_regcmd(rc, 512, m, n, K, inb.dma, wtb.dma,
outb.dma, IZP, WZP, OZP);
if (nw < 0) { printf("tile %d build -1 (m=%u n=%u)\n", t, m, n); return 1; }
for (int i = 0; i < nw; 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, &regb, rc, (unsigned)n * sizeof(uint64_t));
((uint64_t)(uint32_t)bib.dma << 16);
wr(fd, &regb, rc, (unsigned)nw * 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_task task = { .regcmd = (uint32_t)regb.dma, .regcmd_count = (uint32_t)nw };
uint32_t ih[] = { inb.h, wtb.h, bib.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;
@@ -102,21 +113,21 @@ int main(void)
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);
uint8_t *got = malloc((size_t)m * n);
rkt_unpack_output(outb.m, 1, m, n, got);
for (unsigned y = 0; y < m; y++)
memcpy(Y + (size_t)(r + y) * N + cc, got + (size_t)y * n, 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);
rocket_close_bo(fd, inb.h); rocket_close_bo(fd, wtb.h); rocket_close_bo(fd, bib.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);
printf("Y[0]=%u Y[M-1][N-1]=%u | %s (%ld/%ld bad vs 144)\n",
Y[0], Y[(long)M * N - 1],
bad ? "MISMATCH" : "2D-TILED NPU MATMUL CORRECT", bad, (long)M * N);
rocket_close_bo(fd, guard.h);
return bad ? 3 : 0;
}