// rl_outcome_bwd.cu — Outcome aux head backward (single linear layer).
//
// Given dL/dlogits [B, K_CLASSES=3] (produced by rl_outcome_ce),
// compute the linear layer gradients via the chain rule:
//
//   logits[b, k] = Σ_i h_t[b, i] × W[i, k] + b[k]
//
//   dL/dW_per_batch[b, i, k] = h_t[b, i] × dL/dlogits[b, k]
//   dL/db_per_batch[b, k]    = dL/dlogits[b, k]
//   dL/dh_t[b, i]            = Σ_k W[i, k] × dL/dlogits[b, k]
//
// Block layout: 1 block per batch, HIDDEN_DIM=128 threads.
//   * Thread i (i < 128) computes dL/dh_t[b, i] (inner product
//     across K_CLASSES=3 outputs — trivially unrolled).
//   * Thread i also writes grad_W_per_batch[b, i, 0..2] (3 writes).
//   * Threads 0..2 write grad_b_per_batch[b, tid] (thread 3..127 idle
//     for that single-cycle step).
//
// Per `feedback_no_atomicadd`: sole-writer per (b, i, k).
// Per `feedback_cpu_is_read_only`: pure device-side.

#include <stdint.h>

#define HIDDEN_DIM  128
#define K_CLASSES     3

extern "C" __global__ void rl_outcome_bwd(
    const float* __restrict__ h_t,                // [B, HIDDEN_DIM]
    const float* __restrict__ grad_logits,         // [B, K_CLASSES]
    const float* __restrict__ w,                   // [HIDDEN_DIM, K_CLASSES]
    int                       b_size,
    float*       __restrict__ grad_w_per_batch,    // [B, HIDDEN_DIM, K_CLASSES]
    float*       __restrict__ grad_b_per_batch,    // [B, K_CLASSES]
    float*       __restrict__ grad_h_t             // [B, HIDDEN_DIM]
) {
    const int b = blockIdx.x;
    const int i = threadIdx.x;
    if (b >= b_size || i >= HIDDEN_DIM) return;

    // Stage grad_logits into shared — K_CLASSES=3 so only threads 0..2
    // participate, but all 128 threads wait on the sync.
    __shared__ float s_grad_logits[K_CLASSES];
    if (i < K_CLASSES) {
        s_grad_logits[i] = grad_logits[b * K_CLASSES + i];
    }
    __syncthreads();

    const float h_bi = h_t[b * HIDDEN_DIM + i];

    // grad_W_per_batch[b, i, k] = h_bi × s_grad_logits[k]
    const int row_off = (b * HIDDEN_DIM + i) * K_CLASSES;
    #pragma unroll
    for (int k = 0; k < K_CLASSES; ++k) {
        grad_w_per_batch[row_off + k] = h_bi * s_grad_logits[k];
    }

    // grad_h_t[b, i] = Σ_k W[i, k] × s_grad_logits[k]
    float acc = 0.0f;
    #pragma unroll
    for (int k = 0; k < K_CLASSES; ++k) {
        acc += w[i * K_CLASSES + k] * s_grad_logits[k];
    }
    grad_h_t[b * HIDDEN_DIM + i] = acc;

    // grad_b_per_batch[b, k] = s_grad_logits[k] — sole writer per
    // (b, k); thread i writes slot k == i (covers k < 3 only).
    if (i < K_CLASSES) {
        grad_b_per_batch[b * K_CLASSES + i] = s_grad_logits[i];
    }
}
