bigmerge
728207aabf
The original brief targeted engram_activate's O(N*D) cosq prescan, but #109 (this branch) already retires that loop algorithmically (HNSW seed selection + lazy memoized cosine) — GPU-accelerating a loop being deleted isn't real work, so that target was dropped rather than forced. Re-investigated for a genuine remaining GPU-shaped call site (not a manufactured one): HNSW insert's candidate-list distance work is bounded- degree (M=24-48) and sequential/adaptive — too fine-grained for a GPU dispatch to pay off. No O(N^2) pairwise cosine pass exists (dedup only checks the K=8 already-selected seed slots). No concurrent multi-query traffic exists (server.el: the soul's curiosity loop is a single in-process caller). vindex_bench.c's brute_topk — the correctness oracle this same PR adds to validate HNSW recall — is the one real, unforced fit: genuine 1-query-vs-N-vectors, embarrassingly parallel, no adaptivity. Adds: - eg_cosine_batch.metal: batched cosine kernel, single- and multi-query variants, same -2.0 dim-mismatch/zero-norm sentinel as eg_cosine(). - eg_metal_cosine.h/.m: C-callable Objective-C bridge. Lazy one-time device/pipeline init, MTLResourceStorageModeShared buffers, returns false on ANY failure so callers fall back to the scalar CPU loop unconditionally — never partial, never throws. - eg_metal_cosine_stub.c: zero-dependency CPU-only implementation for non-Darwin builds (Linux CI) — same symbols, always returns false, no #ifdef needed at any call site. - build_vindex_bench.sh: one-command build, real bridge + Metal frameworks on Darwin, stub everywhere else. vindex_bench.c: brute_topk_metal / brute_topk_metal_batch call the bridge, falling back to the existing CPU brute_topk on any failure or EL_METAL_COSINE=0. The multi-query batched path exists because the first version (one GPU call per query) measured SLOWER than CPU at N~13.7k — it re-uploaded the full N*D matrix every query. Fixed by uploading the matrix once per query batch. Measured against a real nsbx-sandboxed clone of the live store (never :8742/:7770), 13,671 real embedded nodes, dim=768, 300 real queries: BRUTE-FORCE (CPU): 2.013 ms/query BRUTE-METAL (GPU): 0.117 ms/query (17.2x) id-recall vs CPU oracle: 0.9990 over 300 queries same-rank |Δdist|: max 2.98e-07, mean 7.53e-08 (float32 rounding, not a bug) Synthetic scaling sweep (13k -> 50k nodes, same dim/queries) shows the GPU speedup holding (~11x) as N grows toward the mathematical-foundations doc's 1.3M-node target, with CPU brute-force cost growing linearly as expected. Not wired into engram_activate or the daemon build (nsbx's _build_binary) — vindex_bench is a standalone offline tool, not part of the request-serving binary, so no engram_activate/server-latency claim is made here. The bridge is a reusable primitive (single eg_cosine_batch_metal + batched eg_cosine_batch_metal_multi) other call sites can adopt later without re-deriving any of this. Based on feat/reframe-region-setop (PR #109), not dev directly: the only genuine batch-cosine call site (vindex_bench.c) exists solely on this branch. Flagged explicitly in the PR description as a deliberate deviation from the original "base off dev" instruction.
338 lines
16 KiB
Objective-C
338 lines
16 KiB
Objective-C
/* 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;
|
|
}
|
|
}
|