// ppo_loss_reduce_b.cu — [B] → [1] mean reducer for the two PPO loss
// scalars produced by `ppo_clipped_surrogate_fwd`.
//
// Replaces the prior `atomicAdd(loss_pi, l_pi/B)` pattern. Per
// `feedback_no_atomicadd.md`: block tree-reduce in shared memory only,
// no atomics. The forward kernel writes per-batch values to
// `loss_pi_per_b[B]` and `loss_entropy_per_b[B]` (single writer per
// slot), then this kernel reduces each to a scalar mean.
//
// Computes:
//   loss_pi_out[0]      = (1/B) * Σ_b loss_pi_per_b[b]
//   loss_entropy_out[0] = (1/B) * Σ_b loss_entropy_per_b[b]
//
// Block layout:
//   grid  = (1, 1, 1)
//   block = (BLOCK_X, 1, 1)         (BLOCK_X = 256; power of 2 required)
//   shared = 2 * BLOCK_X * sizeof(float)
//
// Caller passes B = batch size; the kernel grid-strides if B > BLOCK_X
// (so a single launch supports any cluster batch size).

#define BLOCK_X 256

extern "C" __global__ void ppo_loss_reduce_b(
    const float* __restrict__ loss_pi_per_b,        // [B]
    const float* __restrict__ loss_entropy_per_b,   // [B]
    int                       B,
    float*       __restrict__ loss_pi_out,          // [1]  OVERWRITE
    float*       __restrict__ loss_entropy_out      // [1]  OVERWRITE
) {
    __shared__ float s_pi[BLOCK_X];
    __shared__ float s_ent[BLOCK_X];

    const int tid = threadIdx.x;

    // Per-thread grid-stride sum over the [B] axis.
    float local_pi  = 0.0f;
    float local_ent = 0.0f;
    for (int b = tid; b < B; b += BLOCK_X) {
        local_pi  += loss_pi_per_b[b];
        local_ent += loss_entropy_per_b[b];
    }
    s_pi[tid]  = local_pi;
    s_ent[tid] = local_ent;
    __syncthreads();

    // Block tree-reduce in shared memory. BLOCK_X must be a power of 2.
    for (int stride = BLOCK_X / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            s_pi[tid]  += s_pi[tid  + stride];
            s_ent[tid] += s_ent[tid + stride];
        }
        __syncthreads();
    }

    // Single-writer store of the batch-mean diagnostics.
    if (tid == 0) {
        const float b_inv = 1.0f / (float)((B > 0) ? B : 1);
        loss_pi_out[0]      = s_pi[0]  * b_inv;
        loss_entropy_out[0] = s_ent[0] * b_inv;
    }
}


// ─────────────────────────────────────────────────────────────────────
// mean_reduce_b_f32 — single-input variant of the same pattern. Used by
// `dqn_distributional_q_bwd` whose per-batch CE values live in
// `loss_per_batch[B]` (single-writer already); this kernel produces
// the scalar mean for the trainer's diagnostic readout, replacing the
// prior `atomicAdd(loss_out, ce / B)` (Phase C "loud-flagged deferral"
// — see dqn_distributional_q.cu header).
//
// Block layout: identical to `ppo_loss_reduce_b`.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void mean_reduce_b_f32(
    const float* __restrict__ values_per_b,   // [B]
    int                       B,
    float*       __restrict__ mean_out        // [1]  OVERWRITE
) {
    __shared__ float s_val[BLOCK_X];

    const int tid = threadIdx.x;

    float local = 0.0f;
    for (int b = tid; b < B; b += BLOCK_X) {
        local += values_per_b[b];
    }
    s_val[tid] = local;
    __syncthreads();

    for (int stride = BLOCK_X / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            s_val[tid] += s_val[tid + stride];
        }
        __syncthreads();
    }

    if (tid == 0) {
        const float b_inv = 1.0f / (float)((B > 0) ? B : 1);
        mean_out[0] = s_val[0] * b_inv;
    }
}
