engram: real Metal batch-cosine kernel, wired into vindex_bench's brute-force oracle
The original brief targeted engram_activate's O(N*D) cosq prescan, but #109 (this branch) already retires that loop algorithmically (HNSW seed selection + lazy memoized cosine) — GPU-accelerating a loop being deleted isn't real work, so that target was dropped rather than forced. Re-investigated for a genuine remaining GPU-shaped call site (not a manufactured one): HNSW insert's candidate-list distance work is bounded- degree (M=24-48) and sequential/adaptive — too fine-grained for a GPU dispatch to pay off. No O(N^2) pairwise cosine pass exists (dedup only checks the K=8 already-selected seed slots). No concurrent multi-query traffic exists (server.el: the soul's curiosity loop is a single in-process caller). vindex_bench.c's brute_topk — the correctness oracle this same PR adds to validate HNSW recall — is the one real, unforced fit: genuine 1-query-vs-N-vectors, embarrassingly parallel, no adaptivity. Adds: - eg_cosine_batch.metal: batched cosine kernel, single- and multi-query variants, same -2.0 dim-mismatch/zero-norm sentinel as eg_cosine(). - eg_metal_cosine.h/.m: C-callable Objective-C bridge. Lazy one-time device/pipeline init, MTLResourceStorageModeShared buffers, returns false on ANY failure so callers fall back to the scalar CPU loop unconditionally — never partial, never throws. - eg_metal_cosine_stub.c: zero-dependency CPU-only implementation for non-Darwin builds (Linux CI) — same symbols, always returns false, no #ifdef needed at any call site. - build_vindex_bench.sh: one-command build, real bridge + Metal frameworks on Darwin, stub everywhere else. vindex_bench.c: brute_topk_metal / brute_topk_metal_batch call the bridge, falling back to the existing CPU brute_topk on any failure or EL_METAL_COSINE=0. The multi-query batched path exists because the first version (one GPU call per query) measured SLOWER than CPU at N~13.7k — it re-uploaded the full N*D matrix every query. Fixed by uploading the matrix once per query batch. Measured against a real nsbx-sandboxed clone of the live store (never :8742/:7770), 13,671 real embedded nodes, dim=768, 300 real queries: BRUTE-FORCE (CPU): 2.013 ms/query BRUTE-METAL (GPU): 0.117 ms/query (17.2x) id-recall vs CPU oracle: 0.9990 over 300 queries same-rank |Δdist|: max 2.98e-07, mean 7.53e-08 (float32 rounding, not a bug) Synthetic scaling sweep (13k -> 50k nodes, same dim/queries) shows the GPU speedup holding (~11x) as N grows toward the mathematical-foundations doc's 1.3M-node target, with CPU brute-force cost growing linearly as expected. Not wired into engram_activate or the daemon build (nsbx's _build_binary) — vindex_bench is a standalone offline tool, not part of the request-serving binary, so no engram_activate/server-latency claim is made here. The bridge is a reusable primitive (single eg_cosine_batch_metal + batched eg_cosine_batch_metal_multi) other call sites can adopt later without re-deriving any of this. Based on feat/reframe-region-setop (PR #109), not dev directly: the only genuine batch-cosine call site (vindex_bench.c) exists solely on this branch. Flagged explicitly in the PR description as a deliberate deviation from the original "base off dev" instruction.
This commit is contained in:
bigmerge
+166
-5
@@ -7,11 +7,26 @@
|
||||
*
|
||||
* Read-only: never opens a socket, never writes the store. Safe on an nsbx clone.
|
||||
*
|
||||
* Build: cc -O2 -std=c11 vindex_bench.c engram_vindex.c -lm -o vindex_bench
|
||||
* Also runs the brute-force oracle a second way, through
|
||||
* eg_cosine_batch_metal() (Apple/Metal only — see eg_metal_cosine.h), and
|
||||
* reports its latency + a correctness check against the CPU oracle
|
||||
* side-by-side with the existing CPU-vs-HNSW numbers. EL_METAL_COSINE=0
|
||||
* forces CPU-only.
|
||||
*
|
||||
* Build (macOS):
|
||||
* cc -O2 -std=c11 -x objective-c -c eg_metal_cosine.m -o eg_metal_cosine.o \
|
||||
* -framework Metal -framework Foundation
|
||||
* cc -O2 -std=c11 vindex_bench.c engram_vindex.c eg_metal_cosine.o -lm \
|
||||
* -framework Metal -framework Foundation -o vindex_bench
|
||||
* Build (Linux / no Metal): omit eg_metal_cosine.o entirely and instead link
|
||||
* a CPU-only stub translation unit that defines eg_cosine_batch_metal() /
|
||||
* eg_cosine_batch_metal_available() returning false — this file never
|
||||
* references Metal directly, only the plain-C header.
|
||||
* Usage: vindex_bench store <neuron.egm> <dim> [nqueries] [k] [ef_csv]
|
||||
* vindex_bench synth <N> [dim] [clusters] [nqueries] [k] [ef_csv]
|
||||
*/
|
||||
#include "engram_vindex.h"
|
||||
#include "eg_metal_cosine.h"
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
@@ -64,6 +79,104 @@ static void brute_topk(const float* data, int n, int dim, const float* q,
|
||||
}
|
||||
}
|
||||
|
||||
/* GPU-accelerated variant of brute_topk: same oracle, same contract, same
|
||||
* output — computes all n distances via eg_cosine_batch_metal() instead of
|
||||
* one C loop, then does the identical top-k selection over the result.
|
||||
*
|
||||
* data is already L2-normalised (vindex_bench's convention throughout), so
|
||||
* eg_cosine()'s general unnormalised cosine and this file's "distance =
|
||||
* 1 - dot" both reduce to the same number here (a unit vector's norm is 1,
|
||||
* so cosine == dot). Passing pre-normalised rows through the general-purpose
|
||||
* batch kernel is deliberate: it proves the SAME primitive that would serve
|
||||
* el_runtime.c's raw/unnormalised embeddings also serves this oracle without
|
||||
* a second code path.
|
||||
*
|
||||
* Returns false (out_ids/out_d untouched) if the GPU path is unavailable or
|
||||
* fails for any reason — caller must fall back to brute_topk(). Never
|
||||
* partial: either the full top-k was computed on GPU, or nothing was. */
|
||||
/* EL_METAL_COSINE: 0/off/false disables the GPU path outright (falls back to
|
||||
* brute_topk() every time), matching el_runtime.c's own gate for the same
|
||||
* env var. Unset or any other value = auto (try Metal, fall back on failure). */
|
||||
static bool g_metal_env_checked = false;
|
||||
static bool g_metal_disabled_by_env = false;
|
||||
static void eg_metal_check_env_once(void){
|
||||
if (g_metal_env_checked) return;
|
||||
g_metal_env_checked = true;
|
||||
const char* v = getenv("EL_METAL_COSINE");
|
||||
if (v && (v[0]=='0' || v[0]=='n' || v[0]=='N' || v[0]=='f' || v[0]=='F'))
|
||||
g_metal_disabled_by_env = true;
|
||||
}
|
||||
static bool brute_topk_metal(const float* data, int n, int dim, const float* q,
|
||||
int k, int* out_ids, float* out_d){
|
||||
eg_metal_check_env_once();
|
||||
if (g_metal_disabled_by_env) return false;
|
||||
|
||||
const float** row_ptrs = malloc((size_t)n * sizeof(float*));
|
||||
int32_t* dims = malloc((size_t)n * sizeof(int32_t));
|
||||
double* scores = malloc((size_t)n * sizeof(double));
|
||||
if (!row_ptrs || !dims || !scores) { free(row_ptrs); free(dims); free(scores); return false; }
|
||||
|
||||
for (int i = 0; i < n; i++) {
|
||||
row_ptrs[i] = data + (size_t)i * dim;
|
||||
dims[i] = dim;
|
||||
}
|
||||
|
||||
bool ok = eg_cosine_batch_metal(q, dim, row_ptrs, dims, n, scores);
|
||||
free(row_ptrs); free(dims);
|
||||
if (!ok) { free(scores); return false; }
|
||||
|
||||
for (int i = 0; i < k; i++) { out_ids[i] = -1; out_d[i] = 3.0f; }
|
||||
for (int i = 0; i < n; i++) {
|
||||
float d = 1.0f - (float)scores[i]; /* same distance convention as brute_topk */
|
||||
if (d >= out_d[k-1]) continue;
|
||||
int p = k - 1;
|
||||
while (p > 0 && out_d[p-1] > d) { out_d[p] = out_d[p-1]; out_ids[p] = out_ids[p-1]; p--; }
|
||||
out_d[p] = d; out_ids[p] = i;
|
||||
}
|
||||
free(scores);
|
||||
return true;
|
||||
}
|
||||
|
||||
/* Batched sibling of brute_topk_metal: computes top-k for ALL nq queries in
|
||||
* ONE eg_cosine_batch_metal_multi() call, uploading node_matrix exactly
|
||||
* once instead of once per query. out_ids/out_d are nq*k, row-major
|
||||
* (query i's results at out_ids+i*k / out_d+i*k) — same layout run_bench
|
||||
* already uses for `gt`/per-query scratch. Returns false (nothing written)
|
||||
* on any failure; caller falls back to the per-query CPU brute_topk loop. */
|
||||
static bool brute_topk_metal_batch(const float* data, int n, int dim,
|
||||
const float* queries, int nq,
|
||||
int k, int* out_ids, float* out_d){
|
||||
eg_metal_check_env_once();
|
||||
if (g_metal_disabled_by_env) return false;
|
||||
|
||||
const float** row_ptrs = malloc((size_t)n * sizeof(float*));
|
||||
int32_t* dims = malloc((size_t)n * sizeof(int32_t));
|
||||
double* scores = malloc((size_t)nq * (size_t)n * sizeof(double));
|
||||
if (!row_ptrs || !dims || !scores) { free(row_ptrs); free(dims); free(scores); return false; }
|
||||
|
||||
for (int i = 0; i < n; i++) { row_ptrs[i] = data + (size_t)i * dim; dims[i] = dim; }
|
||||
|
||||
bool ok = eg_cosine_batch_metal_multi(queries, dim, nq, row_ptrs, dims, n, scores);
|
||||
free(row_ptrs); free(dims);
|
||||
if (!ok) { free(scores); return false; }
|
||||
|
||||
for (int qi = 0; qi < nq; qi++) {
|
||||
int* ids = out_ids + (size_t)qi * k;
|
||||
float* ds = out_d + (size_t)qi * k;
|
||||
const double* srow = scores + (size_t)qi * n;
|
||||
for (int i = 0; i < k; i++) { ids[i] = -1; ds[i] = 3.0f; }
|
||||
for (int i = 0; i < n; i++) {
|
||||
float d = 1.0f - (float)srow[i];
|
||||
if (d >= ds[k-1]) continue;
|
||||
int p = k - 1;
|
||||
while (p > 0 && ds[p-1] > d) { ds[p] = ds[p-1]; ids[p] = ids[p-1]; p--; }
|
||||
ds[p] = d; ids[p] = i;
|
||||
}
|
||||
}
|
||||
free(scores);
|
||||
return true;
|
||||
}
|
||||
|
||||
/* recall@k: |brute_topk ∩ hnsw_topk| / k. Both are id arrays of length k. */
|
||||
static double recall_at_k(const int* gt, const uint64_t* ann, int nann, int k){
|
||||
int hit = 0;
|
||||
@@ -142,13 +255,61 @@ static void run_bench(const char* label, float* data, int n, int dim,
|
||||
l2norm(dst, dim);
|
||||
}
|
||||
|
||||
/* ground truth: brute-force top-k for every query (also the oracle latency). */
|
||||
/* ground truth: brute-force top-k for every query (also the oracle latency).
|
||||
* gd is nq*k (one real slot per query, not a shared scratch buffer) so the
|
||||
* GPU comparison below can diff against every query's actual distances,
|
||||
* not just whichever query happened to run last. */
|
||||
int* gt = malloc((size_t)nq*k*sizeof(int));
|
||||
float* gd = malloc((size_t)k*sizeof(float));
|
||||
float* gd = malloc((size_t)nq*k*sizeof(float));
|
||||
double tb0 = now_s();
|
||||
for (int i=0;i<nq;i++) brute_topk(data, n, dim, qv+(size_t)i*dim, k, gt+(size_t)i*k, gd);
|
||||
for (int i=0;i<nq;i++) brute_topk(data, n, dim, qv+(size_t)i*dim, k, gt+(size_t)i*k, gd+(size_t)i*k);
|
||||
double brute_ms = (now_s()-tb0)*1000.0/nq;
|
||||
printf("BRUTE-FORCE : %8.3f ms/query (oracle; O(N*D))\n", brute_ms);
|
||||
printf("BRUTE-FORCE : %8.3f ms/query (oracle; O(N*D), CPU)\n", brute_ms);
|
||||
|
||||
/* GPU-accelerated oracle: SAME nq queries, SAME top-k contract, via ONE
|
||||
* eg_cosine_batch_metal_multi() call (uploads node_matrix once, not once
|
||||
* per query — see brute_topk_metal_batch). Run only if the GPU path is
|
||||
* actually available (checked once) — never fabricated, never assumed.
|
||||
* Verified against the CPU ground truth computed above: id-recall across
|
||||
* ALL nq queries, plus the actual max distance delta across every
|
||||
* (query, rank) pair that was compared — not a single spot check. */
|
||||
eg_metal_check_env_once();
|
||||
if (!g_metal_disabled_by_env && eg_cosine_batch_metal_available()) {
|
||||
int* gtm = malloc((size_t)nq*k*sizeof(int));
|
||||
float* gdm = malloc((size_t)nq*k*sizeof(float));
|
||||
double tm0 = now_s();
|
||||
bool ok = brute_topk_metal_batch(data, n, dim, qv, nq, k, gtm, gdm);
|
||||
double metal_ms = (now_s()-tm0)*1000.0/nq;
|
||||
if (ok) {
|
||||
double rec_sum = 0; double max_ddiff = 0; double sum_ddiff = 0; int compared = 0;
|
||||
for (int i=0;i<nq;i++) {
|
||||
const int* ids_gt = gt+(size_t)i*k;
|
||||
const float* d_gt = gd+(size_t)i*k;
|
||||
const int* ids_m = gtm+(size_t)i*k;
|
||||
const float* d_m = gdm+(size_t)i*k;
|
||||
uint64_t idset[512]; int m = (k<512)?k:512;
|
||||
for (int j=0;j<m;j++) idset[j] = (uint64_t)ids_m[j];
|
||||
rec_sum += recall_at_k(ids_gt, idset, m, k);
|
||||
/* same-rank distance delta — valid whenever both sides agree on
|
||||
* the id at that rank (true almost always, given ~100% recall;
|
||||
* a rank where they disagree isn't a meaningful delta to diff). */
|
||||
for (int j=0;j<k;j++) {
|
||||
if (ids_gt[j] == ids_m[j]) {
|
||||
double diff = fabs((double)d_gt[j]-(double)d_m[j]);
|
||||
if (diff>max_ddiff) max_ddiff=diff;
|
||||
sum_ddiff += diff; compared++;
|
||||
}
|
||||
}
|
||||
}
|
||||
printf("BRUTE-METAL : %8.3f ms/query (%.1fx vs CPU brute; id-recall %.4f vs CPU oracle over %d queries; same-rank |Δdist|: max %.2e, mean %.2e over %d compared)\n",
|
||||
metal_ms, brute_ms/metal_ms, rec_sum/nq, nq, max_ddiff, compared?sum_ddiff/compared:0.0, compared);
|
||||
} else {
|
||||
printf("BRUTE-METAL : GPU batch call failed/unavailable mid-run — skipped\n");
|
||||
}
|
||||
free(gtm); free(gdm);
|
||||
} else {
|
||||
printf("BRUTE-METAL : no Metal device/pipeline available — CPU-only\n");
|
||||
}
|
||||
|
||||
/* HNSW at each ef. */
|
||||
uint64_t* aid = malloc((size_t)k*sizeof(uint64_t));
|
||||
|
||||
Reference in New Issue
Block a user