/* eg_cosine_batch.metal — batched cosine similarity, one query vs N node vectors.
 *
 * GPU-shaped counterpart to eg_cosine() in el_runtime.c: same math, same
 * dim-mismatch/zero-norm sentinel (-2.0), applied to N independent rows in
 * parallel instead of one pair at a time in a CPU loop.
 *
 * Semantics MUST match eg_cosine() exactly:
 *   - inputs are raw, UNNORMALIZED vectors (nomic-embed-text magnitudes are
 *     not 1.0) — this kernel computes the full dot/(|a|*|b|) cosine, not a
 *     plain dot product.
 *   - a node whose declared dim differs from the query dim, or whose norm is
 *     zero, scores exactly -2.0 (below any valid cosine in [-1,1]), so a
 *     caller doing `if (score > threshold)` behaves identically whether the
 *     scalar or the batched path filled the array.
 *
 * Precision: Apple GPUs do not support double in Metal Shading Language —
 * everything here is float32. eg_cosine accumulates in CPU double, but its
 * *inputs* are float32 embeddings, so the achievable precision ceiling is
 * bounded by the input data regardless of accumulator width. To keep the
 * float32 reduction from drifting relative to the double-accumulated CPU
 * result across dim=768 terms, each thread accumulates with 4 independent
 * partial sums (unrolled) rather than one running scalar — the same
 * error-reduction trick already used by the CPU brute-force loop in
 * vindex_bench.c. The measured float-vs-double delta is reported in the PR
 * description; this is not assumed to be "close enough" without measurement.
 */
#include <metal_stdlib>
using namespace metal;

/* Per-dispatch invariants. `dim` is the query's dimensionality — the
 * dimensionality every comparable node vector must match. */
struct EgCosineParams {
    uint n;      /* number of node rows */
    uint dim;    /* vector width (both query and node rows are `dim` wide in
                  * the packed buffer; node_dims[] carries each node's REAL
                  * embedded dim for the mismatch check) */
};

/* One thread per node row. node_matrix is n*dim floats, row-major, packed at
 * `dim` stride regardless of a row's real dim (the CPU side zero-pads or
 * skips packing rows that don't match — see eg_cosine_batch_metal in
 * eg_metal_cosine.m for the exact packing contract). node_dims[i] is the
 * node's true emb_dim, used only for the mismatch sentinel — never used to
 * index, since every row is packed at uniform `dim` stride. */
kernel void eg_cosine_batch_kernel(
    device const float*  query        [[buffer(0)]],
    device const float*  node_matrix  [[buffer(1)]],
    device const int*    node_dims    [[buffer(2)]],
    constant EgCosineParams& p        [[buffer(3)]],
    device float*        out_scores   [[buffer(4)]],
    uint gid [[thread_position_in_grid]])
{
    if (gid >= p.n) return;

    if (node_dims[gid] != int(p.dim)) {
        out_scores[gid] = -2.0f;
        return;
    }

    device const float* row = node_matrix + (uint64_t)gid * (uint64_t)p.dim;

    /* 4-way partial accumulation — same shape as vindex_bench.c's brute_topk
     * unroll, done here for float32 accuracy rather than raw throughput. */
    float dot0 = 0.0f, dot1 = 0.0f, dot2 = 0.0f, dot3 = 0.0f;
    float na0  = 0.0f, na1  = 0.0f, na2  = 0.0f, na3  = 0.0f;
    float nb0  = 0.0f, nb1  = 0.0f, nb2  = 0.0f, nb3  = 0.0f;

    uint d = 0;
    uint dim4 = p.dim & ~3u;
    for (; d < dim4; d += 4) {
        float a0 = row[d],   b0 = query[d];
        float a1 = row[d+1], b1 = query[d+1];
        float a2 = row[d+2], b2 = query[d+2];
        float a3 = row[d+3], b3 = query[d+3];
        dot0 += a0*b0; dot1 += a1*b1; dot2 += a2*b2; dot3 += a3*b3;
        na0  += a0*a0; na1  += a1*a1; na2  += a2*a2; na3  += a3*a3;
        nb0  += b0*b0; nb1  += b1*b1; nb2  += b2*b2; nb3  += b3*b3;
    }
    float dot = (dot0 + dot1) + (dot2 + dot3);
    float na  = (na0  + na1)  + (na2  + na3);
    float nb  = (nb0  + nb1)  + (nb2  + nb3);
    for (; d < p.dim; d++) {
        float a = row[d], b = query[d];
        dot += a*b; na += a*a; nb += b*b;
    }

    if (na <= 0.0f || nb <= 0.0f) {
        out_scores[gid] = -2.0f;
        return;
    }
    out_scores[gid] = dot / sqrt(na * nb);
}

/* ── multi-query variant ──────────────────────────────────────────────────
 * Same per-pair math as eg_cosine_batch_kernel, but amortizes ONE upload of
 * node_matrix (the expensive part at real store size — 13k*768 floats is
 * ~42MB) across `nq` queries instead of re-uploading it once per query.
 * Measured need: a naive one-query-at-a-time loop calling the single-query
 * kernel nq times was SLOWER than the CPU oracle at N≈13.7k (re-gather +
 * re-upload dominated the actual compute) — this is the fix, not a
 * hypothetical optimization.
 *
 * 2D grid: x = node index [0,n), y = query index [0,nq). out_scores is
 * nq*n, row-major by query (out_scores[qid*n + nid]). */
struct EgCosineMultiParams { uint n; uint dim; uint nq; };

kernel void eg_cosine_batch_multi_kernel(
    device const float*  queries      [[buffer(0)]],   /* nq*dim */
    device const float*  node_matrix  [[buffer(1)]],   /* n*dim */
    device const int*    node_dims    [[buffer(2)]],   /* n */
    constant EgCosineMultiParams& p   [[buffer(3)]],
    device float*        out_scores   [[buffer(4)]],   /* nq*n */
    uint2 gid [[thread_position_in_grid]])
{
    uint nid = gid.x, qid = gid.y;
    if (nid >= p.n || qid >= p.nq) return;

    uint64_t out_idx = (uint64_t)qid * (uint64_t)p.n + (uint64_t)nid;

    if (node_dims[nid] != int(p.dim)) {
        out_scores[out_idx] = -2.0f;
        return;
    }

    device const float* row   = node_matrix + (uint64_t)nid * (uint64_t)p.dim;
    device const float* query = queries     + (uint64_t)qid * (uint64_t)p.dim;

    float dot0 = 0.0f, dot1 = 0.0f, dot2 = 0.0f, dot3 = 0.0f;
    float na0  = 0.0f, na1  = 0.0f, na2  = 0.0f, na3  = 0.0f;
    float nb0  = 0.0f, nb1  = 0.0f, nb2  = 0.0f, nb3  = 0.0f;

    uint d = 0;
    uint dim4 = p.dim & ~3u;
    for (; d < dim4; d += 4) {
        float a0 = row[d],   b0 = query[d];
        float a1 = row[d+1], b1 = query[d+1];
        float a2 = row[d+2], b2 = query[d+2];
        float a3 = row[d+3], b3 = query[d+3];
        dot0 += a0*b0; dot1 += a1*b1; dot2 += a2*b2; dot3 += a3*b3;
        na0  += a0*a0; na1  += a1*a1; na2  += a2*a2; na3  += a3*a3;
        nb0  += b0*b0; nb1  += b1*b1; nb2  += b2*b2; nb3  += b3*b3;
    }
    float dot = (dot0 + dot1) + (dot2 + dot3);
    float na  = (na0  + na1)  + (na2  + na3);
    float nb  = (nb0  + nb1)  + (nb2  + nb3);
    for (; d < p.dim; d++) {
        float a = row[d], b = query[d];
        dot += a*b; na += a*a; nb += b*b;
    }

    if (na <= 0.0f || nb <= 0.0f) {
        out_scores[out_idx] = -2.0f;
        return;
    }
    out_scores[out_idx] = dot / sqrt(na * nb);
}