/*
 * aux_trunk_backward_kernel.cu — SP14 Layer C Phase C.4 (2026-05-08).
 *
 * Backward pass for the auxiliary trunk's 3-layer Linear→ELU→Linear→ELU→
 * Linear MLP introduced in Phase C.3 (`aux_trunk_forward_kernel.cu`).
 * Receives `dh_s2_aux [B, AUX_HIDDEN_DIM]` from the aux head's backward
 * and propagates through w3/w2/w1 to accumulate parameter gradients
 * `dW{1,2,3}` and bias gradients `db{1,2,3}`.
 *
 * Forward chain (recap from C.3):
 *
 *   x_in [B, ENCODER_OUT_DIM]
 *     → Linear(w1, b1) → h_aux1_pre → ELU → h_aux1 [B, H1]   (saved-fwd)
 *     → Linear(w2, b2) → h_aux2_pre → ELU → h_aux2 [B, H2]   (saved-fwd)
 *     → Linear(w3, b3)               → h_s2_aux [B, AUX_HIDDEN_DIM] (no act)
 *
 * Backward math (in reverse order):
 *
 *   d_logits = dh_s2_aux                                   (output is linear)
 *   db3[j]    = sum_b d_logits[b, j]
 *   dW3[k, j] = sum_b h_aux2[b, k] * d_logits[b, j]
 *   dh_aux2_post[b, k] = sum_j d_logits[b, j] * w3[k, j]
 *   dh_aux2_pre [b, k] = dh_aux2_post[b, k] * elu_bwd_post(h_aux2[b, k])
 *
 *   db2[j]    = sum_b dh_aux2_pre[b, j]
 *   dW2[k, j] = sum_b h_aux1[b, k] * dh_aux2_pre[b, j]
 *   dh_aux1_post[b, k] = sum_j dh_aux2_pre[b, j] * w2[k, j]
 *   dh_aux1_pre [b, k] = dh_aux1_post[b, k] * elu_bwd_post(h_aux1[b, k])
 *
 *   db1[j]    = sum_b dh_aux1_pre[b, j]
 *   dW1[k, j] = sum_b x_in[b, k] * dh_aux1_pre[b, j]
 *
 *   STOP — NO `dx_in` write. The encoder boundary is the stop-gradient
 *   per the locked design decision (see plan §C.4 "Critical: NO
 *   `dx_in_out` parameter"). The aux trunk reshapes its own representation
 *   from x_in but cannot pull on the encoder. Q-loss remains the sole
 *   shaping force on the encoder.
 *
 * Pearls applied:
 *   - feedback_no_atomicadd.md — every dW/db output element is owned by
 *     exactly one block; the block performs a shmem-tree reduction over
 *     the batch dim. No atomicAdd anywhere.
 *   - feedback_no_stubs.md — full SGEMM-style accumulation, no
 *     placeholder body, no missing layer.
 *   - pearl_no_host_branches_in_captured_graph.md — kernel signatures are
 *     pure-device; the wrapper's `launch()` issues the fixed sequence of
 *     kernel launches with no host-side branching.
 *
 * Topology / layout (mirrors C.3 forward; passed at launch as runtime args):
 *   - ENCODER_OUT_DIM = config.shared_h1 (256 in production)
 *   - H1 = AUX_TRUNK_H1 (256)
 *   - H2 = AUX_TRUNK_H2 (128)
 *   - AUX_HIDDEN_DIM = config.shared_h2 (256)
 *
 *   w1 [ENCODER_OUT_DIM, H1]  row-major  → w1[k, j] = w1_ptr[k * H1 + j]
 *   w2 [H1, H2]               row-major  → w2[k, j] = w2_ptr[k * H2 + j]
 *   w3 [H2, AUX_HIDDEN_DIM]   row-major  → w3[k, j] = w3_ptr[k * AUX_HIDDEN_DIM + j]
 *
 * Saved-fwd buffers (POST-ELU values from C.3 forward):
 *   h_aux1 [B, H1]  — Layer-1 post-ELU
 *   h_aux2 [B, H2]  — Layer-2 post-ELU
 *
 * ELU' from POST-activation `y = ELU(x)` per `aux_elu_bwd_from_post` in
 * `aux_heads_kernel.cu`:
 *   if y > 0: x > 0 ⇒ ELU'(x) = 1
 *   else:     x ≤ 0 ⇒ ELU(x) = exp(x) - 1, so ELU'(x) = exp(x) = 1 + y
 *   (the y > 0 branch covers both x = 0 and x > 0 — at x = 0, y = 0 so
 *   the ≤ branch fires and returns 1.0 anyway: ELU is C¹ at the origin.)
 *
 * Kernel sequence (orchestrated by `AuxTrunkBackwardOps::launch`):
 *
 *   1. aux_trunk_bwd_dh_pre — [B] blocks, AUX_TRUNK_BLOCK threads.
 *        in:  dh_s2_aux [B, AUX_HIDDEN_DIM], w3, w2, h_aux1, h_aux2
 *        out: dh_aux2_pre [B, H2], dh_aux1_pre [B, H1]   (scratch)
 *
 *   2. aux_trunk_bwd_dW_reduce — generic dW = sum_b A[b, k] * B[b, j].
 *      Grid: (Krows * Jcols, 1, 1). Block-tree-reduce over batch.
 *      Launched 3× (dW3, dW2, dW1) with different (A, B) input pairs.
 *
 *   3. aux_trunk_bwd_db_reduce — generic db = sum_b B[b, j].
 *      Grid: (J, 1, 1). Block-tree-reduce over batch.
 *      Launched 3× (db3, db2, db1).
 *
 * No layer4 / dx_in kernel exists by design — that's the stop-gradient.
 */

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

#ifndef AUX_TRUNK_BWD_BLOCK
#define AUX_TRUNK_BWD_BLOCK 256
#endif

/* ELU backward from POST-activation y = ELU(x). Mirrors
 * `aux_elu_bwd_from_post` in aux_heads_kernel.cu — kept inline because the
 * common header pre-pended by build.rs only ships generic device fns; the
 * aux-specific ELU helpers live in their own kernel files. */
__device__ __forceinline__ float aux_trunk_elu_bwd_from_post(float y) {
    return (y > 0.0f) ? 1.0f : (1.0f + y);
}

/* =====================================================================
 * aux_trunk_bwd_dh_pre — pre-pass for the parameter-gradient reduces.
 *
 * Per sample b (one block per sample, AUX_TRUNK_BWD_BLOCK threads):
 *   1. Compute dh_aux2_pre[b, :] = (d_logits[b, :] @ w3.T) ⊙ ELU'(h_aux2[b, :]).
 *      Stored to global `dh_aux2_pre_out [B, H2]` for the dW2/db2/dW1
 *      reduces. Cached in shmem `sh_dh2_pre[H2]` for the next layer's
 *      matmul.
 *   2. Compute dh_aux1_pre[b, :] = (sh_dh2_pre @ w2.T) ⊙ ELU'(h_aux1[b, :]).
 *      Stored to global `dh_aux1_pre_out [B, H1]` for dW1/db1.
 *
 * Note: each block computes both layers serially within itself (no inter-
 * block communication required because both layers consume the same
 * sample's data). The shmem `sh_dh2_pre[H2]` lives across layers so
 * Layer-2 matmul reads from shmem rather than a second global-memory
 * round-trip.
 *
 * Block: AUX_TRUNK_BWD_BLOCK = 256 threads. Grid: (B, 1, 1).
 * Shared memory: H2 floats (Layer-2 post-multiply cache).
 *
 * No atomicAdd, no inter-block contention — each (b, k) output cell is
 * owned by exactly one block, written by exactly one thread.
 * ===================================================================== */
extern "C" __global__ void aux_trunk_bwd_dh_pre(
    const float* __restrict__ d_logits,         /* [B, AUX_HIDDEN_DIM] */
    const float* __restrict__ w3,               /* [H2, AUX_HIDDEN_DIM] row-major */
    const float* __restrict__ w2,               /* [H1, H2]              row-major */
    const float* __restrict__ h_aux1,           /* [B, H1] post-ELU */
    const float* __restrict__ h_aux2,           /* [B, H2] post-ELU */
    float* __restrict__ dh_aux2_pre_out,        /* [B, H2] OUT */
    float* __restrict__ dh_aux1_pre_out,        /* [B, H1] OUT */
    int B,
    int H1,
    int H2,
    int AUX_HIDDEN_DIM
) {
    const int b = blockIdx.x;
    if (b >= B) return;
    const int tid = threadIdx.x;

    /* Shared memory: H2-vector cache for Layer-2 post-multiply (so the
     * Layer-1 dh_aux1 matmul reads from shmem). */
    extern __shared__ float smem[];
    float* sh_dh2_pre = smem;                /* [H2] */

    /* ─────────────── Layer 3 backward → dh_aux2_pre ─────────────── */
    /* Each thread owns a slice of H2 via stride loop. Serial dot over
     * AUX_HIDDEN_DIM per lane. */
    {
        const float* dlogits_row = d_logits + (size_t)b * AUX_HIDDEN_DIM;
        for (int k = tid; k < H2; k += blockDim.x) {
            float acc = 0.0f;
            const float* w3_row = w3 + (size_t)k * AUX_HIDDEN_DIM;
            for (int j = 0; j < AUX_HIDDEN_DIM; ++j) {
                acc += dlogits_row[j] * w3_row[j];
            }
            const float h2_post = h_aux2[(size_t)b * H2 + k];
            const float dh2_pre = acc * aux_trunk_elu_bwd_from_post(h2_post);
            sh_dh2_pre[k] = dh2_pre;
            dh_aux2_pre_out[(size_t)b * H2 + k] = dh2_pre;
        }
    }
    __syncthreads();

    /* ─────────────── Layer 2 backward → dh_aux1_pre ─────────────── */
    /* dh_aux1_post[b, k] = sum_j sh_dh2_pre[j] * w2[k, j]
     * dh_aux1_pre [b, k] = dh_aux1_post[b, k] * ELU'(h_aux1[b, k]) */
    for (int k = tid; k < H1; k += blockDim.x) {
        float acc = 0.0f;
        const float* w2_row = w2 + (size_t)k * H2;
        for (int j = 0; j < H2; ++j) {
            acc += sh_dh2_pre[j] * w2_row[j];
        }
        const float h1_post = h_aux1[(size_t)b * H1 + k];
        const float dh1_pre = acc * aux_trunk_elu_bwd_from_post(h1_post);
        dh_aux1_pre_out[(size_t)b * H1 + k] = dh1_pre;
    }

    /* No Layer-1 backward here — that lives in `aux_trunk_bwd_dW_reduce`
     * launched with (A=x_in, B=dh_aux1_pre_out). NO `dx_in` is computed:
     * stop-gradient at the encoder boundary per locked design. */
}

/* =====================================================================
 * aux_trunk_bwd_dW_reduce — generic outer-product accumulation:
 *
 *   dW[k, j] = sum_b A[b, k] * B[b, j]
 *
 * Used 3× per backward step:
 *   - dW3 ←  A=h_aux2 [B, H2],  B=d_logits     [B, AUX_HIDDEN_DIM]
 *   - dW2 ←  A=h_aux1 [B, H1],  B=dh_aux2_pre  [B, H2]
 *   - dW1 ←  A=x_in   [B, ENC], B=dh_aux1_pre  [B, H1]
 *
 * Grid: (Krows * Jcols, 1, 1) — one block per output element.
 * Block: AUX_TRUNK_BWD_BLOCK = 256 threads — block-tree-reduce over batch.
 * Shared memory: AUX_TRUNK_BWD_BLOCK floats.
 *
 * Each block:
 *   1. Threads read strided slices of B samples, accumulate `local +=
 *      A[b, k] * B[b, j]` into per-thread register.
 *   2. Stage local values into shmem.
 *   3. Tree-reduce over shmem.
 *   4. Thread 0 writes final dW[k, j] = smem[0].
 *
 * No atomicAdd: each block owns its (k, j) cell exclusively.
 * ===================================================================== */
extern "C" __global__ void aux_trunk_bwd_dW_reduce(
    const float* __restrict__ A,        /* [B, Krows] row-major (saved-fwd or input) */
    const float* __restrict__ B_grad,   /* [B, Jcols] row-major (gradient) */
    float* __restrict__ dW_out,         /* [Krows, Jcols] row-major OUT */
    int B,
    int Krows,
    int Jcols
) {
    const int kj = blockIdx.x;
    if (kj >= Krows * Jcols) return;

    const int k = kj / Jcols;
    const int j = kj - k * Jcols;
    const int tid = threadIdx.x;
    const int block = blockDim.x;

    extern __shared__ float sdata[];

    float local = 0.0f;
    for (int i = tid; i < B; i += block) {
        local += A[(size_t)i * Krows + k] * B_grad[(size_t)i * Jcols + j];
    }
    sdata[tid] = local;
    __syncthreads();

    /* Power-of-two tree reduction. AUX_TRUNK_BWD_BLOCK is 256 — pure
     * power-of-two, no odd-stride bookkeeping needed. */
    for (int s = block / 2; s > 0; s >>= 1) {
        if (tid < s) sdata[tid] += sdata[tid + s];
        __syncthreads();
    }

    if (tid == 0) {
        dW_out[(size_t)k * Jcols + j] = sdata[0];
    }
}

/* =====================================================================
 * aux_trunk_bwd_db_reduce — generic batch-reduce for bias gradients:
 *
 *   db[j] = sum_b B[b, j]
 *
 * Used 3× per backward step:
 *   - db3 ←  B=d_logits     [B, AUX_HIDDEN_DIM]
 *   - db2 ←  B=dh_aux2_pre  [B, H2]
 *   - db1 ←  B=dh_aux1_pre  [B, H1]
 *
 * Grid: (Jcols, 1, 1) — one block per bias element.
 * Block: AUX_TRUNK_BWD_BLOCK = 256 threads — block-tree-reduce over batch.
 * Shared memory: AUX_TRUNK_BWD_BLOCK floats.
 * ===================================================================== */
extern "C" __global__ void aux_trunk_bwd_db_reduce(
    const float* __restrict__ B_grad,   /* [B, Jcols] row-major */
    float* __restrict__ db_out,         /* [Jcols] OUT */
    int B,
    int Jcols
) {
    const int j = blockIdx.x;
    if (j >= Jcols) return;
    const int tid = threadIdx.x;
    const int block = blockDim.x;

    extern __shared__ float sdata[];

    float local = 0.0f;
    for (int i = tid; i < B; i += block) {
        local += B_grad[(size_t)i * Jcols + j];
    }
    sdata[tid] = local;
    __syncthreads();

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

    if (tid == 0) {
        db_out[j] = sdata[0];
    }
}
