gguf: add release_weight_buffer, keyed to the buffer rather than the load

src/core/gguf_loader.cpp

@@ -7,6 +7,8 @@
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
+#include <map>
+#include <mutex>
 
 #if defined(_WIN32)
 #include <io.h>
@@ -201,8 +203,9 @@ std::vector<uint8_t> kv_u8_array(gguf_context * gctx, const char * key) {
 
 namespace {
 
-// Platform unmap, shared by MappedFile's destructor and free_weights() (the
-// no-copy path transfers the mapping into WeightLoad, which unmaps on free).
+// Platform unmap, shared by MappedFile's destructor and release_weight_buffer()
+// (the no-copy path transfers the mapping to the backend buffer, which unmaps
+// when that buffer is released).
 void core_unmap(void * base, size_t size) {
     if (!base) return;
 #if defined(__EMSCRIPTEN__)
@@ -215,6 +218,44 @@ void core_unmap(void * base, size_t size) {
 #endif
 }
 
+// Which host mapping belongs to which backend buffer.
+//
+// On the no-copy path the backend buffer is a view onto pages the backend does
+// not own: ggml_backend_dev_buffer_from_host_ptr has no deallocator parameter,
+// so freeing the buffer releases the view and nothing else. WeightLoad carries
+// mmap_addr/mmap_len for the caller that keeps the whole struct, but a caller
+// that moves `buf` into its model and lets the WeightLoad die drops the only
+// record of the mapping — and eleven of the models in src/ tear down that way.
+// Keying the region to the buffer instead means the mapping is released by
+// whoever releases the buffer, whichever of the two shapes the caller uses.
+struct mmap_region {
+    void * base = nullptr;
+    size_t size = 0;
+};
+
+std::mutex g_buf_mmap_mu;
+std::map<ggml_backend_buffer_t, mmap_region> g_buf_mmap;
+
+void register_buf_mmap(ggml_backend_buffer_t buf, void * base, size_t size) {
+    std::lock_guard<std::mutex> lk(g_buf_mmap_mu);
+    g_buf_mmap[buf] = { base, size };
+}
+
+// Look the region up and remove the entry in one critical section. A lookup
+// followed by a separate erase would let a second release of the same buffer
+// read the entry before the first erased it and unmap twice; it would also
+// race a concurrent load whose fresh buffer landed on the same address after
+// the free. A default-constructed region means no entry, which is the ordinary
+// case — the copy path maps nothing that outlives the load.
+mmap_region take_buf_mmap(ggml_backend_buffer_t buf) {
+    std::lock_guard<std::mutex> lk(g_buf_mmap_mu);
+    auto it = g_buf_mmap.find(buf);
+    if (it == g_buf_mmap.end()) return mmap_region{};
+    const mmap_region r = it->second;
+    g_buf_mmap.erase(it);
+    return r;
+}
+
 // Read a file slice into a backend tensor. Uses mmap on POSIX; falls back
 // to pread/lseek+read when mmap is unavailable (rare in practice).
 //
@@ -350,7 +391,10 @@ bool load_weights(const char * path, ggml_backend_t backend, const char * model_
                     out.mmap_addr = mf.base;
                     out.mmap_len = mf.size;
                     out.used_mmap = true;
-                    mf.release(); // WeightLoad now owns the mapping
+                    // The buffer owns the mapping from here; the WeightLoad
+                    // fields above are a record of it, not a second owner.
+                    register_buf_mmap(buf, mf.base, mf.size);
+                    mf.release();
                     gguf_free(gctx);
                     return true;
                 }
@@ -630,23 +674,30 @@ bool load_weights_split(const char * path, ggml_backend_t gpu_backend, ggml_back
     return true;
 }
 
+void release_weight_buffer(ggml_backend_buffer_t & buf) {
+    if (!buf) return;
+    // Take the entry before the free, not after: between a free and a later
+    // erase, a concurrent load could receive a new buffer at the same address
+    // and register it, and the erase would then drop a live mapping's record.
+    const mmap_region r = take_buf_mmap(buf);
+    // Free the buffer first. On the no-copy path it is a view onto these pages,
+    // so unmapping them while it is alive would leave the device addressing
+    // unmapped memory.
+    ggml_backend_buffer_free(buf);
+    buf = nullptr;
+    core_unmap(r.base, r.size);
+}
+
 void free_weights(WeightLoad & wl) {
-    if (wl.buf) {
-        ggml_backend_buffer_free(wl.buf); // no-copy buffer doesn't own the pages
-        wl.buf = nullptr;
-    }
-    if (wl.buf_cpu) {
-        ggml_backend_buffer_free(wl.buf_cpu);
-        wl.buf_cpu = nullptr;
-    }
-    for (auto * b : wl.split_bufs) ggml_backend_buffer_free(b);
+    release_weight_buffer(wl.buf);
+    release_weight_buffer(wl.buf_cpu);
+    for (auto * b : wl.split_bufs) release_weight_buffer(b);
     wl.split_bufs.clear();
-    if (wl.mmap_addr) { // unmap after the buffer is freed
-        core_unmap(wl.mmap_addr, wl.mmap_len);
-        wl.mmap_addr = nullptr;
-        wl.mmap_len = 0;
-        wl.used_mmap = false;
-    }
+    // The mapping was released with its buffer above; these fields are the
+    // caller-visible record of the load, so clear them without unmapping again.
+    wl.mmap_addr = nullptr;
+    wl.mmap_len = 0;
+    wl.used_mmap = false;
     if (wl.ctx) {
         ggml_free(wl.ctx);
         wl.ctx = nullptr;

src/core/gguf_loader.h

@@ -167,8 +167,37 @@ using IsGpuTensor = bool (*)(const char * tensor_name, void * user);
 bool load_weights_split(const char * path, ggml_backend_t gpu_backend, ggml_backend_t cpu_backend, IsGpuTensor is_gpu,
                         void * user, const char * model_tag, WeightLoad & out);
 
+// Free a backend buffer that came from one of this loader's weight-loading
+// entry points, releasing the host mmap behind it when there is one.
+//
+// USE THIS INSTEAD OF ggml_backend_buffer_free() FOR ANY BUFFER OBTAINED FROM
+// load_weights() / load_weights_split(), including after the buffer has been
+// moved into a model struct. On the no-copy path
+// (`load_weights(..., try_mmap=true)` on a device advertising
+// buffer_from_host_ptr) the backend buffer is a view onto a host mapping the
+// backend does not own — buffer_from_host_ptr has no deallocator parameter —
+// so freeing the buffer alone leaves the weight file mapped for the life of
+// the process. free_weights() reaches the same mapping through WeightLoad, but
+// only for a caller that still holds the whole struct; a caller that moved
+// `buf` into its model has this entry point and nothing else.
+//
+// Semantics:
+//   * A null handle is a no-op, and the caller's handle is nulled on return,
+//     so a second call cannot double-free or double-unmap.
+//   * A buffer with no recorded mapping is the ordinary case — the copy path
+//     maps nothing that outlives the load — and is released like any other
+//     backend buffer.
+//   * The backend buffer is freed before the region is unmapped, so a
+//     device-side view of the pages never outlives them.
+//
+// This is also the entry point CrispASR's shared `crisp_audio` source calls,
+// so the name and behaviour have to hold in both repos — see the cross-repo
+// contract note above.
+void release_weight_buffer(ggml_backend_buffer_t & buf);
+
 // Free a WeightLoad's resources. Call when the model is being destroyed
-// and the buffer/context are not held elsewhere.
+// and the buffer/context are not held elsewhere. Releases every buffer
+// through release_weight_buffer().
 void free_weights(WeightLoad & wl);
 
 // ---------------------------------------------------------------------------

tests/test_gguf_loader_mmap.cpp

@@ -2,7 +2,8 @@
 // Builds a small GGUF, loads it both ways on the CPU backend (which advertises
 // buffer_from_host_ptr), and asserts the tensors are byte-identical and match
 // the values written. Validates the no-copy path is actually taken and that
-// free_weights() cleans up (incl. the mmap) without crashing.
+// free_weights() releases the mapping — asked of the kernel, not inferred from
+// the call returning.
 
 #include "core/gguf_loader.h"
 #include "core/clean_exit.h"
@@ -13,12 +14,69 @@
 
 #include <cmath>
 #include <cstdio>
+#include <climits>
 #include <cstring>
 #include <string>
 #include <vector>
 
+#if defined(__APPLE__)
+#include <libproc.h>
+#include <sys/proc_info.h>
+#include <unistd.h>
+#elif defined(__linux__)
+#include <fstream>
+#endif
+
 static const char * kPath = "/tmp/crispembed_test_loader_mmap.gguf";
 
+// The kernel names a region by its resolved path, so the comparison has to be
+// made against the same. On macOS /tmp is a symlink to /private/tmp, which is
+// enough on its own to make every region look like a stranger's.
+static std::string resolved_path(const char * path) {
+#if defined(_WIN32)
+    return std::string(path);
+#else
+    char buf[PATH_MAX];
+    return realpath(path, buf) ? std::string(buf) : std::string(path);
+#endif
+}
+
+// How many regions of this process are backed by `path`. The no-copy path keeps
+// the weight file mapped for the buffer's lifetime, so this is the exact way to
+// ask whether the mapping went away — no footprint threshold, nothing to settle.
+// Returns (size_t)-1 where the platform offers no region enumeration, which the
+// caller treats as "cannot assert here".
+static size_t count_regions_backed_by(const std::string & path) {
+#if defined(__APPLE__)
+    // PROC_PIDREGIONPATHINFO returns the region and its backing path in one
+    // record. proc_regionfilename() alone is not usable: asked about the base of
+    // an anonymous region it answers with the file of the next region above,
+    // counting an unrelated neighbour as a mapping of this file.
+    size_t n = 0;
+    uint64_t addr = 0;
+    for (;;) {
+        struct proc_regionwithpathinfo rpi;
+        if (proc_pidinfo(getpid(), PROC_PIDREGIONPATHINFO, addr, &rpi, sizeof(rpi)) != (int)sizeof(rpi)) break;
+        if (rpi.prp_vip.vip_path[0] != '\0' && path == rpi.prp_vip.vip_path) n++;
+        const uint64_t next = rpi.prp_prinfo.pri_address + rpi.prp_prinfo.pri_size;
+        if (next <= addr) break; // no forward progress; stop rather than spin
+        addr = next;
+    }
+    return n;
+#elif defined(__linux__)
+    size_t n = 0;
+    const size_t len = path.size();
+    std::ifstream maps("/proc/self/maps");
+    std::string line;
+    while (std::getline(maps, line))
+        if (line.size() > len && line.compare(line.size() - len, len, path) == 0) n++;
+    return n;
+#else
+    (void)path;
+    return (size_t)-1;
+#endif
+}
+
 static float expected(int i) {
     return sinf((float)i * 0.013f) + 0.5f;
 }
@@ -101,8 +159,31 @@ static int crispembed_test_main() {
         if (!ok) fails++;
     }
 
+    // The mapping is keyed to the backend buffer, so free_weights reaches it
+    // through release_weight_buffer rather than through mw.mmap_addr. Check the
+    // kernel, not the field: a desync between where the region is registered
+    // and where it is taken would leave the file mapped and clear the field
+    // anyway.
+    const std::string kAbs = resolved_path(kPath);
+    const size_t mapped_before = count_regions_backed_by(kAbs);
     core_gguf::free_weights(cw);
-    core_gguf::free_weights(mw); // also unmaps
+    core_gguf::free_weights(mw);
+    if (mapped_before == (size_t)-1) {
+        printf("region enumeration unavailable on this platform — mapping release not asserted\n");
+    } else if (mapped_before < 1) {
+        // Positive control. Without it the zero below would also be satisfied
+        // by a load that never mapped the file in the first place.
+        fprintf(stderr, "FAIL: no-copy load left no mapping to release (found %zu regions)\n", mapped_before);
+        fails++;
+    } else {
+        const size_t mapped_after = count_regions_backed_by(kAbs);
+        printf("  weight file regions: %zu while loaded, %zu after free_weights\n", mapped_before, mapped_after);
+        if (mapped_after != 0) {
+            fprintf(stderr, "FAIL: free_weights left %zu mapping(s) of %s\n", mapped_after, kPath);
+            fails++;
+        }
+    }
+
     ggml_backend_free(backend);
     remove(kPath);