/*
 * grn_kernel.cu — Gated Residual Network forward + backward kernels.
 *
 * Plan 4 Task 2c.1. Additive module — these kernels compile and load
 * but have NO production callers in this commit. Task 2c.3+4 wires
 * them into the trunk encoder/backward path.
 *
 * GRN block formula (per spec §4.E.2):
 *   GRN(x) = LayerNorm(P(x) + GLU(Linear_b(ELU(Linear_a(x)))))
 *
 *   where:
 *     Linear_a:   [hidden, in_dim]
 *     ELU:        x if x >= 0 else exp(x) - 1   (canonical α = 1.0)
 *     Linear_b:   [2*hidden, hidden]            (output 2× for GLU split)
 *     GLU(z):     split z into [a, b] each [hidden]; output = a * sigmoid(b)
 *     P(x):       Linear_residual(x) if dim_in != dim_out, else identity
 *     LayerNorm:  per-sample mean+var; gamma, beta both [hidden]
 *
 * Composition: cuBLAS GEMMs (Linear_a, Linear_b, Linear_residual) +
 * the 6 custom kernels in this file (grn_elu_inplace, grn_glu_forward,
 * grn_residual_layernorm_forward, grn_layernorm_backward,
 * grn_glu_backward, grn_elu_backward).
 *
 * IMPORTANT: the LayerNorm backward in this file is the FULL JACOBIAN,
 * not the simplified element-wise d_x = d_out * gamma * rstd that
 * attention_backward_kernel.cu uses. The simplified version works in
 * attention's residual+LN composition because the residual gradient
 * absorbs the missing terms — for GRN's trunk-encoder backward the
 * full Jacobian is required:
 *
 *   d_out_g[b, d] = d_out[b, d] * gamma[d]
 *   s1[b]         = sum_d(d_out_g[b, :])
 *   s2[b]         = sum_d(d_out_g[b, :] * normed[b, :])
 *   d_x[b, d]     = (1/H) * rstd[b] *
 *                   (H * d_out_g[b, d] - s1[b] - normed[b, d] * s2[b])
 *
 * Layout convention for GRN: tensors are ROW-MAJOR [B, H] (sample-major,
 * matches spec §4.E.2 and dt_layernorm_kernel in dt_kernels.cu). This is
 * intentionally different from attention_kernel.cu's [D, B] col-major —
 * the trunk encoder downstream of GRN uses row-major buffers, so doing
 * the LN per row in row-major avoids a transpose.
 *
 * Pearls applied:
 *   - pearl_cold_path_no_exception_to_gpu_drives.md — every reduction
 *     (s1, s2, d_gamma, d_beta) is GPU-side; no host roundtrip even
 *     in cold-path use.
 *   - feedback_no_atomicadd.md — d_gamma, d_beta accumulation across
 *     batch uses per-block partial reduction + final reduction pass,
 *     not atomicAdd.
 */

#include <cuda_runtime.h>
#include <math_constants.h>

/* =====================================================================
 * grn_elu_inplace — Element-wise ELU activation (canonical α = 1.0).
 *
 *   f(x) = x          if x >= 0
 *        = exp(x) - 1 if x  < 0
 *
 * In-place mutation of `x`. The original value is NOT saved here — the
 * backward path needs the post-activation value (via x_saved) and uses
 * the identity exp(x_pre) = x_post + 1 for x_pre < 0 to avoid extra
 * forward stash. (Equivalently the caller passes the post-activation
 * tensor as x_saved to grn_elu_backward.)
 *
 * Launch:  grid=ceil(n/256), block=256.
 * Shapes:  x is any flat tensor of length n.
 * ===================================================================== */
extern "C" __global__ void grn_elu_inplace(
    float* __restrict__ x,
    int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) return;
    float v = x[idx];
    /* Canonical ELU with α = 1.0; expm1f is numerically preferable to
       exp(v) - 1 for small negative v. */
    x[idx] = (v >= 0.0f) ? v : expm1f(v);
}

/* =====================================================================
 * grn_glu_forward — Gated Linear Unit forward.
 *
 *   in_2h is row-major [B, 2*H]. Per sample b, the row is split into
 *   chunk_a = in_2h[b, 0..H] and chunk_b = in_2h[b, H..2*H]. Output:
 *
 *       out_h[b, d] = chunk_a[b, d] * sigmoid(chunk_b[b, d])
 *
 *   sigmoid(chunk_b) is saved per element for backward — GLU's d/dx of
 *   the sigmoid factor needs both chunk_a and σ(b)·(1-σ(b)) which is
 *   cheaper to recompute from σ(b) than to reload b and re-apply σ.
 *
 * Why GLU: per spec §4.E.2 the GRN's gating mechanism. The sigmoid gate
 * lets the model attenuate its own non-linear branch independently of
 * the residual, which empirically stabilises depth.
 *
 * Launch:  grid=(ceil(B*H/256)), block=256.
 *          Flat 1D dispatch over (B*H) since the split is per element.
 * ===================================================================== */
extern "C" __global__ void grn_glu_forward(
    const float* __restrict__ in_2h,        /* [B, 2*H] row-major */
    int B, int H,
    float* __restrict__ out_h,              /* [B, H] row-major */
    float* __restrict__ save_sigmoid_b      /* [B, H] for backward */
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * H;
    if (idx >= total) return;
    int b = idx / H;
    int d = idx - b * H;
    /* Row stride is 2*H, chunk_a starts at column 0, chunk_b at column H. */
    const float* row = in_2h + (size_t)b * (2 * H);
    float a_val = row[d];
    float b_val = row[H + d];
    /* Numerically stable sigmoid: 1 / (1 + exp(-x)). For very large |b_val|
       this saturates safely to 0/1 in fp32. */
    float sig_b = 1.0f / (1.0f + expf(-b_val));
    save_sigmoid_b[idx] = sig_b;
    out_h[idx] = a_val * sig_b;
}

/* =====================================================================
 * grn_residual_layernorm_forward — Fused residual + LayerNorm.
 *
 * Per spec §4.E.2 GRN: y = LayerNorm(gated + residual). Implemented
 * fused so we never materialize the (gated + residual) intermediate
 * to global memory before LN.
 *
 *   For each sample b:
 *       z[b, d]      = gated_h[b, d] + residual_h[b, d]
 *       mean[b]      = (1/H) * sum_d(z[b, :])
 *       var[b]       = (1/H) * sum_d((z[b, d] - mean[b])^2)
 *       rstd[b]      = 1 / sqrt(var[b] + eps)
 *       normed[b, d] = (z[b, d] - mean[b]) * rstd[b]
 *       y[b, d]      = gamma[d] * normed[b, d] + beta[d]
 *
 * save_mean / save_rstd / save_normed are stashed for backward. Saving
 * `normed` (already affine-naked) means the backward kernel doesn't
 * need to reload (gated, residual) and recompute the centring step —
 * memory is cheap, recompute is not.
 *
 * Reduction strategy (matches dt_layernorm_kernel):
 *   - 1 block per sample (grid.x = B)
 *   - block size = min(H, 256)
 *   - Two-pass: pass 1 = mean (warp shuffle + smem cross-warp); pass 2
 *     = variance using the published mean.
 *   - eps = 1e-5f (matches existing LN kernels in attention/dt).
 *
 * Launch:  grid=(B), block=(min(H, 256)).
 * Shapes:  gated_h, residual_h, y, save_normed are [B, H] row-major.
 *          gamma, beta are [H]. save_mean, save_rstd are [B].
 * ===================================================================== */
extern "C" __global__ void grn_residual_layernorm_forward(
    const float* __restrict__ gated_h,      /* [B, H] */
    const float* __restrict__ residual_h,   /* [B, H] */
    const float* __restrict__ gamma,        /* [H] */
    const float* __restrict__ beta,         /* [H] */
    int B, int H,
    float* __restrict__ y,                  /* [B, H] */
    float* __restrict__ save_mean,          /* [B] */
    float* __restrict__ save_rstd,          /* [B] */
    float* __restrict__ save_normed         /* [B, H] */
) {
    int b = blockIdx.x;
    if (b >= B) return;
    int tid = threadIdx.x;

    const float* g_row = gated_h + (size_t)b * H;
    const float* r_row = residual_h + (size_t)b * H;
    float* y_row       = y + (size_t)b * H;
    float* n_row       = save_normed + (size_t)b * H;

    /* ----- Pass 1: mean of (gated + residual) per row. ----- */
    float local_sum = 0.0f;
    for (int d = tid; d < H; d += blockDim.x) {
        local_sum += g_row[d] + r_row[d];
    }
    /* Warp-level reduction. */
    for (int off = 16; off > 0; off >>= 1)
        local_sum += __shfl_xor_sync(0xFFFFFFFF, local_sum, off);
    /* Cross-warp reduction via shared memory (up to 8 warps for blockDim 256). */
    __shared__ float warp_sums[8];
    int warp_id = tid >> 5;
    int lane    = tid & 31;
    if (lane == 0 && warp_id < 8) warp_sums[warp_id] = local_sum;
    __syncthreads();
    float row_sum;
    if (warp_id == 0) {
        int n_warps = (blockDim.x + 31) >> 5;
        float v = (lane < n_warps) ? warp_sums[lane] : 0.0f;
        for (int off = 16; off > 0; off >>= 1)
            v += __shfl_xor_sync(0xFFFFFFFF, v, off);
        row_sum = v;
    }
    __shared__ float sh_mean;
    if (tid == 0) sh_mean = row_sum / (float)H;
    __syncthreads();
    float mean = sh_mean;

    /* ----- Pass 2: variance using published mean. ----- */
    float local_var = 0.0f;
    for (int d = tid; d < H; d += blockDim.x) {
        float diff = (g_row[d] + r_row[d]) - mean;
        local_var += diff * diff;
    }
    for (int off = 16; off > 0; off >>= 1)
        local_var += __shfl_xor_sync(0xFFFFFFFF, local_var, off);
    if (lane == 0 && warp_id < 8) warp_sums[warp_id] = local_var;
    __syncthreads();
    float row_var;
    if (warp_id == 0) {
        int n_warps = (blockDim.x + 31) >> 5;
        float v = (lane < n_warps) ? warp_sums[lane] : 0.0f;
        for (int off = 16; off > 0; off >>= 1)
            v += __shfl_xor_sync(0xFFFFFFFF, v, off);
        row_var = v;
    }
    __shared__ float sh_rstd;
    if (tid == 0) sh_rstd = rsqrtf(row_var / (float)H + 1e-5f);
    __syncthreads();
    float rstd = sh_rstd;

    if (tid == 0) {
        save_mean[b] = mean;
        save_rstd[b] = rstd;
    }

    /* ----- Pass 3: normalise + affine, write y and save_normed. ----- */
    for (int d = tid; d < H; d += blockDim.x) {
        float z      = g_row[d] + r_row[d];
        float normed = (z - mean) * rstd;
        n_row[d] = normed;
        y_row[d] = gamma[d] * normed + beta[d];
    }
}

/* =====================================================================
 * grn_layernorm_backward_dx — Full-Jacobian LN backward (d_x only).
 *
 * Computes the FULL LayerNorm Jacobian, NOT the simplified
 *     d_x = d_out * gamma * rstd
 * pattern that attn_layer_norm_bwd_dx uses. The simplified form is
 * correct only when there is a residual gradient flowing in parallel
 * that absorbs the missing s1/s2 terms — GRN's trunk-encoder backward
 * has no such absorbing channel, so we need the exact Jacobian:
 *
 *   d_out_g[b, d] = d_y[b, d] * gamma[d]
 *   s1[b]         = sum_d(d_out_g[b, :])
 *   s2[b]         = sum_d(d_out_g[b, :] * save_normed[b, :])
 *   d_x[b, d]     = (1/H) * rstd[b] *
 *                   (H * d_out_g[b, d] - s1[b] - save_normed[b, d] * s2[b])
 *
 * d_x flows backward through both the residual add (so the caller
 * routes the same d_x to BOTH the gated branch and the residual
 * branch — that is just an outer-edge wiring detail handled in 2c.3+4,
 * not here).
 *
 * Reduction strategy: 1 block per sample, two-pass.
 *   - Pass 1: per-thread accumulate (d_out_g_d) and (d_out_g_d * normed_d),
 *     then warp + cross-warp reduce → publish s1, s2 in shared memory.
 *   - Pass 2: per-thread compute d_x[b, d] using s1, s2.
 *
 * Launch:  grid=(B), block=(min(H, 256)).
 * Shapes:  d_y, save_normed, d_x are [B, H] row-major.
 *          gamma is [H]. save_rstd is [B].
 * ===================================================================== */
extern "C" __global__ void grn_layernorm_backward_dx(
    const float* __restrict__ d_y,          /* [B, H] grad in */
    const float* __restrict__ save_normed,  /* [B, H] from forward */
    const float* __restrict__ save_rstd,    /* [B] from forward */
    const float* __restrict__ gamma,        /* [H] */
    int B, int H,
    float* __restrict__ d_x                 /* [B, H] grad out */
) {
    int b = blockIdx.x;
    if (b >= B) return;
    int tid = threadIdx.x;

    const float* dy_row = d_y + (size_t)b * H;
    const float* nm_row = save_normed + (size_t)b * H;
    float*       dx_row = d_x + (size_t)b * H;
    float        rstd   = save_rstd[b];

    /* ----- Pass 1: compute s1 and s2 via two parallel reductions. ----- */
    float local_s1 = 0.0f;
    float local_s2 = 0.0f;
    for (int d = tid; d < H; d += blockDim.x) {
        float dyg = dy_row[d] * gamma[d];   /* d_out_g */
        local_s1 += dyg;
        local_s2 += dyg * nm_row[d];
    }
    /* Warp-level reductions for s1 and s2 in lockstep. */
    for (int off = 16; off > 0; off >>= 1) {
        local_s1 += __shfl_xor_sync(0xFFFFFFFF, local_s1, off);
        local_s2 += __shfl_xor_sync(0xFFFFFFFF, local_s2, off);
    }
    /* Cross-warp reduction via shared memory: 2 × 8 entries (one each for s1, s2). */
    __shared__ float warp_s1[8];
    __shared__ float warp_s2[8];
    int warp_id = tid >> 5;
    int lane    = tid & 31;
    if (lane == 0 && warp_id < 8) {
        warp_s1[warp_id] = local_s1;
        warp_s2[warp_id] = local_s2;
    }
    __syncthreads();
    float row_s1, row_s2;
    if (warp_id == 0) {
        int n_warps = (blockDim.x + 31) >> 5;
        float v1 = (lane < n_warps) ? warp_s1[lane] : 0.0f;
        float v2 = (lane < n_warps) ? warp_s2[lane] : 0.0f;
        for (int off = 16; off > 0; off >>= 1) {
            v1 += __shfl_xor_sync(0xFFFFFFFF, v1, off);
            v2 += __shfl_xor_sync(0xFFFFFFFF, v2, off);
        }
        row_s1 = v1;
        row_s2 = v2;
    }
    __shared__ float sh_s1;
    __shared__ float sh_s2;
    if (tid == 0) { sh_s1 = row_s1; sh_s2 = row_s2; }
    __syncthreads();
    float s1 = sh_s1;
    float s2 = sh_s2;

    /* ----- Pass 2: full-Jacobian d_x. -----
     * d_x[b, d] = (1/H) * rstd * (H * d_out_g[b, d] - s1 - normed[b, d] * s2)
     */
    float inv_H = 1.0f / (float)H;
    for (int d = tid; d < H; d += blockDim.x) {
        float dyg    = dy_row[d] * gamma[d];
        float normed = nm_row[d];
        dx_row[d] = inv_H * rstd * ((float)H * dyg - s1 - normed * s2);
    }
}

/* =====================================================================
 * grn_layernorm_backward_dgamma_dbeta_p1 — Per-block partial reduction.
 *
 * d_gamma[d] = sum_b(d_y[b, d] * save_normed[b, d])    (over batch)
 * d_beta[d]  = sum_b(d_y[b, d])                        (over batch)
 *
 * No atomicAdd per feedback_no_atomicadd.md. Uses two-phase reduction:
 *   Phase 1 (this kernel): grid=(num_blocks_b, H), block=256.
 *           Each block reduces a stripe of B for a single d, writes
 *           gamma_partials[block_b * H + d] and beta_partials[...].
 *   Phase 2 (next kernel): final reduce over num_blocks_b.
 *
 * Layout note: the partials are stored stripe-major
 *     gamma_partials[block_b, d] = block_b * H + d
 * so the phase-2 final reduce iterates with unit stride across blocks
 * for the same d.
 *
 * Shared memory: 2 × blockDim.x floats (one buffer for gamma partial,
 * one for beta partial; they reduce in lockstep).
 *
 * Launch:  grid=(num_blocks_b, H), block=256, smem=2*256*sizeof(float).
 * ===================================================================== */
extern "C" __global__ void grn_layernorm_backward_dgamma_dbeta_p1(
    const float* __restrict__ d_y,              /* [B, H] */
    const float* __restrict__ save_normed,      /* [B, H] */
    int B, int H,
    float* __restrict__ gamma_partials,         /* [num_blocks_b, H] */
    float* __restrict__ beta_partials           /* [num_blocks_b, H] */
) {
    extern __shared__ float sdata[];
    float* sg = sdata;
    float* sb = sdata + blockDim.x;

    int d   = blockIdx.y;       /* hidden dim this block handles */
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + tid;    /* batch index */

    float gv = 0.0f;
    float bv = 0.0f;
    if (gid < B) {
        size_t off = (size_t)gid * H + d;       /* row-major [B, H] */
        float dy_val = d_y[off];
        float nm_val = save_normed[off];
        gv = dy_val * nm_val;
        bv = dy_val;
    }
    sg[tid] = gv;
    sb[tid] = bv;
    __syncthreads();

    /* Smem tree reduction for both partials in lockstep. */
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sg[tid] += sg[tid + s];
            sb[tid] += sb[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        gamma_partials[blockIdx.x * H + d] = sg[0];
        beta_partials[blockIdx.x * H + d]  = sb[0];
    }
}

/* =====================================================================
 * grn_layernorm_backward_dgamma_dbeta_p2 — Final reduce over partials.
 *
 * For each d, sums gamma_partials[0..num_blocks_b, d] → d_gamma[d] and
 * beta_partials[0..num_blocks_b, d] → d_beta[d]. This kernel produces
 * the final batch-summed d_gamma and d_beta.
 *
 * Launch:  grid=(ceil(H/256)), block=256.
 *          One thread per hidden dim; sequential sum over blocks.
 * ===================================================================== */
extern "C" __global__ void grn_layernorm_backward_dgamma_dbeta_p2(
    const float* __restrict__ gamma_partials,   /* [num_blocks_b, H] */
    const float* __restrict__ beta_partials,    /* [num_blocks_b, H] */
    int H, int num_blocks_b,
    float* __restrict__ d_gamma,                /* [H] */
    float* __restrict__ d_beta                  /* [H] */
) {
    int d = blockIdx.x * blockDim.x + threadIdx.x;
    if (d >= H) return;
    float dg = 0.0f;
    float db = 0.0f;
    for (int i = 0; i < num_blocks_b; i++) {
        dg += gamma_partials[(size_t)i * H + d];
        db += beta_partials[(size_t)i * H + d];
    }
    d_gamma[d] = dg;
    d_beta[d]  = db;
}

/* =====================================================================
 * grn_glu_backward — Gated Linear Unit backward.
 *
 * Forward was: out_h[b, d] = chunk_a[b, d] * sigmoid(chunk_b[b, d])
 *
 * Let σ = sigmoid(chunk_b). The forward saved σ in save_sigmoid_b.
 * Standard chain rule gives:
 *
 *   d_chunk_a[b, d] = d_out_h[b, d] * σ
 *   d_chunk_b[b, d] = d_out_h[b, d] * chunk_a * σ * (1 - σ)
 *
 * The forward stored chunk_a as the first H columns of in_2h_saved
 * (the Linear_b output before GLU split). chunk_b lives in the second H
 * columns but is NOT needed for backward — σ already encapsulates it.
 *
 * Output layout: d_in_2h is [B, 2*H] row-major; the first H columns
 * receive d_chunk_a, the second H columns receive d_chunk_b.
 *
 * Launch:  grid=(ceil(B*H/256)), block=256.
 *          Flat 1D dispatch over (B*H); each thread writes both
 *          d_chunk_a[b, d] and d_chunk_b[b, d] for its (b, d).
 * ===================================================================== */
extern "C" __global__ void grn_glu_backward(
    const float* __restrict__ d_out_h,          /* [B, H] */
    const float* __restrict__ in_2h_saved,      /* [B, 2*H] from forward (Linear_b out) */
    const float* __restrict__ save_sigmoid_b,   /* [B, H] */
    int B, int H,
    float* __restrict__ d_in_2h                 /* [B, 2*H] grad out */
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * H;
    if (idx >= total) return;
    int b = idx / H;
    int d = idx - b * H;

    /* Read forward-saved σ and the original chunk_a (column d in the first
       half of row b of in_2h_saved). */
    float sig_b = save_sigmoid_b[idx];
    const float* row_2h = in_2h_saved + (size_t)b * (2 * H);
    float chunk_a = row_2h[d];
    float dout    = d_out_h[idx];

    float d_chunk_a = dout * sig_b;
    /* d/dx of σ(x) = σ(x) * (1 - σ(x)); the inner derivative passes
       through chunk_a unmodified (chunk_a does not depend on chunk_b). */
    float d_chunk_b = dout * chunk_a * sig_b * (1.0f - sig_b);

    /* Scatter to the same row-major [B, 2*H] layout: columns [0..H) for
       d_chunk_a, [H..2H) for d_chunk_b. */
    float* drow_2h = d_in_2h + (size_t)b * (2 * H);
    drow_2h[d]     = d_chunk_a;
    drow_2h[H + d] = d_chunk_b;
}

/* =====================================================================
 * grn_elu_backward — Element-wise ELU backward.
 *
 * Forward was: f(x) = x if x >= 0, else exp(x) - 1.
 * Derivative: f'(x) = 1     if x >= 0
 *                   = exp(x) if x  < 0
 *
 * `x_saved` is the forward POST-activation value (i.e. f(x_pre)). We
 * recover f'(x_pre) without needing the pre-activation:
 *   - If x_saved >= 0, then x_pre >= 0 too → f'(x_pre) = 1.
 *   - If x_saved <  0, then exp(x_pre) = x_saved + 1 (since
 *     x_saved = exp(x_pre) - 1) → f'(x_pre) = x_saved + 1.
 *
 * This avoids a separate pre-activation stash — the post-activation is
 * already on the forward path so we reuse it.
 *
 * Launch:  grid=ceil(n/256), block=256.
 * ===================================================================== */
extern "C" __global__ void grn_elu_backward(
    const float* __restrict__ d_out,
    const float* __restrict__ x_saved,          /* post-activation from forward */
    int n,
    float* __restrict__ d_x
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) return;
    float v  = x_saved[idx];
    /* For x_pre < 0: f(x_pre) = exp(x_pre) - 1 < 0, and f'(x_pre) = exp(x_pre) = v + 1.
       For x_pre >= 0: f(x_pre) = x_pre >= 0, and f'(x_pre) = 1. */
    float deriv = (v >= 0.0f) ? 1.0f : (v + 1.0f);
    d_x[idx] = d_out[idx] * deriv;
}
