/* =====================================================================
 * rl_outcome_fwd.cu -- K=3 trade-outcome classifier forward pass.
 *
 * Linear forward: logits = h_t x W + b
 *
 * Grid  = (B, 1, 1)   -- one block per batch element.
 * Block = (3, 1, 1)    -- one thread per class (K=3).
 *
 * Layout (row-major unless noted):
 *   h_t    : [B, 128]       hidden states from trunk
 *   w      : [128, 3]       column-major (so thread k reads stride-3)
 *   b      : [3]            bias per class
 *   logits : [B, 3]         output logits
 *
 * Each thread k in {0,1,2} computes:
 *   logits[batch*3 + k] = b[k] + sum_{i=0}^{127} h_t[batch*128 + i] * w[i*3 + k]
 *
 * No atomicAdd -- single-thread accumulation per output element.
 * No shared memory needed -- each thread reads the full 128-dim row
 * independently (128 loads is register-friendly at K=3 threads/block).
 * ===================================================================== */

#define HIDDEN_DIM 128
#define K_CLASSES  3

extern "C" __global__ void rl_outcome_fwd(
    const float* __restrict__ h_t,      /* [B, 128]  */
    const float* __restrict__ w,        /* [128, 3] column-major */
    const float* __restrict__ b,        /* [3]       */
    float*       __restrict__ logits,   /* [B, 3]    */
    int b_size
) {
    int batch = blockIdx.x;
    int k     = threadIdx.x;

    if (batch >= b_size || k >= K_CLASSES) return;

    const float* h_row = h_t + batch * HIDDEN_DIM;

    float sum = b[k];
    #pragma unroll 16
    for (int i = 0; i < HIDDEN_DIM; ++i) {
        sum += h_row[i] * w[i * K_CLASSES + k];
    }
    logits[batch * K_CLASSES + k] = sum;
}
