// rl_frd_fwd.cu — Forward-Return-Distribution head forward pass (SP20 P3).
//
// 2-layer MLP per batch entry:
//   hidden[i] = ReLU(Σ_k h_t[b, k] × w1[k, i] + b1[i])    for i in 0..FRD_HIDDEN_DIM
//   logits[j] = Σ_i hidden[i] × w2[i, j] + b2[j]          for j in 0..FRD_OUT_DIM
//
// where FRD_OUT_DIM = FRD_N_HORIZONS × FRD_N_ATOMS = 3 × 21 = 63
// (3 horizons × 21 return-bucket atoms — softmax + CE happen in bwd).
//
// Hidden activations are cached in `hidden_out` for the backward kernel —
// avoids recomputing the ReLU mask + the W1 × h_t product on bwd.
//
// Block layout: 1 block per batch entry, 64 threads.
//   - 64 covers FRD_HIDDEN_DIM=64 in a single pass for the hidden loop
//   - 64 ≥ FRD_OUT_DIM=63 covers the output loop with one idle thread
//   - One shared-mem stage of `hidden[]` between the two phases
//
// Per `feedback_no_atomicadd`: each thread is sole writer of its output slot.
// Per `feedback_cpu_is_read_only`: pure device-side.
// Per `feedback_no_nvrtc`: pre-compiled cubin via build.rs.

#include <stdint.h>

#define HIDDEN_DIM       128
#define FRD_HIDDEN_DIM    64
#define FRD_N_HORIZONS     3
#define FRD_N_ATOMS       21
#define FRD_OUT_DIM      (FRD_N_HORIZONS * FRD_N_ATOMS)  // 63

extern "C" __global__ void rl_frd_fwd(
    const float* __restrict__ h_t,        // [B, HIDDEN_DIM]
    const float* __restrict__ w1,         // [HIDDEN_DIM, FRD_HIDDEN_DIM]
    const float* __restrict__ b1,         // [FRD_HIDDEN_DIM]
    const float* __restrict__ w2,         // [FRD_HIDDEN_DIM, FRD_OUT_DIM]
    const float* __restrict__ b2,         // [FRD_OUT_DIM]
    float*       __restrict__ hidden_out, // [B, FRD_HIDDEN_DIM]  (cached for bwd)
    float*       __restrict__ logits_out, // [B, FRD_OUT_DIM]
    int                       b_size
) {
    const int b = blockIdx.x;
    const int tid = threadIdx.x;
    if (b >= b_size) return;

    __shared__ float hidden[FRD_HIDDEN_DIM];

    // Phase 1: each thread computes one hidden activation (64 threads, 64 slots).
    if (tid < FRD_HIDDEN_DIM) {
        float acc = b1[tid];
        const float* h_row = h_t + b * HIDDEN_DIM;
        for (int k = 0; k < HIDDEN_DIM; ++k) {
            acc += h_row[k] * w1[k * FRD_HIDDEN_DIM + tid];
        }
        hidden[tid] = fmaxf(0.0f, acc);  // ReLU
        hidden_out[b * FRD_HIDDEN_DIM + tid] = hidden[tid];
    }
    __syncthreads();

    // Phase 2: each thread computes one output logit (63 slots, 64 threads → tid=63 idle).
    if (tid < FRD_OUT_DIM) {
        float acc = b2[tid];
        for (int i = 0; i < FRD_HIDDEN_DIM; ++i) {
            acc += hidden[i] * w2[i * FRD_OUT_DIM + tid];
        }
        logits_out[b * FRD_OUT_DIM + tid] = acc;
    }
}
