// grad_norm.cu — GPU-resident grad-norm + clip-scale computation.
//
// Replaces the host-side memcpy_dtoh + sum-of-squares loop in
// `Mamba2AdamW::step_from_buffers`. The previous path did 9× dtoh per
// training step — each a stream sync barrier costing ~1ms+ on the hot
// path. This kernel suite keeps the entire computation on GPU.
//
// Workflow (per training step):
//   1. Zero `grad_norm_sq_d` once via `memset_zeros`.
//   2. For each of N gradient tensors:
//      - launch `grad_norm_sq_phase1` → writes block-tree partials to scratch
//      - launch `grad_norm_sq_phase2` → reduces partials, ACCUMULATES into grad_norm_sq_d
//   3. launch `grad_clip_scale` → reads grad_norm_sq_d, writes
//      `grad_scale_d = min(1, max_norm / sqrt(norm_sq))`
//   4. AdamW kernels read `grad_scale_d` as a device pointer (not host scalar).
//
// All per-block reductions use shared-memory tree reduce per
// `feedback_no_atomicadd.md`. Phase 2 uses a single-thread `+=` for
// the cross-tensor accumulation, which is race-free because phase-2
// launches are sequential in the stream and each launch writes exactly
// one float to the accumulator slot.

#include <cuda_runtime.h>

#define GRAD_NORM_BLOCK 256

/// Phase 1: per-block tree-reduce of `x[]*x[]` → block_partials[blockIdx.x].
extern "C" __global__ void grad_norm_sq_phase1(
    const float* __restrict__ x,
    int n,
    float* __restrict__ block_partials
) {
    __shared__ float smem[GRAD_NORM_BLOCK];
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + tid;
    smem[tid] = (gid < n) ? x[gid] * x[gid] : 0.0f;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    if (tid == 0) block_partials[blockIdx.x] = smem[0];
}

/// Phase 2: single-block tree-reduce of `block_partials[0..n_blocks]`
/// → accumulates into `accum[0]`. `accumulate=1` adds; `accumulate=0`
/// overwrites (use for the first call in a step).
extern "C" __global__ void grad_norm_sq_phase2(
    const float* __restrict__ block_partials,
    int n_blocks,
    float* __restrict__ accum,
    int accumulate
) {
    __shared__ float smem[GRAD_NORM_BLOCK];
    int tid = threadIdx.x;
    float v = 0.0f;
    for (int i = tid; i < n_blocks; i += blockDim.x) {
        v += block_partials[i];
    }
    smem[tid] = v;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    if (tid == 0) {
        float prev = (accumulate != 0) ? accum[0] : 0.0f;
        accum[0] = prev + smem[0];
    }
}

/// Final clamp: grad_scale = min(1, max_norm / sqrt(norm_sq + eps)).
/// One-thread kernel — writes a single float to `grad_scale_out[0]`.
/// AdamW kernels read this device pointer instead of a host scalar.
extern "C" __global__ void grad_clip_scale(
    const float* __restrict__ norm_sq,
    float max_norm,
    float* __restrict__ grad_scale_out
) {
    if (threadIdx.x != 0 || blockIdx.x != 0) return;
    const float eps = 1e-12f;
    float norm = sqrtf(norm_sq[0] + eps);
    grad_scale_out[0] = (norm > max_norm) ? (max_norm / norm) : 1.0f;
}