diff --git a/crates/ml/src/cuda_pipeline/gpu_statistics.rs b/crates/ml/src/cuda_pipeline/gpu_statistics.rs new file mode 100644 index 000000000..45f6b7597 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/gpu_statistics.rs @@ -0,0 +1,165 @@ +#![allow(unsafe_code)] + +//! GPU batch statistics computation via parallel reduction. +//! +//! Replaces per-sample CPU accumulation of reward mean, variance, +//! drawdown, PnL, and win rate with a single CUDA kernel launch +//! and 40-byte readback. Derived statistics (mean, variance) are +//! computed on the host from raw sums for numerical correctness. + +use candle_core::cuda_backend::cudarc; +use candle_core::Device; +use cudarc::driver::{CudaFunction, CudaSlice, LaunchConfig, PushKernelArg}; +use cudarc::nvrtc::Ptx; +use std::sync::OnceLock; + +use crate::MLError; + +static STATISTICS_PTX: OnceLock> = OnceLock::new(); + +fn compile_statistics_ptx() -> Result { + let kernel_src = include_str!("statistics_kernel.cu"); + cudarc::nvrtc::compile_ptx(kernel_src) + .map_err(|e| format!("batch_statistics CUDA kernel compilation failed: {e}")) +} + +/// Aggregated batch statistics from GPU parallel reduction. +#[derive(Debug, Clone)] +pub struct BatchStatistics { + pub mean_reward: f32, + pub reward_variance: f32, + pub max_drawdown: f32, + pub total_pnl: f32, + pub win_rate: f32, + pub total_count: f32, + pub reward_sum: f32, + pub reward_sq_sum: f32, + pub mean_position_size: f32, +} + +/// GPU batch statistics computer. +pub struct GpuStatistics { + kernel_func: CudaFunction, + output_buf: CudaSlice, + device: Device, +} + +impl GpuStatistics { + /// Create a new GPU statistics computer. + pub fn new(device: &Device) -> Result { + let cuda_dev = match device { + Device::Cuda(ref dev) => dev, + _ => return Err(MLError::ModelError("GpuStatistics requires CUDA".into())), + }; + + let ptx_result = STATISTICS_PTX.get_or_init(compile_statistics_ptx); + let ptx = ptx_result.as_ref().map_err(|e| { + MLError::ModelError(format!("statistics PTX: {e}")) + })?; + + let stream = cuda_dev.cuda_stream(); + let context = stream.context(); + let module = context.load_module(ptx.clone()).map_err(|e| { + MLError::ModelError(format!("statistics module load: {e}")) + })?; + let kernel_func = module.load_function("batch_statistics").map_err(|e| { + MLError::ModelError(format!("batch_statistics function load: {e}")) + })?; + + let output_buf = stream.alloc_zeros::(10).map_err(|e| { + MLError::ModelError(format!("alloc output_buf: {e}")) + })?; + + Ok(Self { kernel_func, output_buf, device: device.clone() }) + } + + /// Compute batch statistics from GPU-resident buffers. + /// + /// Single kernel launch + 40-byte readback. Mean, variance, and + /// drawdown are computed on the host from raw sums for correctness. + pub fn compute( + &mut self, + rewards: &CudaSlice, + portfolio_values: &CudaSlice, + positions: &CudaSlice, + n: usize, + ) -> Result { + if n == 0 { + return Ok(BatchStatistics { + mean_reward: 0.0, + reward_variance: 0.0, + max_drawdown: 0.0, + total_pnl: 0.0, + win_rate: 0.0, + total_count: 0.0, + reward_sum: 0.0, + reward_sq_sum: 0.0, + mean_position_size: 0.0, + }); + } + + let cuda_dev = match &self.device { + Device::Cuda(ref dev) => dev, + _ => return Err(MLError::ModelError("Not CUDA".into())), + }; + let stream = cuda_dev.cuda_stream(); + + // Zero output buffer + stream.memset_zeros(&mut self.output_buf).map_err(|e| { + MLError::ModelError(format!("memset output_buf: {e}")) + })?; + + let n_i32 = n as i32; + // Single block — grid-stride handles N > 256 + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + + unsafe { + stream + .launch_builder(&self.kernel_func) + .arg(rewards) + .arg(portfolio_values) + .arg(positions) + .arg(&mut self.output_buf) + .arg(&n_i32) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("batch_statistics launch: {e}")))?; + } + + // Single 40-byte readback + let mut host = [0.0_f32; 10]; + stream.memcpy_dtoh(&self.output_buf, &mut host).map_err(|e| { + MLError::ModelError(format!("statistics readback: {e}")) + })?; + + let reward_sum = host[0]; + let reward_sq_sum = host[1]; + let global_min_pv = host[2]; + let global_max_pv = host[3]; + let win_count = host[4]; + let total_count = host[5].max(1.0); + let position_sum = host[7]; + + // Compute derived statistics on host (mathematically correct) + let mean = reward_sum / total_count; + let variance = (reward_sq_sum / total_count) - mean * mean; + let max_drawdown = (global_max_pv - global_min_pv).max(0.0); + let win_rate = win_count / total_count; + let mean_position = position_sum / total_count; + + Ok(BatchStatistics { + mean_reward: mean, + reward_variance: variance.max(0.0), + max_drawdown, + total_pnl: reward_sum, + win_rate, + total_count, + reward_sum, + reward_sq_sum, + mean_position_size: mean_position, + }) + } +} diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index c21106b64..622bff6f3 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -20,6 +20,10 @@ pub mod gpu_weights; pub mod gpu_experience_collector; #[cfg(feature = "cuda")] pub mod gpu_ppo_collector; +#[cfg(feature = "cuda")] +pub mod gpu_action_selector; +#[cfg(feature = "cuda")] +pub mod gpu_statistics; // gpu_replay_buffer moved to ml-dqn crate /// Maximum bytes allowed for a single GPU upload (2 GB safety limit).