/*
 * vsn_feature_selection_kernel.cu — TFT Variable Selection Network gating.
 *
 * Plan 4 Task 1B-i. Additive module — these kernels compile and load
 * but have NO production callers in this commit. Tasks 1B-iii / 1B-iv
 * wire them into the pre-encoder forward + backward path.
 *
 * VSN composition (per spec §4.E.1):
 *
 *   for each of the SL_NUM_FEATURE_GROUPS=6 feature groups g:
 *       v_g          = Linear_2_g(ReLU(Linear_1_g(state[..., group_g])))   // scalar importance
 *   logits           = concat(v_0, v_1, …, v_5)                            // [num_groups]
 *   mask             = softmax(logits)                                      // [num_groups]
 *   gated_state[..., i] = state[..., i] * mask[group_of(i)]
 *
 * The per-group `Linear_1` / `ReLU` / `Linear_2` MLPs are caller-side
 * cuBLAS GEMMs + an existing ReLU primitive (mirroring `gpu_grn.rs`'s
 * contract). This file owns ONLY the softmax-over-groups + per-feature
 * gate-multiply step (forward) and its symmetric backward.
 *
 * Layout convention: row-major `[B, state_dim_padded]` for tensors,
 * row-major `[B, num_groups]` for the mask. Matches `state_layout.cuh`
 * + `dt_layernorm_kernel`'s sample-major convention.
 *
 * Pearls applied:
 *   - feedback_no_atomicadd.md — d_dot[g] reductions use per-block shmem
 *     tree reduce; no atomicAdd across batch / features.
 *   - pearl_cold_path_no_exception_to_gpu_drives.md — every reduction
 *     (softmax max+sum, d_dot, inner_product) is GPU-side; no host
 *     roundtrip in any commit that wires this module.
 *
 * Indices outside any group's `[group_begins[g], group_ends[g])` range —
 * i.e. the 7-element padding tail past `SL_PADDING_START` — are PASSED
 * THROUGH unchanged on forward (`gated_state[..., i] = state[..., i]`)
 * and on backward (`d_state[..., i] = d_gated_state[..., i]`). This keeps
 * the padding bytes byte-identical to the raw input regardless of how
 * downstream consumers handle the 8-alignment tail (matches the contract
 * established by `assemble_state` in `state_layout.cuh`).
 */

#include <cuda_runtime.h>
#include <math_constants.h>
#include "state_layout.cuh"

/* Look up which group a feature index belongs to. Linear scan over 6
 * groups; returns -1 if the index is outside every group's half-open
 * `[begin, end)` range (the 7-element padding tail).
 *
 * The runtime cost is constant — `SL_NUM_FEATURE_GROUPS` is a compile-time
 * macro and the loop unrolls for `num_groups <= 6`. Caller passes the
 * runtime `num_groups` for signature stability with the host launcher
 * (which mirrors `gpu_grn.rs`'s style of taking dimensions as scalar
 * arguments instead of relying on macro values for graph-capture
 * friendliness).
 */
__device__ __forceinline__ int vsn_group_of(
    int feature_idx,
    int num_groups,
    const int* __restrict__ group_begins,
    const int* __restrict__ group_ends
) {
    #pragma unroll
    for (int g = 0; g < SL_NUM_FEATURE_GROUPS; ++g) {
        if (g >= num_groups) break;
        if (feature_idx >= group_begins[g] && feature_idx < group_ends[g]) {
            return g;
        }
    }
    return -1;
}

/* =====================================================================
 * vsn_softmax_and_gate_forward — softmax-over-groups + per-feature gate.
 *
 * For each sample b (one block per sample):
 *   1. Read `num_groups` logits into shared memory.
 *   2. Numerically-stable softmax across the 6 logits → `mask[b, g]`.
 *   3. Per-feature gate-multiply the [B, state_dim_padded] row, with
 *      passthrough for indices outside any group (padding tail).
 *
 * `mask` is saved for backward (the softmax Jacobian and the per-feature
 * d_state both need it). Caller allocates `mask [B, num_groups]`.
 *
 * Block size: 256 threads. Grid: B blocks.
 * Shared memory: num_groups floats (24 bytes for num_groups=6).
 *
 * Launch:  grid=(B), block=(256), smem=num_groups*sizeof(float).
 * ===================================================================== */
extern "C" __global__ void vsn_softmax_and_gate_forward(
    const float* __restrict__ state_in,         /* [B, state_dim_padded] row-major */
    const float* __restrict__ logits,           /* [B, num_groups] row-major */
    int B,
    int state_dim_padded,
    int num_groups,
    const int* __restrict__ group_begins,       /* [num_groups] */
    const int* __restrict__ group_ends,         /* [num_groups] */
    float* __restrict__ mask,                   /* [B, num_groups] OUT (saved) */
    float* __restrict__ gated_state             /* [B, state_dim_padded] OUT */
) {
    int b = blockIdx.x;
    if (b >= B) return;
    int tid   = threadIdx.x;
    int block = blockDim.x;

    extern __shared__ float smem[];             /* [num_groups] — softmax scratch */

    /* ----- Step 1: pull the 6 logits into shmem (one thread per group). ----- */
    if (tid < num_groups) {
        smem[tid] = logits[(size_t)b * num_groups + tid];
    }
    __syncthreads();

    /* ----- Step 2: numerically-stable softmax across num_groups elements. -----
     * Done by thread 0 only since num_groups is small (6) — a tree reduction
     * over 6 lanes wastes more cycles in shfl bookkeeping than a serial
     * pass over 6 floats. The result is broadcast to the whole block via
     * shmem (smem already holds logits → reuse for mask). */
    if (tid == 0) {
        float max_l = smem[0];
        for (int g = 1; g < num_groups; ++g) {
            float v = smem[g];
            if (v > max_l) max_l = v;
        }
        float sum_e = 0.0f;
        for (int g = 0; g < num_groups; ++g) {
            float e = expf(smem[g] - max_l);
            smem[g] = e;     /* stash unnormalised exp */
            sum_e += e;
        }
        /* Normalise. The host caller's invariant (`num_groups > 0` and
         * `logits` finite) keeps `sum_e > 0`; no defensive `(sum_e > 0) ? ...`
         * mask — masking would silently produce uniform mask if the invariant
         * ever broke; instead the NaN from divide-by-zero propagates loudly. */
        float inv_sum = 1.0f / sum_e;
        for (int g = 0; g < num_groups; ++g) {
            float m = smem[g] * inv_sum;
            smem[g] = m;
            mask[(size_t)b * num_groups + g] = m;
        }
    }
    __syncthreads();

    /* ----- Step 3: per-feature gate-multiply or passthrough. -----
     * Each thread strides across the row. group_of() returns -1 for indices
     * outside any group (the 7-element padding tail past SL_PADDING_START),
     * which we pass through unchanged. */
    const float* in_row  = state_in    + (size_t)b * state_dim_padded;
    float*       out_row = gated_state + (size_t)b * state_dim_padded;
    for (int i = tid; i < state_dim_padded; i += block) {
        int g = vsn_group_of(i, num_groups, group_begins, group_ends);
        float v = in_row[i];
        out_row[i] = (g >= 0) ? (v * smem[g]) : v;
    }
}

/* =====================================================================
 * vsn_softmax_and_gate_backward — full softmax-over-groups Jacobian.
 *
 * Forward was:
 *   mask[b, :]      = softmax(logits[b, :])
 *   gated[b, i]     = state[b, i] * mask[b, group_of(i)]   (or passthrough)
 *
 * Backward computes BOTH `d_logits [B, num_groups]` and
 * `d_state [B, state_dim_padded]` from `d_gated_state` + saved `mask`
 * + saved `state` (we need state to compute the dot product that drives
 * the softmax Jacobian).
 *
 * Steps per sample b (one block):
 *   1. Per-group dot: `d_dot[g] = sum_{i in group_g} d_gated[b, i] * state[b, i]`.
 *      Each thread accumulates partials across its strided slice; in-block
 *      shmem tree reduction collapses to one `d_dot[g]` per group.
 *   2. Softmax Jacobian: `inner = sum_h(mask[b, h] * d_dot[h])`; then
 *      `d_logit[g] = mask[b, g] * (d_dot[g] - inner)`.
 *   3. Per-feature d_state: `d_state[b, i] = d_gated[b, i] * mask[b, group_of(i)]`
 *      with passthrough on indices outside any group.
 *
 * Block size: 256 threads. Grid: B blocks.
 *
 * Shared memory layout:
 *   smem[0 ..    num_groups)              — per-group `d_dot` (final)
 *   smem[num_groups .. 2*num_groups)      — per-group `mask` (cached for inner)
 *   smem[2*num_groups .. 2*num_groups+block) — reduction scratch (one float per thread)
 *
 * Launch:  grid=(B), block=(256),
 *          smem=(2*num_groups + block) * sizeof(float).
 * ===================================================================== */
extern "C" __global__ void vsn_softmax_and_gate_backward(
    const float* __restrict__ d_gated_state,    /* [B, state_dim_padded] IN */
    const float* __restrict__ state,            /* [B, state_dim_padded] saved */
    const float* __restrict__ mask,             /* [B, num_groups] saved */
    int B,
    int state_dim_padded,
    int num_groups,
    const int* __restrict__ group_begins,       /* [num_groups] */
    const int* __restrict__ group_ends,         /* [num_groups] */
    float* __restrict__ d_logits,               /* [B, num_groups] OUT */
    float* __restrict__ d_state                 /* [B, state_dim_padded] OUT */
) {
    int b = blockIdx.x;
    if (b >= B) return;
    int tid   = threadIdx.x;
    int block = blockDim.x;

    extern __shared__ float smem[];
    float* sh_d_dot = smem;                          /* [num_groups] */
    float* sh_mask  = smem + num_groups;             /* [num_groups] */
    float* sh_red   = smem + 2 * num_groups;         /* [block]      */

    const float* dg_row = d_gated_state + (size_t)b * state_dim_padded;
    const float* st_row = state         + (size_t)b * state_dim_padded;
    const float* mk_row = mask          + (size_t)b * num_groups;

    /* Cache mask in shmem — every thread reads it many times below. */
    if (tid < num_groups) {
        sh_mask[tid] = mk_row[tid];
    }
    __syncthreads();

    /* ----- Step 1: per-group d_dot via in-block shmem tree reduction. -----
     * One pass per group. Each thread accumulates a partial across its strided
     * slice of [group_begins[g], group_ends[g]); then a standard tree reduce
     * collapses to thread 0's lane. No atomicAdd. */
    for (int g = 0; g < num_groups; ++g) {
        int gb_idx = group_begins[g];
        int ge_idx = group_ends[g];

        float local = 0.0f;
        for (int i = gb_idx + tid; i < ge_idx; i += block) {
            local += dg_row[i] * st_row[i];
        }
        sh_red[tid] = local;
        __syncthreads();

        for (int s = block / 2; s > 0; s >>= 1) {
            if (tid < s) sh_red[tid] += sh_red[tid + s];
            __syncthreads();
        }
        if (tid == 0) sh_d_dot[g] = sh_red[0];
        __syncthreads();
    }

    /* ----- Step 2: softmax Jacobian. ----- */
    if (tid == 0) {
        float inner = 0.0f;
        for (int h = 0; h < num_groups; ++h) {
            inner += sh_mask[h] * sh_d_dot[h];
        }
        for (int g = 0; g < num_groups; ++g) {
            d_logits[(size_t)b * num_groups + g] = sh_mask[g] * (sh_d_dot[g] - inner);
        }
    }
    /* No __syncthreads() needed before Step 3 — Step 3 only reads
     * `sh_mask`, which thread 0 didn't modify. */

    /* ----- Step 3: per-feature d_state. ----- */
    float* dx_row = d_state + (size_t)b * state_dim_padded;
    for (int i = tid; i < state_dim_padded; i += block) {
        int g = vsn_group_of(i, num_groups, group_begins, group_ends);
        float dg = dg_row[i];
        dx_row[i] = (g >= 0) ? (dg * sh_mask[g]) : dg;
    }
}
