/**
 * Attention backward element-wise cuBLAS kernels.
 *
 * Standalone element-wise kernels for the cuBLAS attention backward path.
 * Q/K/V/O projections use cublasLtMatmul (batched GEMM) for weight gradients.
 * These kernels handle the non-GEMM parts: SDP backward, LayerNorm backward,
 * bias gradient reduction, gradient norm computation, and Adam optimizer.
 *
 * Zero atomicAdd — all cross-sample gradient reduction is handled by cuBLAS
 * GEMM (dW = dY @ X^T) with alpha=1/B for mean reduction.
 *
 * ALL tensors use native f32.
 */

/* =====================================================================
 * ATTENTION GRADIENT NORM — TWO-PHASE DETERMINISTIC REDUCTION
 *
 * Phase 1: per-block partial sum-of-squares → block_sums[blockIdx.x]
 * Phase 2: single-block sequential sum → norm_out[0] (as L2 norm)
 *
 * Replaces the old attn_grad_norm_kernel which used atomicAdd.
 * ===================================================================== */
extern "C" __global__ void attn_grad_norm_phase1(
    const float* __restrict__ grads,
    float* __restrict__ block_sums,
    int total_params
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;

    float partial = 0.0f;
    for (int i = tid; i < total_params; i += gridDim.x * blockDim.x) {
        float g = grads[i];
        partial += g * g;
    }

    __shared__ float shared[256];
    shared[threadIdx.x] = partial;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < (unsigned int)s)
            shared[threadIdx.x] += shared[threadIdx.x + s];
        __syncthreads();
    }

    if (threadIdx.x == 0)
        block_sums[blockIdx.x] = shared[0];
}

extern "C" __global__ void attn_grad_norm_phase2(
    const float* __restrict__ block_sums,
    float* __restrict__ norm_out,
    int num_blocks
) {
    /* Sequential sum → sqrt → L2 norm. Grid=(1), Block=(1). */
    float sum = 0.0f;
    for (int i = 0; i < num_blocks; i++)
        sum += block_sums[i];
    norm_out[0] = sqrtf(fmaxf(sum, 0.0f));
}

/* =====================================================================
 * ATTENTION ADAM UPDATE KERNEL
 *
 * AdamW with gradient clipping for attention parameters.
 * Same algorithm as iqn_adam_kernel.
 *
 * Launch: grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1)
 * ===================================================================== */
extern "C" __global__ void attn_adam_kernel(
    float* __restrict__ params,
    float* __restrict__ grads,
    float* __restrict__ m,
    float* __restrict__ v,
    const float* __restrict__ norm,   /* [1] gradient L2 norm */
    float lr, float beta1, float beta2, float eps,
    float weight_decay, float max_grad_norm,
    const int* __restrict__ adam_t_buf,
    int total_params,
    /* SP4 Layer B: ISV-driven |p| bound (`ISV[WEIGHT_BOUND[Attn]].max(EPS_CLAMP_FLOOR)`)
     * — same pattern wired into all five Adam kernels. Used by both
     * GpuAttention and GpuTlob (TLOB shares this kernel via the same cubin
     * and co-lives in the Attn group taxonomy). 0=disabled (debug only). */
    float weight_clamp_max_abs,
    /* SP4 Task A14 (2026-04-30): Pearl C engagement counter via warp+
     * block tree-reduce (no atomicAdd). GpuAttention + GpuTlob both pass
     * `engage_buf_offset = ParamGroup::Attn × MAX_BLOCKS_PER_ADAM` so
     * Pearl C tracks the kernel-level engagement rate (TLOB and core
     * attention share the same SP4 group). `engage_buf_offset < 0`
     * skips the writeback. */
    int* __restrict__ nan_flags_buf,
    int diag_slot,
    int* __restrict__ clamp_engage_per_block_buf,
    int engage_buf_offset
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    bool active = (tid < total_params);
    int local_engage = 0;

    if (active) {
        /* Skip NaN/Inf gradients — leave param, momentum, variance unchanged. */
        float g_check = (float)grads[tid];
        if (!isfinite(g_check)) {
            grads[tid] = 0.0f;
        } else {
            /* Gradient clipping */
            float grad_norm = norm[0];
            float clip_scale = (grad_norm > max_grad_norm && grad_norm > 0.0f)
                ? max_grad_norm / grad_norm : 1.0f;
            float g = grads[tid] * clip_scale;

            /* AdamW: weight decay applied to params directly */
            params[tid] = params[tid] * (1.0f - lr * weight_decay);

            /* Moment updates */
            float m_val = beta1 * m[tid] + (1.0f - beta1) * g;
            float v_val = beta2 * v[tid] + (1.0f - beta2) * g * g;
            m[tid] = m_val;
            v[tid] = v_val;

            /* Bias correction */
            int adam_t = adam_t_buf[0];
            float f_t = (float)adam_t;
            float bc1 = 1.0f - powf(beta1, f_t);
            float bc2 = 1.0f - powf(beta2, f_t);
            float m_hat = m_val / (bc1 + 1e-12f);
            float v_hat = v_val / (bc2 + 1e-12f);

            /* Parameter update */
            float p = params[tid] - lr * m_hat / (sqrtf(v_hat) + eps);

            /* SP3 Mech 9 + SP4 Task A14: post-Adam |p| clamp +
             * engagement counter — last step before writeback. */
            if (weight_clamp_max_abs > 0.0f) {
                if (fabsf(p) > weight_clamp_max_abs) {
                    if (nan_flags_buf != nullptr && diag_slot >= 0) {
                        nan_flags_buf[diag_slot] = 1;
                    }
                    local_engage = 1;
                }
                p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs);
            }

            params[tid] = p;

            /* Zero gradient for next step */
            grads[tid] = 0.0f;
        }
    }

    /* SP4 Task A14: warp + block tree-reduce of `local_engage`. */
    int e = local_engage;
    for (int offset = 16; offset > 0; offset >>= 1)
        e += __shfl_xor_sync(0xFFFFFFFF, e, offset);

    __shared__ int warp_engage[8]; /* blockDim.x = 256 ⇒ 8 warps */
    int warp_id   = threadIdx.x / 32;
    int warp_lane = threadIdx.x % 32;
    if (warp_lane == 0) warp_engage[warp_id] = e;
    __syncthreads();

    if (warp_id == 0) {
        int val = (warp_lane < blockDim.x / 32) ? warp_engage[warp_lane] : 0;
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        if (warp_lane == 0 && engage_buf_offset >= 0 && clamp_engage_per_block_buf != nullptr) {
            clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x] = val;
        }
    }
}

/* =====================================================================
 * ATTENTION ELEMENT-WISE CUBIN KERNELS
 *
 * Standalone element-wise kernels for batched cuBLAS attention path.
 * All tensors use column-major layout [D, B] (feature-major stride).
 * ===================================================================== */

/* =====================================================================
 * attn_sdp_fwd — Scaled dot-product attention forward (feature-level)
 *
 * Feature attention on h_s2[D, B] where D=256, B=16384, num_heads=4,
 * head_dim=64. NOT sequence attention — each sample is independent.
 * Per head: score = dot(Q_h, K_h) / sqrt(head_dim), gate = sigmoid(score),
 * attn_out_h = V_h * gate.
 *
 * Layout: Q, K, V, attn_out are [D, B] (column-major).
 * save_scores is [num_heads, B].
 * Grid: (B), Block: (num_heads or power-of-two >= num_heads).
 * ===================================================================== */
extern "C" __global__
void attn_sdp_fwd(
    const float* __restrict__ Q,
    const float* __restrict__ K,
    const float* __restrict__ V,
    float* __restrict__ attn_out,
    float* __restrict__ save_scores,
    int B, int D, int num_heads
) {
    int sample = blockIdx.x;
    if (sample >= B) return;
    int head_dim = D / num_heads;
    float scale = 1.0f / sqrtf((float)head_dim);

    for (int h = threadIdx.x; h < num_heads; h += blockDim.x) {
        int off = h * head_dim;
        float score = 0.0f;
        for (int d = 0; d < head_dim; d++) {
            int idx = (off + d) * B + sample;
            score += Q[idx] * K[idx];
        }
        score *= scale;
        float gate = 1.0f / (1.0f + expf(-score));
        save_scores[h * B + sample] = gate;
        for (int d = 0; d < head_dim; d++) {
            int idx = (off + d) * B + sample;
            attn_out[idx] = V[idx] * gate;
        }
    }
}

/* =====================================================================
 * attn_sdp_bwd — Scaled dot-product attention backward
 *
 * Backpropagates through sigmoid-gated feature attention.
 * Computes d_Q, d_K, d_V given d_attn_out and saved sigmoid scores.
 *
 * Layout: all [D, B] except scores which is [num_heads, B].
 * Grid: (B), Block: (num_heads or power-of-two >= num_heads).
 * ===================================================================== */
extern "C" __global__
void attn_sdp_bwd(
    const float* __restrict__ d_attn_out,
    const float* __restrict__ Q,
    const float* __restrict__ K,
    const float* __restrict__ V,
    const float* __restrict__ scores,
    float* __restrict__ d_Q,
    float* __restrict__ d_K,
    float* __restrict__ d_V,
    int B, int D, int num_heads
) {
    int sample = blockIdx.x;
    if (sample >= B) return;
    int head_dim = D / num_heads;
    float scale = 1.0f / sqrtf((float)head_dim);

    for (int h = threadIdx.x; h < num_heads; h += blockDim.x) {
        int off = h * head_dim;
        float gate = scores[h * B + sample];
        float dgate = 0.0f;
        for (int d = 0; d < head_dim; d++) {
            int idx = (off + d) * B + sample;
            d_V[idx] = d_attn_out[idx] * gate;
            dgate += d_attn_out[idx] * V[idx];
        }
        float dscore = dgate * gate * (1.0f - gate) * scale;
        for (int d = 0; d < head_dim; d++) {
            int idx = (off + d) * B + sample;
            d_Q[idx] = dscore * K[idx];
            d_K[idx] = dscore * Q[idx];
        }
    }
}

/* =====================================================================
 * attn_layer_norm_fwd — Per-sample LayerNorm with residual connection
 *
 * Computes: out = gamma * ((x + residual - mean) / std) + beta
 * Saves mean and rstd per sample for backward pass.
 *
 * Layout: x, residual, out are [D, B]. gamma, beta are [D].
 * save_mean, save_rstd are [B].
 * Grid: (B), Block: (1) — serial over D per sample.
 * ===================================================================== */
extern "C" __global__
void attn_layer_norm_fwd(
    const float* __restrict__ x,
    const float* __restrict__ residual,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float* __restrict__ out,
    float* __restrict__ save_mean,
    float* __restrict__ save_rstd,
    int D, int B
) {
    int sample = blockIdx.x;
    if (sample >= B) return;
    float mean = 0.0f;
    for (int d = 0; d < D; d++) mean += x[d * B + sample] + residual[d * B + sample];
    mean /= (float)D;
    float var = 0.0f;
    for (int d = 0; d < D; d++) {
        float v = x[d * B + sample] + residual[d * B + sample] - mean;
        var += v * v;
    }
    var /= (float)D;
    float rstd = 1.0f / sqrtf(var + 1e-5f);
    save_mean[sample] = mean;
    save_rstd[sample] = rstd;
    for (int d = 0; d < D; d++) {
        float normed = (x[d * B + sample] + residual[d * B + sample] - mean) * rstd;
        out[d * B + sample] = gamma[d] * normed + beta[d];
    }
}

/* =====================================================================
 * attn_layer_norm_bwd — LayerNorm backward with residual connection
 *
 * Backpropagates through per-sample LayerNorm + residual.
 * Computes d_x (same gradient as d_residual due to residual add),
 * d_gamma, and d_beta accumulated over the batch dimension.
 *
 * Layout: d_out, x, residual, d_x are [D, B]. gamma, d_gamma, d_beta
 * are [D]. save_mean, save_rstd are [B].
 * Grid: (ceil(D/blockDim.x)), Block: (256).
 * ===================================================================== */
/* attn_layer_norm_bwd_dx — Batch-parallel d_x computation
 * d_x[d, b] = d_out[d, b] * gamma[d] * rstd[b]
 * Grid: ceil(D*B/256), Block: 256. Layout: [D, B] col-major.
 */
extern "C" __global__
void attn_layer_norm_bwd_dx(
    const float* __restrict__ d_out,
    const float* __restrict__ gamma,
    const float* __restrict__ save_rstd,
    float* __restrict__ d_x,
    int D, int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int N = D * B;
    if (idx >= N) return;
    int d = idx % D;
    int b = idx / D;
    d_x[idx] = d_out[idx] * gamma[d] * save_rstd[b];
}

/* attn_layer_norm_bwd_dgamma_p1 — Phase 1: block-level reduce for d_gamma AND d_beta
 * d_gamma_partial = d_out * normed, d_beta_partial = d_out
 * Grid: (num_blocks, D), Block: 256, shared: 512*sizeof(float) (2 × 256)
 */
extern "C" __global__
void attn_layer_norm_bwd_dgamma_p1(
    const float* __restrict__ d_out,
    const float* __restrict__ x,
    const float* __restrict__ residual,
    const float* __restrict__ save_mean,
    const float* __restrict__ save_rstd,
    float* __restrict__ gamma_partials,  /* [num_blocks, D] */
    float* __restrict__ beta_partials,   /* [num_blocks, D] */
    int D, int B
) {
    extern __shared__ float sdata[];  /* [512]: first 256 for gamma, next 256 for beta */
    float* sg = sdata;
    float* sb = sdata + blockDim.x;
    int d   = blockIdx.y;
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + tid;
    float gv = 0.0f, bv = 0.0f;
    if (gid < B) {
        float mean = save_mean[gid];
        float rstd = save_rstd[gid];
        float normed = (x[d * B + gid] + residual[d * B + gid] - mean) * rstd;
        float dout_val = d_out[d * B + gid];
        gv = dout_val * normed;
        bv = dout_val;
    }
    sg[tid] = gv;
    sb[tid] = bv;
    __syncthreads();
    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 * D + d] = sg[0];
        beta_partials[blockIdx.x * D + d] = sb[0];
    }
}

/* attn_layer_norm_bwd_dgamma_p2 — Phase 2: final reduce for d_gamma + d_beta
 * Grid: ceil(D/256), Block: 256
 */
extern "C" __global__
void attn_layer_norm_bwd_dgamma_p2(
    const float* __restrict__ gamma_partials,
    const float* __restrict__ beta_partials,
    float* __restrict__ d_gamma,
    float* __restrict__ d_beta,
    int D, int num_blocks
) {
    int d = blockIdx.x * blockDim.x + threadIdx.x;
    if (d >= D) return;
    float dg = 0.0f, db = 0.0f;
    for (int i = 0; i < num_blocks; i++) {
        dg += gamma_partials[i * D + d];
        db += beta_partials[i * D + d];
    }
    d_gamma[d] = dg;
    d_beta[d] = db;
}

/* =====================================================================
 * attn_concat_input — Concatenate states[D,B] + ofi_embed[embed_dim,B]
 *   into a single col-major [D+embed_dim, B] buffer.
 *
 * states is col-major [D, B]: element (d, b) at d + b*D.
 * ofi_embed is col-major [embed_dim, B]: element (d, b) at d + b*embed_dim.
 * output is col-major [D+embed_dim, B]: element (d, b) at d + b*(D+embed_dim).
 *
 * Grid: (ceil(B/256), D+embed_dim), Block: 256.
 * ===================================================================== */
extern "C" __global__
void attn_concat_input(
    const float* __restrict__ states,       /* [D, B] col-major */
    const float* __restrict__ ofi_embed,    /* [embed_dim, B] col-major */
    float* __restrict__ output,             /* [D+embed_dim, B] col-major */
    int B, int D, int embed_dim
) {
    int d = blockIdx.y;        /* which row dimension */
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;
    int out_width = D + embed_dim;
    float val;
    if (d < D) {
        val = states[d + (long long)b * D];
    } else {
        int ed = d - D;
        if (ed >= embed_dim) return;
        val = ofi_embed[ed + (long long)b * embed_dim];
    }
    output[d + (long long)b * out_width] = val;
}

/* =====================================================================
 * extract_attn_ofi_grad — Extract OFI embed gradient from d_input_scratch
 *
 * d_input_scratch is col-major [input_dim, B] where input_dim = D + embed_dim.
 * The OFI embed gradient occupies the last embed_dim rows (rows D..D+embed_dim-1).
 * This kernel extracts them into a contiguous [embed_dim, B] col-major buffer.
 *
 *   d_ofi_out[d, b] = d_input[(D + d) + b * input_dim]
 *
 * Grid: (ceil(B/256), embed_dim), Block: 256.
 * ===================================================================== */
extern "C" __global__
void extract_attn_ofi_grad(
    const float* __restrict__ d_input,    /* [input_dim, B] col-major */
    float* __restrict__ d_ofi_out,        /* [embed_dim, B] col-major output */
    int B, int D, int embed_dim
) {
    int d = blockIdx.y;        /* which embed dimension */
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B || d >= embed_dim) return;
    int input_dim = D + embed_dim;
    d_ofi_out[d + b * embed_dim] = d_input[(D + d) + (long long)b * input_dim];
}

/* =====================================================================
 * attn_bias_add — In-place bias addition for attention projections
 *
 * Adds bias vector (length out_dim) to a flat [out_dim, B] activation.
 * idx % out_dim selects the bias element for each output position.
 *
 * Grid: (ceil(N/blockDim.x)), Block: (256). N = out_dim * B.
 * ===================================================================== */
extern "C" __global__
void attn_bias_add(
    float* __restrict__ x,
    const float* __restrict__ bias,
    int out_dim, int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) { x[idx] += bias[idx % out_dim]; }
}

/* =====================================================================
 * attn_bias_grad_reduce_p1 — Phase 1: block-level shared-memory reduce
 *
 * Reduces d_pre [out_dim, B] (col-major, cuBLAS output) over B.
 * Grid: (num_blocks, out_dim), Block: 256, shared: 256*sizeof(float).
 * ===================================================================== */
extern "C" __global__
void attn_bias_grad_reduce_p1(
    const float* __restrict__ d_pre,
    float* __restrict__ partials,
    int out_dim, int B
) {
    extern __shared__ float sdata[];
    int j   = blockIdx.y;
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + tid;
    float val = (gid < B) ? d_pre[j + (long)gid * out_dim] : 0.0f;
    sdata[tid] = val;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) sdata[tid] += sdata[tid + s];
        __syncthreads();
    }
    if (tid == 0) partials[blockIdx.x * out_dim + j] = sdata[0];
}

/* =====================================================================
 * attn_bias_grad_reduce_p2 — Phase 2: warp-shuffle final reduce + scale
 *
 * Grid: ceil(out_dim/256), Block: 256.
 * ===================================================================== */
extern "C" __global__
void attn_bias_grad_reduce_p2(
    const float* __restrict__ partials,
    float* __restrict__ d_bias,
    int out_dim, int num_blocks, float inv_batch
) {
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    if (j >= out_dim) return;
    float sum = 0.0f;
    for (int i = 0; i < num_blocks; i++) sum += partials[i * out_dim + j];
    d_bias[j] = sum * inv_batch;
}
