313 lines
13 KiB
C
313 lines
13 KiB
C
/* test_vindex.c — build + RUN gate for the M8 HNSW vector index.
|
|
*
|
|
* Covers: recall@10 vs brute-force oracle, brute-force-vs-index speedup,
|
|
* correctness edge cases (k>N, identical vectors, self-query, zero vector),
|
|
* determinism (seeded PRNG → identical graphs), and vindex_build_from_store
|
|
* over a real engram_store on-disk file.
|
|
*
|
|
* Pure C11; links engram_vindex.c + engram_store.c; -lm. ASan/UBSan clean.
|
|
*/
|
|
#include "engram_vindex.h"
|
|
#include "engram_store.h"
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
#include <stdint.h>
|
|
#include <time.h>
|
|
#include <unistd.h>
|
|
|
|
#define DIM 768
|
|
|
|
static int g_fail = 0;
|
|
/* VINDEX_QUICK=1 shrinks the two large builds so the ASan/UBSan pass (which runs
|
|
* ~5-10x slower) stays fast — memory-safety is size-independent. The perf numbers
|
|
* (recall gate + speedup) come from the un-sanitized, full-size pass. */
|
|
static int g_quick = 0;
|
|
static int envint(const char* k, int dflt){ const char* s=getenv(k); return s?atoi(s):dflt; }
|
|
#define CHECK(cond, msg) do{ if(!(cond)){ printf(" FAIL: %s\n", msg); g_fail=1; } else { printf(" ok: %s\n", msg); } }while(0)
|
|
|
|
/* deterministic test PRNG (splitmix64) */
|
|
static uint64_t rng_state = 0xABCDEF0123456789ULL;
|
|
static uint64_t xrng(uint64_t* s){
|
|
uint64_t z=(*s+=0x9E3779B97F4A7C15ULL);
|
|
z=(z^(z>>30))*0xBF58476D1CE4E5B9ULL; z=(z^(z>>27))*0x94D049BB133111EBULL;
|
|
return z^(z>>31);
|
|
}
|
|
static float frand(uint64_t* s){ return (float)((xrng(s)>>11)*(1.0/9007199254740992.0)) - 0.5f; }
|
|
|
|
static double now_s(void){
|
|
struct timespec t; clock_gettime(CLOCK_MONOTONIC,&t);
|
|
return t.tv_sec + t.tv_nsec*1e-9;
|
|
}
|
|
|
|
/* fill vec[N*DIM]: mostly random, some clustered groups (center + small noise). */
|
|
static void gen_vectors(float* v, int N, uint64_t seed){
|
|
uint64_t s = seed;
|
|
int clustered = N/5; /* last fifth is clustered */
|
|
int ncenters = 20;
|
|
float* centers = (float*)malloc((size_t)ncenters*DIM*sizeof(float));
|
|
for (int c=0;c<ncenters;c++) for(int d=0;d<DIM;d++) centers[c*DIM+d]=frand(&s);
|
|
for (int i=0;i<N;i++){
|
|
if (i < N-clustered){
|
|
for (int d=0;d<DIM;d++) v[i*DIM+d]=frand(&s);
|
|
} else {
|
|
int c = (int)(xrng(&s)%ncenters);
|
|
for (int d=0;d<DIM;d++) v[i*DIM+d]=centers[c*DIM+d] + 0.05f*frand(&s);
|
|
}
|
|
}
|
|
free(centers);
|
|
}
|
|
|
|
static float cosdist(const float* a, const float* b){
|
|
double da=0,db=0,dot=0;
|
|
for(int i=0;i<DIM;i++){ da+=(double)a[i]*a[i]; db+=(double)b[i]*b[i]; dot+=(double)a[i]*b[i]; }
|
|
if (da<=0||db<=0) return 1.0f;
|
|
return (float)(1.0 - dot/(sqrt(da)*sqrt(db)));
|
|
}
|
|
|
|
/* brute-force top-k node ids into ids[k] (ascending distance). */
|
|
static void brute_topk(const float* v, int N, const float* q, int k, int* ids){
|
|
float* bd = (float*)malloc((size_t)k*sizeof(float));
|
|
for (int i=0;i<k;i++){ ids[i]=-1; bd[i]=1e30f; }
|
|
for (int i=0;i<N;i++){
|
|
float d = cosdist(q, v+(size_t)i*DIM);
|
|
if (d < bd[k-1]){
|
|
int p=k-1;
|
|
while (p>0 && bd[p-1]>d){ bd[p]=bd[p-1]; ids[p]=ids[p-1]; p--; }
|
|
bd[p]=d; ids[p]=i;
|
|
}
|
|
}
|
|
free(bd);
|
|
}
|
|
|
|
/* ── Test 1: recall@10 vs brute force + latency/recall tradeoff ────────────── */
|
|
static void test_recall(void){
|
|
int N=envint("VINDEX_N_RECALL", g_quick?1500:5000), Q=200, K=10;
|
|
printf("\n== Test 1: recall@10 vs brute force (N=%d, DIM=768) ==\n", N);
|
|
float* v = (float*)malloc((size_t)N*DIM*sizeof(float));
|
|
gen_vectors(v, N, 111);
|
|
|
|
double t0=now_s();
|
|
VIndex* ix = vindex_create(DIM, VINDEX_DEFAULT_M, VINDEX_DEFAULT_EF_CONSTRUCTION);
|
|
for (int i=0;i<N;i++) vindex_insert(ix, (uint64_t)i, v+(size_t)i*DIM);
|
|
double build_s = now_s()-t0;
|
|
printf(" build: %d vectors in %.2fs (M=%d, ef_construction=%d)\n",
|
|
N, build_s, VINDEX_DEFAULT_M, VINDEX_DEFAULT_EF_CONSTRUCTION);
|
|
|
|
/* queries: half random, half near a real vector (perturbed). */
|
|
float* qs = (float*)malloc((size_t)Q*DIM*sizeof(float));
|
|
uint64_t s=999;
|
|
for (int i=0;i<Q;i++){
|
|
if (i<Q/2) for(int d=0;d<DIM;d++) qs[i*DIM+d]=frand(&s);
|
|
else { int base=(int)(xrng(&s)%N); for(int d=0;d<DIM;d++) qs[i*DIM+d]=v[base*DIM+d]+0.03f*frand(&s); }
|
|
}
|
|
|
|
/* oracle */
|
|
int* oracle = (int*)malloc((size_t)Q*K*sizeof(int));
|
|
for (int i=0;i<Q;i++) brute_topk(v, N, qs+(size_t)i*DIM, K, oracle+(size_t)i*K);
|
|
|
|
int efs[] = { 10, 32, 64, 128 };
|
|
for (int e=0;e<4;e++){
|
|
int ef=efs[e];
|
|
uint64_t ids[64]; float dd[64];
|
|
int hits=0;
|
|
double qt0=now_s();
|
|
for (int i=0;i<Q;i++){
|
|
int n=vindex_search(ix, qs+(size_t)i*DIM, K, ef, ids, dd);
|
|
for (int a=0;a<n;a++) for(int b=0;b<K;b++) if((int)ids[a]==oracle[i*K+b]){ hits++; break; }
|
|
}
|
|
double qs_ms = (now_s()-qt0)*1000.0/Q;
|
|
double recall = (double)hits/(Q*K);
|
|
printf(" ef_search=%-4d recall@10=%.4f latency=%.3f ms/query\n", ef, recall, qs_ms);
|
|
if (ef==VINDEX_DEFAULT_EF_SEARCH && !g_quick)
|
|
CHECK(recall >= 0.90, "recall@10 >= 0.90 at default ef_search=128");
|
|
}
|
|
free(oracle); free(qs); free(v); vindex_free(ix);
|
|
}
|
|
|
|
/* ── Test 2: speedup vs brute force ───────────────────────────────────────── */
|
|
static void speedup_at(int N){
|
|
int Q=100, K=10;
|
|
float* v=(float*)malloc((size_t)N*DIM*sizeof(float));
|
|
gen_vectors(v,N,222);
|
|
VIndex* ix=vindex_create(DIM,16,200);
|
|
double bt0=now_s();
|
|
for(int i=0;i<N;i++) vindex_insert(ix,(uint64_t)i,v+(size_t)i*DIM);
|
|
printf(" N=%d build=%.2fs\n", N, now_s()-bt0);
|
|
|
|
float* qs=(float*)malloc((size_t)Q*DIM*sizeof(float));
|
|
uint64_t s=333; for(int i=0;i<Q*DIM;i++) qs[i]=frand(&s);
|
|
|
|
/* brute force */
|
|
int scratch[16];
|
|
double b0=now_s();
|
|
for(int i=0;i<Q;i++) brute_topk(v,N,qs+(size_t)i*DIM,K,scratch);
|
|
double bf=(now_s()-b0)/Q;
|
|
|
|
/* index */
|
|
uint64_t ids[16]; float dd[16];
|
|
double i0=now_s();
|
|
for(int i=0;i<Q;i++) vindex_search(ix,qs+(size_t)i*DIM,K,64,ids,dd);
|
|
double iq=(now_s()-i0)/Q;
|
|
|
|
printf(" N=%d brute=%.4f ms/q index=%.4f ms/q speedup=%.1fx\n",
|
|
N, bf*1000, iq*1000, bf/iq);
|
|
CHECK(iq < bf, "index query faster than brute force");
|
|
free(qs); free(v); vindex_free(ix);
|
|
}
|
|
static void test_speedup(void){
|
|
printf("\n== Test 2: brute-force vs index speedup ==\n");
|
|
speedup_at(g_quick?2000:5000);
|
|
speedup_at(envint("VINDEX_N_BIG", g_quick?3000:20000));
|
|
}
|
|
|
|
/* ── Test 3: edge cases ───────────────────────────────────────────────────── */
|
|
static void test_edges(void){
|
|
printf("\n== Test 3: correctness edge cases ==\n");
|
|
/* k larger than node count */
|
|
{
|
|
VIndex* ix=vindex_create(DIM,16,200);
|
|
float vec[DIM]; uint64_t s=1;
|
|
for(int i=0;i<3;i++){ for(int d=0;d<DIM;d++) vec[d]=frand(&s); vindex_insert(ix,(uint64_t)i,vec); }
|
|
uint64_t ids[50]; float dd[50];
|
|
int n=vindex_search(ix, vec, 50, 64, ids, dd);
|
|
CHECK(n==3, "k > node count returns exactly node-count results");
|
|
vindex_free(ix);
|
|
}
|
|
/* duplicate / identical vectors */
|
|
{
|
|
VIndex* ix=vindex_create(DIM,16,200);
|
|
float a[DIM]; uint64_t s=2; for(int d=0;d<DIM;d++) a[d]=frand(&s);
|
|
for(int i=0;i<10;i++) vindex_insert(ix,(uint64_t)i,a); /* all identical */
|
|
float b[DIM]; for(int d=0;d<DIM;d++) b[d]=frand(&s);
|
|
vindex_insert(ix,100,b);
|
|
uint64_t ids[5]; float dd[5];
|
|
int n=vindex_search(ix,a,5,64,ids,dd);
|
|
CHECK(n==5, "identical-vector index returns k results");
|
|
CHECK(dd[0] < 1e-4f, "top-1 distance ~0 for a duplicated vector");
|
|
vindex_free(ix);
|
|
}
|
|
/* query equal to an indexed vector returns itself as top-1, dist ~0 */
|
|
{
|
|
VIndex* ix=vindex_create(DIM,16,200);
|
|
int N=500; float* v=(float*)malloc((size_t)N*DIM*sizeof(float)); gen_vectors(v,N,7);
|
|
for(int i=0;i<N;i++) vindex_insert(ix,(uint64_t)(1000+i),v+(size_t)i*DIM);
|
|
int probe=137;
|
|
uint64_t ids[3]; float dd[3];
|
|
int n=vindex_search(ix, v+(size_t)probe*DIM, 3, 64, ids, dd);
|
|
CHECK(n>=1 && ids[0]==(uint64_t)(1000+probe), "self-query returns itself as top-1");
|
|
CHECK(dd[0] < 1e-4f, "self-query top-1 distance ~0");
|
|
free(v); vindex_free(ix);
|
|
}
|
|
/* zero vector: no NaN, handled */
|
|
{
|
|
VIndex* ix=vindex_create(DIM,16,200);
|
|
float z[DIM]; memset(z,0,sizeof z);
|
|
float a[DIM]; uint64_t s=3; for(int d=0;d<DIM;d++) a[d]=frand(&s);
|
|
vindex_insert(ix,0,z); vindex_insert(ix,1,a);
|
|
uint64_t ids[2]; float dd[2];
|
|
int n=vindex_search(ix, z, 2, 64, ids, dd); /* zero query */
|
|
int nan=0; for(int i=0;i<n;i++) if(isnan(dd[i])||isinf(dd[i])) nan=1;
|
|
CHECK(n>=1 && !nan, "zero vector query produces no NaN/Inf");
|
|
n=vindex_search(ix, a, 2, 64, ids, dd); /* zero indexed */
|
|
nan=0; for(int i=0;i<n;i++) if(isnan(dd[i])||isinf(dd[i])) nan=1;
|
|
CHECK(!nan, "indexed zero vector produces no NaN/Inf");
|
|
vindex_free(ix);
|
|
}
|
|
}
|
|
|
|
/* ── Test 4: determinism ──────────────────────────────────────────────────── */
|
|
static void test_determinism(void){
|
|
printf("\n== Test 4: determinism (seeded PRNG → identical results) ==\n");
|
|
int N=1500;
|
|
float* v=(float*)malloc((size_t)N*DIM*sizeof(float)); gen_vectors(v,N,55);
|
|
uint64_t ids1[10],ids2[10]; float d1[10],d2[10];
|
|
int identical=1;
|
|
for (int build=0; build<2; build++){
|
|
VIndex* ix=vindex_create(DIM,16,200);
|
|
for(int i=0;i<N;i++) vindex_insert(ix,(uint64_t)i,v+(size_t)i*DIM);
|
|
/* probe several queries */
|
|
for (int q=0;q<20;q++){
|
|
uint64_t* ida = build? ids2 : ids1; float* da = build? d2 : d1;
|
|
vindex_search(ix, v+(size_t)(q*37%N)*DIM, 10, 64, ida, da);
|
|
if (build==1){
|
|
/* re-run build-0 query stored? simpler: compare within-run below */
|
|
}
|
|
}
|
|
vindex_free(ix);
|
|
}
|
|
/* Proper comparison: run two fresh builds, same single query. */
|
|
identical=1;
|
|
for (int q=0;q<25;q++){
|
|
int qi=(q*61)%N;
|
|
VIndex* a=vindex_create(DIM,16,200); for(int i=0;i<N;i++) vindex_insert(a,(uint64_t)i,v+(size_t)i*DIM);
|
|
VIndex* b=vindex_create(DIM,16,200); for(int i=0;i<N;i++) vindex_insert(b,(uint64_t)i,v+(size_t)i*DIM);
|
|
int na=vindex_search(a, v+(size_t)qi*DIM,10,64,ids1,d1);
|
|
int nb=vindex_search(b, v+(size_t)qi*DIM,10,64,ids2,d2);
|
|
if (na!=nb) identical=0;
|
|
for(int i=0;i<na;i++) if(ids1[i]!=ids2[i] || d1[i]!=d2[i]) identical=0;
|
|
vindex_free(a); vindex_free(b);
|
|
}
|
|
CHECK(identical, "two independent builds give byte-identical query results");
|
|
free(v);
|
|
}
|
|
|
|
/* ── Test 5: build_from_store ─────────────────────────────────────────────── */
|
|
static void test_build_from_store(void){
|
|
printf("\n== Test 5: vindex_build_from_store over a real engram_store ==\n");
|
|
char path[256];
|
|
snprintf(path,sizeof path,"/tmp/vindex_test_store_%d.engram",(int)getpid());
|
|
unlink(path);
|
|
EngramPagedStore* st = store_create(path);
|
|
if (!st){ printf(" FAIL: store_create\n"); g_fail=1; return; }
|
|
|
|
int N=300;
|
|
float* v=(float*)malloc((size_t)N*DIM*sizeof(float)); gen_vectors(v,N,88);
|
|
for (int i=0;i<N;i++){
|
|
StoreNode n; memset(&n,0,sizeof n);
|
|
char id[32]; snprintf(id,sizeof id,"node-%d",i);
|
|
n.id=id; n.content="x"; n.node_type="concept"; n.tier="Semantic";
|
|
n.emb = v+(size_t)i*DIM; n.emb_dim=DIM;
|
|
if (store_put_node(st,&n)!=0){ printf(" FAIL: put_node %d\n",i); g_fail=1; }
|
|
}
|
|
/* a node WITHOUT an emb — must be skipped by build_from_store. */
|
|
{ StoreNode n; memset(&n,0,sizeof n); n.id=(char*)"no-emb"; n.content="y"; n.node_type="concept"; n.tier="Semantic";
|
|
store_put_node(st,&n); }
|
|
store_close(st);
|
|
|
|
VIndex* ix = vindex_create(DIM,16,200);
|
|
char** ids=NULL; int nids=0;
|
|
int ins = vindex_build_from_store(ix, path, &ids, &nids);
|
|
printf(" build_from_store inserted %d vectors (expected %d; 1 emb-less skipped)\n", ins, N);
|
|
CHECK(ins==N, "build_from_store inserts exactly the emb'd nodes");
|
|
CHECK((size_t)ins==vindex_size(ix), "index size matches insert count");
|
|
|
|
/* query with a known vector → must return its own node id as top-1. */
|
|
int probe=42;
|
|
uint64_t rids[5]; float dd[5];
|
|
int n=vindex_search(ix, v+(size_t)probe*DIM, 5, 64, rids, dd);
|
|
int correct = (n>=1 && rids[0]<(uint64_t)nids && strcmp(ids[rids[0]], "node-42")==0);
|
|
printf(" query for node-42's vector → top-1 id=%s dist=%.5f\n",
|
|
(n>=1 && rids[0]<(uint64_t)nids)? ids[rids[0]] : "?", n?dd[0]:-1);
|
|
CHECK(correct, "build_from_store query resolves to the right node id");
|
|
CHECK(n>=1 && dd[0]<1e-4f, "top-1 distance ~0 for exact stored vector");
|
|
|
|
for (int i=0;i<nids;i++) free(ids[i]);
|
|
free(ids); free(v); vindex_free(ix); unlink(path);
|
|
}
|
|
|
|
int main(void){
|
|
(void)rng_state;
|
|
g_quick = envint("VINDEX_QUICK", 0);
|
|
printf("=== engram_vindex (HNSW) test suite ===%s\n", g_quick?" [QUICK]":"");
|
|
test_recall();
|
|
test_speedup();
|
|
test_edges();
|
|
test_determinism();
|
|
test_build_from_store();
|
|
printf("\n=== %s ===\n", g_fail? "FAILURES PRESENT" : "ALL TESTS PASSED");
|
|
return g_fail;
|
|
}
|