/* =====================================================================
 * rl_spectral_norm.cu — power iteration for max singular value estimation
 * + optional weight rescaling for spectral normalisation.
 *
 * One launch per weight matrix.
 * Grid=(1), Block=(min(rows, 256)).
 *
 * Persistent v_buffer across steps provides warm-start: one power
 * iteration per step is sufficient for tracking slowly-evolving weight
 * matrices under Adam/SGD updates.
 *
 * Discipline:
 *   - No atomicAdd (feedback_no_atomicadd.md) — block-tree reduce.
 *   - GPU-only (feedback_cpu_is_read_only.md).
 *   - Pre-compiled cubin (feedback_no_nvrtc.md).
 *   - ISV-driven sigma_max (feedback_isv_for_adaptive_bounds.md).
 * ===================================================================== */

#include <cuda_runtime.h>

#define SPECTRAL_NORM_MAX_INDEX 559

extern "C" __global__ void rl_spectral_norm(
    float* __restrict__ W,          /* [rows x cols] IN/OUT (rescaled if sigma > max) */
    float* __restrict__ v_buffer,   /* [cols] persistent across steps               */
    const float* __restrict__ isv,  /* ISV signals (device-mapped pinned)            */
    int rows,
    int cols
)
{
    if (blockIdx.x != 0) return;

    extern __shared__ float smem[];     /* >= blockDim.x floats */
    const int tid   = threadIdx.x;
    const int block = blockDim.x;

    /* ------------------------------------------------------------------
     * Step 1: Load v from v_buffer into shared memory.
     *
     * shared v lives at smem[0..cols). Threads stride-load then sync so
     * the entire vector is visible to all threads for the matrix-vector
     * products below.
     * ------------------------------------------------------------------ */
    float* sh_v = smem;
    for (int c = tid; c < cols; c += block)
        sh_v[c] = v_buffer[c];
    __syncthreads();

    /* ------------------------------------------------------------------
     * Step 2: u = W * v   (each thread handles one or more rows via
     *                       grid-stride loop over rows)
     *
     * u is stored in the global v_buffer area temporarily — we allocate
     * a register per-thread for u[tid] and then need the norm.
     * Since rows can exceed blockDim, each thread accumulates its row's
     * dot product in registers then we tree-reduce ||u||.
     * ------------------------------------------------------------------ */
    float* sh_reduce = &smem[cols];  /* scratch for tree-reduce: blockDim.x floats */

    /* We need u[rows] but rows can be > block. We do the norm accumulation
     * in-register: each thread holds the sum-of-squares for its row(s). */
    float u_norm_sq_local = 0.0f;

    /* We also need u values for step 3. Since rows <= typical weight
     * matrix size, we process in two passes:
     *   Pass A: compute ||u|| (accumulate u_i^2)
     *   Pass B: normalize u, then compute v_new.
     *
     * To avoid storing the full u vector in shared memory (could be
     * thousands of rows), we recompute u_i in pass B from W and v —
     * two GEMV passes is the standard power-iteration approach used by
     * the batched spectral norm kernel in dqn_utility_kernels.cu. */

    /* Pass A: accumulate ||u||^2 */
    for (int r = tid; r < rows; r += block) {
        float dot = 0.0f;
        for (int c = 0; c < cols; c++)
            dot += W[(size_t)r * cols + c] * sh_v[c];
        u_norm_sq_local += dot * dot;
    }

    /* Tree-reduce ||u||^2 */
    sh_reduce[tid] = u_norm_sq_local;
    __syncthreads();
    for (int s = block / 2; s > 0; s >>= 1) {
        if (tid < s) sh_reduce[tid] += sh_reduce[tid + s];
        __syncthreads();
    }
    float u_norm = sqrtf(sh_reduce[0] + 1e-12f);

    /* ------------------------------------------------------------------
     * Step 3: Normalize u (in-register on recompute), compute
     *         v_new = W^T * u_hat.
     *
     * u_hat[r] = (W[r,:] . v) / ||u||. Each thread handles a subset
     * of columns for v_new, looping over all rows to accumulate the
     * transpose product.
     *
     * We first need ALL u_hat values accessible. Since rows can exceed
     * shared-memory capacity, we store u_hat into the v_buffer
     * temporarily (it has >= max(rows, cols) elements? No — it has
     * exactly `cols` elements). So we use a different strategy:
     *
     * For each column c (stride over block), loop over rows r, compute
     * u_hat[r] on the fly, accumulate W[r,c]*u_hat[r].
     *
     * This is O(rows*cols/block) = same cost as a GEMV.
     * ------------------------------------------------------------------ */
    float inv_u_norm = 1.0f / (u_norm + 1e-12f);

    for (int c = tid; c < cols; c += block) {
        float accum = 0.0f;
        for (int r = 0; r < rows; r++) {
            /* Recompute u_hat[r] = dot(W[r,:], v) / ||u|| */
            float dot = 0.0f;
            for (int cc = 0; cc < cols; cc++)
                dot += W[(size_t)r * cols + cc] * sh_v[cc];
            float u_hat_r = dot * inv_u_norm;
            accum += W[(size_t)r * cols + c] * u_hat_r;
        }
        v_buffer[c] = accum;
    }
    __syncthreads();

    /* ------------------------------------------------------------------
     * Step 4: Normalize v_new and store back to v_buffer.
     * ------------------------------------------------------------------ */
    float v_norm_sq_local = 0.0f;
    for (int c = tid; c < cols; c += block) {
        float vc = v_buffer[c];
        v_norm_sq_local += vc * vc;
    }
    sh_reduce[tid] = v_norm_sq_local;
    __syncthreads();
    for (int s = block / 2; s > 0; s >>= 1) {
        if (tid < s) sh_reduce[tid] += sh_reduce[tid + s];
        __syncthreads();
    }
    float v_norm = sqrtf(sh_reduce[0] + 1e-12f);
    float inv_v_norm = 1.0f / (v_norm + 1e-12f);

    for (int c = tid; c < cols; c += block)
        v_buffer[c] *= inv_v_norm;
    __syncthreads();

    /* Reload normalized v_new into shared memory for sigma computation. */
    for (int c = tid; c < cols; c += block)
        sh_v[c] = v_buffer[c];
    __syncthreads();

    /* ------------------------------------------------------------------
     * Step 5: sigma = ||W * v_new||
     *
     * Recompute u with the updated v to get the singular value estimate.
     * ------------------------------------------------------------------ */
    float sigma_sq_local = 0.0f;
    for (int r = tid; r < rows; r += block) {
        float dot = 0.0f;
        for (int c = 0; c < cols; c++)
            dot += W[(size_t)r * cols + c] * sh_v[c];
        sigma_sq_local += dot * dot;
    }
    sh_reduce[tid] = sigma_sq_local;
    __syncthreads();
    for (int s = block / 2; s > 0; s >>= 1) {
        if (tid < s) sh_reduce[tid] += sh_reduce[tid + s];
        __syncthreads();
    }
    float sigma = sqrtf(sh_reduce[0] + 1e-12f);

    /* ------------------------------------------------------------------
     * Step 6: Conditional rescaling.
     *
     * max_sigma from ISV; default 3.0 if ISV slot reads 0.0 (cold-start
     * sentinel per pearl_hardcoded_bootstrap_target_collision.md).
     * ------------------------------------------------------------------ */
    float max_sigma = isv[SPECTRAL_NORM_MAX_INDEX];
    if (max_sigma < 1e-6f) max_sigma = 3.0f;

    if (sigma > max_sigma) {
        float scale = max_sigma / sigma;
        int n_total = rows * cols;
        for (int i = tid; i < n_total; i += block)
            W[i] *= scale;
    }
}
