foxhunt/crates/ml-alpha/src/data/gpu_dataset.rs

//! GPU-resident dataset: all pre-converted snapshots + labels on device.
//!
//! Eliminates ALL per-step CPU data loading work. At init, the host
//! pre-converts every `Mbp10Snapshot` → `Mbp10RawInput`, concatenates
//! across files, and uploads as a single contiguous device buffer. A
//! GPU kernel (`gpu_sample_and_gather`) does random sampling + window
//! gathering per step — zero host work, zero host↔device transfers in
//! the training hot path.
//!
//! Memory budget (45M snapshots):
//!   snapshots: 45M × 216B = 9.7 GB
//!   FRD labels: 45M × 3 × 4B = 0.54 GB
//!   Total: ~10.2 GB on L40S (48GB), leaving 34+ GB free.
//!
//! Per `feedback_no_htod_htoh_only_mapped_pinned.md`: the init-time
//! upload uses cudarc's `htod_copy` (a one-time bulk transfer, not in
//! the captured-graph hot path). The per-step path is pure device-to-
//! device (kernel reads from one device buffer, writes to another).

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
use cudarc::driver::sys::CUstream;
use ml_core::device::MlDevice;

use crate::rl::common::FRD_N_HORIZONS;
use crate::trainer::perception::SoaBufferPtrs;
use crate::trainer::raw_launch::{RawArgs, raw_launch};

const GPU_SAMPLE_GATHER_CUBIN: &[u8] = include_bytes!(
    concat!(env!("OUT_DIR"), "/gpu_sample_and_gather.cubin")
);

/// All pre-converted snapshots + metadata resident on GPU. Created once
/// at init by [`super::loader::MultiHorizonLoader::upload_to_gpu`].
pub struct GpuDataset {
    /// All snapshots concatenated across files, on GPU.
    /// Byte layout: `[total_snaps]` of `Mbp10Raw` (216 bytes each).
    /// The kernel reinterprets this as `Mbp10Raw*`.
    pub snapshots_d: CudaSlice<u8>,
    /// Per-file start offset into `snapshots_d` (in snapshot units, not
    /// bytes). `file_offsets_d[f]` is the index of the first snapshot
    /// from file `f` in the flat `snapshots_d` array.
    pub file_offsets_d: CudaSlice<i32>,
    /// Per-file snapshot count. `file_sizes_d[f]` is the number of
    /// snapshots in file `f`.
    pub file_sizes_d: CudaSlice<i32>,
    /// FRD labels, row-major `[FRD_N_HORIZONS, total_snaps]`.
    /// Sentinel -1 at right-edge positions.
    pub frd_labels_d: CudaSlice<i32>,
    /// PRNG state per batch element `[B]`. Seeded at init, advanced by
    /// the kernel each step.
    pub prng_d: CudaSlice<u32>,

    // ── Supervised BCE labels + aux outcome labels ─────────────────────
    // Uploaded at init alongside snapshots so `gpu_gather_bce_labels`
    // can write per-step labels directly into the perception trainer's
    // mapped-pinned staging buffers — closing the supervised-loss gap
    // that left the encoder BCE-blind on the GPU data path.

    /// BCE direction labels, row-major `[N_HORIZONS, total_snaps]`.
    /// NaN at right-edge positions (matches host-side `labels_full`).
    pub labels_d: CudaSlice<f32>,
    /// Candidate-B prof-binary long labels `[N_HORIZONS, total_snaps]`.
    pub outcome_prof_long_d: CudaSlice<f32>,
    /// Candidate-B prof-binary short labels `[N_HORIZONS, total_snaps]`.
    pub outcome_prof_short_d: CudaSlice<f32>,
    /// Candidate-B σ-normalised size long targets `[N_HORIZONS, total_snaps]`.
    pub outcome_size_long_d: CudaSlice<f32>,
    /// Candidate-B σ-normalised size short targets `[N_HORIZONS, total_snaps]`.
    pub outcome_size_short_d: CudaSlice<f32>,
    /// Kendall σ per horizon `[N_HORIZONS, total_snaps]`.
    pub sigma_k_d: CudaSlice<f32>,
    /// Per-file positive-class fraction `[n_files, 2 * N_HORIZONS]`.
    /// Indexed as `file_idx * (2 * N_HORIZONS) + h` (long) or
    /// `file_idx * (2 * N_HORIZONS) + N_HORIZONS + h` (short).
    pub pos_fraction_d: CudaSlice<f32>,

    /// Number of loaded files.
    pub n_files: usize,
    /// Total snapshots across all files.
    pub total_snapshots: usize,
    /// Maximum forward horizon in snapshot units (max of
    /// `cfg.horizons.iter().max()`), needed by the sampling kernel to
    /// avoid right-edge overrun.
    pub max_horizon: usize,
}

/// GPU-resident data loader: dispatches the `gpu_sample_and_gather`
/// kernel family to fill SoA buffers each step with zero CPU work.
///
/// Usage pattern per training step (mirrors the existing
/// `step_with_lobsim` sequence):
///   1. `sample_and_gather(dataset, soa, ...)` — samples B random
///      (file, anchor) pairs AND gathers the s_t window into SoA.
///   2. Caller runs `forward_encoder_from_device()` → h_t. But wait:
///      the integrated trainer needs h_{t+1} FIRST. So the actual
///      sequence is:
///        a. `sample(dataset, ...)` — samples only, writes (file_offset, anchor) to device.
///        b. `gather_next_window(dataset, soa, ...)` — gathers s_{t+1}.
///        c. `forward_encoder_from_device()` → copy h_t → h_tp1.
///        d. `gather_current_window(dataset, soa, ...)` — gathers s_t.
///        e. `forward_encoder_from_device()` → h_t.
///        f. `gather_frd_labels(dataset, ...)` — gathers FRD labels.
pub struct GpuDataLoader {
    /// Kept alive to anchor the `raw_stream` pointer lifetime.
    _stream: Arc<CudaStream>,
    raw_stream: CUstream,
    _module: Arc<CudaModule>,
    sample_gather_fn: CudaFunction,
    gather_next_fn: CudaFunction,
    gather_current_fn: CudaFunction,
    gather_frd_labels_fn: CudaFunction,
    gather_bce_labels_fn: CudaFunction,
    gather_pos_fraction_fn: CudaFunction,
    /// Per-batch sampled file offsets — written by the sampling kernel,
    /// read by `gather_next_window` and `gather_frd_labels` so they
    /// use the same (file, anchor) pair. `[B]`.
    sample_file_offset_d: CudaSlice<i32>,
    /// Per-batch sampled anchors. `[B]`.
    sample_anchor_d: CudaSlice<i32>,
    /// Per-batch sampled file index — written by the sampling kernel,
    /// read by `gather_pos_fraction` to index into `pos_fraction_d`.
    sample_file_idx_d: CudaSlice<i32>,
}

impl GpuDataLoader {
    pub fn new(dev: &MlDevice, batch_size: usize) -> Result<Self> {
        let stream = dev.cuda_stream().context("GpuDataLoader: CUDA stream")?.clone();
        let raw_stream = stream.cu_stream();
        let ctx = dev.cuda_context().context("GpuDataLoader: CUDA context")?;
        let module = ctx
            .load_cubin(GPU_SAMPLE_GATHER_CUBIN.to_vec())
            .context("load gpu_sample_and_gather cubin")?;
        let sample_gather_fn = module
            .load_function("gpu_sample_and_gather")
            .context("load gpu_sample_and_gather function")?;
        let gather_next_fn = module
            .load_function("gpu_gather_next")
            .context("load gpu_gather_next function")?;
        let gather_current_fn = module
            .load_function("gpu_gather_current")
            .context("load gpu_gather_current function")?;
        let gather_frd_labels_fn = module
            .load_function("gpu_gather_frd_labels")
            .context("load gpu_gather_frd_labels function")?;
        let gather_bce_labels_fn = module
            .load_function("gpu_gather_bce_labels")
            .context("load gpu_gather_bce_labels function")?;
        let gather_pos_fraction_fn = module
            .load_function("gpu_gather_pos_fraction")
            .context("load gpu_gather_pos_fraction function")?;

        let sample_file_offset_d = stream
            .alloc_zeros::<i32>(batch_size)
            .context("alloc sample_file_offset_d")?;
        let sample_anchor_d = stream
            .alloc_zeros::<i32>(batch_size)
            .context("alloc sample_anchor_d")?;
        let sample_file_idx_d = stream
            .alloc_zeros::<i32>(batch_size)
            .context("alloc sample_file_idx_d")?;

        Ok(Self {
            _stream: stream,
            raw_stream,
            _module: module,
            sample_gather_fn,
            gather_next_fn,
            gather_current_fn,
            gather_frd_labels_fn,
            gather_bce_labels_fn,
            gather_pos_fraction_fn,
            sample_file_offset_d,
            sample_anchor_d,
            sample_file_idx_d,
        })
    }

    /// Sample B random (file, anchor) pairs AND gather the s_t window
    /// into the perception trainer's SoA buffers. The sampled
    /// (file_offset, anchor) values are saved to device memory for
    /// subsequent `gather_next_window` and `gather_frd_labels` calls.
    pub fn sample_and_gather(
        &mut self,
        dataset: &GpuDataset,
        soa: &SoaBufferPtrs,
        seq_len: usize,
        batch_size: usize,
        label_outs: [u64; 5],
    ) -> Result<()> {
        let mut args = RawArgs::new();
        args.push_ptr(dataset.snapshots_d.raw_ptr());
        args.push_ptr(dataset.file_offsets_d.raw_ptr());
        args.push_ptr(dataset.file_sizes_d.raw_ptr());
        args.push_ptr(dataset.prng_d.raw_ptr());
        args.push_i32(dataset.n_files as i32);
        args.push_i32(seq_len as i32);
        args.push_i32(dataset.max_horizon as i32);
        args.push_ptr(soa.bid_px);
        args.push_ptr(soa.bid_sz);
        args.push_ptr(soa.ask_px);
        args.push_ptr(soa.ask_sz);
        args.push_ptr(soa.regime);
        args.push_ptr(soa.prev_mid);
        args.push_ptr(soa.trade_signed_vol);
        args.push_ptr(soa.trade_count);
        args.push_ptr(soa.ts_ns);
        args.push_ptr(soa.prev_ts_ns);
        args.push_ptr(self.sample_file_offset_d.raw_ptr());
        args.push_ptr(self.sample_anchor_d.raw_ptr());
        args.push_ptr(self.sample_file_idx_d.raw_ptr());
        // Fused label gather arguments
        args.push_i32(dataset.total_snapshots as i32);
        args.push_ptr(dataset.labels_d.raw_ptr());
        args.push_ptr(dataset.outcome_prof_long_d.raw_ptr());
        args.push_ptr(dataset.outcome_prof_short_d.raw_ptr());
        args.push_ptr(dataset.outcome_size_long_d.raw_ptr());
        args.push_ptr(dataset.outcome_size_short_d.raw_ptr());
        args.push_ptr(label_outs[0]);
        args.push_ptr(label_outs[1]);
        args.push_ptr(label_outs[2]);
        args.push_ptr(label_outs[3]);
        args.push_ptr(label_outs[4]);
        args.push_i32(batch_size as i32);
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.sample_gather_fn.cu_function(),
                (batch_size as u32, 1, 1),
                (seq_len as u32, 1, 1),
                0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            )
            .map_err(|e| anyhow::anyhow!("gpu_sample_and_gather: {:?}", e))?;
        }
        Ok(())
    }

    /// Gather the s_{t+1} window (anchor+1) into the specified SoA
    /// buffers. Uses the (file_offset, anchor) sampled by the most
    /// recent `sample_and_gather` call.
    pub fn gather_next_window(
        &mut self,
        dataset: &GpuDataset,
        soa: &SoaBufferPtrs,
        seq_len: usize,
        batch_size: usize,
    ) -> Result<()> {
        let mut args = RawArgs::new();
        args.push_ptr(dataset.snapshots_d.raw_ptr());
        args.push_ptr(self.sample_file_offset_d.raw_ptr());
        args.push_ptr(self.sample_anchor_d.raw_ptr());
        args.push_i32(seq_len as i32);
        args.push_ptr(soa.bid_px);
        args.push_ptr(soa.bid_sz);
        args.push_ptr(soa.ask_px);
        args.push_ptr(soa.ask_sz);
        args.push_ptr(soa.regime);
        args.push_ptr(soa.prev_mid);
        args.push_ptr(soa.trade_signed_vol);
        args.push_ptr(soa.trade_count);
        args.push_ptr(soa.ts_ns);
        args.push_ptr(soa.prev_ts_ns);
        args.push_i32(batch_size as i32);
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.gather_next_fn.cu_function(),
                (batch_size as u32, 1, 1),
                (seq_len as u32, 1, 1),
                0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            )
            .map_err(|e| anyhow::anyhow!("gpu_gather_next: {:?}", e))?;
        }
        Ok(())
    }

    /// Re-gather the s_t window (anchor+0) into the SoA buffers. Uses
    /// the saved (file_offset, anchor) from the most recent
    /// `sample_and_gather` call — no PRNG re-sampling. Called after
    /// `gather_next_window` overwrites the SoA with s_{t+1} and the
    /// encoder forward on s_{t+1} has completed, to restore s_t in the
    /// SoA for the second encoder forward.
    pub fn gather_current_window(
        &mut self,
        dataset: &GpuDataset,
        soa: &SoaBufferPtrs,
        seq_len: usize,
        batch_size: usize,
    ) -> Result<()> {
        let mut args = RawArgs::new();
        args.push_ptr(dataset.snapshots_d.raw_ptr());
        args.push_ptr(self.sample_file_offset_d.raw_ptr());
        args.push_ptr(self.sample_anchor_d.raw_ptr());
        args.push_i32(seq_len as i32);
        args.push_ptr(soa.bid_px);
        args.push_ptr(soa.bid_sz);
        args.push_ptr(soa.ask_px);
        args.push_ptr(soa.ask_sz);
        args.push_ptr(soa.regime);
        args.push_ptr(soa.prev_mid);
        args.push_ptr(soa.trade_signed_vol);
        args.push_ptr(soa.trade_count);
        args.push_ptr(soa.ts_ns);
        args.push_ptr(soa.prev_ts_ns);
        args.push_i32(batch_size as i32);
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.gather_current_fn.cu_function(),
                (batch_size as u32, 1, 1),
                (seq_len as u32, 1, 1),
                0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            )
            .map_err(|e| anyhow::anyhow!("gpu_gather_current: {:?}", e))?;
        }
        Ok(())
    }

    /// Gather FRD labels at the newest snapshot (anchor + seq_len - 1)
    /// of the s_t window into the output buffer. Uses the (file_offset,
    /// anchor) sampled by the most recent `sample_and_gather` call.
    pub fn gather_frd_labels(
        &mut self,
        dataset: &GpuDataset,
        seq_len: usize,
        batch_size: usize,
        frd_labels_out_ptr: u64,
    ) -> Result<()> {
        let mut args = RawArgs::new();
        args.push_ptr(dataset.frd_labels_d.raw_ptr());
        args.push_ptr(dataset.file_offsets_d.raw_ptr());
        args.push_ptr(self.sample_file_offset_d.raw_ptr());
        args.push_ptr(self.sample_anchor_d.raw_ptr());
        args.push_i32(seq_len as i32);
        args.push_i32(FRD_N_HORIZONS as i32);
        args.push_i32(dataset.total_snapshots as i32);
        args.push_ptr(frd_labels_out_ptr);
        args.push_i32(batch_size as i32);
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.gather_frd_labels_fn.cu_function(),
                (batch_size as u32, 1, 1),
                (FRD_N_HORIZONS as u32, 1, 1),
                0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            )
            .map_err(|e| anyhow::anyhow!("gpu_gather_frd_labels: {:?}", e))?;
        }
        Ok(())
    }

    /// Gather per-horizon float labels for the sampled window into a
    /// mapped-pinned staging buffer. The kernel writes the `[K, B, n_h]`
    /// layout that `dispatch_train_step`'s DtoD expects. The caller
    /// passes the `dev_ptr` of the destination mapped-pinned buffer (e.g.
    /// `perception.stg_labels.dev_ptr`).
    ///
    /// `all_labels_ptr` is the raw device pointer to one of the 6
    /// per-horizon label slices on `GpuDataset` (e.g. `labels_d`,
    /// `outcome_prof_long_d`). Caller selects which array to gather.
    pub fn gather_bce_labels(
        &mut self,
        all_labels_ptr: u64,
        total_snapshots: usize,
        n_horizons: usize,
        seq_len: usize,
        batch_size: usize,
        labels_out_ptr: u64,
    ) -> Result<()> {
        let mut args = RawArgs::new();
        args.push_ptr(all_labels_ptr);
        args.push_ptr(self.sample_file_offset_d.raw_ptr());
        args.push_ptr(self.sample_anchor_d.raw_ptr());
        args.push_i32(seq_len as i32);
        args.push_i32(n_horizons as i32);
        args.push_i32(total_snapshots as i32);
        args.push_ptr(labels_out_ptr);
        args.push_i32(batch_size as i32);
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.gather_bce_labels_fn.cu_function(),
                (batch_size as u32, 1, 1),
                (seq_len as u32, 1, 1),
                0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            )
            .map_err(|e| anyhow::anyhow!("gpu_gather_bce_labels: {:?}", e))?;
        }
        Ok(())
    }

    /// Gather per-file `pos_fraction` for batch element 0's sampled file
    /// into the output buffer `[2 * N_HORIZONS]`. The trainer uses this
    /// to compute BCE `pos_weight = (1 - p) / p` per direction/horizon.
    pub fn gather_pos_fraction(
        &mut self,
        dataset: &GpuDataset,
        n_horizons: usize,
        batch_size: usize,
        pos_fraction_out_ptr: u64,
    ) -> Result<()> {
        let n_slots = 2 * n_horizons;
        let mut args = RawArgs::new();
        args.push_ptr(dataset.pos_fraction_d.raw_ptr());
        args.push_ptr(self.sample_file_idx_d.raw_ptr());
        args.push_i32(n_horizons as i32);
        args.push_ptr(pos_fraction_out_ptr);
        args.push_i32(batch_size as i32);
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.gather_pos_fraction_fn.cu_function(),
                (1, 1, 1),
                (n_slots as u32, 1, 1),
                0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            )
            .map_err(|e| anyhow::anyhow!("gpu_gather_pos_fraction: {:?}", e))?;
        }
        Ok(())
    }
}