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

/* rl_curriculum_weights — GPU-side E8 difficulty-weighted curriculum computation.
 *
 * Replaces the host-side `enrichment::compute_curriculum_weights` + manual
 * ISV write loop with a single device-resident kernel that reads per-segment
 * Sharpe from ISV input slots, computes z-score-normalised softmax difficulty
 * weights, and writes:
 *   - per-segment weights to ISV[CURRICULUM_WEIGHT_0_INDEX..+n_segments]
 *   - scalar concentration (1 - entropy/log(n)) to ISV[CURRICULUM_CONCENTRATION_INDEX]
 *
 * Run once per epoch (not per step) — launched after the host writes per-segment
 * Sharpe into ISV[CURRICULUM_SHARPE_0_INDEX..+n_segments].
 *
 * Grid  = (1, 1, 1)
 * Block = (min(n_segments, 256), 1, 1)
 * Shared memory = 2 * MAX_SEGMENTS * sizeof(float)
 *
 * Per `feedback_no_atomicadd.md`: reductions use block tree-reduce only.
 * Per `feedback_no_nvrtc.md`: pre-compiled cubin only (build.rs).
 */

/* ISV slot indices — literals mirror canonical Rust constants in
 * sp21_isv_slots.rs (CURRICULUM_WEIGHT_0_INDEX=528, CURRICULUM_CONCENTRATION_INDEX=527)
 * and the new CURRICULUM_SHARPE_0_INDEX=554 in e8_curriculum_isv_slots.rs.
 * Trainer constructor pins names to these indices via compile-time assertions. */
#define CURRICULUM_SHARPE_0_INDEX       554
#define CURRICULUM_WEIGHT_0_INDEX       528
#define CURRICULUM_CONCENTRATION_INDEX  527
#define MAX_SEGMENTS                     16

/* Block tree-reduce on shmem[0..bdim], total sum → shmem[0].
 * Caller must __syncthreads() BEFORE (writes visible) and AFTER
 * (safe to reuse buffer). Per feedback_no_atomicadd.md. */
__device__ __forceinline__ void curriculum_tree_reduce(
    float* shmem,
    int    tid,
    int    bdim)
{
    for (int s = bdim / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shmem[tid] += shmem[tid + s];
        }
        __syncthreads();
    }
}

/* Block tree-reduce for max: shmem[0] holds the block-wide max.
 * Same sync discipline as curriculum_tree_reduce. */
__device__ __forceinline__ void curriculum_tree_reduce_max(
    float* shmem,
    int    tid,
    int    bdim)
{
    for (int s = bdim / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shmem[tid] = fmaxf(shmem[tid], shmem[tid + s]);
        }
        __syncthreads();
    }
}

extern "C" __global__ void rl_curriculum_weights(
    float* __restrict__ isv,
    int                 n_segments   /* actual segment count, <= MAX_SEGMENTS */
)
{
    int tid  = threadIdx.x;
    int bdim = blockDim.x;

    /* Guard: threads beyond n_segments are inactive for data access
     * but participate in tree-reduce sync barriers. */
    const int active = (tid < n_segments) ? 1 : 0;

    /* Shared memory layout: two buffers of bdim floats each.
     * buf_a: working buffer for reductions (mean, std, sum_exp).
     * buf_b: secondary buffer for numerically-stable softmax (max finding). */
    extern __shared__ float shmem[];
    float* buf_a = shmem;
    float* buf_b = shmem + bdim;

    /* ---------------------------------------------------------------
     * Pass 1: Load per-segment Sharpe from ISV and compute mean.
     * --------------------------------------------------------------- */
    float sharpe_i = 0.0f;
    if (active) {
        sharpe_i = isv[CURRICULUM_SHARPE_0_INDEX + tid];
    }

    /* Sum of Sharpe values via tree-reduce → mean. */
    buf_a[tid] = sharpe_i;
    __syncthreads();
    curriculum_tree_reduce(buf_a, tid, bdim);
    /* buf_a[0] = sum of all sharpe values (inactive threads contributed 0). */

    const float sharpe_mean = buf_a[0] / fmaxf((float)n_segments, 1.0f);

    /* ---------------------------------------------------------------
     * Pass 2: Variance (for z-score std).
     * --------------------------------------------------------------- */
    float diff = active ? (sharpe_i - sharpe_mean) : 0.0f;
    buf_a[tid] = diff * diff;
    __syncthreads();
    curriculum_tree_reduce(buf_a, tid, bdim);

    const float sharpe_var = buf_a[0] / fmaxf((float)n_segments, 1.0f);
    const float sharpe_std = fmaxf(sqrtf(sharpe_var), 1e-6f);

    /* ---------------------------------------------------------------
     * Pass 3: z-score normalised difficulty → exp weights.
     *
     * Negate so HARDER segments (lower Sharpe) get larger weights.
     * z-clamp to [-3, 3] keeps exp range well-conditioned [0.05, 20].
     *
     * Numerically-stable softmax: subtract max(z) before exp to
     * prevent overflow. With z in [-3, 3] overflow is unlikely, but
     * the pattern costs one extra tree-reduce and is strictly correct.
     * --------------------------------------------------------------- */
    float z_i = 0.0f;
    if (active) {
        z_i = -(sharpe_i - sharpe_mean) / sharpe_std;
        z_i = fmaxf(-3.0f, fminf(3.0f, z_i));
    }

    /* Find max z for numerically-stable softmax. */
    buf_b[tid] = active ? z_i : -1e30f;
    __syncthreads();
    curriculum_tree_reduce_max(buf_b, tid, bdim);
    const float z_max = buf_b[0];

    /* exp(z_i - z_max) for numerical stability. */
    float exp_i = active ? expf(z_i - z_max) : 0.0f;

    /* Sum of exp weights via tree-reduce. */
    buf_a[tid] = exp_i;
    __syncthreads();
    curriculum_tree_reduce(buf_a, tid, bdim);
    const float sum_exp = buf_a[0];

    /* ---------------------------------------------------------------
     * Pass 4: Normalise weights (softmax) and write to ISV.
     * --------------------------------------------------------------- */
    float weight_i = 0.0f;
    if (active && sum_exp > 1e-12f) {
        weight_i = exp_i / sum_exp;
    } else if (active) {
        /* Degenerate case: uniform fallback. */
        weight_i = 1.0f / (float)n_segments;
    }

    if (active) {
        isv[CURRICULUM_WEIGHT_0_INDEX + tid] = weight_i;
    }

    /* ---------------------------------------------------------------
     * Pass 5: Concentration scalar = 1 - entropy / log(n_segments).
     *
     * entropy = -sum(w_i * log(w_i)) with 0*log(0) := 0.
     * Normalised entropy ∈ [0, 1]; concentration = 1 - normalised.
     *   concentration = 0 → perfectly uniform (all segments equal)
     *   concentration = 1 → single segment dominates
     *
     * Thread 0 writes the scalar to ISV[CURRICULUM_CONCENTRATION_INDEX].
     * --------------------------------------------------------------- */
    float h_i = 0.0f;
    if (active && weight_i > 1e-12f) {
        h_i = -weight_i * logf(weight_i);
    }

    buf_a[tid] = h_i;
    __syncthreads();
    curriculum_tree_reduce(buf_a, tid, bdim);

    if (tid == 0) {
        const float entropy = buf_a[0];
        const float log_n   = logf(fmaxf((float)n_segments, 2.0f));
        const float norm_entropy = fminf(entropy / log_n, 1.0f);
        const float concentration = 1.0f - norm_entropy;
        isv[CURRICULUM_CONCENTRATION_INDEX] = fmaxf(0.0f, fminf(1.0f, concentration));
    }
}