/* 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 #import #include "eg_cosine_batch_strategy.h" #include #include /* ── BEGIN embedded shader source (keep in sync with eg_cosine_batch.metal) ── */ static const char* kEgCosineBatchMetalSrc = "#include \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 g_device = nil; static id g_queue = nil; static id g_pipeline = nil; /* single-query kernel */ static id 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 dev = MTLCreateSystemDefaultDevice(); if (!dev) return false; id q = [dev newCommandQueue]; if (!q) return false; NSError* err = nil; NSString* src = [NSString stringWithUTF8String:kEgCosineBatchMetalSrc]; MTLCompileOptions* opts = [MTLCompileOptions new]; id lib = [dev newLibraryWithSource:src options:opts error:&err]; if (!lib) return false; id fn = [lib newFunctionWithName:@"eg_cosine_batch_kernel"]; if (!fn) return false; id pipe = [dev newComputePipelineStateWithFunction:fn error:&err]; if (!pipe) return false; id fnMulti = [lib newFunctionWithName:@"eg_cosine_batch_multi_kernel"]; if (!fnMulti) return false; id 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 bufQuery = [g_device newBufferWithBytes:query length:dim * sizeof(float) options:MTLResourceStorageModeShared]; id bufMatrix = [g_device newBufferWithBytes:matrix length:nu * dim * sizeof(float) options:MTLResourceStorageModeShared]; id bufDims = [g_device newBufferWithBytes:dims_i32 length:nu * sizeof(int32_t) options:MTLResourceStorageModeShared]; EgCosineParamsC params = { (uint32_t)nu, (uint32_t)dim }; id bufParams = [g_device newBufferWithBytes:¶ms length:sizeof(params) options:MTLResourceStorageModeShared]; id bufOut = [g_device newBufferWithLength:nu * sizeof(float) options:MTLResourceStorageModeShared]; free(matrix); free(dims_i32); if (!bufQuery || !bufMatrix || !bufDims || !bufParams || !bufOut) return false; id cmd = [g_queue commandBuffer]; if (!cmd) return false; id 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 bufMatrix = [g_device newBufferWithBytes:matrix length:nu * dim * sizeof(float) options:MTLResourceStorageModeShared]; id bufDims = [g_device newBufferWithBytes:dims_i32 length:nu * sizeof(int32_t) options:MTLResourceStorageModeShared]; id bufQueries = [g_device newBufferWithBytes:queries length:nqu * dim * sizeof(float) options:MTLResourceStorageModeShared]; EgCosineMultiParamsC params = { (uint32_t)nu, (uint32_t)dim, (uint32_t)nqu }; id bufParams = [g_device newBufferWithBytes:¶ms length:sizeof(params) options:MTLResourceStorageModeShared]; id bufOut = [g_device newBufferWithLength:nqu * nu * sizeof(float) options:MTLResourceStorageModeShared]; free(matrix); free(dims_i32); if (!bufMatrix || !bufDims || !bufQueries || !bufParams || !bufOut) return false; id cmd = [g_queue commandBuffer]; if (!cmd) return false; id 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; }