npu-probe: backport mmap-leak fix + K>8192 backstop from ggml-rocket

bo_free() (munmap+close) fixes the per-tile mmap leak in rkt_npu_matmul; add
rocket_munmap_bo to librocket. rkt_matmul rejects K>8192 (int8 K-limit). Keeps
the canonical primitive in sync with the ggml-rocket backend copy in rk-llama.cpp.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EWpfhDgYNA21tETDP9ueBE

userspace/npu-probe/librocket.c

@@ -50,6 +50,12 @@ void *rocket_mmap_bo(int fd, uint64_t mmap_offset, uint32_t size)
 		    mmap_offset);
 }
 
+void rocket_munmap_bo(void *map, uint32_t size)
+{
+	if (map && map != MAP_FAILED)
+		munmap(map, size);
+}
+
 int rocket_prep_bo(int fd, uint32_t handle, int64_t timeout_ns)
 {
 	struct drm_rocket_prep_bo prep_bo;

userspace/npu-probe/librocket.h

@@ -48,6 +48,10 @@ int rocket_create_bo(int fd, uint32_t size, uint32_t *handle_out,
  */
 void *rocket_mmap_bo(int fd, uint64_t mmap_offset, uint32_t size);
 
+/* rocket_munmap_bo - undo rocket_mmap_bo (the VMA pins the GEM object; GEM_CLOSE
+ * alone does NOT free it, so every mmap must be paired with a munmap). */
+void rocket_munmap_bo(void *map, uint32_t size);
+
 /**
  * rocket_prep_bo - take CPU ownership of a BO (wait for NPU, sync caches)
  * @fd:         rocket device fd

userspace/npu-probe/rkt_matmul.c

@@ -95,6 +95,13 @@ int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity,
 {
 	if (M == 0 || N == 0 || K == 0)
 		return -1;
+	/* Hard backstop: the RK3588 int8 matmul K-limit is ~8192 (verified: correct
+	 * through 8192, catastrophic garbage at 11008). Weights exceeding CBUF are
+	 * streamed (reuse_weights_cbuf=0, emitted below) and are correct up to that
+	 * K limit, so gate on K, not on bank count. Backend supports_op also caps
+	 * this; the check here protects direct callers. */
+	if (K > 8192)
+		return -1;
 
 	struct op o = {
 		.stride = 1,

userspace/npu-probe/rkt_npu_matmul.c

@@ -14,6 +14,17 @@
 #include <stdlib.h>
 #include <string.h>
 #include <sys/mman.h>
+#include <time.h>
+
+/* rocket_prep_bo timeout is an ABSOLUTE CLOCK_MONOTONIC deadline. Build one
+ * rel_ns in the future, for the job-completion wait so a hung NPU job returns
+ * an error instead of blocking forever (the driver watchdog is ~5s). */
+static int64_t rkt_deadline_ns(int64_t rel_ns)
+{
+	struct timespec t;
+	clock_gettime(CLOCK_MONOTONIC, &t);
+	return (int64_t)t.tv_sec * 1000000000LL + (int64_t)t.tv_nsec + rel_ns;
+}
 
 int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity,
 				   uint32_t M, uint32_t N, uint32_t K,
@@ -45,6 +56,18 @@ static void bo_write(int fd, struct rkt_bo *b, const void *src, unsigned n)
 	rocket_fini_bo(fd, b->h);
 }
 
+/* munmap + GEM_CLOSE a BO (safe on a zero-initialised struct). GEM_CLOSE alone
+ * leaks the mapping — the VMA keeps the object pinned until process exit. */
+static void bo_free(int fd, struct rkt_bo *b)
+{
+	if (b->map && b->map != MAP_FAILED)
+		rocket_munmap_bo(b->map, b->sz);
+	if (b->h)
+		rocket_close_bo(fd, b->h);
+	b->map = NULL;
+	b->h = 0;
+}
+
 int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc,
 		   const int32_t *bias, uint32_t M, uint32_t N, uint32_t K,
 		   uint8_t izp, uint8_t wzp, uint8_t ozp,
@@ -97,8 +120,11 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc,
 		    bo_alloc(fd, &w, rkt_packed_weights_size(1, 1, K, n)) ||
 		    bo_alloc(fd, &b, n * sizeof(int32_t)) ||
 		    bo_alloc(fd, &o, rkt_raw_output_size(1, m, n)) ||
-		    bo_alloc(fd, &reg, 0x1000))
+		    bo_alloc(fd, &reg, 0x1000)) {
+			bo_free(fd, &in); bo_free(fd, &w); bo_free(fd, &b);
+			bo_free(fd, &o); bo_free(fd, &reg);
 			goto out;
+		}
 		bo_write(fd, &in, ipk, rkt_raw_input_size(1, m, K));
 		bo_write(fd, &w, wpk, rkt_packed_weights_size(1, 1, K, n));
 		bo_write(fd, &b, bpk, n * sizeof(int32_t));
@@ -108,9 +134,9 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc,
 			rc, 4096, m, n, K, in.dma, w.dma, o.dma, izp, wzp, ozp,
 			in_scale, wt_scale, out_scales ? out_scales[cc] : out_scale, b.dma);
 		if (nw < 0) {
-			rocket_close_bo(fd, in.h); rocket_close_bo(fd, w.h);
-			rocket_close_bo(fd, b.h); rocket_close_bo(fd, o.h);
-			rocket_close_bo(fd, reg.h);
+			bo_free(fd, &in); bo_free(fd, &w);
+			bo_free(fd, &b); bo_free(fd, &o);
+			bo_free(fd, &reg);
 			goto out;
 		}
 		bo_write(fd, &reg, rc, (unsigned)nw * sizeof(uint64_t));
@@ -129,7 +155,7 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc,
 		job.out_bo_handle_count = 1;
 
 		int sr = rocket_submit(fd, &job, 1);
-		int wr = sr ? -1 : rocket_prep_bo(fd, o.h, INT64_MAX);
+		int wr = sr ? -1 : rocket_prep_bo(fd, o.h, rkt_deadline_ns(6000000000LL));
 		if (!sr && !wr) {
 			uint8_t *got = malloc((size_t)m * n);
 			if (got) {
@@ -140,9 +166,9 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc,
 				free(got);
 			}
 		}
-		rocket_close_bo(fd, in.h); rocket_close_bo(fd, w.h);
-		rocket_close_bo(fd, b.h); rocket_close_bo(fd, o.h);
-		rocket_close_bo(fd, reg.h);
+		bo_free(fd, &in); bo_free(fd, &w);
+		bo_free(fd, &b); bo_free(fd, &o);
+		bo_free(fd, &reg);
 		if (sr || wr)
 			goto out;
 	}