Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 728207aabf |
Executable
+25
@@ -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"
|
||||
@@ -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,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 <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 — 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 */
|
||||
@@ -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 <Foundation/Foundation.h>
|
||||
#import <Metal/Metal.h>
|
||||
#include "eg_metal_cosine.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;
|
||||
}
|
||||
}
|
||||
|
||||
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<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;
|
||||
}
|
||||
}
|
||||
|
||||
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<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;
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
+166
-5
@@ -7,11 +7,26 @@
|
||||
*
|
||||
* Read-only: never opens a socket, never writes the store. Safe on an nsbx clone.
|
||||
*
|
||||
* Build: cc -O2 -std=c11 vindex_bench.c engram_vindex.c -lm -o vindex_bench
|
||||
* Also runs the brute-force oracle a second way, through
|
||||
* eg_cosine_batch_metal() (Apple/Metal only — see eg_metal_cosine.h), and
|
||||
* reports its latency + a correctness check against the CPU oracle
|
||||
* side-by-side with the existing CPU-vs-HNSW numbers. EL_METAL_COSINE=0
|
||||
* forces CPU-only.
|
||||
*
|
||||
* Build (macOS):
|
||||
* cc -O2 -std=c11 -x objective-c -c eg_metal_cosine.m -o eg_metal_cosine.o \
|
||||
* -framework Metal -framework Foundation
|
||||
* cc -O2 -std=c11 vindex_bench.c engram_vindex.c eg_metal_cosine.o -lm \
|
||||
* -framework Metal -framework Foundation -o vindex_bench
|
||||
* Build (Linux / no Metal): omit eg_metal_cosine.o entirely and instead link
|
||||
* a CPU-only stub translation unit that defines eg_cosine_batch_metal() /
|
||||
* eg_cosine_batch_metal_available() returning false — this file never
|
||||
* references Metal directly, only the plain-C header.
|
||||
* Usage: vindex_bench store <neuron.egm> <dim> [nqueries] [k] [ef_csv]
|
||||
* vindex_bench synth <N> [dim] [clusters] [nqueries] [k] [ef_csv]
|
||||
*/
|
||||
#include "engram_vindex.h"
|
||||
#include "eg_metal_cosine.h"
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
@@ -64,6 +79,104 @@ static void brute_topk(const float* data, int n, int dim, const float* q,
|
||||
}
|
||||
}
|
||||
|
||||
/* GPU-accelerated variant of brute_topk: same oracle, same contract, same
|
||||
* output — computes all n distances via eg_cosine_batch_metal() instead of
|
||||
* one C loop, then does the identical top-k selection over the result.
|
||||
*
|
||||
* data is already L2-normalised (vindex_bench's convention throughout), so
|
||||
* eg_cosine()'s general unnormalised cosine and this file's "distance =
|
||||
* 1 - dot" both reduce to the same number here (a unit vector's norm is 1,
|
||||
* so cosine == dot). Passing pre-normalised rows through the general-purpose
|
||||
* batch kernel is deliberate: it proves the SAME primitive that would serve
|
||||
* el_runtime.c's raw/unnormalised embeddings also serves this oracle without
|
||||
* a second code path.
|
||||
*
|
||||
* Returns false (out_ids/out_d untouched) if the GPU path is unavailable or
|
||||
* fails for any reason — caller must fall back to brute_topk(). Never
|
||||
* partial: either the full top-k was computed on GPU, or nothing was. */
|
||||
/* EL_METAL_COSINE: 0/off/false disables the GPU path outright (falls back to
|
||||
* brute_topk() every time), matching el_runtime.c's own gate for the same
|
||||
* env var. Unset or any other value = auto (try Metal, fall back on failure). */
|
||||
static bool g_metal_env_checked = false;
|
||||
static bool g_metal_disabled_by_env = false;
|
||||
static void eg_metal_check_env_once(void){
|
||||
if (g_metal_env_checked) return;
|
||||
g_metal_env_checked = true;
|
||||
const char* v = getenv("EL_METAL_COSINE");
|
||||
if (v && (v[0]=='0' || v[0]=='n' || v[0]=='N' || v[0]=='f' || v[0]=='F'))
|
||||
g_metal_disabled_by_env = true;
|
||||
}
|
||||
static bool brute_topk_metal(const float* data, int n, int dim, const float* q,
|
||||
int k, int* out_ids, float* out_d){
|
||||
eg_metal_check_env_once();
|
||||
if (g_metal_disabled_by_env) return false;
|
||||
|
||||
const float** row_ptrs = malloc((size_t)n * sizeof(float*));
|
||||
int32_t* dims = malloc((size_t)n * sizeof(int32_t));
|
||||
double* scores = malloc((size_t)n * sizeof(double));
|
||||
if (!row_ptrs || !dims || !scores) { free(row_ptrs); free(dims); free(scores); return false; }
|
||||
|
||||
for (int i = 0; i < n; i++) {
|
||||
row_ptrs[i] = data + (size_t)i * dim;
|
||||
dims[i] = dim;
|
||||
}
|
||||
|
||||
bool ok = eg_cosine_batch_metal(q, dim, row_ptrs, dims, n, scores);
|
||||
free(row_ptrs); free(dims);
|
||||
if (!ok) { free(scores); return false; }
|
||||
|
||||
for (int i = 0; i < k; i++) { out_ids[i] = -1; out_d[i] = 3.0f; }
|
||||
for (int i = 0; i < n; i++) {
|
||||
float d = 1.0f - (float)scores[i]; /* same distance convention as brute_topk */
|
||||
if (d >= out_d[k-1]) continue;
|
||||
int p = k - 1;
|
||||
while (p > 0 && out_d[p-1] > d) { out_d[p] = out_d[p-1]; out_ids[p] = out_ids[p-1]; p--; }
|
||||
out_d[p] = d; out_ids[p] = i;
|
||||
}
|
||||
free(scores);
|
||||
return true;
|
||||
}
|
||||
|
||||
/* Batched sibling of brute_topk_metal: computes top-k for ALL nq queries in
|
||||
* ONE eg_cosine_batch_metal_multi() call, uploading node_matrix exactly
|
||||
* once instead of once per query. out_ids/out_d are nq*k, row-major
|
||||
* (query i's results at out_ids+i*k / out_d+i*k) — same layout run_bench
|
||||
* already uses for `gt`/per-query scratch. Returns false (nothing written)
|
||||
* on any failure; caller falls back to the per-query CPU brute_topk loop. */
|
||||
static bool brute_topk_metal_batch(const float* data, int n, int dim,
|
||||
const float* queries, int nq,
|
||||
int k, int* out_ids, float* out_d){
|
||||
eg_metal_check_env_once();
|
||||
if (g_metal_disabled_by_env) return false;
|
||||
|
||||
const float** row_ptrs = malloc((size_t)n * sizeof(float*));
|
||||
int32_t* dims = malloc((size_t)n * sizeof(int32_t));
|
||||
double* scores = malloc((size_t)nq * (size_t)n * sizeof(double));
|
||||
if (!row_ptrs || !dims || !scores) { free(row_ptrs); free(dims); free(scores); return false; }
|
||||
|
||||
for (int i = 0; i < n; i++) { row_ptrs[i] = data + (size_t)i * dim; dims[i] = dim; }
|
||||
|
||||
bool ok = eg_cosine_batch_metal_multi(queries, dim, nq, row_ptrs, dims, n, scores);
|
||||
free(row_ptrs); free(dims);
|
||||
if (!ok) { free(scores); return false; }
|
||||
|
||||
for (int qi = 0; qi < nq; qi++) {
|
||||
int* ids = out_ids + (size_t)qi * k;
|
||||
float* ds = out_d + (size_t)qi * k;
|
||||
const double* srow = scores + (size_t)qi * n;
|
||||
for (int i = 0; i < k; i++) { ids[i] = -1; ds[i] = 3.0f; }
|
||||
for (int i = 0; i < n; i++) {
|
||||
float d = 1.0f - (float)srow[i];
|
||||
if (d >= ds[k-1]) continue;
|
||||
int p = k - 1;
|
||||
while (p > 0 && ds[p-1] > d) { ds[p] = ds[p-1]; ids[p] = ids[p-1]; p--; }
|
||||
ds[p] = d; ids[p] = i;
|
||||
}
|
||||
}
|
||||
free(scores);
|
||||
return true;
|
||||
}
|
||||
|
||||
/* recall@k: |brute_topk ∩ hnsw_topk| / k. Both are id arrays of length k. */
|
||||
static double recall_at_k(const int* gt, const uint64_t* ann, int nann, int k){
|
||||
int hit = 0;
|
||||
@@ -142,13 +255,61 @@ static void run_bench(const char* label, float* data, int n, int dim,
|
||||
l2norm(dst, dim);
|
||||
}
|
||||
|
||||
/* ground truth: brute-force top-k for every query (also the oracle latency). */
|
||||
/* ground truth: brute-force top-k for every query (also the oracle latency).
|
||||
* gd is nq*k (one real slot per query, not a shared scratch buffer) so the
|
||||
* GPU comparison below can diff against every query's actual distances,
|
||||
* not just whichever query happened to run last. */
|
||||
int* gt = malloc((size_t)nq*k*sizeof(int));
|
||||
float* gd = malloc((size_t)k*sizeof(float));
|
||||
float* gd = malloc((size_t)nq*k*sizeof(float));
|
||||
double tb0 = now_s();
|
||||
for (int i=0;i<nq;i++) brute_topk(data, n, dim, qv+(size_t)i*dim, k, gt+(size_t)i*k, gd);
|
||||
for (int i=0;i<nq;i++) brute_topk(data, n, dim, qv+(size_t)i*dim, k, gt+(size_t)i*k, gd+(size_t)i*k);
|
||||
double brute_ms = (now_s()-tb0)*1000.0/nq;
|
||||
printf("BRUTE-FORCE : %8.3f ms/query (oracle; O(N*D))\n", brute_ms);
|
||||
printf("BRUTE-FORCE : %8.3f ms/query (oracle; O(N*D), CPU)\n", brute_ms);
|
||||
|
||||
/* GPU-accelerated oracle: SAME nq queries, SAME top-k contract, via ONE
|
||||
* eg_cosine_batch_metal_multi() call (uploads node_matrix once, not once
|
||||
* per query — see brute_topk_metal_batch). Run only if the GPU path is
|
||||
* actually available (checked once) — never fabricated, never assumed.
|
||||
* Verified against the CPU ground truth computed above: id-recall across
|
||||
* ALL nq queries, plus the actual max distance delta across every
|
||||
* (query, rank) pair that was compared — not a single spot check. */
|
||||
eg_metal_check_env_once();
|
||||
if (!g_metal_disabled_by_env && eg_cosine_batch_metal_available()) {
|
||||
int* gtm = malloc((size_t)nq*k*sizeof(int));
|
||||
float* gdm = malloc((size_t)nq*k*sizeof(float));
|
||||
double tm0 = now_s();
|
||||
bool ok = brute_topk_metal_batch(data, n, dim, qv, nq, k, gtm, gdm);
|
||||
double metal_ms = (now_s()-tm0)*1000.0/nq;
|
||||
if (ok) {
|
||||
double rec_sum = 0; double max_ddiff = 0; double sum_ddiff = 0; int compared = 0;
|
||||
for (int i=0;i<nq;i++) {
|
||||
const int* ids_gt = gt+(size_t)i*k;
|
||||
const float* d_gt = gd+(size_t)i*k;
|
||||
const int* ids_m = gtm+(size_t)i*k;
|
||||
const float* d_m = gdm+(size_t)i*k;
|
||||
uint64_t idset[512]; int m = (k<512)?k:512;
|
||||
for (int j=0;j<m;j++) idset[j] = (uint64_t)ids_m[j];
|
||||
rec_sum += recall_at_k(ids_gt, idset, m, k);
|
||||
/* same-rank distance delta — valid whenever both sides agree on
|
||||
* the id at that rank (true almost always, given ~100% recall;
|
||||
* a rank where they disagree isn't a meaningful delta to diff). */
|
||||
for (int j=0;j<k;j++) {
|
||||
if (ids_gt[j] == ids_m[j]) {
|
||||
double diff = fabs((double)d_gt[j]-(double)d_m[j]);
|
||||
if (diff>max_ddiff) max_ddiff=diff;
|
||||
sum_ddiff += diff; compared++;
|
||||
}
|
||||
}
|
||||
}
|
||||
printf("BRUTE-METAL : %8.3f ms/query (%.1fx vs CPU brute; id-recall %.4f vs CPU oracle over %d queries; same-rank |Δdist|: max %.2e, mean %.2e over %d compared)\n",
|
||||
metal_ms, brute_ms/metal_ms, rec_sum/nq, nq, max_ddiff, compared?sum_ddiff/compared:0.0, compared);
|
||||
} else {
|
||||
printf("BRUTE-METAL : GPU batch call failed/unavailable mid-run — skipped\n");
|
||||
}
|
||||
free(gtm); free(gdm);
|
||||
} else {
|
||||
printf("BRUTE-METAL : no Metal device/pipeline available — CPU-only\n");
|
||||
}
|
||||
|
||||
/* HNSW at each ef. */
|
||||
uint64_t* aid = malloc((size_t)k*sizeof(uint64_t));
|
||||
|
||||
Reference in New Issue
Block a user