engram: batch-cosine Adapter/Strategy/Factory over ggml, supersedes hand-rolled PR #114

Stop hand-rolling GPU kernels for batch cosine similarity — use ggml (the MIT-licensed compute library underneath llama.cpp, installed standalone via Homebrew) as the preferred backend, without ripping out PR #114's carefully-verified hand-rolled Metal shader. Structure: one stable public adapter (eg_cosine_batch.h, zero #ifdef at call sites) backed by three selectable concrete Strategies behind an internal vtable (eg_cosine_batch_strategy.h) chosen by a Factory (eg_cosine_batch.c): - eg_cosine_batch_strategy_ggml.c — NEW. ggml + dynamically-loaded Metal backend plugin (ggml_backend_load_all_from_path + ggml_mul_mat for the batched dot product), gather/scatter around the -2.0 sentinel contract. - eg_cosine_batch_strategy_metal_hand.m — PR #114's original hand-rolled Metal shader bridge, preserved almost verbatim, now one strategy among several rather than the only option. eg_cosine_batch.metal kept byte-identical to the original. - eg_cosine_batch_strategy_cpu.c — universal always-false fallback (direct descendant of PR #114's eg_metal_cosine_stub.c). Selection: EL_COSINE_BATCH_STRATEGY=ggml|metal|cpu|auto (default: ggml first, then hand-rolled Metal, then CPU — first available wins), plus back-compat EL_METAL_COSINE=0 to disable every GPU-backed strategy. build_vindex_bench.sh compiles all three strategies on Darwin, CPU-fallback-only elsewhere. vindex_bench.c now reports BRUTE-GGML and BRUTE-METAL side by side against the same CPU oracle, on the same dataset, in one run (real numbers vs. real store snapshot in the PR body).
2026-08-15 17:16:41 -05:00
parent 2555e363a6
commit b3f410fc91
9 changed files with 1372 additions and 5 deletions
@@ -0,0 +1,306 @@
+/* eg_cosine_batch_strategy_ggml.c — the GGML Strategy, and the preferred
+ * default whenever it is available (see the factory's selection order in
+ * eg_cosine_batch.c).
+ *
+ * WHY: directive from Will Anderson — stop hand-rolling GPU kernels, use a
+ * real, proven, permissively-licensed library instead. ggml (the compute
+ * library underneath llama.cpp, MIT licensed) is already installed on this
+ * machine as a standalone Homebrew package (`brew info ggml`), independent
+ * of llama.cpp itself. This file is a genuinely bounded COMPUTE UTILITY —
+ * batch cosine-similarity math — analogous to a VBD Accessor calling out to
+ * infrastructure. It is explicitly NOT the engram's reasoning/persistence
+ * core; using ggml here does not cross the "own the core" line, because
+ * batch cosine math is infrastructure, not the graph traversal / activation
+ * spreading / "thinking" that IS the core and stays 100% own-code.
+ *
+ * ── The real API shape (verified against the installed headers + a
+ *    standalone probe program, not assumed from memory of other tensor
+ *    libraries) ──────────────────────────────────────────────────────────
+ *
+ * ggml ships its CPU and Metal implementations as DYNAMICALLY LOADED PLUGIN
+ * .so files (confirmed by nm: `ggml_backend_metal_init` is NOT an exported
+ * symbol of libggml.dylib/libggml-base.dylib — it exists ONLY inside
+ * libggml-metal.so under $(brew --prefix ggml)/libexec/). You cannot link
+ * `-lggml-metal`; you must go through ggml's backend REGISTRY:
+ *
+ *   1. ggml_backend_load_all_from_path(dir) — dlopen()s every backend plugin
+ *      .so found in `dir` and registers its device(s). We point this at
+ *      $(brew --prefix ggml)/libexec (resolved once, at build+init time; see
+ *      eg_ggml_backend_dir() below) rather than relying on
+ *      ggml_backend_load_all()'s own default search heuristics, which are
+ *      tuned for an installed llama.cpp-style app bundle layout, not an
+ *      arbitrary `cc`-built binary invoked from an arbitrary cwd — the exact
+ *      same "must not silently fall back to CPU for reasons that have
+ *      nothing to do with GPU availability" concern PR #114's hand-rolled
+ *      bridge already documented for its own embedded-shader-source choice.
+ *   2. ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU) — find the
+ *      registered Metal device.
+ *   3. ggml_backend_dev_init(dev, NULL) — get a live ggml_backend_t.
+ *   4. Build a tiny ggml_context (no_alloc=true; it holds only tensor
+ *      metadata, not data), declare 2D F32 tensors, ggml_mul_mat(nodes,
+ *      query) — ggml's documented convention: A is [k cols, n rows], B is
+ *      [k cols, m rows] (transposed internally), result is [n cols, m rows]
+ *      — i.e. mul_mat(node_matrix[dim,n], query_matrix[dim,nq]) yields
+ *      out[n,nq] where out[j*n+i] = dot(node_i, query_j). A row-major
+ *      (dim,n) node matrix and a row-major (dim,nq) query matrix is EXACTLY
+ *      the packed layout the hand-rolled Metal kernel already used — one
+ *      matmul replaces the whole per-row dot-product loop.
+ *   5. ggml_backend_alloc_ctx_tensors(ctx, backend) to actually allocate
+ *      device buffers for those tensors, ggml_backend_tensor_set() to upload,
+ *      ggml_backend_graph_compute() to run, ggml_backend_tensor_get() to
+ *      read back.
+ *
+ * This exact sequence was verified end-to-end in a standalone probe (build
+ * it yourself: see the PR description) against a plain-C CPU dot product —
+ * bit-for-bit correct within float rounding. Real numbers against the
+ * el_runtime.c CPU oracle are reported in the PR body via vindex_bench.
+ *
+ * ggml_mul_mat only computes the raw dot products — it has no notion of
+ * "cosine" or of this codebase's -2.0 dim-mismatch/null/zero-norm sentinel.
+ * Per the adapter's directive: gather only VALID, uniform-dim rows into the
+ * packed matrix sent to the GPU (skipping null/mismatched rows entirely,
+ * rather than the hand-rolled kernel's zero-pad-and-sentinel-in-shader
+ * approach), then scatter -2.0 back for every row that was excluded — same
+ * gather/scatter contract eg_cosine_batch.h documents. Norms (||node||,
+ * ||query||) are computed on the CPU host in the same pass that already
+ * touches every element to gather/convert — essentially free — using the
+ * same 4-way-partial-sum accumulation the hand-rolled kernel and the CPU
+ * oracle both use, so the float32 error profile stays comparable across all
+ * three strategies. Only the O(n*dim*nq) dot-product matmul — the actual
+ * expensive part — is offloaded to the GPU.
+ *
+ * Precision: ggml_mul_mat on F32 x F32 inputs computes in F32 on the Metal
+ * backend (verified: no GGML_PREC_F16 default path applies to F32 inputs;
+ * see ggml_mul_mat_set_prec in ggml.h, which exists specifically to raise
+ * precision for lower-than-F32 inputs — ours are already F32 throughout).
+ * The measured delta vs the CPU double-accumulated oracle is reported
+ * honestly in the PR body (vindex_bench's BRUTE-GGML line), not assumed.
+ */
+#include "eg_cosine_batch_strategy.h"
+
+#include <ggml.h>
+#include <ggml-backend.h>
+#include <ggml-alloc.h>
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <math.h>
+
+/* ── lazy, one-time backend init, cached ─────────────────────────────────── */
+static bool g_init_attempted = false;
+static bool g_init_ok = false;
+static ggml_backend_t g_backend = NULL;
+
+/* Where to look for the dynamically-loaded backend plugin .so files.
+ * EL_GGML_BACKEND_PATH overrides for non-standard installs; otherwise we try
+ * the Homebrew opt-prefix symlink (stable across ggml point-version bumps —
+ * $(brew --prefix ggml)/libexec — confirmed to exist and contain
+ * libggml-metal.so / libggml-cpu-*.so / libggml-blas.so on this machine),
+ * falling back to ggml's own default search (ggml_backend_load_all()) in
+ * case a different install layout (e.g. a from-source build with a
+ * standard-prefix install) makes that succeed instead. */
+static const char* eg_ggml_backend_dir(void) {
+    const char* s = getenv("EL_GGML_BACKEND_PATH");
+    if (s && *s) return s;
+    return "/opt/homebrew/opt/ggml/libexec";
+}
+
+static bool eg_ggml_ensure_init(void) {
+    if (g_init_attempted) return g_init_ok;
+    g_init_attempted = true;
+
+    const char* dir = eg_ggml_backend_dir();
+    /* dlopen every backend plugin .so in `dir` and register its device(s).
+     * Never throws; a missing/empty directory just means no devices get
+     * registered and the lookup below fails cleanly. */
+    ggml_backend_load_all_from_path(dir);
+
+    ggml_backend_dev_t dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU);
+    if (!dev) {
+        /* Fall back to ggml's own default search heuristics only if the
+         * explicit path above found nothing — avoids double-registering the
+         * same plugins (ggml does not dedupe two different paths that
+         * happen to resolve to the same files, e.g. our stable opt-prefix
+         * symlink vs. its own Cellar-relative guess) in the common case
+         * where the explicit path already worked. */
+        ggml_backend_load_all();
+        dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU);
+    }
+    if (!dev) return false;
+
+    ggml_backend_t backend = ggml_backend_dev_init(dev, NULL);
+    if (!backend) return false;
+
+    g_backend = backend;
+    g_init_ok = true;
+    return true;
+}
+
+static bool ggml_strategy_available(void) {
+    return eg_ggml_ensure_init();
+}
+
+/* ── shared core: gather valid rows + norms, matmul, scatter ────────────── */
+
+/* 4-way partial-sum squared-norm accumulation over `dim` floats — same shape
+ * as eg_cosine_batch.metal's per-thread accumulation and vindex_bench.c's
+ * CPU brute_topk unroll, kept consistent on purpose so the float32 error
+ * profile is comparable across all three strategies. */
+static float eg_norm_sq_f32(const float* v, int32_t dim) {
+    float s0 = 0, s1 = 0, s2 = 0, s3 = 0;
+    int32_t d = 0, dim4 = dim & ~3;
+    for (; d < dim4; d += 4) {
+        s0 += v[d] * v[d]; s1 += v[d+1] * v[d+1];
+        s2 += v[d+2] * v[d+2]; s3 += v[d+3] * v[d+3];
+    }
+    float s = (s0 + s1) + (s2 + s3);
+    for (; d < dim; d++) s += v[d] * v[d];
+    return s;
+}
+
+/* Runs one ggml_mul_mat(node_matrix[dim,n_valid], query_matrix[dim,nq]) and
+ * combines it with CPU-computed norms into cosine scores, scattering into
+ * out_scores at ORIGINAL (ungathered) indices. out_scores must already be
+ * fully sized for n*nq (or n for the single-query case, nq=1) — every entry
+ * gets written (valid rows get a real cosine, invalid rows get -2.0), so
+ * this never leaves a partial result. Returns false only on a genuine
+ * failure (alloc, compute) at which point out_scores is left as whatever a
+ * caller-supplied scratch buffer already contained — callers here always
+ * pass a fresh buffer they discard on false, matching the adapter contract
+ * of "on failure, out_scores is treated as untouched" from the caller's
+ * point of view. */
+static bool eg_ggml_run(const float* queries, int32_t qdim, int32_t nq,
+                         const float* const* node_ptrs, const int32_t* node_dims,
+                         int32_t n, double* out_scores) {
+    if (!queries || qdim <= 0 || nq <= 0 || !node_ptrs || !node_dims || n <= 0 || !out_scores)
+        return false;
+    if (!eg_ggml_ensure_init()) return false;
+
+    /* Pass 1 (CPU): gather valid rows (non-NULL ptr, dim == qdim) into a
+     * packed (dim, n_valid) row-major matrix, remembering the original index
+     * of each packed row, and compute each valid row's squared norm in the
+     * same pass. Rows excluded here get -2.0 scattered for every query
+     * below without ever touching the GPU. */
+    int32_t* valid_orig = (int32_t*)malloc((size_t)n * sizeof(int32_t));
+    float* node_norm_sq  = (float*)malloc((size_t)n * sizeof(float)); /* indexed by packed position */
+    float* node_matrix   = NULL;
+    if (!valid_orig || !node_norm_sq) { free(valid_orig); free(node_norm_sq); return false; }
+
+    int32_t n_valid = 0;
+    for (int32_t i = 0; i < n; i++) {
+        if (node_ptrs[i] && node_dims[i] == qdim) n_valid++;
+    }
+
+    if (n_valid > 0) {
+        node_matrix = (float*)malloc((size_t)n_valid * (size_t)qdim * sizeof(float));
+        if (!node_matrix) { free(valid_orig); free(node_norm_sq); return false; }
+        int32_t w = 0;
+        for (int32_t i = 0; i < n; i++) {
+            if (!node_ptrs[i] || node_dims[i] != qdim) continue;
+            memcpy(node_matrix + (size_t)w * qdim, node_ptrs[i], (size_t)qdim * sizeof(float));
+            node_norm_sq[w] = eg_norm_sq_f32(node_ptrs[i], qdim);
+            valid_orig[w] = i;
+            w++;
+        }
+    }
+
+    /* Query norms — nq is typically small (1 or the size of one batch of
+     * comparison queries), so this loop is cheap regardless. */
+    float* q_norm_sq = (float*)malloc((size_t)nq * sizeof(float));
+    if (!q_norm_sq) { free(valid_orig); free(node_norm_sq); free(node_matrix); return false; }
+    for (int32_t j = 0; j < nq; j++) q_norm_sq[j] = eg_norm_sq_f32(queries + (size_t)j * qdim, qdim);
+
+    /* Nothing valid to compare against: every output is -2.0. Still a fully
+     * and correctly populated result — no GPU dispatch was needed to know
+     * that. */
+    if (n_valid == 0) {
+        for (size_t k = 0; k < (size_t)n * (size_t)nq; k++) out_scores[k] = -2.0;
+        free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq);
+        return true;
+    }
+
+    /* Pass 2 (GPU via ggml): dot = mul_mat(node_matrix[dim,n_valid],
+     * queries[dim,nq]) -> dot[n_valid, nq], dot[j*n_valid+i] = dot(node_i,query_j). */
+    struct ggml_init_params gp = {
+        .mem_size   = ggml_tensor_overhead() * 8 + ggml_graph_overhead(),
+        .mem_buffer = NULL,
+        .no_alloc   = true,
+    };
+    struct ggml_context* ctx = ggml_init(gp);
+    if (!ctx) { free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq); return false; }
+
+    struct ggml_tensor* t_nodes = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, qdim, n_valid);
+    struct ggml_tensor* t_query = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, qdim, nq);
+    if (!t_nodes || !t_query) { ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq); return false; }
+    struct ggml_tensor* t_dot = ggml_mul_mat(ctx, t_nodes, t_query);
+    if (!t_dot) { ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq); return false; }
+
+    struct ggml_backend_buffer* buf = ggml_backend_alloc_ctx_tensors(ctx, g_backend);
+    if (!buf) { ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq); return false; }
+
+    ggml_backend_tensor_set(t_nodes, node_matrix, 0, (size_t)n_valid * qdim * sizeof(float));
+    ggml_backend_tensor_set(t_query, queries, 0, (size_t)nq * qdim * sizeof(float));
+    free(node_matrix); /* uploaded; the packed CPU copy is no longer needed */
+
+    struct ggml_cgraph* gf = ggml_new_graph(ctx);
+    if (!gf) { ggml_backend_buffer_free(buf); ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(q_norm_sq); return false; }
+    ggml_build_forward_expand(gf, t_dot);
+    enum ggml_status st = ggml_backend_graph_compute(g_backend, gf);
+    if (st != GGML_STATUS_SUCCESS) {
+        ggml_backend_buffer_free(buf); ggml_free(ctx);
+        free(valid_orig); free(node_norm_sq); free(q_norm_sq);
+        return false;
+    }
+
+    float* dot = (float*)malloc(ggml_nbytes(t_dot));
+    if (!dot) { ggml_backend_buffer_free(buf); ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(q_norm_sq); return false; }
+    ggml_backend_tensor_get(t_dot, dot, 0, ggml_nbytes(t_dot));
+
+    /* Pass 3 (CPU): combine dot/(||a||*||b||) per (query,node) pair, scatter
+     * into out_scores at ORIGINAL node indices; every excluded row gets
+     * -2.0 for every query. out_scores is fully populated either way. */
+    for (int32_t j = 0; j < nq; j++) {
+        double* orow = out_scores + (size_t)j * n;
+        for (int32_t i = 0; i < n; i++) orow[i] = -2.0; /* default: excluded */
+        for (int32_t w = 0; w < n_valid; w++) {
+            float na = node_norm_sq[w], nb = q_norm_sq[j];
+            int32_t oi = valid_orig[w];
+            if (na <= 0.0f || nb <= 0.0f) { orow[oi] = -2.0; continue; }
+            float d = dot[(size_t)j * n_valid + w];
+            orow[oi] = (double)(d / sqrtf(na * nb));
+        }
+    }
+
+    free(dot);
+    ggml_backend_buffer_free(buf);
+    ggml_free(ctx);
+    free(valid_orig); free(node_norm_sq); free(q_norm_sq);
+    return true;
+}
+
+static bool ggml_strategy_batch(const float* query, int32_t qdim,
+                                 const float* const* node_ptrs, const int32_t* node_dims,
+                                 int32_t n, double* out_scores) {
+    if (!query || qdim <= 0 || !node_ptrs || !node_dims || n <= 0 || !out_scores) return false;
+    /* out_scores here is n doubles (nq=1); eg_ggml_run writes n*nq = n of
+     * them, laid out identically to the single-query contract. */
+    return eg_ggml_run(query, qdim, 1, node_ptrs, node_dims, n, out_scores);
+}
+
+static bool ggml_strategy_batch_multi(const float* queries, int32_t qdim, int32_t nq,
+                                       const float* const* node_ptrs, const int32_t* node_dims,
+                                       int32_t n, double* out_scores) {
+    return eg_ggml_run(queries, qdim, nq, node_ptrs, node_dims, n, out_scores);
+}
+
+static const EgCosineBatchStrategy g_ggml_strategy = {
+    .name = "ggml",
+    .available = ggml_strategy_available,
+    .batch = ggml_strategy_batch,
+    .batch_multi = ggml_strategy_batch_multi,
+};
+
+const EgCosineBatchStrategy* eg_cosine_batch_strategy_ggml(void) {
+    return &g_ggml_strategy;
+}