/*
 * aux_trunk_forward_kernel.cu — SP14 Layer C Phase C.3 (2026-05-08)
 *
 * 3-layer MLP forward pass for the auxiliary head's separate trunk:
 *
 *   x_in [B, ENCODER_OUT_DIM]
 *     → Linear(w1, b1)         → [B, H1]   (H1 = AUX_TRUNK_H1 = 256)
 *     → ELU
 *     → Linear(w2, b2)         → [B, H2]   (H2 = AUX_TRUNK_H2 = 128)
 *     → ELU
 *     → Linear(w3, b3)         → [B, AUX_HIDDEN_DIM]
 *   h_s2_aux [B, AUX_HIDDEN_DIM]
 *
 * No activation on the final output — the downstream aux head applies its
 * own activation. The encoder output `x_in` flows through stop-gradient at
 * the encoder boundary in C.4 (backward) so this trunk shapes a
 * representation suited to next-bar prediction without pulling on the
 * encoder. Aux head's own w1/b1/w2/b2 train independently from the trunk
 * weights here.
 *
 * Topology (production config — verified C.2):
 *   ENCODER_OUT_DIM = config.shared_h1 = 256
 *   AUX_HIDDEN_DIM  = config.shared_h2 = 256  (matches existing aux head input dim)
 *   H1 = AUX_TRUNK_H1 = 256
 *   H2 = AUX_TRUNK_H2 = 128
 *
 * Saved-for-backward outputs:
 *   - h_aux1_out [B, H1]:  POST-ELU Layer-1 hidden activations
 *   - h_aux2_out [B, H2]:  POST-ELU Layer-2 hidden activations
 *   - h_s2_aux_out [B, AUX_HIDDEN_DIM]:  final linear output (no activation)
 *
 * Backward (C.4) re-reads h_aux1 + h_aux2 to derive ELU' from the post-
 * activation values via the standard `f'(x) = 1 if y > 0 else 1 + y`
 * identity (mirrors `aux_elu_bwd_from_post` in `aux_heads_kernel.cu`).
 *
 * Pearls applied:
 *   - feedback_no_atomicadd.md — forward writes only per-(b, j) outputs;
 *     no contention across blocks. Caller-side cuBLAS is NOT used here
 *     because mirroring `aux_next_bar_forward`'s thread-per-output style
 *     keeps the kernel self-contained for graph capture.
 *   - pearl_no_host_branches_in_captured_graph.md — kernel signature is
 *     pure-device; no host-side dispatch inside the captured region.
 *   - feedback_no_stubs.md — full SGEMM math, no return-zero, no
 *     placeholder body.
 *   - common_device_functions.cuh provides standard helpers (this kernel
 *     gets the common header prepended at build time per `build.rs`).
 *
 * Layout convention: row-major `[B, …]` for activations.
 *   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]
 *
 * ELU activation: f(x) = x if x > 0 else (exp(x) - 1).
 *
 * Block: `AUX_TRUNK_BLOCK = 256` threads. Grid: `(B, 1, 1)` blocks.
 * Shared memory: `(H1 + H2) * sizeof(float)` bytes (post-ELU caches for
 * Layer 2 + Layer 3).
 */

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

#ifndef AUX_TRUNK_BLOCK
#define AUX_TRUNK_BLOCK 256
#endif

/* ELU helper — branchless on `x > 0.f`. Mirrors `aux_elu_fwd` in
 * aux_heads_kernel.cu (kept inline rather than #include'd because the
 * common header pre-pended by build.rs only ships the small generic
 * device fns; the aux-specific ELU helpers live in their own kernel
 * file). */
__device__ __forceinline__ float aux_trunk_elu_fwd(float x) {
    return (x > 0.0f) ? x : (expf(x) - 1.0f);
}

/* =====================================================================
 * aux_trunk_forward — 3-layer Linear→ELU→Linear→ELU→Linear MLP.
 *
 * Per sample b (one block per sample, AUX_TRUNK_BLOCK threads):
 *   1. h_aux1[k] = ELU( b1[k] + sum_j w1[j, k] * x_in[b, j] )    for k ∈ [0, H1)
 *   2. h_aux2[k] = ELU( b2[k] + sum_j w2[j, k] * h_aux1[j] )      for k ∈ [0, H2)
 *   3. h_s2_aux[k] = b3[k] + sum_j w3[j, k] * h_aux2[j]           for k ∈ [0, AUX_HIDDEN_DIM)
 *
 * Hidden layer outputs (post-ELU) are saved to global memory for backward.
 * Final layer output has NO activation — downstream aux head applies its
 * own activation pipeline.
 *
 * Each thread within the block owns a slice of the layer's output dim via
 * stride-loop. The matmul over the input dim is serial within each
 * thread (mirrors `aux_next_bar_forward` — input dims are small enough
 * that block-cooperative reduce wastes more cycles in shfl bookkeeping
 * than the serial pass). The shared-memory `sh_h1` / `sh_h2` caches let
 * the next layer's matmul read post-ELU values without a second global-
 * memory pass.
 *
 * Block: AUX_TRUNK_BLOCK threads. Grid: (B, 1, 1). Shared memory:
 * (H1 + H2) floats — Layer-1 cache for Layer-2 matmul + Layer-2 cache
 * for Layer-3 matmul.
 * ===================================================================== */
extern "C" __global__ void aux_trunk_forward(
    const float* __restrict__ x_in,          /* [B, ENCODER_OUT_DIM] row-major */
    const float* __restrict__ w1,            /* [ENCODER_OUT_DIM, H1] row-major */
    const float* __restrict__ b1,            /* [H1] */
    const float* __restrict__ w2,            /* [H1, H2] row-major */
    const float* __restrict__ b2,            /* [H2] */
    const float* __restrict__ w3,            /* [H2, AUX_HIDDEN_DIM] row-major */
    const float* __restrict__ b3,            /* [AUX_HIDDEN_DIM] */
    float* __restrict__ h_aux1_out,          /* [B, H1] saved for backward */
    float* __restrict__ h_aux2_out,          /* [B, H2] saved for backward */
    float* __restrict__ h_s2_aux_out,        /* [B, AUX_HIDDEN_DIM] final */
    int B,
    int ENCODER_OUT_DIM,
    int H1,
    int H2,
    int AUX_HIDDEN_DIM
) {
    const int b = blockIdx.x;
    if (b >= B) return;
    const int tid = threadIdx.x;

    /* Shared memory: post-ELU caches for Layer 1 + Layer 2.
     * Total: (H1 + H2) floats. Allocator-side smem size is computed in
     * the Rust wrapper based on the runtime H1/H2 values. */
    extern __shared__ float smem[];
    float* sh_h1 = smem;             /* [H1] */
    float* sh_h2 = smem + H1;        /* [H2] */

    /* ─────────────── Layer 1: Linear(w1, b1) → ELU ─────────────── */
    /* Each thread computes one (or more, via stride loop) hidden lanes.
     * Serial dot over ENCODER_OUT_DIM per lane — input dim is small
     * enough that block-cooperative reduce is wasteful. */
    {
        const float* x_row = x_in + (size_t)b * ENCODER_OUT_DIM;
        for (int k = tid; k < H1; k += blockDim.x) {
            float acc = b1[k];
            for (int j = 0; j < ENCODER_OUT_DIM; ++j) {
                acc += x_row[j] * w1[(size_t)j * H1 + k];
            }
            const float h_post = aux_trunk_elu_fwd(acc);
            sh_h1[k] = h_post;
            h_aux1_out[(size_t)b * H1 + k] = h_post;
        }
    }
    __syncthreads();

    /* ─────────────── Layer 2: Linear(w2, b2) → ELU ─────────────── */
    /* Read post-ELU Layer-1 from shmem (no second global-memory pass). */
    for (int k = tid; k < H2; k += blockDim.x) {
        float acc = b2[k];
        for (int j = 0; j < H1; ++j) {
            acc += sh_h1[j] * w2[(size_t)j * H2 + k];
        }
        const float h_post = aux_trunk_elu_fwd(acc);
        sh_h2[k] = h_post;
        h_aux2_out[(size_t)b * H2 + k] = h_post;
    }
    __syncthreads();

    /* ─────────────── Layer 3: Linear(w3, b3) — NO activation ─────────────── */
    /* Read post-ELU Layer-2 from shmem. Output goes directly to global —
     * downstream aux head's first FC applies its own activation. */
    {
        float* out_row = h_s2_aux_out + (size_t)b * AUX_HIDDEN_DIM;
        for (int k = tid; k < AUX_HIDDEN_DIM; k += blockDim.x) {
            float acc = b3[k];
            for (int j = 0; j < H2; ++j) {
                acc += sh_h2[j] * w3[(size_t)j * AUX_HIDDEN_DIM + k];
            }
            out_row[k] = acc;
        }
    }
}
