diff --git a/lang/runtime/build_vindex_bench.sh b/lang/runtime/build_vindex_bench.sh new file mode 100755 index 0000000..37cbd5f --- /dev/null +++ b/lang/runtime/build_vindex_bench.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# build_vindex_bench.sh — build the vindex_bench oracle/proof harness, with the +# real Metal batch-cosine bridge on Darwin and a zero-dependency CPU stub +# everywhere else. Mirrors the two-step recipe documented in vindex_bench.c's +# own header comment; this script exists so that recipe is one command, not a +# copy-pasted paragraph. +# +# Usage: ./build_vindex_bench.sh [output_path] +set -euo pipefail +cd "$(dirname "$0")" +OUT="${1:-./vindex_bench}" +CC="${CC:-cc}" + +if [ "$(uname -s)" = "Darwin" ]; then + echo "== Darwin: building with the real Metal bridge ==" + "$CC" -O2 -std=c11 -x objective-c -c eg_metal_cosine.m -o /tmp/eg_metal_cosine.o \ + -framework Metal -framework Foundation + "$CC" -O2 -std=c11 -w vindex_bench.c engram_vindex.c /tmp/eg_metal_cosine.o -lm \ + -framework Metal -framework Foundation -o "$OUT" +else + echo "== non-Darwin: building with the CPU-only stub (no Metal) ==" + "$CC" -O2 -std=c11 -w vindex_bench.c engram_vindex.c eg_metal_cosine_stub.c -lm -o "$OUT" +fi + +echo "built: $OUT" diff --git a/lang/runtime/eg_cosine_batch.metal b/lang/runtime/eg_cosine_batch.metal new file mode 100644 index 0000000..b6868c6 --- /dev/null +++ b/lang/runtime/eg_cosine_batch.metal @@ -0,0 +1,156 @@ +/* eg_cosine_batch.metal — batched cosine similarity, one query vs N node vectors. + * + * GPU-shaped counterpart to eg_cosine() in el_runtime.c: same math, same + * dim-mismatch/zero-norm sentinel (-2.0), applied to N independent rows in + * parallel instead of one pair at a time in a CPU loop. + * + * Semantics MUST match eg_cosine() exactly: + * - inputs are raw, UNNORMALIZED vectors (nomic-embed-text magnitudes are + * not 1.0) — this kernel computes the full dot/(|a|*|b|) cosine, not a + * plain dot product. + * - a node whose declared dim differs from the query dim, or whose norm is + * zero, scores exactly -2.0 (below any valid cosine in [-1,1]), so a + * caller doing `if (score > threshold)` behaves identically whether the + * scalar or the batched path filled the array. + * + * Precision: Apple GPUs do not support double in Metal Shading Language — + * everything here is float32. eg_cosine accumulates in CPU double, but its + * *inputs* are float32 embeddings, so the achievable precision ceiling is + * bounded by the input data regardless of accumulator width. To keep the + * float32 reduction from drifting relative to the double-accumulated CPU + * result across dim=768 terms, each thread accumulates with 4 independent + * partial sums (unrolled) rather than one running scalar — the same + * error-reduction trick already used by the CPU brute-force loop in + * vindex_bench.c. The measured float-vs-double delta is reported in the PR + * description; this is not assumed to be "close enough" without measurement. + */ +#include +using namespace metal; + +/* Per-dispatch invariants. `dim` is the query's dimensionality — the + * dimensionality every comparable node vector must match. */ +struct EgCosineParams { + uint n; /* number of node rows */ + uint dim; /* vector width (both query and node rows are `dim` wide in + * the packed buffer; node_dims[] carries each node's REAL + * embedded dim for the mismatch check) */ +}; + +/* One thread per node row. node_matrix is n*dim floats, row-major, packed at + * `dim` stride regardless of a row's real dim (the CPU side zero-pads or + * skips packing rows that don't match — see eg_cosine_batch_metal in + * eg_metal_cosine.m for the exact packing contract). node_dims[i] is the + * node's true emb_dim, used only for the mismatch sentinel — never used to + * index, since every row is packed at uniform `dim` stride. */ +kernel void eg_cosine_batch_kernel( + device const float* query [[buffer(0)]], + device const float* node_matrix [[buffer(1)]], + device const int* node_dims [[buffer(2)]], + constant EgCosineParams& p [[buffer(3)]], + device float* out_scores [[buffer(4)]], + uint gid [[thread_position_in_grid]]) +{ + if (gid >= p.n) return; + + if (node_dims[gid] != int(p.dim)) { + out_scores[gid] = -2.0f; + return; + } + + device const float* row = node_matrix + (uint64_t)gid * (uint64_t)p.dim; + + /* 4-way partial accumulation — same shape as vindex_bench.c's brute_topk + * unroll, done here for float32 accuracy rather than raw throughput. */ + float dot0 = 0.0f, dot1 = 0.0f, dot2 = 0.0f, dot3 = 0.0f; + float na0 = 0.0f, na1 = 0.0f, na2 = 0.0f, na3 = 0.0f; + float nb0 = 0.0f, nb1 = 0.0f, nb2 = 0.0f, nb3 = 0.0f; + + uint d = 0; + uint dim4 = p.dim & ~3u; + for (; d < dim4; d += 4) { + float a0 = row[d], b0 = query[d]; + float a1 = row[d+1], b1 = query[d+1]; + float a2 = row[d+2], b2 = query[d+2]; + float a3 = row[d+3], b3 = query[d+3]; + dot0 += a0*b0; dot1 += a1*b1; dot2 += a2*b2; dot3 += a3*b3; + na0 += a0*a0; na1 += a1*a1; na2 += a2*a2; na3 += a3*a3; + nb0 += b0*b0; nb1 += b1*b1; nb2 += b2*b2; nb3 += b3*b3; + } + float dot = (dot0 + dot1) + (dot2 + dot3); + float na = (na0 + na1) + (na2 + na3); + float nb = (nb0 + nb1) + (nb2 + nb3); + for (; d < p.dim; d++) { + float a = row[d], b = query[d]; + dot += a*b; na += a*a; nb += b*b; + } + + if (na <= 0.0f || nb <= 0.0f) { + out_scores[gid] = -2.0f; + return; + } + out_scores[gid] = dot / sqrt(na * nb); +} + +/* ── multi-query variant ────────────────────────────────────────────────── + * Same per-pair math as eg_cosine_batch_kernel, but amortizes ONE upload of + * node_matrix (the expensive part at real store size — 13k*768 floats is + * ~42MB) across `nq` queries instead of re-uploading it once per query. + * Measured need: a naive one-query-at-a-time loop calling the single-query + * kernel nq times was SLOWER than the CPU oracle at N≈13.7k (re-gather + + * re-upload dominated the actual compute) — this is the fix, not a + * hypothetical optimization. + * + * 2D grid: x = node index [0,n), y = query index [0,nq). out_scores is + * nq*n, row-major by query (out_scores[qid*n + nid]). */ +struct EgCosineMultiParams { uint n; uint dim; uint nq; }; + +kernel void eg_cosine_batch_multi_kernel( + device const float* queries [[buffer(0)]], /* nq*dim */ + device const float* node_matrix [[buffer(1)]], /* n*dim */ + device const int* node_dims [[buffer(2)]], /* n */ + constant EgCosineMultiParams& p [[buffer(3)]], + device float* out_scores [[buffer(4)]], /* nq*n */ + uint2 gid [[thread_position_in_grid]]) +{ + uint nid = gid.x, qid = gid.y; + if (nid >= p.n || qid >= p.nq) return; + + uint64_t out_idx = (uint64_t)qid * (uint64_t)p.n + (uint64_t)nid; + + if (node_dims[nid] != int(p.dim)) { + out_scores[out_idx] = -2.0f; + return; + } + + device const float* row = node_matrix + (uint64_t)nid * (uint64_t)p.dim; + device const float* query = queries + (uint64_t)qid * (uint64_t)p.dim; + + float dot0 = 0.0f, dot1 = 0.0f, dot2 = 0.0f, dot3 = 0.0f; + float na0 = 0.0f, na1 = 0.0f, na2 = 0.0f, na3 = 0.0f; + float nb0 = 0.0f, nb1 = 0.0f, nb2 = 0.0f, nb3 = 0.0f; + + uint d = 0; + uint dim4 = p.dim & ~3u; + for (; d < dim4; d += 4) { + float a0 = row[d], b0 = query[d]; + float a1 = row[d+1], b1 = query[d+1]; + float a2 = row[d+2], b2 = query[d+2]; + float a3 = row[d+3], b3 = query[d+3]; + dot0 += a0*b0; dot1 += a1*b1; dot2 += a2*b2; dot3 += a3*b3; + na0 += a0*a0; na1 += a1*a1; na2 += a2*a2; na3 += a3*a3; + nb0 += b0*b0; nb1 += b1*b1; nb2 += b2*b2; nb3 += b3*b3; + } + float dot = (dot0 + dot1) + (dot2 + dot3); + float na = (na0 + na1) + (na2 + na3); + float nb = (nb0 + nb1) + (nb2 + nb3); + for (; d < p.dim; d++) { + float a = row[d], b = query[d]; + dot += a*b; na += a*a; nb += b*b; + } + + if (na <= 0.0f || nb <= 0.0f) { + out_scores[out_idx] = -2.0f; + return; + } + out_scores[out_idx] = dot / sqrt(na * nb); +} diff --git a/lang/runtime/eg_metal_cosine.h b/lang/runtime/eg_metal_cosine.h new file mode 100644 index 0000000..fc625a9 --- /dev/null +++ b/lang/runtime/eg_metal_cosine.h @@ -0,0 +1,85 @@ +/* eg_metal_cosine.h — C-callable bridge to the Metal batched-cosine kernel. + * + * Plain C11 header, safe to #include from el_runtime.c / vindex_bench.c on + * every platform. The implementation (eg_metal_cosine.m) only exists on + * Apple builds; on any other platform (or if Metal init fails for any + * reason at all — no supported GPU, shader compile error, OOM, sandboxing, + * whatever) eg_cosine_batch_metal() returns false and writes nothing, and + * the caller MUST fall back to its existing scalar per-node loop + * unconditionally. This function must never be allowed to crash or hang + * the engram. + */ +#ifndef EG_METAL_COSINE_H +#define EG_METAL_COSINE_H + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* Batched cosine similarity: one query vector against `n` node vectors. + * + * query — qdim floats, the query embedding. Raw/unnormalized. + * qdim — query dimensionality (e.g. 768 for nomic-embed-text). + * node_ptrs — array of n pointers, node_ptrs[i] pointing at a (possibly + * differently-owned, possibly NULL) float vector for node i. + * NOT required to be contiguous — this function performs the + * gather into a packed row-major matrix internally, exactly + * mirroring how EngramNode.emb is one malloc per node. + * node_dims — array of n ints, node_dims[i] = that node's real emb_dim + * (0 or mismatched vs qdim ⇒ that node scores -2.0, matching + * eg_cosine's null/dim-mismatch/zero-norm sentinel exactly). + * n — number of nodes. + * out_scores — caller-owned array of n doubles; out_scores[i] is filled + * with the cosine similarity of node i against query, or + * -2.0 for a null/dim-mismatched/zero-norm node — bit-for-bit + * the same contract as eg_cosine(node_ptrs[i], query, qdim). + * + * Returns true iff the GPU path ran and out_scores was fully populated. + * Returns false (out_scores left untouched) on ANY failure or unavailability + * — no Metal-capable device, shader compile failure, allocation failure, + * n<=0, qdim<=0, null query/node_ptrs/node_dims/out_scores. Never partial: + * either every element of out_scores was written, or none were. + */ +bool eg_cosine_batch_metal(const float* query, int32_t qdim, + const float* const* node_ptrs, + const int32_t* node_dims, + int32_t n, + double* out_scores); + +/* True iff a Metal device + compiled pipeline is available right now (cheap + * after the first call — cached). Purely informational (e.g. for a startup + * log line or /api/stats field); callers should still treat a false return + * from eg_cosine_batch_metal itself as the authoritative fallback signal. */ +bool eg_cosine_batch_metal_available(void); + +/* Multi-query batched cosine: nq query vectors against the SAME n node + * vectors, in one call. Uploads node_matrix once and reuses it for every + * query, instead of nq separate eg_cosine_batch_metal() calls each paying + * the full gather+upload cost — measured necessary: at N≈13.7k/dim=768, + * repeating the single-query call per query was slower than the CPU + * baseline; batching queries together is what makes the GPU path a real win + * at this shape. Use this whenever multiple queries will run against an + * unchanged (or rarely-changing) node population; use the single-query + * function above for a genuinely one-off comparison. + * + * queries — nq*qdim floats, row-major (query i at queries+i*qdim). + * out_scores — caller-owned nq*n doubles, row-major + * (out_scores[i*n+j] = cosine(queries[i], node j)), same + * -2.0 sentinel semantics as eg_cosine_batch_metal. + * + * Returns true iff the GPU path ran and out_scores was fully populated + * (all nq*n entries); false (untouched) on any failure/unavailability. */ +bool eg_cosine_batch_metal_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); + +#ifdef __cplusplus +} +#endif + +#endif /* EG_METAL_COSINE_H */ diff --git a/lang/runtime/eg_metal_cosine.m b/lang/runtime/eg_metal_cosine.m new file mode 100644 index 0000000..94ef383 --- /dev/null +++ b/lang/runtime/eg_metal_cosine.m @@ -0,0 +1,337 @@ +/* eg_metal_cosine.m — Objective-C bridge exposing the Metal batched-cosine + * kernel (eg_cosine_batch.metal) as a plain C function. + * + * Apple-only (Metal has no other platform). This file is excluded from the + * build entirely on non-Darwin — see tools/neuron-sandbox/nsbx's + * _build_binary(), which only compiles/links this file and adds + * -framework Metal -framework Foundation when `uname` is Darwin. On Linux + * (the CI runner), eg_metal_cosine.h is still included by callers, but + * eg_cosine_batch_metal() is never linked in from this file — callers must + * always be prepared for the "GPU path unavailable" fallback, which is also + * exactly what happens here on Apple hardware with no usable GPU. + * + * Design: + * - Device/queue/pipeline are created lazily, once, and cached in static + * globals — every call after the first only allocates buffers + submits. + * - The Metal shader source is embedded as a C string literal (kMetalSrc + * below) rather than loaded from a file at runtime or shipped as a + * precompiled .metallib. Chosen over newLibraryWithFile: /a .metallib + * because the engram binary can be invoked from an arbitrary working + * directory (launchd job, nsbx sandbox, CI) and a file-path shader would + * be one relocation away from silently falling back to CPU for reasons + * that have nothing to do with Metal availability. Embedding costs one + * runtime shader compile (~tens of ms) on first use, amortized over the + * process lifetime, in exchange for a genuinely self-contained binary. + * Source of truth for review/tooling is eg_cosine_batch.metal — this + * string MUST be kept byte-identical to that file (a comment marks both + * ends of the copy). + * - Buffers use MTLResourceStorageModeShared: on Apple Silicon's unified + * memory, CPU and GPU read the same physical pages, so filling a buffer + * is a plain memcpy and there is no separate "upload" step. + * - ANY failure at ANY step (no device, pipeline compile error, buffer + * allocation failure, bad args) returns false and leaves out_scores + * untouched. This function is called from the request-handling hot path + * of a long-lived daemon — it must never throw, crash, or hang it. + */ +#import +#import +#include "eg_metal_cosine.h" +#include +#include + +/* ── BEGIN embedded shader source (keep in sync with eg_cosine_batch.metal) ── */ +static const char* kEgCosineBatchMetalSrc = +"#include \n" +"using namespace metal;\n" +"struct EgCosineParams { uint n; uint dim; };\n" +"kernel void eg_cosine_batch_kernel(\n" +" device const float* query [[buffer(0)]],\n" +" device const float* node_matrix [[buffer(1)]],\n" +" device const int* node_dims [[buffer(2)]],\n" +" constant EgCosineParams& p [[buffer(3)]],\n" +" device float* out_scores [[buffer(4)]],\n" +" uint gid [[thread_position_in_grid]])\n" +"{\n" +" if (gid >= p.n) return;\n" +" if (node_dims[gid] != int(p.dim)) { out_scores[gid] = -2.0f; return; }\n" +" device const float* row = node_matrix + (uint64_t)gid * (uint64_t)p.dim;\n" +" float dot0 = 0.0f, dot1 = 0.0f, dot2 = 0.0f, dot3 = 0.0f;\n" +" float na0 = 0.0f, na1 = 0.0f, na2 = 0.0f, na3 = 0.0f;\n" +" float nb0 = 0.0f, nb1 = 0.0f, nb2 = 0.0f, nb3 = 0.0f;\n" +" uint d = 0;\n" +" uint dim4 = p.dim & ~3u;\n" +" for (; d < dim4; d += 4) {\n" +" float a0 = row[d], b0 = query[d];\n" +" float a1 = row[d+1], b1 = query[d+1];\n" +" float a2 = row[d+2], b2 = query[d+2];\n" +" float a3 = row[d+3], b3 = query[d+3];\n" +" dot0 += a0*b0; dot1 += a1*b1; dot2 += a2*b2; dot3 += a3*b3;\n" +" na0 += a0*a0; na1 += a1*a1; na2 += a2*a2; na3 += a3*a3;\n" +" nb0 += b0*b0; nb1 += b1*b1; nb2 += b2*b2; nb3 += b3*b3;\n" +" }\n" +" float dot = (dot0 + dot1) + (dot2 + dot3);\n" +" float na = (na0 + na1) + (na2 + na3);\n" +" float nb = (nb0 + nb1) + (nb2 + nb3);\n" +" for (; d < p.dim; d++) {\n" +" float a = row[d], b = query[d];\n" +" dot += a*b; na += a*a; nb += b*b;\n" +" }\n" +" if (na <= 0.0f || nb <= 0.0f) { out_scores[gid] = -2.0f; return; }\n" +" out_scores[gid] = dot / sqrt(na * nb);\n" +"}\n" +"struct EgCosineMultiParams { uint n; uint dim; uint nq; };\n" +"kernel void eg_cosine_batch_multi_kernel(\n" +" device const float* queries [[buffer(0)]],\n" +" device const float* node_matrix [[buffer(1)]],\n" +" device const int* node_dims [[buffer(2)]],\n" +" constant EgCosineMultiParams& p [[buffer(3)]],\n" +" device float* out_scores [[buffer(4)]],\n" +" uint2 gid [[thread_position_in_grid]])\n" +"{\n" +" uint nid = gid.x, qid = gid.y;\n" +" if (nid >= p.n || qid >= p.nq) return;\n" +" uint64_t out_idx = (uint64_t)qid * (uint64_t)p.n + (uint64_t)nid;\n" +" if (node_dims[nid] != int(p.dim)) { out_scores[out_idx] = -2.0f; return; }\n" +" device const float* row = node_matrix + (uint64_t)nid * (uint64_t)p.dim;\n" +" device const float* query = queries + (uint64_t)qid * (uint64_t)p.dim;\n" +" float dot0 = 0.0f, dot1 = 0.0f, dot2 = 0.0f, dot3 = 0.0f;\n" +" float na0 = 0.0f, na1 = 0.0f, na2 = 0.0f, na3 = 0.0f;\n" +" float nb0 = 0.0f, nb1 = 0.0f, nb2 = 0.0f, nb3 = 0.0f;\n" +" uint d = 0;\n" +" uint dim4 = p.dim & ~3u;\n" +" for (; d < dim4; d += 4) {\n" +" float a0 = row[d], b0 = query[d];\n" +" float a1 = row[d+1], b1 = query[d+1];\n" +" float a2 = row[d+2], b2 = query[d+2];\n" +" float a3 = row[d+3], b3 = query[d+3];\n" +" dot0 += a0*b0; dot1 += a1*b1; dot2 += a2*b2; dot3 += a3*b3;\n" +" na0 += a0*a0; na1 += a1*a1; na2 += a2*a2; na3 += a3*a3;\n" +" nb0 += b0*b0; nb1 += b1*b1; nb2 += b2*b2; nb3 += b3*b3;\n" +" }\n" +" float dot = (dot0 + dot1) + (dot2 + dot3);\n" +" float na = (na0 + na1) + (na2 + na3);\n" +" float nb = (nb0 + nb1) + (nb2 + nb3);\n" +" for (; d < p.dim; d++) {\n" +" float a = row[d], b = query[d];\n" +" dot += a*b; na += a*a; nb += b*b;\n" +" }\n" +" if (na <= 0.0f || nb <= 0.0f) { out_scores[out_idx] = -2.0f; return; }\n" +" out_scores[out_idx] = dot / sqrt(na * nb);\n" +"}\n"; +/* ── END embedded shader source ── */ + +typedef struct EgCosineParamsC { uint32_t n; uint32_t dim; } EgCosineParamsC; +typedef struct EgCosineMultiParamsC { uint32_t n; uint32_t dim; uint32_t nq; } EgCosineMultiParamsC; + +static id g_device = nil; +static id g_queue = nil; +static id g_pipeline = nil; /* single-query kernel */ +static id g_pipeline_multi = nil; /* multi-query kernel */ +static bool g_init_attempted = false; +static bool g_init_ok = false; + +/* Lazy, one-time setup. Never throws — every Metal call here is the + * "returns nil/NSError on failure" flavor, not an exception-throwing one. */ +static bool eg_metal_ensure_init(void) { + if (g_init_attempted) return g_init_ok; + g_init_attempted = true; + + @autoreleasepool { + id dev = MTLCreateSystemDefaultDevice(); + if (!dev) return false; + + id q = [dev newCommandQueue]; + if (!q) return false; + + NSError* err = nil; + NSString* src = [NSString stringWithUTF8String:kEgCosineBatchMetalSrc]; + MTLCompileOptions* opts = [MTLCompileOptions new]; + id lib = [dev newLibraryWithSource:src options:opts error:&err]; + if (!lib) return false; + + id fn = [lib newFunctionWithName:@"eg_cosine_batch_kernel"]; + if (!fn) return false; + id pipe = [dev newComputePipelineStateWithFunction:fn error:&err]; + if (!pipe) return false; + + id fnMulti = [lib newFunctionWithName:@"eg_cosine_batch_multi_kernel"]; + if (!fnMulti) return false; + id pipeMulti = [dev newComputePipelineStateWithFunction:fnMulti error:&err]; + if (!pipeMulti) return false; + + g_device = dev; + g_queue = q; + g_pipeline = pipe; + g_pipeline_multi = pipeMulti; + g_init_ok = true; + return true; + } +} + +bool eg_cosine_batch_metal_available(void) { + return eg_metal_ensure_init(); +} + +bool eg_cosine_batch_metal(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; + if (!eg_metal_ensure_init()) return false; + + @autoreleasepool { + const size_t dim = (size_t)qdim; + const size_t nu = (size_t)n; + + /* Gather into a packed row-major matrix — EngramNode.emb is one + * malloc per node, not a contiguous array, so this copy is + * unavoidable regardless of backend. Rows whose real dim doesn't + * match qdim are zero-filled (harmless: the kernel sentinels them + * via node_dims before ever reading the row). */ + float* matrix = (float*)calloc(nu * dim, sizeof(float)); + int32_t* dims_i32 = (int32_t*)malloc(nu * sizeof(int32_t)); + if (!matrix || !dims_i32) { free(matrix); free(dims_i32); return false; } + + for (size_t i = 0; i < nu; i++) { + dims_i32[i] = node_dims[i]; + if (node_ptrs[i] && node_dims[i] == qdim) { + memcpy(matrix + i * dim, node_ptrs[i], dim * sizeof(float)); + } + /* else: leave zero-filled; node_dims[i] != qdim (or missing) + * makes the kernel sentinel it to -2.0 without reading the row. */ + } + + id bufQuery = [g_device newBufferWithBytes:query + length:dim * sizeof(float) + options:MTLResourceStorageModeShared]; + id bufMatrix = [g_device newBufferWithBytes:matrix + length:nu * dim * sizeof(float) + options:MTLResourceStorageModeShared]; + id bufDims = [g_device newBufferWithBytes:dims_i32 + length:nu * sizeof(int32_t) + options:MTLResourceStorageModeShared]; + EgCosineParamsC params = { (uint32_t)nu, (uint32_t)dim }; + id bufParams = [g_device newBufferWithBytes:¶ms + length:sizeof(params) + options:MTLResourceStorageModeShared]; + id bufOut = [g_device newBufferWithLength:nu * sizeof(float) + options:MTLResourceStorageModeShared]; + + free(matrix); free(dims_i32); + + if (!bufQuery || !bufMatrix || !bufDims || !bufParams || !bufOut) return false; + + id cmd = [g_queue commandBuffer]; + if (!cmd) return false; + id enc = [cmd computeCommandEncoder]; + if (!enc) return false; + + [enc setComputePipelineState:g_pipeline]; + [enc setBuffer:bufQuery offset:0 atIndex:0]; + [enc setBuffer:bufMatrix offset:0 atIndex:1]; + [enc setBuffer:bufDims offset:0 atIndex:2]; + [enc setBuffer:bufParams offset:0 atIndex:3]; + [enc setBuffer:bufOut offset:0 atIndex:4]; + + NSUInteger tgSize = g_pipeline.maxTotalThreadsPerThreadgroup; + if (tgSize > 256) tgSize = 256; + if (tgSize < 1) tgSize = 1; + MTLSize gridSize = MTLSizeMake(nu, 1, 1); + MTLSize threadgroupSize = MTLSizeMake(tgSize, 1, 1); + [enc dispatchThreads:gridSize threadsPerThreadgroup:threadgroupSize]; + [enc endEncoding]; + + [cmd commit]; + [cmd waitUntilCompleted]; + + if (cmd.status != MTLCommandBufferStatusCompleted) return false; + + const float* results = (const float*)bufOut.contents; + if (!results) return false; + for (size_t i = 0; i < nu; i++) out_scores[i] = (double)results[i]; + return true; + } +} + +bool eg_cosine_batch_metal_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) { + if (!queries || qdim <= 0 || nq <= 0 || !node_ptrs || !node_dims || n <= 0 || !out_scores) return false; + if (!eg_metal_ensure_init()) return false; + + @autoreleasepool { + const size_t dim = (size_t)qdim; + const size_t nu = (size_t)n; + const size_t nqu = (size_t)nq; + + float* matrix = (float*)calloc(nu * dim, sizeof(float)); + int32_t* dims_i32 = (int32_t*)malloc(nu * sizeof(int32_t)); + if (!matrix || !dims_i32) { free(matrix); free(dims_i32); return false; } + + for (size_t i = 0; i < nu; i++) { + dims_i32[i] = node_dims[i]; + if (node_ptrs[i] && node_dims[i] == qdim) { + memcpy(matrix + i * dim, node_ptrs[i], dim * sizeof(float)); + } + } + + /* This is the ONE upload of node_matrix for the whole nq-query batch — + * the fix for the measured re-upload-per-query slowdown. */ + id bufMatrix = [g_device newBufferWithBytes:matrix + length:nu * dim * sizeof(float) + options:MTLResourceStorageModeShared]; + id bufDims = [g_device newBufferWithBytes:dims_i32 + length:nu * sizeof(int32_t) + options:MTLResourceStorageModeShared]; + id bufQueries = [g_device newBufferWithBytes:queries + length:nqu * dim * sizeof(float) + options:MTLResourceStorageModeShared]; + EgCosineMultiParamsC params = { (uint32_t)nu, (uint32_t)dim, (uint32_t)nqu }; + id bufParams = [g_device newBufferWithBytes:¶ms + length:sizeof(params) + options:MTLResourceStorageModeShared]; + id bufOut = [g_device newBufferWithLength:nqu * nu * sizeof(float) + options:MTLResourceStorageModeShared]; + + free(matrix); free(dims_i32); + + if (!bufMatrix || !bufDims || !bufQueries || !bufParams || !bufOut) return false; + + id cmd = [g_queue commandBuffer]; + if (!cmd) return false; + id enc = [cmd computeCommandEncoder]; + if (!enc) return false; + + [enc setComputePipelineState:g_pipeline_multi]; + [enc setBuffer:bufQueries offset:0 atIndex:0]; + [enc setBuffer:bufMatrix offset:0 atIndex:1]; + [enc setBuffer:bufDims offset:0 atIndex:2]; + [enc setBuffer:bufParams offset:0 atIndex:3]; + [enc setBuffer:bufOut offset:0 atIndex:4]; + + /* 2D dispatch: x over nodes, y over queries. Threadgroup width picked + * from the pipeline's own limit, height fixed at 1 — nq is typically + * small (tens to low hundreds) relative to n (thousands+), so tiling + * the wide axis (n) is what matters for occupancy. */ + NSUInteger tgWidth = g_pipeline_multi.maxTotalThreadsPerThreadgroup; + if (tgWidth > 256) tgWidth = 256; + if (tgWidth < 1) tgWidth = 1; + MTLSize gridSize = MTLSizeMake(nu, nqu, 1); + MTLSize threadgroupSize = MTLSizeMake(tgWidth, 1, 1); + [enc dispatchThreads:gridSize threadsPerThreadgroup:threadgroupSize]; + [enc endEncoding]; + + [cmd commit]; + [cmd waitUntilCompleted]; + + if (cmd.status != MTLCommandBufferStatusCompleted) return false; + + const float* results = (const float*)bufOut.contents; + if (!results) return false; + for (size_t i = 0; i < nqu * nu; i++) out_scores[i] = (double)results[i]; + return true; + } +} diff --git a/lang/runtime/eg_metal_cosine_stub.c b/lang/runtime/eg_metal_cosine_stub.c new file mode 100644 index 0000000..7d7af4e --- /dev/null +++ b/lang/runtime/eg_metal_cosine_stub.c @@ -0,0 +1,36 @@ +/* eg_metal_cosine_stub.c — plain-C, zero-dependency implementation of the + * eg_metal_cosine.h contract for platforms without Metal (Linux CI, or any + * build that simply chooses not to link the real Objective-C bridge). + * + * Always returns false / unavailable. Callers already treat that as "fall + * back to the CPU path" unconditionally — this file exists so that exactly + * one of {eg_metal_cosine.m, eg_metal_cosine_stub.c} is linked per build, + * selected by the build script (Darwin → the real bridge + Metal frameworks; + * everything else → this stub, no framework flags, no Objective-C compiler + * needed), and el_runtime.c / vindex_bench.c never need an #ifdef to call + * eg_cosine_batch_metal() — the symbol always exists, its behavior is what + * varies by platform. + */ +#include "eg_metal_cosine.h" + +bool eg_cosine_batch_metal_available(void) { + return false; +} + +bool eg_cosine_batch_metal(const float* query, int32_t qdim, + const float* const* node_ptrs, + const int32_t* node_dims, + int32_t n, + double* out_scores) { + (void)query; (void)qdim; (void)node_ptrs; (void)node_dims; (void)n; (void)out_scores; + return false; +} + +bool eg_cosine_batch_metal_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) { + (void)queries; (void)qdim; (void)nq; (void)node_ptrs; (void)node_dims; (void)n; (void)out_scores; + return false; +} diff --git a/lang/runtime/vindex_bench.c b/lang/runtime/vindex_bench.c index ef0c482..5c7669a 100644 --- a/lang/runtime/vindex_bench.c +++ b/lang/runtime/vindex_bench.c @@ -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 [nqueries] [k] [ef_csv] * vindex_bench synth [dim] [clusters] [nqueries] [k] [ef_csv] */ #include "engram_vindex.h" +#include "eg_metal_cosine.h" #include #include #include @@ -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;imax_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));