engram: batch-cosine Adapter/Strategy/Factory over ggml, supersedes hand-rolled PR #114
El SDK CI - dev / build-and-test (pull_request) Failing after 4m29s
El SDK CI - dev / build-and-test (pull_request) Failing after 4m29s
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).
This commit is contained in:
Executable
+51
@@ -0,0 +1,51 @@
|
||||
#!/bin/bash
|
||||
# build_vindex_bench.sh — build the vindex_bench oracle/proof harness, with
|
||||
# the real ggml + hand-rolled-Metal batch-cosine strategies on Darwin and a
|
||||
# zero-dependency CPU-only 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.
|
||||
#
|
||||
# Darwin build links FOUR strategy translation units:
|
||||
# eg_cosine_batch.c — the Factory (always)
|
||||
# eg_cosine_batch_strategy_cpu.c — universal fallback (always)
|
||||
# eg_cosine_batch_strategy_ggml.c — ggml + dynamic Metal backend plugin
|
||||
# eg_cosine_batch_strategy_metal_hand.m — PR #114's original hand-rolled
|
||||
# Metal shader, preserved as one
|
||||
# selectable strategy
|
||||
# plus -DEG_HAVE_STRATEGY_GGML -DEG_HAVE_STRATEGY_METAL_HAND so the Factory
|
||||
# (and vindex_bench.c's own direct strategy comparison) knows both exist.
|
||||
#
|
||||
# ggml is resolved via `brew --prefix ggml` when available (portable across
|
||||
# Intel /usr/local and Apple Silicon /opt/homebrew installs), falling back to
|
||||
# /opt/homebrew if brew isn't on PATH. Override with GGML_PREFIX=... env var.
|
||||
#
|
||||
# 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
|
||||
GGML_PREFIX="${GGML_PREFIX:-$(brew --prefix ggml 2>/dev/null || echo /opt/homebrew)}"
|
||||
echo "== Darwin: building with the ggml + hand-rolled-Metal strategies (ggml prefix: $GGML_PREFIX) =="
|
||||
|
||||
"$CC" -O2 -std=c11 -x objective-c \
|
||||
-c eg_cosine_batch_strategy_metal_hand.m -o /tmp/eg_cosine_batch_strategy_metal_hand.o \
|
||||
-framework Metal -framework Foundation
|
||||
|
||||
"$CC" -O2 -std=c11 -DEG_HAVE_STRATEGY_GGML -DEG_HAVE_STRATEGY_METAL_HAND -w \
|
||||
-I"$GGML_PREFIX/include" \
|
||||
vindex_bench.c engram_vindex.c \
|
||||
eg_cosine_batch.c eg_cosine_batch_strategy_cpu.c eg_cosine_batch_strategy_ggml.c \
|
||||
/tmp/eg_cosine_batch_strategy_metal_hand.o \
|
||||
-L"$GGML_PREFIX/lib" -lggml -lggml-base \
|
||||
-Wl,-rpath,"$GGML_PREFIX/lib" \
|
||||
-lm -framework Metal -framework Foundation -o "$OUT"
|
||||
else
|
||||
echo "== non-Darwin: building with the CPU-only fallback strategy (no ggml, no Metal) =="
|
||||
"$CC" -O2 -std=c11 -w vindex_bench.c engram_vindex.c \
|
||||
eg_cosine_batch.c eg_cosine_batch_strategy_cpu.c \
|
||||
-lm -o "$OUT"
|
||||
fi
|
||||
|
||||
echo "built: $OUT"
|
||||
@@ -0,0 +1,121 @@
|
||||
/* eg_cosine_batch.c — the Factory. Implements the stable public interface
|
||||
* declared in eg_cosine_batch.h by selecting ONE concrete
|
||||
* EgCosineBatchStrategy (eg_cosine_batch_strategy.h) and dispatching every
|
||||
* call to it. This is the ONLY file that branches on EG_HAVE_STRATEGY_*
|
||||
* (build-time: which strategy .c/.m files were actually compiled in for
|
||||
* this platform) — call sites never see those macros.
|
||||
*
|
||||
* Selection is lazy (first call) and cached — mirrors the lazy-init caching
|
||||
* every individual strategy already does internally, so there is no added
|
||||
* per-call cost after the first.
|
||||
*
|
||||
* Selection mechanism (env var + build-time + runtime capability probe, all
|
||||
* three, exactly as directed):
|
||||
* - BUILD-TIME decides which strategies exist to choose from at all: a
|
||||
* Darwin build compiles+links the ggml strategy and the hand-rolled
|
||||
* Metal strategy (EG_HAVE_STRATEGY_GGML / EG_HAVE_STRATEGY_METAL_HAND
|
||||
* both defined); a non-Darwin build compiles neither, matching PR #114's
|
||||
* original Linux behavior exactly (CPU-fallback only, no Objective-C
|
||||
* compiler or Metal frameworks required).
|
||||
* - RUNTIME CAPABILITY PROBE: each candidate strategy's own available()
|
||||
* does the real, cheap-after-first-call check (device present, backend
|
||||
* plugin loaded, pipeline compiled) — never assumed from build-time
|
||||
* alone. A build that HAS the ggml strategy compiled in but is running
|
||||
* on hardware/software where it can't actually initialize (backend
|
||||
* plugin missing, no GPU) correctly falls through to the next candidate.
|
||||
* - ENV VAR gives explicit, debuggable override for either axis:
|
||||
* EL_COSINE_BATCH_STRATEGY = "ggml" | "metal" | "cpu" | unset/"auto"
|
||||
* forces a specific strategy (falling back to cpu if the forced one
|
||||
* isn't actually available), or leaves the default auto-preference
|
||||
* order in place.
|
||||
* EL_METAL_COSINE = 0/n/N/f/F (back-compat with PR #114's vindex_bench
|
||||
* gate) disables ALL GPU-backed strategies outright, same as before.
|
||||
*
|
||||
* DEFAULT preference order when nothing is forced: ggml, then hand-rolled
|
||||
* Metal, then CPU fallback — first candidate whose available() reports true
|
||||
* wins. This is what makes "stop hand-rolling GPU kernels, use ggml" real
|
||||
* rather than nominal: ggml is what actually runs by default on this
|
||||
* machine today (see the PR body for the measured numbers backing that).
|
||||
*/
|
||||
#include "eg_cosine_batch.h"
|
||||
#include "eg_cosine_batch_strategy.h"
|
||||
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
|
||||
static bool g_selected = false;
|
||||
static const EgCosineBatchStrategy* g_active = NULL;
|
||||
|
||||
static bool eg_env_truthy_off(const char* v) {
|
||||
return v && (v[0]=='0' || v[0]=='n' || v[0]=='N' || v[0]=='f' || v[0]=='F');
|
||||
}
|
||||
|
||||
static const EgCosineBatchStrategy* eg_select_strategy(void) {
|
||||
if (g_selected) return g_active;
|
||||
g_selected = true;
|
||||
|
||||
const EgCosineBatchStrategy* cpu = eg_cosine_batch_strategy_cpu();
|
||||
const char* force = getenv("EL_COSINE_BATCH_STRATEGY");
|
||||
const char* legacy_off = getenv("EL_METAL_COSINE");
|
||||
|
||||
if (eg_env_truthy_off(legacy_off)) { g_active = cpu; return g_active; }
|
||||
|
||||
if (force && strcmp(force, "cpu") == 0) { g_active = cpu; return g_active; }
|
||||
|
||||
if (force && strcmp(force, "ggml") == 0) {
|
||||
#ifdef EG_HAVE_STRATEGY_GGML
|
||||
const EgCosineBatchStrategy* s = eg_cosine_batch_strategy_ggml();
|
||||
if (s->available()) { g_active = s; return g_active; }
|
||||
#endif
|
||||
g_active = cpu; return g_active;
|
||||
}
|
||||
|
||||
if (force && strcmp(force, "metal") == 0) {
|
||||
#ifdef EG_HAVE_STRATEGY_METAL_HAND
|
||||
const EgCosineBatchStrategy* s = eg_cosine_batch_strategy_metal_hand();
|
||||
if (s->available()) { g_active = s; return g_active; }
|
||||
#endif
|
||||
g_active = cpu; return g_active;
|
||||
}
|
||||
|
||||
/* auto (unset, or any other value): ggml -> metal-hand -> cpu, first
|
||||
* available wins. */
|
||||
#ifdef EG_HAVE_STRATEGY_GGML
|
||||
{
|
||||
const EgCosineBatchStrategy* s = eg_cosine_batch_strategy_ggml();
|
||||
if (s->available()) { g_active = s; return g_active; }
|
||||
}
|
||||
#endif
|
||||
#ifdef EG_HAVE_STRATEGY_METAL_HAND
|
||||
{
|
||||
const EgCosineBatchStrategy* s = eg_cosine_batch_strategy_metal_hand();
|
||||
if (s->available()) { g_active = s; return g_active; }
|
||||
}
|
||||
#endif
|
||||
g_active = cpu;
|
||||
return g_active;
|
||||
}
|
||||
|
||||
bool eg_cosine_batch_available(void) {
|
||||
return eg_select_strategy()->available();
|
||||
}
|
||||
|
||||
const char* eg_cosine_batch_strategy_name(void) {
|
||||
return eg_select_strategy()->name;
|
||||
}
|
||||
|
||||
bool eg_cosine_batch(const float* query, int32_t qdim,
|
||||
const float* const* node_ptrs,
|
||||
const int32_t* node_dims,
|
||||
int32_t n,
|
||||
double* out_scores) {
|
||||
return eg_select_strategy()->batch(query, qdim, node_ptrs, node_dims, n, out_scores);
|
||||
}
|
||||
|
||||
bool eg_cosine_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_select_strategy()->batch_multi(queries, qdim, nq, node_ptrs, node_dims, n, out_scores);
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
/* eg_cosine_batch.h — stable Adapter interface over batch-cosine-similarity
|
||||
* BACKEND STRATEGIES. This header is the ONE thing call sites (el_runtime.c,
|
||||
* vindex_bench.c, ...) talk to. Plain C11, safe to #include on every
|
||||
* platform — the symbols declared here always exist and always link,
|
||||
* regardless of what backend actually runs underneath. Zero #ifdef at call
|
||||
* sites: which concrete strategy executes (ggml/Metal, hand-rolled Metal, or
|
||||
* the always-false CPU fallback) is resolved once, lazily, inside
|
||||
* eg_cosine_batch.c's factory — see eg_cosine_batch_strategy.h for that.
|
||||
*
|
||||
* This supersedes eg_metal_cosine.h (PR #114's single hand-rolled-Metal-only
|
||||
* bridge). The contract is UNCHANGED — same shapes, same sentinel, same
|
||||
* never-partial guarantee, same "caller must always be prepared to fall back
|
||||
* to its own scalar per-node loop" rule — only the name changed, because the
|
||||
* thing behind it is no longer "the Metal bridge," it is "whichever batch-
|
||||
* cosine strategy the factory picked." eg_metal_cosine.h's original doc
|
||||
* comments (byte-for-byte, this file is the direct descendant) are preserved
|
||||
* below since they remain the precise spec any strategy must honor.
|
||||
*
|
||||
* On ANY failure at ANY step — no compute device, compile/init error, alloc
|
||||
* failure, bad args — every function here returns false and writes nothing.
|
||||
* Out-params are either fully populated or left completely untouched, never
|
||||
* partial. Callers MUST always be prepared to fall back to their own scalar
|
||||
* per-node CPU loop unconditionally. These functions must never crash, throw,
|
||||
* or hang the calling process — several call sites run inside a long-lived
|
||||
* daemon's request-handling hot path.
|
||||
*/
|
||||
#ifndef EG_COSINE_BATCH_H
|
||||
#define EG_COSINE_BATCH_H
|
||||
|
||||
#include <stdint.h>
|
||||
#include <stdbool.h>
|
||||
|
||||
#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 — every strategy 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 a real backend strategy ran and out_scores was fully
|
||||
* populated. Returns false (out_scores left untouched) on ANY failure or
|
||||
* unavailability — no compute device, compile/init failure, allocation
|
||||
* failure, n<=0, qdim<=0, null query/node_ptrs/node_dims/out_scores.
|
||||
*/
|
||||
bool eg_cosine_batch(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 real (non-CPU-fallback) strategy is available right now (cheap
|
||||
* after the first call — cached). Purely informational (e.g. a startup log
|
||||
* line or /api/stats field); callers should still treat a false return from
|
||||
* eg_cosine_batch()/eg_cosine_batch_multi() itself as the authoritative
|
||||
* fallback signal, not this function. */
|
||||
bool eg_cosine_batch_available(void);
|
||||
|
||||
/* Which concrete strategy is currently selected — "ggml", "metal-hand",
|
||||
* or "cpu-fallback". Purely informational/diagnostic, same spirit as
|
||||
* eg_cosine_batch_available(). Never NULL. */
|
||||
const char* eg_cosine_batch_strategy_name(void);
|
||||
|
||||
/* Multi-query batched cosine: nq query vectors against the SAME n node
|
||||
* vectors, in one call. A real strategy uploads/prepares the node population
|
||||
* ONCE and reuses it for every query, instead of nq separate
|
||||
* eg_cosine_batch() calls each paying the full gather+upload cost — PR #114
|
||||
* measured this 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 a GPU-backed path a real win at this shape. Use this
|
||||
* whenever multiple queries will run against an unchanged (or
|
||||
* rarely-changing) node population; use eg_cosine_batch() 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().
|
||||
*
|
||||
* Returns true iff a real strategy ran and out_scores was fully populated
|
||||
* (all nq*n entries); false (untouched) on any failure/unavailability. */
|
||||
bool eg_cosine_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);
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif /* EG_COSINE_BATCH_H */
|
||||
@@ -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 <metal_stdlib>
|
||||
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);
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
/* eg_cosine_batch_strategy.h — internal Strategy interface, NOT for call
|
||||
* sites (they use eg_cosine_batch.h). Only eg_cosine_batch.c's factory and
|
||||
* the concrete strategy implementation files include this.
|
||||
*
|
||||
* Each concrete strategy exposes exactly one getter returning a pointer to a
|
||||
* static, immutable EgCosineBatchStrategy vtable. Which getters actually
|
||||
* exist as linkable symbols is a BUILD-TIME concern (decided by
|
||||
* build_vindex_bench.sh / the engram daemon's own build, via which .c/.m
|
||||
* files get compiled per platform) gated by the EG_HAVE_STRATEGY_* macros
|
||||
* below — the factory in eg_cosine_batch.c is the ONLY place that branches
|
||||
* on those macros. Call sites never see them; that's the whole point of the
|
||||
* Adapter in eg_cosine_batch.h.
|
||||
*
|
||||
* Three concrete strategies exist:
|
||||
* eg_cosine_batch_strategy_ggml() — ggml + dynamically-loaded Metal
|
||||
* backend plugin. Darwin only in
|
||||
* this build; the default
|
||||
* preferred strategy wherever
|
||||
* available. EG_HAVE_STRATEGY_GGML.
|
||||
* eg_cosine_batch_strategy_metal_hand() — the original hand-rolled Metal
|
||||
* compute shader from PR #114
|
||||
* (eg_cosine_batch.metal),
|
||||
* preserved verbatim as a
|
||||
* selectable fallback strategy,
|
||||
* not deleted. Darwin only.
|
||||
* EG_HAVE_STRATEGY_METAL_HAND.
|
||||
* eg_cosine_batch_strategy_cpu() — universal always-false
|
||||
* fallback. Always compiled, on
|
||||
* every platform; this is what a
|
||||
* non-Darwin build links
|
||||
* exclusively (matching PR #114's
|
||||
* eg_metal_cosine_stub.c), and
|
||||
* what any platform falls back
|
||||
* to when no real strategy is
|
||||
* available at runtime.
|
||||
*/
|
||||
#ifndef EG_COSINE_BATCH_STRATEGY_H
|
||||
#define EG_COSINE_BATCH_STRATEGY_H
|
||||
|
||||
#include <stdint.h>
|
||||
#include <stdbool.h>
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
typedef struct EgCosineBatchStrategy {
|
||||
/* Stable, short, lowercase-hyphenated identifier — what
|
||||
* eg_cosine_batch_strategy_name() surfaces. Never NULL. */
|
||||
const char* name;
|
||||
|
||||
/* Cheap after the first call (lazy init, cached internally). Must never
|
||||
* throw/crash/hang — mirrors eg_cosine_batch_available()'s contract. */
|
||||
bool (*available)(void);
|
||||
|
||||
/* Same shape/contract as eg_cosine_batch() in eg_cosine_batch.h. */
|
||||
bool (*batch)(const float* query, int32_t qdim,
|
||||
const float* const* node_ptrs, const int32_t* node_dims,
|
||||
int32_t n, double* out_scores);
|
||||
|
||||
/* Same shape/contract as eg_cosine_batch_multi() in eg_cosine_batch.h. */
|
||||
bool (*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);
|
||||
} EgCosineBatchStrategy;
|
||||
|
||||
#ifdef EG_HAVE_STRATEGY_GGML
|
||||
const EgCosineBatchStrategy* eg_cosine_batch_strategy_ggml(void);
|
||||
#endif
|
||||
|
||||
#ifdef EG_HAVE_STRATEGY_METAL_HAND
|
||||
const EgCosineBatchStrategy* eg_cosine_batch_strategy_metal_hand(void);
|
||||
#endif
|
||||
|
||||
/* Always declared/linked, on every platform/build. */
|
||||
const EgCosineBatchStrategy* eg_cosine_batch_strategy_cpu(void);
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif /* EG_COSINE_BATCH_STRATEGY_H */
|
||||
@@ -0,0 +1,45 @@
|
||||
/* eg_cosine_batch_strategy_cpu.c — plain-C, zero-dependency universal
|
||||
* fallback strategy. Always returns false / unavailable. Direct descendant
|
||||
* of PR #114's eg_metal_cosine_stub.c, generalized from "the Metal stub" to
|
||||
* "the strategy vtable's universal fallback entry" now that multiple real
|
||||
* strategies can exist.
|
||||
*
|
||||
* Always compiled, on every platform. On Darwin builds it is the last-resort
|
||||
* strategy the factory falls back to when neither ggml nor the hand-rolled
|
||||
* Metal strategy is available at runtime (no device, compile failure, ...).
|
||||
* On non-Darwin builds it is the ONLY strategy compiled in at all — no
|
||||
* Objective-C, no Metal frameworks, no ggml/Metal backend plugin — so
|
||||
* eg_cosine_batch()/eg_cosine_batch_multi() always return false there and
|
||||
* every call site's existing CPU fallback runs unconditionally, exactly as
|
||||
* before this PR.
|
||||
*/
|
||||
#include "eg_cosine_batch_strategy.h"
|
||||
|
||||
static bool cpu_available(void) {
|
||||
return false;
|
||||
}
|
||||
|
||||
static bool cpu_batch(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;
|
||||
}
|
||||
|
||||
static bool cpu_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) {
|
||||
(void)queries; (void)qdim; (void)nq; (void)node_ptrs; (void)node_dims; (void)n; (void)out_scores;
|
||||
return false;
|
||||
}
|
||||
|
||||
static const EgCosineBatchStrategy g_cpu_strategy = {
|
||||
.name = "cpu-fallback",
|
||||
.available = cpu_available,
|
||||
.batch = cpu_batch,
|
||||
.batch_multi = cpu_batch_multi,
|
||||
};
|
||||
|
||||
const EgCosineBatchStrategy* eg_cosine_batch_strategy_cpu(void) {
|
||||
return &g_cpu_strategy;
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
@@ -0,0 +1,358 @@
|
||||
/* eg_cosine_batch_strategy_metal_hand.m — the HAND-ROLLED-METAL Strategy.
|
||||
*
|
||||
* This is PR #114's original Objective-C bridge (formerly eg_metal_cosine.m)
|
||||
* exposing the hand-written Metal compute shader (eg_cosine_batch.metal) as
|
||||
* one concrete EgCosineBatchStrategy. It is preserved here almost verbatim —
|
||||
* real, carefully verified work, not discarded — now living behind the
|
||||
* Adapter/Strategy/Factory restructuring (see eg_cosine_batch.h and
|
||||
* eg_cosine_batch_strategy.h) alongside the new ggml-Metal strategy
|
||||
* (eg_cosine_batch_strategy_ggml.c) and the universal CPU fallback
|
||||
* (eg_cosine_batch_strategy_cpu.c). The factory in eg_cosine_batch.c prefers
|
||||
* ggml by default when both are available; this strategy remains selectable
|
||||
* via EL_COSINE_BATCH_STRATEGY=metal, and is what the factory falls back to
|
||||
* if ggml's backend plugin fails to load/init for any reason.
|
||||
*
|
||||
* Apple-only (Metal has no other platform). This file is excluded from the
|
||||
* build entirely on non-Darwin — see build_vindex_bench.sh, which only
|
||||
* compiles/links this file and defines EG_HAVE_STRATEGY_METAL_HAND when
|
||||
* `uname` is Darwin. On Linux the factory never sees this strategy at all —
|
||||
* callers must always be prepared for the "no real strategy available"
|
||||
* fallback via the CPU strategy, which is also exactly what happens here on
|
||||
* Apple hardware with no usable GPU.
|
||||
*
|
||||
* Design (unchanged from PR #114):
|
||||
* - 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 <Foundation/Foundation.h>
|
||||
#import <Metal/Metal.h>
|
||||
#include "eg_cosine_batch_strategy.h"
|
||||
#include <string.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
/* ── BEGIN embedded shader source (keep in sync with eg_cosine_batch.metal) ── */
|
||||
static const char* kEgCosineBatchMetalSrc =
|
||||
"#include <metal_stdlib>\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<MTLDevice> g_device = nil;
|
||||
static id<MTLCommandQueue> g_queue = nil;
|
||||
static id<MTLComputePipelineState> g_pipeline = nil; /* single-query kernel */
|
||||
static id<MTLComputePipelineState> 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<MTLDevice> dev = MTLCreateSystemDefaultDevice();
|
||||
if (!dev) return false;
|
||||
|
||||
id<MTLCommandQueue> q = [dev newCommandQueue];
|
||||
if (!q) return false;
|
||||
|
||||
NSError* err = nil;
|
||||
NSString* src = [NSString stringWithUTF8String:kEgCosineBatchMetalSrc];
|
||||
MTLCompileOptions* opts = [MTLCompileOptions new];
|
||||
id<MTLLibrary> lib = [dev newLibraryWithSource:src options:opts error:&err];
|
||||
if (!lib) return false;
|
||||
|
||||
id<MTLFunction> fn = [lib newFunctionWithName:@"eg_cosine_batch_kernel"];
|
||||
if (!fn) return false;
|
||||
id<MTLComputePipelineState> pipe = [dev newComputePipelineStateWithFunction:fn error:&err];
|
||||
if (!pipe) return false;
|
||||
|
||||
id<MTLFunction> fnMulti = [lib newFunctionWithName:@"eg_cosine_batch_multi_kernel"];
|
||||
if (!fnMulti) return false;
|
||||
id<MTLComputePipelineState> 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;
|
||||
}
|
||||
}
|
||||
|
||||
static bool mh_available(void) {
|
||||
return eg_metal_ensure_init();
|
||||
}
|
||||
|
||||
static bool mh_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;
|
||||
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<MTLBuffer> bufQuery = [g_device newBufferWithBytes:query
|
||||
length:dim * sizeof(float)
|
||||
options:MTLResourceStorageModeShared];
|
||||
id<MTLBuffer> bufMatrix = [g_device newBufferWithBytes:matrix
|
||||
length:nu * dim * sizeof(float)
|
||||
options:MTLResourceStorageModeShared];
|
||||
id<MTLBuffer> bufDims = [g_device newBufferWithBytes:dims_i32
|
||||
length:nu * sizeof(int32_t)
|
||||
options:MTLResourceStorageModeShared];
|
||||
EgCosineParamsC params = { (uint32_t)nu, (uint32_t)dim };
|
||||
id<MTLBuffer> bufParams = [g_device newBufferWithBytes:¶ms
|
||||
length:sizeof(params)
|
||||
options:MTLResourceStorageModeShared];
|
||||
id<MTLBuffer> bufOut = [g_device newBufferWithLength:nu * sizeof(float)
|
||||
options:MTLResourceStorageModeShared];
|
||||
|
||||
free(matrix); free(dims_i32);
|
||||
|
||||
if (!bufQuery || !bufMatrix || !bufDims || !bufParams || !bufOut) return false;
|
||||
|
||||
id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
|
||||
if (!cmd) return false;
|
||||
id<MTLComputeCommandEncoder> 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;
|
||||
}
|
||||
}
|
||||
|
||||
static bool mh_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) {
|
||||
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<MTLBuffer> bufMatrix = [g_device newBufferWithBytes:matrix
|
||||
length:nu * dim * sizeof(float)
|
||||
options:MTLResourceStorageModeShared];
|
||||
id<MTLBuffer> bufDims = [g_device newBufferWithBytes:dims_i32
|
||||
length:nu * sizeof(int32_t)
|
||||
options:MTLResourceStorageModeShared];
|
||||
id<MTLBuffer> bufQueries = [g_device newBufferWithBytes:queries
|
||||
length:nqu * dim * sizeof(float)
|
||||
options:MTLResourceStorageModeShared];
|
||||
EgCosineMultiParamsC params = { (uint32_t)nu, (uint32_t)dim, (uint32_t)nqu };
|
||||
id<MTLBuffer> bufParams = [g_device newBufferWithBytes:¶ms
|
||||
length:sizeof(params)
|
||||
options:MTLResourceStorageModeShared];
|
||||
id<MTLBuffer> bufOut = [g_device newBufferWithLength:nqu * nu * sizeof(float)
|
||||
options:MTLResourceStorageModeShared];
|
||||
|
||||
free(matrix); free(dims_i32);
|
||||
|
||||
if (!bufMatrix || !bufDims || !bufQueries || !bufParams || !bufOut) return false;
|
||||
|
||||
id<MTLCommandBuffer> cmd = [g_queue commandBuffer];
|
||||
if (!cmd) return false;
|
||||
id<MTLComputeCommandEncoder> 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;
|
||||
}
|
||||
}
|
||||
|
||||
static const EgCosineBatchStrategy g_metal_hand_strategy = {
|
||||
.name = "metal-hand",
|
||||
.available = mh_available,
|
||||
.batch = mh_batch,
|
||||
.batch_multi = mh_batch_multi,
|
||||
};
|
||||
|
||||
const EgCosineBatchStrategy* eg_cosine_batch_strategy_metal_hand(void) {
|
||||
return &g_metal_hand_strategy;
|
||||
}
|
||||
+147
-5
@@ -7,11 +7,29 @@
|
||||
*
|
||||
* 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 (and third) way, through the
|
||||
* batch-cosine Strategies behind eg_cosine_batch_strategy.h — the ggml
|
||||
* strategy and the hand-rolled-Metal strategy (Apple/Metal only; see
|
||||
* eg_cosine_batch.h/eg_cosine_batch_strategy.h) — and reports each one's
|
||||
* latency + a correctness check against the CPU oracle side-by-side with the
|
||||
* existing CPU-vs-HNSW numbers. This harness deliberately reaches past the
|
||||
* single-selection Factory (eg_cosine_batch.c) to instantiate every
|
||||
* compiled-in strategy directly, so it can compare all of them against the
|
||||
* SAME dataset in one run — that is the harness's whole job; a real call
|
||||
* site (el_runtime.c) never does this, it only ever calls the plain
|
||||
* eg_cosine_batch()/eg_cosine_batch_multi() adapter functions.
|
||||
* EL_METAL_COSINE=0 forces CPU-only (skips every strategy comparison).
|
||||
*
|
||||
* Build (macOS, ggml + hand-rolled Metal): see build_vindex_bench.sh.
|
||||
* Build (Linux / no Metal): omit every eg_cosine_batch_strategy_*.{c,m} file
|
||||
* except eg_cosine_batch_strategy_cpu.c — this file never references
|
||||
* ggml/Metal directly except through the plain-C strategy header, guarded
|
||||
* by the same EG_HAVE_STRATEGY_* build macros the Factory itself uses.
|
||||
* 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_cosine_batch_strategy.h"
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
@@ -74,6 +92,106 @@ static double recall_at_k(const int* gt, const uint64_t* ann, int nann, int k){
|
||||
return (double)hit / (double)k;
|
||||
}
|
||||
|
||||
/* EL_METAL_COSINE: 0/off/false disables EVERY strategy comparison outright
|
||||
* (falls back to brute_topk() only), matching el_runtime.c's own gate for
|
||||
* the same env var (back-compat name kept from PR #114; it now gates all
|
||||
* GPU-backed strategies, not just the hand-rolled Metal one). Unset or any
|
||||
* other value = try every compiled-in strategy, report each that's
|
||||
* available, skip (without failing the run) any that isn't. */
|
||||
static bool g_strategy_env_checked = false;
|
||||
static bool g_strategy_disabled_by_env = false;
|
||||
static void eg_strategy_check_env_once(void){
|
||||
if (g_strategy_env_checked) return;
|
||||
g_strategy_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_strategy_disabled_by_env = true;
|
||||
}
|
||||
|
||||
/* Batched sibling of brute_topk, generalized over ANY EgCosineBatchStrategy:
|
||||
* computes top-k for ALL nq queries in ONE strategy->batch_multi() call,
|
||||
* uploading/preparing the node population 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). Returns false (nothing written) on any
|
||||
* failure/unavailability; caller treats that as "skip this strategy in the
|
||||
* report", never as a hard error. */
|
||||
static bool batch_topk_strategy(const EgCosineBatchStrategy* strat,
|
||||
const float* data, int n, int dim,
|
||||
const float* queries, int nq,
|
||||
int k, int* out_ids, float* out_d){
|
||||
if (!strat || !strat->available()) 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 = strat->batch_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]; /* same distance convention as brute_topk */
|
||||
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;
|
||||
}
|
||||
|
||||
/* Runs batch_topk_strategy for one named strategy over ALL nq queries, diffs
|
||||
* against the CPU ground truth (gt/gd, both nq*k), and prints a report line
|
||||
* in the same shape PR #114 established for BRUTE-METAL — id-recall over
|
||||
* every query plus the actual max/mean same-rank distance delta across
|
||||
* every (query,rank) pair that was compared, never fabricated or assumed. */
|
||||
static void report_strategy_vs_oracle(const char* label, const EgCosineBatchStrategy* strat,
|
||||
const float* data, int n, int dim,
|
||||
const float* qv, int nq, int k,
|
||||
const int* gt, const float* gd, double brute_ms){
|
||||
if (g_strategy_disabled_by_env) { printf("%-13s: disabled via EL_METAL_COSINE\n", label); return; }
|
||||
if (!strat || !strat->available()) { printf("%-13s: not available on this build/host — skipped\n", label); return; }
|
||||
|
||||
int* gtm = malloc((size_t)nq*k*sizeof(int));
|
||||
float* gdm = malloc((size_t)nq*k*sizeof(float));
|
||||
double tm0 = now_s();
|
||||
bool ok = batch_topk_strategy(strat, data, n, dim, qv, nq, k, gtm, gdm);
|
||||
double strat_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);
|
||||
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("%-13s: %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",
|
||||
label, strat_ms, brute_ms/strat_ms, rec_sum/nq, nq, max_ddiff, compared?sum_ddiff/compared:0.0, compared);
|
||||
} else {
|
||||
printf("%-13s: batch call failed mid-run — skipped\n", label);
|
||||
}
|
||||
free(gtm); free(gdm);
|
||||
}
|
||||
|
||||
/* Parse "64,128,256" into an int array; returns count. */
|
||||
static int parse_csv(const char* s, int* out, int maxo){
|
||||
int n=0; if(!s||!*s) return 0;
|
||||
@@ -142,13 +260,37 @@ 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
|
||||
* strategy comparisons 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-backed oracles: SAME nq queries, SAME top-k contract, via each
|
||||
* compiled-in Strategy's batch_multi() (uploads/prepares the node
|
||||
* population once, not once per query). Run only for strategies that
|
||||
* are actually available (checked internally) — never fabricated, never
|
||||
* assumed. Verified against the CPU ground truth computed above:
|
||||
* id-recall across ALL nq queries, plus the actual max/mean distance
|
||||
* delta across every (query,rank) pair that was compared. */
|
||||
eg_strategy_check_env_once();
|
||||
#ifdef EG_HAVE_STRATEGY_GGML
|
||||
report_strategy_vs_oracle("BRUTE-GGML", eg_cosine_batch_strategy_ggml(),
|
||||
data, n, dim, qv, nq, k, gt, gd, brute_ms);
|
||||
#else
|
||||
printf("BRUTE-GGML : strategy not compiled into this build\n");
|
||||
#endif
|
||||
#ifdef EG_HAVE_STRATEGY_METAL_HAND
|
||||
report_strategy_vs_oracle("BRUTE-METAL", eg_cosine_batch_strategy_metal_hand(),
|
||||
data, n, dim, qv, nq, k, gt, gd, brute_ms);
|
||||
#else
|
||||
printf("BRUTE-METAL : strategy not compiled into this build\n");
|
||||
#endif
|
||||
|
||||
/* HNSW at each ef. */
|
||||
uint64_t* aid = malloc((size_t)k*sizeof(uint64_t));
|
||||
|
||||
Reference in New Issue
Block a user