From a2244675b126371cf80c000cee2b32459fb9f461 Mon Sep 17 00:00:00 2001 From: Markus Fritsche Date: Tue, 14 Jul 2026 11:32:13 +0200 Subject: [PATCH] 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 Claude-Session: https://claude.ai/code/session_01EWpfhDgYNA21tETDP9ueBE --- userspace/npu-probe/librocket.c | 6 ++++ userspace/npu-probe/librocket.h | 4 +++ userspace/npu-probe/rkt_matmul.c | 7 +++++ userspace/npu-probe/rkt_npu_matmul.c | 42 ++++++++++++++++++++++------ 4 files changed, 51 insertions(+), 8 deletions(-) diff --git a/userspace/npu-probe/librocket.c b/userspace/npu-probe/librocket.c index 1eb2707..91368be 100644 --- a/userspace/npu-probe/librocket.c +++ b/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; diff --git a/userspace/npu-probe/librocket.h b/userspace/npu-probe/librocket.h index 7427d90..6a60212 100644 --- a/userspace/npu-probe/librocket.h +++ b/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 diff --git a/userspace/npu-probe/rkt_matmul.c b/userspace/npu-probe/rkt_matmul.c index d570fe5..d3da360 100644 --- a/userspace/npu-probe/rkt_matmul.c +++ b/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, diff --git a/userspace/npu-probe/rkt_npu_matmul.c b/userspace/npu-probe/rkt_npu_matmul.c index 84759a3..5b53adc 100644 --- a/userspace/npu-probe/rkt_npu_matmul.c +++ b/userspace/npu-probe/rkt_npu_matmul.c @@ -14,6 +14,17 @@ #include #include #include +#include + +/* 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, ®, 0x1000)) + bo_alloc(fd, ®, 0x1000)) { + bo_free(fd, &in); bo_free(fd, &w); bo_free(fd, &b); + bo_free(fd, &o); bo_free(fd, ®); 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, ®); goto out; } bo_write(fd, ®, 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, ®); if (sr || wr) goto out; }