bigmerge
b3f410fc91
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).
359 lines
17 KiB
Objective-C
359 lines
17 KiB
Objective-C
/* 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;
|
|
}
|