/* 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 #import #include "eg_metal_cosine.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; } } 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 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; } } 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 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; } }