feat(cuda): add GpuStatistics wrapper with host-side mean/variance computation

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

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<Result<Ptx, String>> = OnceLock::new();
+
+fn compile_statistics_ptx() -> Result<Ptx, String> {
+    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<f32>,
+    device: Device,
+}
+
+impl GpuStatistics {
+    /// Create a new GPU statistics computer.
+    pub fn new(device: &Device) -> Result<Self, MLError> {
+        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::<f32>(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<f32>,
+        portfolio_values: &CudaSlice<f32>,
+        positions: &CudaSlice<f32>,
+        n: usize,
+    ) -> Result<BatchStatistics, MLError> {
+        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,
+        })
+    }
+}

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).