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

// gpu_sample_and_gather.cu — GPU-resident batch sampler + AoS→SoA scatter
//
// Replaces the CPU-side `next_sequence_pair` + `build_sequence_at` +
// `extend_from_slice` + mapped-pinned memcpy hot path. At b=256, K=32,
// the CPU version spent ~16ms/step on 512 build_sequence_at calls with
// 7,700 heap allocations while GPU compute was 0.62ms. This kernel
// samples B random windows and gathers snapshots directly from a
// pre-uploaded GPU-resident dataset — zero CPU work per step.
//
// Architecture:
//   * At init, the host pre-converts ALL Mbp10Snapshot → Mbp10RawInput
//     and uploads as a single flat CudaSlice<u8> (reinterpreted as
//     Mbp10Raw on device). File offsets + sizes are uploaded as i32
//     arrays. Labels/regime are uploaded similarly.
//   * Per step, this kernel runs with Grid=(B,1,1), Block=(K,1,1):
//     - Thread 0 of each block uses device-side xorshift32 PRNG to
//       sample file_idx and anchor within that file.
//     - All K threads gather all_snapshots[file_offset + anchor + k]
//       and scatter to SoA output buffers (same layout as
//       snapshot_aos_to_soa.cu).
//
// Per feedback_no_atomicadd: block-local shared memory for sampling,
// no atomics. Per feedback_no_nvrtc: pre-compiled cubin via build.rs.

#define BOOK_LEVELS 10
#define REGIME_DIM  6
#define N_HORIZONS  3
#define N_LABEL_SOURCES 5

// Must match #[repr(C)] Mbp10RawInput in snap_features.rs (216 bytes).
struct __align__(8) Mbp10Raw {
    float bid_px[BOOK_LEVELS];   // offset   0, 40 bytes
    float bid_sz[BOOK_LEVELS];   // offset  40, 40 bytes
    float ask_px[BOOK_LEVELS];   // offset  80, 40 bytes
    float ask_sz[BOOK_LEVELS];   // offset 120, 40 bytes
    float prev_mid;              // offset 160, 4 bytes
    float trade_signed_vol;      // offset 164, 4 bytes
    unsigned int trade_count;    // offset 168, 4 bytes
    // 4 bytes padding for u64 alignment
    unsigned long long ts_ns;    // offset 176, 8 bytes
    unsigned long long prev_ts_ns; // offset 184, 8 bytes
    float regime[REGIME_DIM];    // offset 192, 24 bytes
    // total: 216 bytes
};

// Marsaglia xorshift32 — minimal state, sufficient quality for
// stochastic sampling. Each batch element carries its own seed to
// avoid inter-block contention.
__device__ __forceinline__ unsigned int xorshift32(unsigned int s) {
    s ^= s << 13;
    s ^= s >> 17;
    s ^= s << 5;
    return s;
}

// ── Primary kernel: sample + gather + SoA scatter ────────────────────
//
// Grid  = (B, 1, 1)   — one block per batch element
// Block = (K, 1, 1)   — one thread per sequence position
//
// Thread 0 samples (file_idx, anchor) via xorshift32. All K threads
// then read their assigned snapshot from the flat GPU array and write
// to SoA outputs at position [b * K + k].

extern "C" __global__ void gpu_sample_and_gather(
    const Mbp10Raw* __restrict__ all_snapshots,  // [total_snaps]
    const int*      __restrict__ file_offsets,    // [n_files]
    const int*      __restrict__ file_sizes,      // [n_files]
    unsigned int*   __restrict__ prng_state,      // [B]
    int n_files,
    int seq_len,
    int max_horizon,
    // SoA outputs (same layout as snapshot_aos_to_soa):
    float*          __restrict__ bid_px_soa,      // [B*K * BOOK_LEVELS]
    float*          __restrict__ bid_sz_soa,
    float*          __restrict__ ask_px_soa,
    float*          __restrict__ ask_sz_soa,
    float*          __restrict__ regime_soa,       // [B*K * REGIME_DIM]
    float*          __restrict__ prev_mid_soa,     // [B*K]
    float*          __restrict__ tsv_soa,          // [B*K]
    int*            __restrict__ tc_soa,           // [B*K]
    long long*      __restrict__ ts_ns_soa,        // [B*K]
    long long*      __restrict__ prev_ts_ns_soa,   // [B*K]
    // Global-memory outputs for the sampled (file_offset, anchor, file_idx)
    // per batch element. Read by gpu_gather_next and gpu_gather_frd_labels
    // to reuse the same sampling decision.
    int*            __restrict__ out_file_offset,   // [B]
    int*            __restrict__ out_anchor,         // [B]
    int*            __restrict__ out_file_idx,       // [B]
    // Fused label gather — 5 label sources (BCE, prof_long, prof_short,
    // size_long, size_short), each [N_HORIZONS × total_snaps]. Reads in
    // the same pass as snapshot gather to avoid 5 separate random-access
    // kernel launches that consumed 95.6% of GPU time.
    int             total_snaps,
    const float*    __restrict__ labels_src_0,      // [N_HORIZONS × total_snaps]
    const float*    __restrict__ labels_src_1,
    const float*    __restrict__ labels_src_2,
    const float*    __restrict__ labels_src_3,
    const float*    __restrict__ labels_src_4,
    float*          __restrict__ labels_out_0,      // [K, B, N_HORIZONS]
    float*          __restrict__ labels_out_1,
    float*          __restrict__ labels_out_2,
    float*          __restrict__ labels_out_3,
    float*          __restrict__ labels_out_4,
    int B
) {
    int b = blockIdx.x;
    int k = threadIdx.x;
    if (b >= B || k >= seq_len) return;

    // ── Thread 0: sample file + anchor via device-side PRNG ──────────
    __shared__ int s_file_offset;
    __shared__ int s_anchor;

    if (k == 0) {
        unsigned int seed = prng_state[b];
        // Sample file index (uniform across files).
        seed = xorshift32(seed);
        int file_idx = (int)(seed % (unsigned int)n_files);
        // Sample anchor within file. Must leave room for seq_len +
        // max_horizon snapshots after the anchor (+ 1 for the s_{t+1}
        // window gathered by gpu_gather_next).
        int fsize = file_sizes[file_idx];
        int usable = fsize - seq_len - max_horizon - 1;
        seed = xorshift32(seed);
        int anchor;
        if (usable > 0) {
            anchor = (int)(seed % (unsigned int)usable);
        } else {
            // Degenerate file too short — clamp to 0 (defensive).
            anchor = 0;
        }
        s_file_offset = file_offsets[file_idx];
        s_anchor = anchor;
        // Write to global memory so downstream kernels can read the
        // same (file_offset, anchor, file_idx) without re-sampling.
        out_file_offset[b] = s_file_offset;
        out_anchor[b] = anchor;
        out_file_idx[b] = file_idx;
        prng_state[b] = seed;
    }
    __syncthreads();

    // ── All K threads: gather + SoA scatter ──────────────────────────
    int global_idx = s_file_offset + s_anchor + k;
    const Mbp10Raw& snap = all_snapshots[global_idx];
    int n = b * seq_len + k;  // output position in [B*K] flat layout

    int base_book   = n * BOOK_LEVELS;
    int base_regime = n * REGIME_DIM;

    #pragma unroll
    for (int i = 0; i < BOOK_LEVELS; i++) {
        bid_px_soa[base_book + i] = snap.bid_px[i];
        bid_sz_soa[base_book + i] = snap.bid_sz[i];
        ask_px_soa[base_book + i] = snap.ask_px[i];
        ask_sz_soa[base_book + i] = snap.ask_sz[i];
    }
    #pragma unroll
    for (int i = 0; i < REGIME_DIM; i++) {
        regime_soa[base_regime + i] = snap.regime[i];
    }
    prev_mid_soa[n]  = snap.prev_mid;
    tsv_soa[n]       = snap.trade_signed_vol;
    tc_soa[n]        = (int)snap.trade_count;
    ts_ns_soa[n]     = (long long)snap.ts_ns;
    prev_ts_ns_soa[n] = (long long)snap.prev_ts_ns;

    // Fused label gather — same global_idx, zero extra random access.
    // Output layout: [K, B, N_HORIZONS] row-major.
    int lbl_base = (k * B + b) * N_HORIZONS;
    const float* srcs[N_LABEL_SOURCES] = {
        labels_src_0, labels_src_1, labels_src_2, labels_src_3, labels_src_4
    };
    float* outs[N_LABEL_SOURCES] = {
        labels_out_0, labels_out_1, labels_out_2, labels_out_3, labels_out_4
    };
    #pragma unroll
    for (int s = 0; s < N_LABEL_SOURCES; s++) {
        #pragma unroll
        for (int h = 0; h < N_HORIZONS; h++) {
            outs[s][lbl_base + h] = srcs[s][h * total_snaps + global_idx];
        }
    }
}

// ── Label gather kernel (STANDALONE — kept for backward compat) ─────
//
// Runs AFTER gpu_sample_and_gather. Gathers per-horizon FRD labels at
// the anchor position (rightmost K position = newest snapshot in the
// window, matching the supervised label semantics).
//
// Grid  = (B, 1, 1)
// Block = (FRD_N_HORIZONS, 1, 1)  — typically 3 threads
//
// For RL training, the FRD labels are the primary use case. BCE labels
// and outcome labels are handled by the supervised path and are not
// needed per-step in the RL loop.

extern "C" __global__ void gpu_gather_frd_labels(
    const int*      __restrict__ all_frd_labels,   // [n_frd_horizons * total_snaps]
    const int*      __restrict__ file_offsets,      // [n_files] (same as above)
    const int*      __restrict__ sample_file_offset,// [B] — file_offset chosen by gpu_sample_and_gather
    const int*      __restrict__ sample_anchor,     // [B] — anchor chosen by gpu_sample_and_gather
    int             seq_len,
    int             n_frd_horizons,
    int             total_snaps,
    int*            __restrict__ frd_labels_out,    // [B * n_frd_horizons]
    int B
) {
    int b = blockIdx.x;
    int h = threadIdx.x;
    if (b >= B || h >= n_frd_horizons) return;

    // The FRD label is at the NEWEST snapshot in the window = anchor + seq_len - 1.
    int newest_idx = sample_file_offset[b] + sample_anchor[b] + seq_len - 1;
    int label_idx = h * total_snaps + newest_idx;  // row-major [n_frd_horizons, total_snaps]
    frd_labels_out[b * n_frd_horizons + h] = all_frd_labels[label_idx];
}

// ── Pair gather: s_{t+1} window for Bellman targets ──────────────────
//
// The RL trainer needs (s_t, s_{t+1}) consecutive windows. This kernel
// gathers the SECOND window at anchor+1 into a separate set of SoA
// buffers, reusing the same (file_offset, anchor) from the primary
// gpu_sample_and_gather call. Avoids a second sampling pass.
//
// Grid  = (B, 1, 1)
// Block = (K, 1, 1)

extern "C" __global__ void gpu_gather_next(
    const Mbp10Raw* __restrict__ all_snapshots,    // [total_snaps]
    const int*      __restrict__ sample_file_offset,// [B]
    const int*      __restrict__ sample_anchor,     // [B]
    int             seq_len,
    // SoA outputs for s_{t+1}:
    float*          __restrict__ bid_px_soa,
    float*          __restrict__ bid_sz_soa,
    float*          __restrict__ ask_px_soa,
    float*          __restrict__ ask_sz_soa,
    float*          __restrict__ regime_soa,
    float*          __restrict__ prev_mid_soa,
    float*          __restrict__ tsv_soa,
    int*            __restrict__ tc_soa,
    long long*      __restrict__ ts_ns_soa,
    long long*      __restrict__ prev_ts_ns_soa,
    int B
) {
    int b = blockIdx.x;
    int k = threadIdx.x;
    if (b >= B || k >= seq_len) return;

    // anchor + 1 for the next-step window.
    int global_idx = sample_file_offset[b] + sample_anchor[b] + 1 + k;
    const Mbp10Raw& snap = all_snapshots[global_idx];
    int n = b * seq_len + k;

    int base_book   = n * BOOK_LEVELS;
    int base_regime = n * REGIME_DIM;

    #pragma unroll
    for (int i = 0; i < BOOK_LEVELS; i++) {
        bid_px_soa[base_book + i] = snap.bid_px[i];
        bid_sz_soa[base_book + i] = snap.bid_sz[i];
        ask_px_soa[base_book + i] = snap.ask_px[i];
        ask_sz_soa[base_book + i] = snap.ask_sz[i];
    }
    #pragma unroll
    for (int i = 0; i < REGIME_DIM; i++) {
        regime_soa[base_regime + i] = snap.regime[i];
    }
    prev_mid_soa[n]  = snap.prev_mid;
    tsv_soa[n]       = snap.trade_signed_vol;
    tc_soa[n]        = (int)snap.trade_count;
    ts_ns_soa[n]     = (long long)snap.ts_ns;
    prev_ts_ns_soa[n] = (long long)snap.prev_ts_ns;
}

// ── Re-gather s_t: current window at anchor+0 ──────────────────────
//
// After gpu_sample_and_gather samples an anchor and gpu_gather_next
// overwrites the SoA with s_{t+1}, this kernel re-populates the SoA
// with the original s_t window. Uses the saved (file_offset, anchor)
// from gpu_sample_and_gather — no PRNG re-sampling. Identical to
// gpu_gather_next except the global_idx offset is +0 instead of +1.
//
// Grid  = (B, 1, 1)
// Block = (K, 1, 1)

extern "C" __global__ void gpu_gather_current(
    const Mbp10Raw* __restrict__ all_snapshots,
    const int*      __restrict__ sample_file_offset,
    const int*      __restrict__ sample_anchor,
    int             seq_len,
    float*          __restrict__ bid_px_soa,
    float*          __restrict__ bid_sz_soa,
    float*          __restrict__ ask_px_soa,
    float*          __restrict__ ask_sz_soa,
    float*          __restrict__ regime_soa,
    float*          __restrict__ prev_mid_soa,
    float*          __restrict__ tsv_soa,
    int*            __restrict__ tc_soa,
    long long*      __restrict__ ts_ns_soa,
    long long*      __restrict__ prev_ts_ns_soa,
    int B
) {
    int b = blockIdx.x;
    int k = threadIdx.x;
    if (b >= B || k >= seq_len) return;

    int global_idx = sample_file_offset[b] + sample_anchor[b] + k;
    const Mbp10Raw& snap = all_snapshots[global_idx];
    int n = b * seq_len + k;

    int base_book   = n * BOOK_LEVELS;
    int base_regime = n * REGIME_DIM;

    #pragma unroll
    for (int i = 0; i < BOOK_LEVELS; i++) {
        bid_px_soa[base_book + i] = snap.bid_px[i];
        bid_sz_soa[base_book + i] = snap.bid_sz[i];
        ask_px_soa[base_book + i] = snap.ask_px[i];
        ask_sz_soa[base_book + i] = snap.ask_sz[i];
    }
    #pragma unroll
    for (int i = 0; i < REGIME_DIM; i++) {
        regime_soa[base_regime + i] = snap.regime[i];
    }
    prev_mid_soa[n]  = snap.prev_mid;
    tsv_soa[n]       = snap.trade_signed_vol;
    tc_soa[n]        = (int)snap.trade_count;
    ts_ns_soa[n]     = (long long)snap.ts_ns;
    prev_ts_ns_soa[n] = (long long)snap.prev_ts_ns;
}

// ── BCE / aux label gather kernel ───────────────────────────────────
//
// Gathers per-horizon float labels at each of the K window positions
// (anchor + k, for k in [0, K)) into the output buffer with the
// same [K, B, N_H] row-major layout that step_batched's mapped-pinned
// staging uses. This unified kernel serves all 6 per-horizon float
// label arrays (labels, outcome_prof_{long,short}, outcome_size_{long,
// short}, sigma_k) — the caller selects which by passing the
// appropriate `all_labels` source pointer from GpuDataset.
//
// Grid  = (B, 1, 1)
// Block = (K, 1, 1) where K = seq_len
//
// The output layout is:
//   out[k * B * n_horizons + b * n_horizons + h] = all_labels[h * total_snaps + global_idx]
// matching the staging fill loop in perception.rs step_batched.

extern "C" __global__ void gpu_gather_bce_labels(
    const float*    __restrict__ all_labels,          // [n_horizons, total_snaps] row-major
    const int*      __restrict__ sample_file_offset,  // [B]
    const int*      __restrict__ sample_anchor,       // [B]
    int             seq_len,
    int             n_horizons,
    int             total_snaps,
    float*          __restrict__ labels_out,           // [K, B, n_horizons] row-major
    int B
) {
    int b = blockIdx.x;
    int k = threadIdx.x;
    if (b >= B || k >= seq_len) return;

    int global_idx = sample_file_offset[b] + sample_anchor[b] + k;
    int out_base = (k * B + b) * n_horizons;

    for (int h = 0; h < n_horizons; h++) {
        labels_out[out_base + h] = all_labels[h * total_snaps + global_idx];
    }
}

// ── Per-file pos_fraction gather kernel ─────────────────────────────
//
// Reads the per-file positive-class fraction for batch element 0's
// sampled file and writes a single pos_fraction vector. The trainer
// uses per-file pos_fraction as a coarse class-balance signal; all
// batch elements within a step use the same pw value (batch 0's file).
//
// Grid  = (1, 1, 1)
// Block = (2 * N_HORIZONS, 1, 1) — one thread per pos_fraction slot

extern "C" __global__ void gpu_gather_pos_fraction(
    const float*    __restrict__ all_pos_fraction,    // [n_files, 2 * n_horizons]
    const int*      __restrict__ sample_file_idx,     // [B]
    int             n_horizons,
    float*          __restrict__ pos_fraction_out,    // [2 * n_horizons]
    int B
) {
    int slot = threadIdx.x;
    int n_slots = 2 * n_horizons;
    if (slot >= n_slots) return;

    // Use batch 0's file index as the representative.
    int file_idx = sample_file_idx[0];
    pos_fraction_out[slot] = all_pos_fraction[file_idx * n_slots + slot];
}