foxhunt/crates/ml-alpha/cuda/rl_per_sample.cu

/* =====================================================================
 * rl_per_sample.cu — GPU-resident PER: proportional sampling via sum-tree
 *
 * Grid=(b_size), Block=(1). One thread per sample.
 *
 * Each thread samples a leaf from the priority sum-tree using stratified
 * sampling (segment per thread) with xorshift32 PRNG, then gathers the
 * transition data from replay storage.
 *
 * ISV reads: none (alpha used only at priority-update time)
 * ===================================================================== */

#define HIDDEN_DIM             128
#define SCALARS_PER_TRANSITION 7

/* xorshift32 PRNG — same pattern used throughout the codebase */
__device__ __forceinline__ unsigned int xorshift32(unsigned int* state) {
    unsigned int x = *state;
    x ^= x << 13;
    x ^= x >> 17;
    x ^= x << 5;
    *state = x;
    return x;
}

extern "C" __global__ void rl_per_sample(
    const float* __restrict__ priority_tree,     /* [2 * capacity] */
    const float* __restrict__ replay_h_t,        /* [capacity, 128] */
    const float* __restrict__ replay_h_tp1,      /* [capacity, 128] */
    const float* __restrict__ replay_scalars,    /* [capacity, 7] */
    const unsigned int* __restrict__ replay_len, /* [1] */
    const float* __restrict__ isv,
    unsigned int* __restrict__ prng_state,       /* [B] */
    /* Outputs */
    float* __restrict__ sampled_h_t,             /* [B, 128] */
    float* __restrict__ sampled_h_tp1,           /* [B, 128] */
    float* __restrict__ sampled_rewards,         /* [B] */
    float* __restrict__ sampled_dones,           /* [B] */
    float* __restrict__ sampled_log_pi_old,      /* [B] */
    float* __restrict__ sampled_n_step_gammas,   /* [B] */
    int*   __restrict__ sampled_actions,         /* [B] */
    unsigned int* __restrict__ sample_indices,   /* [B] */
    int b_size,
    int capacity
)
{
    const int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= b_size) return;

    const unsigned int len = replay_len[0];
    const float total_priority = priority_tree[1]; /* root */

    /* Guard: empty or zero-priority replay */
    if (total_priority < 1e-9f || len == 0) {
        for (int i = 0; i < HIDDEN_DIM; ++i) {
            sampled_h_t[b * HIDDEN_DIM + i] = 0.0f;
            sampled_h_tp1[b * HIDDEN_DIM + i] = 0.0f;
        }
        sampled_rewards[b] = 0.0f;
        sampled_dones[b] = 0.0f;
        sampled_log_pi_old[b] = 0.0f;
        sampled_n_step_gammas[b] = 0.0f;
        sampled_actions[b] = 0;
        sample_indices[b] = 0;
        return;
    }

    /* ── Seed PRNG if cold ────────────────────────────────────────────── */
    if (prng_state[b] == 0) {
        prng_state[b] = (unsigned int)(b + 1) * 2654435761u;
    }

    /* ── Stratified sampling: draw u in [b*segment, (b+1)*segment) ────── */
    const float segment = total_priority / (float)b_size;
    unsigned int rng = xorshift32(&prng_state[b]);
    const float u_frac = (float)(rng & 0x00FFFFFFu) / (float)0x01000000u; /* [0, 1) */
    float u = ((float)b + u_frac) * segment;

    /* Clamp to avoid floating-point overshoot */
    if (u >= total_priority) u = total_priority - 1e-6f;
    if (u < 0.0f) u = 0.0f;

    /* ── Walk tree top-down ───────────────────────────────────────────── */
    int idx = 1;
    while (idx < capacity) {
        const int left = 2 * idx;
        const float left_val = priority_tree[left];
        if (u <= left_val) {
            idx = left;
        } else {
            u -= left_val;
            idx = left + 1;
        }
    }
    const int leaf = idx - capacity;

    /* Clamp leaf to valid range */
    const int safe_leaf = (leaf >= 0 && leaf < (int)len) ? leaf : 0;
    sample_indices[b] = (unsigned int)safe_leaf;

    /* ── Gather h_t ──────────────────────────────────────────────────── */
    for (int i = 0; i < HIDDEN_DIM; ++i) {
        sampled_h_t[b * HIDDEN_DIM + i] = replay_h_t[safe_leaf * HIDDEN_DIM + i];
    }

    /* ── Gather h_tp1 ────────────────────────────────────────────────── */
    for (int i = 0; i < HIDDEN_DIM; ++i) {
        sampled_h_tp1[b * HIDDEN_DIM + i] = replay_h_tp1[safe_leaf * HIDDEN_DIM + i];
    }

    /* ── Unpack scalars ──────────────────────────────────────────────── */
    const int sc_base = safe_leaf * SCALARS_PER_TRANSITION;
    sampled_actions[b]        = (int)replay_scalars[sc_base + 0];
    sampled_rewards[b]        = replay_scalars[sc_base + 1]; /* n-step scaled return */
    sampled_dones[b]          = replay_scalars[sc_base + 3];
    sampled_log_pi_old[b]     = replay_scalars[sc_base + 4];
    sampled_n_step_gammas[b]  = replay_scalars[sc_base + 5];
}