/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * topk_logits.cu — device-side exact top-k over the LM-head logits
 *
 * Target: sm_37 (K80, Kepler)
 * Compile: nvcc -ptx -arch=sm_37 topk_logits.cu -o topk_logits.ptx
 *
 * Why: on a small-BAR K80 host the host link runs at ~0.4 GB/s, so copying
 * the full 248k-vocab logits (1 MB) D2H costs ~4 ms per decoded token — the
 * single largest non-compute slice of the token budget. Greedy decode only
 * needs the argmax and the serving sampler (top_k=40) only needs the top 40
 * (id, logit) pairs, so we reduce on the die and copy back k pairs instead.
 *
 * Exactness: selection is ITERATIVE ARGMAX — k rounds, each finding the
 * global max of the remaining logits with the comparator
 *     (val_a > val_b) || (val_a == val_b && idx_a < idx_b)
 * which reproduces, value-for-value and tie-for-tie, the host's serial
 *     for (j...) if (logits[j] > logits[best]) best = j;
 * scan (lowest index wins ties) applied k times with previous winners
 * removed. No float arithmetic is performed on the logits — only compares —
 * so the returned values are the bit-exact LM-head outputs.
 *
 * Round structure (driven by fc_topk_logits_ctx in fc_cuda.c):
 *   1. fc_topk_partial_argmax — grid of blocks, grid-stride scan, one
 *      (val, idx) candidate per block.
 *   2. fc_topk_finish — single block reduces the per-block candidates,
 *      records the round winner in out_val/out_idx[slot], and (for k > 1)
 *      overwrites logits[winner] with -inf so the next round skips it.
 *   3. fc_topk_restore — after the last round, writes the recorded winner
 *      values back so the logits buffer is left intact.
 *
 * PTX symbols emitted (for embedding into ../cuda/fc_cuda_ptx.h):
 *   FC_PTX_TOPK_LOGITS -> kernels
 *     fc_topk_partial_argmax, fc_topk_finish, fc_topk_restore
 *
 * Kepler notes: pure fp32, shared-memory tree reductions (no __shfl).
 */

extern "C" {

#define FC_TOPK_NEG_INF __int_as_float(0xff800000)

/* Comparator: does candidate (av, ai) beat incumbent (bv, bi)?
 * Strictly greater value wins; equal value -> lower index wins. */
__device__ __forceinline__ int fc_topk_beats(float av, int ai,
                                             float bv, int bi)
{
    return (av > bv) || (av == bv && ai < bi);
}

/*
 * Stage 1: per-block argmax candidates.
 *
 * x:    [n] fp32 logits (read-only this stage)
 * pval: [gridDim.x] fp32 — per-block best value (output)
 * pidx: [gridDim.x] int32 — per-block best index (output)
 *
 * Launched with gridDim.x blocks (host caps at FC_TOPK_NPARTS), blockDim.x
 * threads, grid-stride loop. Shared memory: blockDim.x * (4 + 4) bytes.
 */
__global__ void fc_topk_partial_argmax(const float *x, int n,
                                       float *pval, int *pidx)
{
    extern __shared__ float sh[];
    float *s_val = sh;
    int   *s_idx = (int *)(sh + blockDim.x);

    int tid    = threadIdx.x;
    int stride = gridDim.x * blockDim.x;

    /* Thread-local scan in increasing index order: strict '>' keeps the
     * first-seen (lowest) index on ties, same as the host loop. */
    float best_v = FC_TOPK_NEG_INF;
    int   best_i = 0x7fffffff;
    for (int i = blockIdx.x * blockDim.x + tid; i < n; i += stride) {
        float v = x[i];
        if (v > best_v) { best_v = v; best_i = i; }
    }

    s_val[tid] = best_v;
    s_idx[tid] = best_i;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s &&
            fc_topk_beats(s_val[tid + s], s_idx[tid + s],
                          s_val[tid],     s_idx[tid])) {
            s_val[tid] = s_val[tid + s];
            s_idx[tid] = s_idx[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        pval[blockIdx.x] = s_val[0];
        pidx[blockIdx.x] = s_idx[0];
    }
}

/*
 * Stage 2: reduce the per-block candidates to the round winner.
 *
 * x:       [n] fp32 logits — winner slot overwritten with -inf when clobber
 * pval/pidx: stage-1 outputs, nparts entries
 * out_val: [k] fp32 — winner value written at [slot]
 * out_idx: [k] int32 — winner index written at [slot]
 * clobber: nonzero -> x[winner] = -inf (skip it in later rounds)
 *
 * Launched with 1 block, blockDim.x >= nparts threads.
 * Shared memory: blockDim.x * (4 + 4) bytes.
 */
__global__ void fc_topk_finish(float *x, const float *pval, const int *pidx,
                               int nparts, float *out_val, int *out_idx,
                               int slot, int clobber)
{
    extern __shared__ float sh[];
    float *s_val = sh;
    int   *s_idx = (int *)(sh + blockDim.x);

    int tid = threadIdx.x;
    if (tid < nparts) {
        s_val[tid] = pval[tid];
        s_idx[tid] = pidx[tid];
    } else {
        s_val[tid] = FC_TOPK_NEG_INF;
        s_idx[tid] = 0x7fffffff;
    }
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s &&
            fc_topk_beats(s_val[tid + s], s_idx[tid + s],
                          s_val[tid],     s_idx[tid])) {
            s_val[tid] = s_val[tid + s];
            s_idx[tid] = s_idx[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        out_val[slot] = s_val[0];
        out_idx[slot] = s_idx[0];
        if (clobber)
            x[s_idx[0]] = FC_TOPK_NEG_INF;
    }
}

/*
 * Stage 3: undo the -inf clobbers so the logits buffer stays intact.
 * Launched with 1 block, k threads.
 */
__global__ void fc_topk_restore(float *x, const float *out_val,
                                const int *out_idx, int k)
{
    int tid = threadIdx.x;
    if (tid < k)
        x[out_idx[tid]] = out_val[tid];
}

/* ============================================================
 * Batched (multi-slot) top-k (L8): the SAME iterative-argmax rounds run
 * for K slot rows of a [K][n] logits matrix IN PLACE — no per-slot D2D
 * copy, one sync and one K*k-pair D2H for the whole batch instead of a
 * serialized sync-per-user loop. Per row the scan order, comparator and
 * reduction shapes are identical to the single-row kernels above, so each
 * row's (id, val) pairs are bit-exact to fc_topk_logits run on that row.
 *
 * Layouts: x row s at x + s*n; stage-1 candidates at pval/pidx
 * [s*nparts + part]; winners at out_val/out_idx[s*k + round].
 * ============================================================ */

/* Stage 1: per-(block, slot) argmax candidates. Grid (nparts, nslots). */
__global__ void fc_topk_partial_argmax_slots(const float *x, int n,
                                             float *pval, int *pidx)
{
    extern __shared__ float sh[];
    float *s_val = sh;
    int   *s_idx = (int *)(sh + blockDim.x);

    const float *row = x + (size_t)blockIdx.y * n;
    int tid    = threadIdx.x;
    int stride = gridDim.x * blockDim.x;

    float best_v = FC_TOPK_NEG_INF;
    int   best_i = 0x7fffffff;
    for (int i = blockIdx.x * blockDim.x + tid; i < n; i += stride) {
        float v = row[i];
        if (v > best_v) { best_v = v; best_i = i; }
    }

    s_val[tid] = best_v;
    s_idx[tid] = best_i;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s &&
            fc_topk_beats(s_val[tid + s], s_idx[tid + s],
                          s_val[tid],     s_idx[tid])) {
            s_val[tid] = s_val[tid + s];
            s_idx[tid] = s_idx[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        pval[(size_t)blockIdx.y * gridDim.x + blockIdx.x] = s_val[0];
        pidx[(size_t)blockIdx.y * gridDim.x + blockIdx.x] = s_idx[0];
    }
}

/* Stage 2: per-slot round winner. Grid (1, nslots), blockDim >= nparts. */
__global__ void fc_topk_finish_slots(float *x, const float *pval,
                                     const int *pidx, int n, int nparts,
                                     float *out_val, int *out_idx, int k,
                                     int slot, int clobber)
{
    extern __shared__ float sh[];
    float *s_val = sh;
    int   *s_idx = (int *)(sh + blockDim.x);

    const int u = blockIdx.y;
    const float *pv = pval + (size_t)u * nparts;
    const int   *pi = pidx + (size_t)u * nparts;

    int tid = threadIdx.x;
    if (tid < nparts) {
        s_val[tid] = pv[tid];
        s_idx[tid] = pi[tid];
    } else {
        s_val[tid] = FC_TOPK_NEG_INF;
        s_idx[tid] = 0x7fffffff;
    }
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s &&
            fc_topk_beats(s_val[tid + s], s_idx[tid + s],
                          s_val[tid],     s_idx[tid])) {
            s_val[tid] = s_val[tid + s];
            s_idx[tid] = s_idx[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        out_val[(size_t)u * k + slot] = s_val[0];
        out_idx[(size_t)u * k + slot] = s_idx[0];
        if (clobber)
            x[(size_t)u * n + s_idx[0]] = FC_TOPK_NEG_INF;
    }
}

/* Stage 3: per-slot clobber restore. Grid (nslots), k threads. */
__global__ void fc_topk_restore_slots(float *x, const float *out_val,
                                      const int *out_idx, int n, int k)
{
    const int u = blockIdx.x;
    int tid = threadIdx.x;
    if (tid < k)
        x[(size_t)u * n + out_idx[(size_t)u * k + tid]] =
            out_val[(size_t)u * k + tid];
}

} /* extern "C" */