Rosenblatt: arbitrary-size matmul on the NPU via tiling (U6 core primitive)

gemmtest.c computes Y[M][N]=X[M][K]*W[K][N] for M larger than one CBUF tile by
splitting M via rkt_gemm_plan and running each chunk as a verified single-op
matmul (fresh BOs/tile, weights packed once), assembling the output. M=4096
via 4x1024-row tiles -> output matches CPU exactly (0/32768 wrong).

Key HW limit found: a single rocket op writes correct output only up to
M~1024 rows (K=16); M=2048 silently drops rows past ~971 (CBUF input capacity),
so tile_m<=1024 is the safe per-op bound. Also: fresh BOs per tile are needed
(reusing in/out/reg BOs across jobs corrupted later tiles). This is the tiling
orchestration U6/ggml will use; N-tiling (weight-column) + weight-reuse
optimization are next.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

userspace/npu-probe/.gitignore

@@ -7,3 +7,5 @@ verify/golden.diff
 
 hwtest
 openprobe
+
+gemmtest

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 <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/mman.h>
+
+#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, &regb, 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, &regb, 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;
+}