/* =====================================================================
 * rl_per_push_ring.cu — Coalesced n-step ring write kernel
 *
 * Writes the current step's hidden state, scaled reward, raw reward,
 * and done flag into a per-batch n-step ring buffer. The ring write
 * for h_t is fully coalesced: 128 threads each write one float,
 * producing a single 512-byte transaction per batch element.
 *
 * This is the first of two kernels splitting the old rl_per_push
 * (Grid=1, Block=B, 124us/call). The second kernel (rl_per_flush)
 * handles the actual PER priority computation and replay-buffer push
 * for batches whose flush_flag is set.
 *
 * Launch config: Grid=(b_size, 1, 1), Block=(128, 1, 1).
 * One block per batch element. 128 threads = HIDDEN_DIM.
 *
 * Memory access pattern:
 *   h_t read:  coalesced 512B (128 threads x 4B, contiguous in b)
 *   h_t write: coalesced 512B (128 threads x 4B, ring slot contiguous)
 *   scalars:   thread-0 only, 3 floats (negligible)
 *   ring meta: thread-0 only, 2 uint32 reads + 2 writes
 *
 * No atomicAdd. No shared memory. No nvrtc (pre-compiled cubin).
 * ===================================================================== */

#define HIDDEN_DIM      128
#define N_STEP_MAX      16
#define RL_GAMMA_INDEX  400
#define RL_N_STEP_INDEX 403

extern "C" __global__ void rl_per_push_ring(
    const float* __restrict__ h_t_current,      /* [B, 128]                      */
    const float* __restrict__ rewards_scaled,   /* [B]                           */
    const float* __restrict__ rewards_raw,      /* [B]                           */
    const float* __restrict__ dones,            /* [B]                           */
    const float* __restrict__ log_pi_old,       /* [B]                           */
    const int*   __restrict__ actions,          /* [B]                           */
    const float* __restrict__ isv,              /* ISV vector (shared, read-only)*/
    /* N-step ring buffers (per batch element) */
    float*        __restrict__ nstep_scalars,   /* [B, N_STEP_MAX, 3]            */
    float*        __restrict__ nstep_h_t,       /* [B, N_STEP_MAX, 128]          */
    unsigned int* __restrict__ nstep_write_idx, /* [B]                           */
    unsigned int* __restrict__ nstep_count,     /* [B]                           */
    /* Output: flush decision per batch element */
    int*          __restrict__ flush_flags,     /* [B]                           */
    int b_size
)
{
    const int b   = blockIdx.x;
    const int tid = threadIdx.x;

    if (b >= b_size) return;

    /* ------------------------------------------------------------------
     * Read ring metadata (same value for every thread in the block).
     * No shared memory needed — all 128 threads read the same two scalars
     * from global memory; the L1 cache serves subsequent reads from the
     * same cache line at zero cost.
     * ------------------------------------------------------------------ */
    const unsigned int ring_idx = nstep_write_idx[b];
    const unsigned int count    = nstep_count[b];

    /* ------------------------------------------------------------------
     * ISV-driven n_step: clamp to [1, N_STEP_MAX].
     * Thread 0 reads the ISV slot; other threads don't need the value
     * until the flush decision, which only thread 0 computes.
     * ------------------------------------------------------------------ */
    int n_step = N_STEP_MAX; /* default if ISV read is out of range */
    if (tid == 0) {
        float n_step_f = isv[RL_N_STEP_INDEX];
        int n_step_i = (int)n_step_f;
        n_step = (n_step_i < 1) ? 1 : ((n_step_i > N_STEP_MAX) ? N_STEP_MAX : n_step_i);
    }

    /* ------------------------------------------------------------------
     * Coalesced h_t copy: each thread writes one float.
     * Source: h_t_current[b * HIDDEN_DIM + tid]  (contiguous per batch)
     * Dest:   nstep_h_t[(b * N_STEP_MAX + ring_idx) * HIDDEN_DIM + tid]
     *
     * 128 threads x 4 bytes = 512 bytes = one coalesced transaction.
     * ------------------------------------------------------------------ */
    const long long src_offset = (long long)b * HIDDEN_DIM + tid;
    const long long dst_offset = ((long long)b * N_STEP_MAX + ring_idx) * HIDDEN_DIM + tid;
    nstep_h_t[dst_offset] = h_t_current[src_offset];

    /* ------------------------------------------------------------------
     * Thread 0: write scalar ring entries and advance ring pointer.
     * ------------------------------------------------------------------ */
    if (tid == 0) {
        /* Scalar layout: [rewards_scaled, rewards_raw, dones] per ring slot */
        const long long scalar_base = ((long long)b * N_STEP_MAX + ring_idx) * 3;
        nstep_scalars[scalar_base + 0] = rewards_scaled[b];
        nstep_scalars[scalar_base + 1] = rewards_raw[b];
        nstep_scalars[scalar_base + 2] = dones[b];

        /* Advance ring write pointer (modular) */
        nstep_write_idx[b] = (ring_idx + 1) % N_STEP_MAX;

        /* Increment count (saturates at N_STEP_MAX — ring is full) */
        const unsigned int new_count = (count < (unsigned int)N_STEP_MAX)
                                       ? count + 1
                                       : (unsigned int)N_STEP_MAX;
        nstep_count[b] = new_count;

        /* Flush decision: ring is full OR episode terminated */
        const int should_flush = ((int)new_count >= n_step) || (dones[b] > 0.5f);
        flush_flags[b] = should_flush ? 1 : 0;
    }
}
