// rl_frd_layer1_bwd.cu — FRD head backward stage 3 (SP20 P3 F.3c).
//
// Given dL/dhidden [B, FRD_HIDDEN_DIM] (from rl_frd_layer2_bwd) and
// the cached post-ReLU hidden [B, FRD_HIDDEN_DIM] (from rl_frd_fwd),
// compute the layer-1 gradients via the chain rule with ReLU mask:
//
//   pre_hidden[b, i] = Σ_k h_t[b, k] × W1[k, i] + b1[i]
//   hidden[b, i]     = ReLU(pre_hidden[b, i])
//
//   dL/dpre_hidden[b, i] = dL/dhidden[b, i] × 1{hidden[b, i] > 0}
//   dL/dW1[k, i]     = Σ_b h_t[b, k] × dL/dpre_hidden[b, i]    ← reduce over B
//   dL/db1[i]        = Σ_b dL/dpre_hidden[b, i]                 ← reduce over B
//   dL/dh_t[b, k]    = Σ_i W1[k, i] × dL/dpre_hidden[b, i]      ← per-batch overwrite
//
// Per-batch scratch:
//   grad_W1_per_batch[b, k, i] = h_t[b, k] × dL/dpre_hidden[b, i]
//   grad_b1_per_batch[b, i]    = dL/dpre_hidden[b, i]
//
// Block layout: 1 block per batch, HIDDEN_DIM=128 threads.
//   * Phase 0: stage post-ReLU `hidden` (64 slots) + dL/dhidden into
//     shared, apply ReLU mask in-place to produce dL/dpre_hidden
//     (only first 64 threads work; remaining 64 idle).
//   * Phase 1: each thread k (k < 128) is the sole writer of:
//     - grad_W1_per_batch[b, k, 0..64] (64 writes per thread)
//     - grad_h_t[b, k] (one Σ_i across 64 outputs)
//   * Phase 2: thread i (i < 64) writes grad_b1_per_batch[b, i].
//
// Per `feedback_no_atomicadd`: per-(b, k, i) sole-writer pattern.
// Per `feedback_cpu_is_read_only`: pure device-side.

#include <stdint.h>

#define HIDDEN_DIM       128
#define FRD_HIDDEN_DIM    64

extern "C" __global__ void rl_frd_layer1_bwd(
    const float* __restrict__ h_t,             // [B, HIDDEN_DIM]
    const float* __restrict__ hidden,          // [B, FRD_HIDDEN_DIM] (post-ReLU cache)
    const float* __restrict__ grad_hidden,     // [B, FRD_HIDDEN_DIM] (from layer2_bwd)
    const float* __restrict__ w1,              // [HIDDEN_DIM, FRD_HIDDEN_DIM]
    int                       b_size,
    float*       __restrict__ grad_w1_per_batch, // [B, HIDDEN_DIM, FRD_HIDDEN_DIM]
    float*       __restrict__ grad_b1_per_batch, // [B, FRD_HIDDEN_DIM]
    float*       __restrict__ grad_h_t           // [B, HIDDEN_DIM]
) {
    const int b = blockIdx.x;
    const int k = threadIdx.x;
    if (b >= b_size || k >= HIDDEN_DIM) return;

    // Stage dL/dpre_hidden in shared — apply ReLU mask using the
    // cached post-ReLU hidden (hidden > 0 ⇒ mask=1, else mask=0).
    __shared__ float s_grad_pre[FRD_HIDDEN_DIM];
    if (k < FRD_HIDDEN_DIM) {
        const int idx = b * FRD_HIDDEN_DIM + k;
        const float h_val = hidden[idx];
        const float g_val = grad_hidden[idx];
        s_grad_pre[k] = (h_val > 0.0f) ? g_val : 0.0f;
        // grad_b1_per_batch[b, i] = dL/dpre_hidden[b, i] — same value.
        grad_b1_per_batch[idx] = s_grad_pre[k];
    }
    __syncthreads();

    const float h_t_bk = h_t[b * HIDDEN_DIM + k];

    // grad_W1_per_batch[b, k, i] = h_t_bk × s_grad_pre[i]
    const int row_off = (b * HIDDEN_DIM + k) * FRD_HIDDEN_DIM;
    #pragma unroll
    for (int i = 0; i < FRD_HIDDEN_DIM; ++i) {
        grad_w1_per_batch[row_off + i] = h_t_bk * s_grad_pre[i];
    }

    // grad_h_t[b, k] = Σ_i W1[k, i] × s_grad_pre[i]
    float acc = 0.0f;
    #pragma unroll
    for (int i = 0; i < FRD_HIDDEN_DIM; ++i) {
        acc += w1[k * FRD_HIDDEN_DIM + i] * s_grad_pre[i];
    }
    grad_h_t[b * HIDDEN_DIM + k] = acc;
}
