eef58e5c6a
- MultiGpuConfig::detect() probes CUDA ordinals 0..8, returns None on single-GPU/CPU setups - shard_indices() splits dataset across devices with remainder handling - NcclGradientSync (behind `nccl` feature flag) for all-reduce averaging - Feature: `nccl = ["cuda"]` — requires NCCL library on system Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
998 lines
36 KiB
Rust
998 lines
36 KiB
Rust
//! GPU Pipeline: Pre-uploaded training data for DQN/PPO trainers
|
||
//!
|
||
//! Eliminates per-bar heap allocations by uploading all training data
|
||
//! as GPU tensors at epoch start. Both DQN and PPO trainers index
|
||
//! into these pre-uploaded tensors instead of creating per-step
|
||
//! `Tensor::from_vec()` copies.
|
||
|
||
use candle_core::{Device, Tensor};
|
||
use crate::MLError;
|
||
|
||
pub mod double_buffer;
|
||
pub mod multi_gpu;
|
||
pub mod prefetch;
|
||
|
||
#[cfg(feature = "cuda")]
|
||
pub mod gpu_portfolio;
|
||
#[cfg(feature = "cuda")]
|
||
pub mod gpu_weights;
|
||
#[cfg(feature = "cuda")]
|
||
pub mod gpu_experience_collector;
|
||
#[cfg(feature = "cuda")]
|
||
pub mod gpu_ppo_collector;
|
||
|
||
/// Maximum bytes allowed for a single GPU upload (2 GB safety limit).
|
||
const MAX_UPLOAD_BYTES: usize = 2 * 1024 * 1024 * 1024;
|
||
|
||
/// Round dimension up to nearest multiple of 8 for tensor core alignment.
|
||
///
|
||
/// Tensor cores on Ampere (A100, L4, L40S) and Hopper (H100) require
|
||
/// matrix dimensions divisible by 8 (FP16/BF16) for optimal throughput.
|
||
/// This pads dimensions like state_dim=54 → 56 to enable tensor core paths.
|
||
pub fn align_to_tensor_cores(dim: usize) -> usize {
|
||
(dim + 7) & !7
|
||
}
|
||
|
||
/// Estimate VRAM usage in bytes for pre-uploaded f32 data.
|
||
pub fn estimate_vram_bytes(num_elements: usize) -> usize {
|
||
num_elements * std::mem::size_of::<f32>()
|
||
}
|
||
|
||
/// Compute optimal (grid_dim, block_dim) for a 1-D kernel launch.
|
||
///
|
||
/// Selects the largest block size that divides 32 (warp size) and does not
|
||
/// exceed the GPU's `max_threads_per_block`. On data-centre GPUs with
|
||
/// high SM counts (H100: 132 SMs, A100: 108) the larger block sizes
|
||
/// improve occupancy vs. the default 32.
|
||
///
|
||
/// # Arguments
|
||
/// * `n_items` - Total work items (e.g. number of episodes)
|
||
/// * `max_threads_per_block` - Device limit (typically 1024). Pass 0 for default (256).
|
||
///
|
||
/// # Returns
|
||
/// `(grid_x, block_x)` tuple suitable for `LaunchConfig`.
|
||
pub fn optimal_launch_dims(n_items: u32, max_threads_per_block: u32) -> (u32, u32) {
|
||
let max_tpb = if max_threads_per_block == 0 { 256 } else { max_threads_per_block };
|
||
|
||
// Clamp to a warp-multiple. For experience-collection kernels each thread
|
||
// runs one full episode, so very large blocks aren't needed. Cap at 256 to
|
||
// keep register pressure manageable while still improving over 32.
|
||
let block = if n_items >= 256 && max_tpb >= 256 {
|
||
256
|
||
} else if n_items >= 128 && max_tpb >= 128 {
|
||
128
|
||
} else if n_items >= 64 && max_tpb >= 64 {
|
||
64
|
||
} else {
|
||
32
|
||
};
|
||
|
||
let grid = n_items.saturating_add(block - 1) / block;
|
||
(grid, block)
|
||
}
|
||
|
||
/// Pre-uploaded GPU training data for DQN trainer.
|
||
///
|
||
/// Holds market features [N, 51] and target prices [N, 4] as GPU tensors.
|
||
/// Portfolio features (3 dims) are computed per-bar by the trainer and
|
||
/// concatenated on-device via `Tensor::cat`.
|
||
#[derive(Debug)]
|
||
pub struct DqnGpuData {
|
||
/// Market features tensor [num_bars, 51] on GPU (f32)
|
||
pub features: Tensor,
|
||
/// Target prices tensor [num_bars, 4] on GPU (f32)
|
||
pub targets: Tensor,
|
||
/// Number of training bars
|
||
pub num_bars: usize,
|
||
/// Feature dimension (51)
|
||
pub feature_dim: usize,
|
||
}
|
||
|
||
impl DqnGpuData {
|
||
/// Upload DQN training data to GPU as f32 tensors.
|
||
///
|
||
/// Converts `[f64; 51]` features and `Vec<f64>` targets to f32,
|
||
/// flattens into contiguous arrays, and uploads once.
|
||
pub fn upload(
|
||
data: &[([f64; 51], Vec<f64>)],
|
||
device: &Device,
|
||
) -> Result<Self, MLError> {
|
||
let num_bars = data.len();
|
||
if num_bars == 0 {
|
||
return Err(MLError::ModelError("Empty training data".to_owned()));
|
||
}
|
||
|
||
let feature_dim = 51;
|
||
let target_dim = 4;
|
||
|
||
let estimated_bytes = estimate_vram_bytes(num_bars * (feature_dim + target_dim));
|
||
if estimated_bytes > MAX_UPLOAD_BYTES {
|
||
return Err(MLError::ModelError(format!(
|
||
"Training data too large for GPU upload: {:.1} GB > 2.0 GB limit ({} bars)",
|
||
estimated_bytes as f64 / 1_073_741_824.0,
|
||
num_bars,
|
||
)));
|
||
}
|
||
|
||
let mut flat_features = Vec::with_capacity(num_bars * feature_dim);
|
||
for (features, _) in data {
|
||
for &v in features.iter() {
|
||
flat_features.push(v as f32);
|
||
}
|
||
}
|
||
|
||
let mut flat_targets = Vec::with_capacity(num_bars * target_dim);
|
||
for (_, targets) in data {
|
||
for i in 0..target_dim {
|
||
flat_targets.push(targets.get(i).copied().unwrap_or(0.0) as f32);
|
||
}
|
||
}
|
||
|
||
let features = Tensor::from_vec(flat_features, (num_bars, feature_dim), device)
|
||
.map_err(|e| MLError::ModelError(format!("GPU feature upload failed: {e}")))?;
|
||
|
||
let targets = Tensor::from_vec(flat_targets, (num_bars, target_dim), device)
|
||
.map_err(|e| MLError::ModelError(format!("GPU target upload failed: {e}")))?;
|
||
|
||
Ok(Self {
|
||
features,
|
||
targets,
|
||
num_bars,
|
||
feature_dim,
|
||
})
|
||
}
|
||
|
||
/// Estimated VRAM usage in bytes.
|
||
pub fn vram_bytes(&self) -> usize {
|
||
estimate_vram_bytes(self.num_bars * (self.feature_dim + 4))
|
||
}
|
||
|
||
/// Get market features for a single bar as a [1, 51] tensor slice (zero-copy on GPU).
|
||
pub fn bar_features(&self, bar_idx: usize) -> Result<Tensor, MLError> {
|
||
self.features
|
||
.narrow(0, bar_idx, 1)
|
||
.map_err(|e| MLError::ModelError(format!("Bar {bar_idx} feature slice failed: {e}")))
|
||
}
|
||
|
||
/// Get market features for a batch of bars as a [batch_size, 51] tensor slice.
|
||
pub fn batch_features(&self, start: usize, count: usize) -> Result<Tensor, MLError> {
|
||
let count = count.min(self.num_bars.saturating_sub(start));
|
||
self.features
|
||
.narrow(0, start, count)
|
||
.map_err(|e| MLError::ModelError(format!("Batch feature slice [{start}..+{count}] failed: {e}")))
|
||
}
|
||
|
||
/// Get target prices for a single bar as a [1, 4] tensor slice.
|
||
pub fn bar_targets(&self, bar_idx: usize) -> Result<Tensor, MLError> {
|
||
self.targets
|
||
.narrow(0, bar_idx, 1)
|
||
.map_err(|e| MLError::ModelError(format!("Bar {bar_idx} target slice failed: {e}")))
|
||
}
|
||
|
||
/// Get target prices for a batch as [batch_size, 4] tensor slice.
|
||
pub fn batch_targets(&self, start: usize, count: usize) -> Result<Tensor, MLError> {
|
||
let count = count.min(self.num_bars.saturating_sub(start));
|
||
self.targets
|
||
.narrow(0, start, count)
|
||
.map_err(|e| MLError::ModelError(format!("Batch target slice [{start}..+{count}] failed: {e}")))
|
||
}
|
||
|
||
/// Get raw f32 target values for a bar (for CPU-side portfolio/barrier calculations).
|
||
///
|
||
/// Returns [current_close_preproc, next_close_preproc, current_close_raw, next_close_raw].
|
||
pub fn bar_target_values(&self, bar_idx: usize) -> Result<[f32; 4], MLError> {
|
||
let slice = self.bar_targets(bar_idx)?
|
||
.flatten_all()
|
||
.map_err(|e| MLError::ModelError(format!("Target flatten failed: {e}")))?
|
||
.to_vec1::<f32>()
|
||
.map_err(|e| MLError::ModelError(format!("Target to_vec1 failed: {e}")))?;
|
||
Ok([
|
||
slice.get(0).copied().unwrap_or(0.0),
|
||
slice.get(1).copied().unwrap_or(0.0),
|
||
slice.get(2).copied().unwrap_or(0.0),
|
||
slice.get(3).copied().unwrap_or(0.0),
|
||
])
|
||
}
|
||
|
||
/// Build a complete 54-dim state tensor by concatenating pre-uploaded market features
|
||
/// with per-bar portfolio features on-device.
|
||
pub fn build_state_tensor(
|
||
&self,
|
||
bar_idx: usize,
|
||
portfolio_features: &[f32; 3],
|
||
device: &Device,
|
||
) -> Result<Tensor, MLError> {
|
||
let market = self.bar_features(bar_idx)?;
|
||
let portfolio = Tensor::from_vec(
|
||
portfolio_features.to_vec(),
|
||
(1, 3),
|
||
device,
|
||
).map_err(|e| MLError::ModelError(format!("Portfolio tensor failed: {e}")))?;
|
||
|
||
Tensor::cat(&[&market, &portfolio], 1)
|
||
.map_err(|e| MLError::ModelError(format!("State cat failed: {e}")))
|
||
}
|
||
|
||
/// Build a batch of complete 54-dim state tensors on GPU.
|
||
///
|
||
/// Concatenates pre-uploaded [batch_size, 51] market features with
|
||
/// broadcast [1, 3] portfolio features to produce [batch_size, 54].
|
||
/// Eliminates ~770 Vec allocations per batch vs feature_vector_to_state().
|
||
///
|
||
/// # Arguments
|
||
/// * `start` - First bar index into pre-uploaded features
|
||
/// * `count` - Number of bars in batch
|
||
/// * `portfolio_features` - 3 f32 values: [normalized_value, normalized_position, spread]
|
||
/// * `device` - Target device (must match self.features device)
|
||
pub fn build_batch_states(
|
||
&self,
|
||
start: usize,
|
||
count: usize,
|
||
portfolio_features: &[f32; 3],
|
||
device: &Device,
|
||
) -> Result<Tensor, MLError> {
|
||
let count = count.min(self.num_bars.saturating_sub(start));
|
||
if count == 0 {
|
||
return Err(MLError::ModelError("Empty batch for state construction".to_owned()));
|
||
}
|
||
|
||
// [count, 51] zero-copy slice from pre-uploaded GPU tensor
|
||
let market = self.features
|
||
.narrow(0, start, count)
|
||
.map_err(|e| MLError::ModelError(format!("Batch feature slice failed: {e}")))?;
|
||
|
||
// [1, 3] portfolio features — single small upload
|
||
let portfolio = Tensor::from_vec(
|
||
portfolio_features.to_vec(),
|
||
(1, 3),
|
||
device,
|
||
).map_err(|e| MLError::ModelError(format!("Portfolio tensor failed: {e}")))?;
|
||
|
||
// Broadcast [1, 3] → [count, 3] then concatenate with [count, 51] → [count, 54]
|
||
let portfolio_broadcast = portfolio
|
||
.broadcast_as((count, 3))
|
||
.map_err(|e| MLError::ModelError(format!("Portfolio broadcast failed: {e}")))?;
|
||
|
||
Tensor::cat(&[&market, &portfolio_broadcast], 1)
|
||
.map_err(|e| MLError::ModelError(format!("State cat failed: {e}")))
|
||
}
|
||
}
|
||
|
||
/// Reusable CPU staging buffers for GPU upload.
|
||
///
|
||
/// Pre-allocates flat `Vec<f32>` staging arrays sized for the maximum
|
||
/// expected fold. Reusing these across walk-forward folds avoids
|
||
/// repeated heap allocation/deallocation, reducing GC pauses and
|
||
/// memory fragmentation — especially important on GPUs with limited
|
||
/// VRAM (L4 24 GB) or large datasets.
|
||
///
|
||
/// # Usage
|
||
/// ```ignore
|
||
/// let mut pool = GpuBufferPool::new(300_000, 51, 4);
|
||
/// for fold in folds {
|
||
/// let gpu_data = pool.upload_dqn(&fold_data, &device)?;
|
||
/// // ... train on gpu_data ...
|
||
/// }
|
||
/// ```
|
||
#[derive(Debug)]
|
||
pub struct GpuBufferPool {
|
||
/// Staging buffer for features [max_bars * feature_dim]
|
||
feature_buf: Vec<f32>,
|
||
/// Staging buffer for targets [max_bars * target_dim]
|
||
target_buf: Vec<f32>,
|
||
/// Maximum bars this pool can handle
|
||
pub max_bars: usize,
|
||
/// Feature dimension (51 for DQN)
|
||
pub feature_dim: usize,
|
||
/// Target dimension (4 for DQN)
|
||
pub target_dim: usize,
|
||
}
|
||
|
||
impl GpuBufferPool {
|
||
/// Create a pool pre-allocated for `max_bars` training bars.
|
||
pub fn new(max_bars: usize, feature_dim: usize, target_dim: usize) -> Self {
|
||
Self {
|
||
feature_buf: vec![0.0_f32; max_bars * feature_dim],
|
||
target_buf: vec![0.0_f32; max_bars * target_dim],
|
||
max_bars,
|
||
feature_dim,
|
||
target_dim,
|
||
}
|
||
}
|
||
|
||
/// Upload DQN data to GPU, reusing the pre-allocated staging buffers.
|
||
///
|
||
/// Data is copied into the staging `Vec<f32>` (no heap alloc if within
|
||
/// `max_bars`), then uploaded to GPU in a single `Tensor::from_vec` call.
|
||
pub fn upload_dqn(
|
||
&mut self,
|
||
data: &[([f64; 51], Vec<f64>)],
|
||
device: &Device,
|
||
) -> Result<DqnGpuData, MLError> {
|
||
let num_bars = data.len();
|
||
if num_bars == 0 {
|
||
return Err(MLError::ModelError("Empty training data".to_owned()));
|
||
}
|
||
|
||
let feat_len = num_bars * self.feature_dim;
|
||
let targ_len = num_bars * self.target_dim;
|
||
|
||
// Grow staging buffers only if the fold exceeds the pre-allocated size.
|
||
if feat_len > self.feature_buf.len() {
|
||
self.feature_buf.resize(feat_len, 0.0);
|
||
}
|
||
if targ_len > self.target_buf.len() {
|
||
self.target_buf.resize(targ_len, 0.0);
|
||
}
|
||
|
||
// Fill staging buffers (zero-alloc copy).
|
||
for (i, (features, targets)) in data.iter().enumerate() {
|
||
let f_start = i * self.feature_dim;
|
||
for (j, &v) in features.iter().enumerate() {
|
||
self.feature_buf[f_start + j] = v as f32;
|
||
}
|
||
|
||
let t_start = i * self.target_dim;
|
||
for j in 0..self.target_dim {
|
||
self.target_buf[t_start + j] = targets.get(j).copied().unwrap_or(0.0) as f32;
|
||
}
|
||
}
|
||
|
||
// Upload the used slice to GPU.
|
||
let features = Tensor::from_vec(
|
||
self.feature_buf[..feat_len].to_vec(),
|
||
(num_bars, self.feature_dim),
|
||
device,
|
||
)
|
||
.map_err(|e| MLError::ModelError(format!("GPU feature upload failed: {e}")))?;
|
||
|
||
let targets = Tensor::from_vec(
|
||
self.target_buf[..targ_len].to_vec(),
|
||
(num_bars, self.target_dim),
|
||
device,
|
||
)
|
||
.map_err(|e| MLError::ModelError(format!("GPU target upload failed: {e}")))?;
|
||
|
||
Ok(DqnGpuData {
|
||
features,
|
||
targets,
|
||
num_bars,
|
||
feature_dim: self.feature_dim,
|
||
})
|
||
}
|
||
|
||
/// Total bytes of pre-allocated CPU staging memory.
|
||
pub fn staging_bytes(&self) -> usize {
|
||
(self.feature_buf.len() + self.target_buf.len()) * std::mem::size_of::<f32>()
|
||
}
|
||
}
|
||
|
||
/// Pre-uploaded GPU training data for PPO trainer.
|
||
///
|
||
/// Holds all rollout states as a single [N, state_dim] GPU tensor.
|
||
#[derive(Debug)]
|
||
pub struct PpoGpuData {
|
||
/// All states tensor [num_steps, state_dim] on GPU (f32)
|
||
pub states: Tensor,
|
||
/// Number of training steps
|
||
pub num_steps: usize,
|
||
/// State dimension
|
||
pub state_dim: usize,
|
||
}
|
||
|
||
impl PpoGpuData {
|
||
/// Upload PPO market data to GPU as a single contiguous tensor.
|
||
pub fn upload(market_data: &[Vec<f32>], device: &Device) -> Result<Self, MLError> {
|
||
let num_steps = market_data.len();
|
||
if num_steps == 0 {
|
||
return Err(MLError::ModelError("Empty market data".to_owned()));
|
||
}
|
||
|
||
let state_dim = market_data[0].len();
|
||
|
||
let estimated_bytes = estimate_vram_bytes(num_steps * state_dim);
|
||
if estimated_bytes > MAX_UPLOAD_BYTES {
|
||
return Err(MLError::ModelError(format!(
|
||
"Market data too large for GPU upload: {:.1} GB > 2.0 GB limit ({} steps × {} dims)",
|
||
estimated_bytes as f64 / 1_073_741_824.0,
|
||
num_steps,
|
||
state_dim,
|
||
)));
|
||
}
|
||
|
||
let mut flat_states = Vec::with_capacity(num_steps * state_dim);
|
||
for state in market_data {
|
||
if state.len() != state_dim {
|
||
return Err(MLError::ModelError(format!(
|
||
"Inconsistent state dim: expected {state_dim}, got {}",
|
||
state.len()
|
||
)));
|
||
}
|
||
flat_states.extend_from_slice(state);
|
||
}
|
||
|
||
let states = Tensor::from_vec(flat_states, (num_steps, state_dim), device)
|
||
.map_err(|e| MLError::ModelError(format!("GPU state upload failed: {e}")))?;
|
||
|
||
Ok(Self {
|
||
states,
|
||
num_steps,
|
||
state_dim,
|
||
})
|
||
}
|
||
|
||
/// Estimated VRAM usage in bytes.
|
||
pub fn vram_bytes(&self) -> usize {
|
||
estimate_vram_bytes(self.num_steps * self.state_dim)
|
||
}
|
||
|
||
/// Get a single step's state as a [1, state_dim] tensor slice (zero-copy on GPU).
|
||
pub fn step_state(&self, step_idx: usize) -> Result<Tensor, MLError> {
|
||
self.states
|
||
.narrow(0, step_idx, 1)
|
||
.map_err(|e| MLError::ModelError(format!("Step {step_idx} state slice failed: {e}")))
|
||
}
|
||
|
||
/// Get a batch of states as [batch_size, state_dim] tensor slice.
|
||
pub fn batch_states(&self, start: usize, count: usize) -> Result<Tensor, MLError> {
|
||
let count = count.min(self.num_steps.saturating_sub(start));
|
||
self.states
|
||
.narrow(0, start, count)
|
||
.map_err(|e| MLError::ModelError(format!("Batch state slice [{start}..+{count}] failed: {e}")))
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_dqn_gpu_data_upload_cpu() {
|
||
let data: Vec<([f64; 51], Vec<f64>)> = (0..100)
|
||
.map(|i| {
|
||
let features = [i as f64 * 0.01; 51];
|
||
let targets = vec![100.0 + i as f64, 101.0 + i as f64, 100.5, 101.5];
|
||
(features, targets)
|
||
})
|
||
.collect();
|
||
|
||
let gpu_data = DqnGpuData::upload(&data, &Device::Cpu).unwrap();
|
||
assert_eq!(gpu_data.num_bars, 100);
|
||
assert_eq!(gpu_data.feature_dim, 51);
|
||
|
||
let bar0 = gpu_data.bar_features(0).unwrap();
|
||
assert_eq!(bar0.dims(), &[1, 51]);
|
||
|
||
let batch = gpu_data.batch_features(0, 10).unwrap();
|
||
assert_eq!(batch.dims(), &[10, 51]);
|
||
|
||
let targets = gpu_data.bar_target_values(0).unwrap();
|
||
assert!((targets[0] - 100.0).abs() < 0.01);
|
||
}
|
||
|
||
#[test]
|
||
fn test_dqn_build_state_tensor() {
|
||
let data: Vec<([f64; 51], Vec<f64>)> = vec![
|
||
([1.0; 51], vec![100.0, 101.0, 100.5, 101.5]),
|
||
];
|
||
|
||
let gpu_data = DqnGpuData::upload(&data, &Device::Cpu).unwrap();
|
||
let portfolio = [0.95_f32, 0.5, 0.0001];
|
||
let state = gpu_data.build_state_tensor(0, &portfolio, &Device::Cpu).unwrap();
|
||
assert_eq!(state.dims(), &[1, 54]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_tensor_core_alignment() {
|
||
assert_eq!(align_to_tensor_cores(54), 56);
|
||
assert_eq!(align_to_tensor_cores(64), 64); // already aligned
|
||
assert_eq!(align_to_tensor_cores(51), 56); // pad up
|
||
assert_eq!(align_to_tensor_cores(1), 8);
|
||
assert_eq!(align_to_tensor_cores(8), 8);
|
||
assert_eq!(align_to_tensor_cores(0), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_optimal_launch_dims() {
|
||
// Small: 16 episodes → block=32, grid=1
|
||
let (g, b) = optimal_launch_dims(16, 0);
|
||
assert_eq!(b, 32);
|
||
assert_eq!(g, 1);
|
||
|
||
// Medium: 128 episodes → block=128, grid=1
|
||
let (g, b) = optimal_launch_dims(128, 0);
|
||
assert_eq!(b, 128);
|
||
assert_eq!(g, 1);
|
||
|
||
// Large: 256 episodes → block=256, grid=1
|
||
let (g, b) = optimal_launch_dims(256, 0);
|
||
assert_eq!(b, 256);
|
||
assert_eq!(g, 1);
|
||
|
||
// Very large: 512 episodes → block=256, grid=2
|
||
let (g, b) = optimal_launch_dims(512, 0);
|
||
assert_eq!(b, 256);
|
||
assert_eq!(g, 2);
|
||
|
||
// Constrained GPU: max_tpb=64
|
||
let (g, b) = optimal_launch_dims(256, 64);
|
||
assert_eq!(b, 64);
|
||
assert_eq!(g, 4);
|
||
|
||
// Edge: 1 episode
|
||
let (g, b) = optimal_launch_dims(1, 0);
|
||
assert_eq!(b, 32);
|
||
assert_eq!(g, 1);
|
||
}
|
||
|
||
#[test]
|
||
fn test_ppo_gpu_data_upload_cpu() {
|
||
let market_data: Vec<Vec<f32>> = (0..200)
|
||
.map(|i| vec![i as f32 * 0.01; 54])
|
||
.collect();
|
||
|
||
let gpu_data = PpoGpuData::upload(&market_data, &Device::Cpu).unwrap();
|
||
assert_eq!(gpu_data.num_steps, 200);
|
||
assert_eq!(gpu_data.state_dim, 54);
|
||
|
||
let step0 = gpu_data.step_state(0).unwrap();
|
||
assert_eq!(step0.dims(), &[1, 54]);
|
||
|
||
let batch = gpu_data.batch_states(10, 20).unwrap();
|
||
assert_eq!(batch.dims(), &[20, 54]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_ppo_gpu_data_inconsistent_dims() {
|
||
let market_data = vec![
|
||
vec![1.0_f32; 54],
|
||
vec![2.0_f32; 53],
|
||
];
|
||
|
||
let result = PpoGpuData::upload(&market_data, &Device::Cpu);
|
||
assert!(result.is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn test_dqn_gpu_data_empty() {
|
||
let data: Vec<([f64; 51], Vec<f64>)> = vec![];
|
||
let result = DqnGpuData::upload(&data, &Device::Cpu);
|
||
assert!(result.is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn test_ppo_gpu_data_empty() {
|
||
let data: Vec<Vec<f32>> = vec![];
|
||
let result = PpoGpuData::upload(&data, &Device::Cpu);
|
||
assert!(result.is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn test_vram_estimation() {
|
||
// 300K bars × 55 features × 4 bytes = ~66 MB
|
||
let bytes = estimate_vram_bytes(300_000 * 55);
|
||
assert!(bytes < 100_000_000); // Under 100 MB
|
||
assert!(bytes > 60_000_000); // Over 60 MB
|
||
}
|
||
|
||
#[test]
|
||
fn test_dqn_boundary_access() {
|
||
let data: Vec<([f64; 51], Vec<f64>)> = vec![
|
||
([0.5; 51], vec![100.0, 101.0, 100.5, 101.5]),
|
||
];
|
||
let gpu_data = DqnGpuData::upload(&data, &Device::Cpu).unwrap();
|
||
|
||
assert!(gpu_data.bar_features(0).is_ok());
|
||
assert!(gpu_data.bar_features(1).is_err());
|
||
|
||
let batch = gpu_data.batch_features(0, 100).unwrap();
|
||
assert_eq!(batch.dims(), &[1, 51]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_dqn_build_batch_states() {
|
||
let data: Vec<([f64; 51], Vec<f64>)> = (0..100)
|
||
.map(|i| {
|
||
let features = [i as f64 * 0.01; 51];
|
||
let targets = vec![100.0, 101.0, 100.5, 101.5];
|
||
(features, targets)
|
||
})
|
||
.collect();
|
||
|
||
let gpu_data = DqnGpuData::upload(&data, &Device::Cpu).unwrap();
|
||
let portfolio = [0.95_f32, 0.5, 0.0001];
|
||
|
||
// Build batch of 32 states starting at bar 10
|
||
let batch = gpu_data.build_batch_states(10, 32, &portfolio, &Device::Cpu).unwrap();
|
||
assert_eq!(batch.dims(), &[32, 54]); // 51 market + 3 portfolio
|
||
|
||
// Verify portfolio features are broadcast correctly
|
||
let row0 = batch.narrow(0, 0, 1).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap();
|
||
assert!((row0[51] - 0.95).abs() < 1e-5, "portfolio_value");
|
||
assert!((row0[52] - 0.5).abs() < 1e-5, "position");
|
||
assert!((row0[53] - 0.0001).abs() < 1e-5, "spread");
|
||
|
||
let row31 = batch.narrow(0, 31, 1).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap();
|
||
assert!((row31[51] - 0.95).abs() < 1e-5, "portfolio_value row31");
|
||
|
||
// Verify market features differ between rows
|
||
assert!((row0[0] - row31[0]).abs() > 1e-6, "market features should differ");
|
||
}
|
||
|
||
#[test]
|
||
fn test_dqn_build_batch_states_clamping() {
|
||
let data: Vec<([f64; 51], Vec<f64>)> = vec![
|
||
([1.0; 51], vec![100.0, 101.0, 100.5, 101.5]),
|
||
([2.0; 51], vec![100.0, 101.0, 100.5, 101.5]),
|
||
];
|
||
let gpu_data = DqnGpuData::upload(&data, &Device::Cpu).unwrap();
|
||
let portfolio = [1.0_f32, 0.0, 0.0001];
|
||
|
||
// Request more bars than available — should clamp
|
||
let batch = gpu_data.build_batch_states(0, 100, &portfolio, &Device::Cpu).unwrap();
|
||
assert_eq!(batch.dims(), &[2, 54]); // Clamped to 2 bars
|
||
|
||
// Start beyond range — should error
|
||
let err = gpu_data.build_batch_states(5, 10, &portfolio, &Device::Cpu);
|
||
assert!(err.is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn test_gpu_buffer_pool_basic() {
|
||
let mut pool = GpuBufferPool::new(1000, 51, 4);
|
||
assert_eq!(pool.max_bars, 1000);
|
||
assert_eq!(pool.staging_bytes(), (1000 * 51 + 1000 * 4) * 4);
|
||
|
||
let data: Vec<([f64; 51], Vec<f64>)> = (0..50)
|
||
.map(|i| ([i as f64 * 0.01; 51], vec![100.0, 101.0, 100.5, 101.5]))
|
||
.collect();
|
||
|
||
let gpu_data = pool.upload_dqn(&data, &Device::Cpu).unwrap();
|
||
assert_eq!(gpu_data.num_bars, 50);
|
||
assert_eq!(gpu_data.feature_dim, 51);
|
||
}
|
||
|
||
#[test]
|
||
fn test_gpu_buffer_pool_reuse() {
|
||
let mut pool = GpuBufferPool::new(500, 51, 4);
|
||
|
||
// First upload
|
||
let data1: Vec<([f64; 51], Vec<f64>)> = (0..100)
|
||
.map(|i| ([i as f64 * 0.01; 51], vec![100.0, 101.0, 100.5, 101.5]))
|
||
.collect();
|
||
let g1 = pool.upload_dqn(&data1, &Device::Cpu).unwrap();
|
||
assert_eq!(g1.num_bars, 100);
|
||
|
||
// Second upload — reuses the same staging buffers
|
||
let data2: Vec<([f64; 51], Vec<f64>)> = (0..200)
|
||
.map(|i| ([i as f64 * 0.02; 51], vec![200.0, 201.0, 200.5, 201.5]))
|
||
.collect();
|
||
let g2 = pool.upload_dqn(&data2, &Device::Cpu).unwrap();
|
||
assert_eq!(g2.num_bars, 200);
|
||
|
||
// Verify data integrity: g2 should have data2's values
|
||
let t = g2.bar_target_values(0).unwrap();
|
||
assert!((t[0] - 200.0).abs() < 0.01);
|
||
}
|
||
|
||
#[test]
|
||
fn test_gpu_buffer_pool_grow() {
|
||
let mut pool = GpuBufferPool::new(10, 51, 4);
|
||
|
||
// Upload more bars than pre-allocated — should grow
|
||
let data: Vec<([f64; 51], Vec<f64>)> = (0..50)
|
||
.map(|i| ([i as f64; 51], vec![1.0, 2.0, 3.0, 4.0]))
|
||
.collect();
|
||
let gpu_data = pool.upload_dqn(&data, &Device::Cpu).unwrap();
|
||
assert_eq!(gpu_data.num_bars, 50);
|
||
// Staging buffers should have grown
|
||
assert!(pool.feature_buf.len() >= 50 * 51);
|
||
}
|
||
|
||
#[test]
|
||
fn test_gpu_buffer_pool_empty() {
|
||
let mut pool = GpuBufferPool::new(100, 51, 4);
|
||
let data: Vec<([f64; 51], Vec<f64>)> = vec![];
|
||
assert!(pool.upload_dqn(&data, &Device::Cpu).is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn test_kernel_source_contains_entry_point() {
|
||
let src = include_str!("dqn_experience_kernel.cu");
|
||
assert!(
|
||
src.contains("dqn_full_experience_kernel"),
|
||
"Kernel source must contain the entry-point function name"
|
||
);
|
||
assert!(
|
||
src.contains("extern \"C\""),
|
||
"Kernel source must declare extern C linkage"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_common_header_contains_shared_functions() {
|
||
let src = include_str!("common_device_functions.cuh");
|
||
assert!(src.contains("gpu_random"), "Missing gpu_random");
|
||
assert!(src.contains("matvec_leaky_relu"), "Missing matvec_leaky_relu");
|
||
assert!(src.contains("action_to_exposure"), "Missing action_to_exposure");
|
||
assert!(src.contains("barrier_init"), "Missing barrier_init");
|
||
assert!(src.contains("diversity_entropy"), "Missing diversity_entropy");
|
||
assert!(src.contains("curiosity_inference"), "Missing curiosity_inference");
|
||
assert!(!src.contains("q_forward_dueling"), "DQN-specific function leaked into common header");
|
||
assert!(!src.contains("dqn_full_experience_kernel"), "DQN kernel leaked into common header");
|
||
}
|
||
|
||
#[cfg(feature = "cuda")]
|
||
#[test]
|
||
fn test_experience_config_defaults() {
|
||
use super::gpu_experience_collector::ExperienceCollectorConfig;
|
||
|
||
let cfg = ExperienceCollectorConfig::default();
|
||
assert_eq!(cfg.n_episodes, 128);
|
||
assert_eq!(cfg.timesteps_per_episode, 500);
|
||
assert_eq!(cfg.total_experiences(), 128 * 500);
|
||
assert!((cfg.epsilon - 0.1).abs() < f32::EPSILON);
|
||
assert!((cfg.gamma - 0.99).abs() < f32::EPSILON);
|
||
assert!((cfg.barrier_profit_mult - 1.02).abs() < f32::EPSILON);
|
||
assert!((cfg.barrier_loss_mult - 0.98).abs() < f32::EPSILON);
|
||
}
|
||
|
||
#[test]
|
||
fn test_ppo_kernel_source_contains_all_functions() {
|
||
let src = include_str!("ppo_experience_kernel.cu");
|
||
assert!(src.contains("ppo_actor_forward"), "Missing ppo_actor_forward");
|
||
assert!(src.contains("softmax_sample"), "Missing softmax_sample");
|
||
assert!(src.contains("ppo_critic_forward"), "Missing ppo_critic_forward");
|
||
assert!(src.contains("compute_gae_backward"), "Missing compute_gae_backward");
|
||
assert!(src.contains("ppo_full_experience_kernel"), "Missing kernel entry point");
|
||
assert!(src.contains("extern \"C\""), "Missing extern C linkage");
|
||
}
|
||
|
||
#[test]
|
||
fn test_ppo_kernel_source_concatenation() {
|
||
let common = include_str!("common_device_functions.cuh");
|
||
let kernel = include_str!("ppo_experience_kernel.cu");
|
||
let full = format!("{}\n{}", common, kernel);
|
||
assert!(full.contains("gpu_random"));
|
||
assert!(full.contains("matvec_leaky_relu"));
|
||
assert!(full.contains("ppo_full_experience_kernel"));
|
||
assert!(full.contains("compute_gae_backward"));
|
||
assert!(full.contains("softmax_sample"));
|
||
assert!(!full.contains("q_forward_dueling"));
|
||
}
|
||
|
||
#[cfg(feature = "cuda")]
|
||
#[test]
|
||
fn test_ppo_collector_config_defaults() {
|
||
use super::gpu_ppo_collector::PpoCollectorConfig;
|
||
let cfg = PpoCollectorConfig::default();
|
||
assert_eq!(cfg.n_episodes, 128);
|
||
assert_eq!(cfg.timesteps_per_episode, 500);
|
||
assert_eq!(cfg.total_experiences(), 64_000);
|
||
assert!((cfg.gamma - 0.99).abs() < f32::EPSILON);
|
||
assert!((cfg.gae_lambda - 0.95).abs() < f32::EPSILON);
|
||
}
|
||
|
||
// ── Phase 3: Batch struct layout & index-math validation ──────────
|
||
|
||
#[cfg(feature = "cuda")]
|
||
#[test]
|
||
fn test_experience_batch_index_math() {
|
||
use super::gpu_experience_collector::ExperienceBatch;
|
||
|
||
let n_episodes = 4;
|
||
let timesteps = 10;
|
||
let state_dim = 54;
|
||
let total = n_episodes * timesteps;
|
||
|
||
// Build synthetic batch matching kernel output layout
|
||
let states: Vec<f32> = (0..total * state_dim)
|
||
.map(|i| i as f32 * 0.001)
|
||
.collect();
|
||
let actions: Vec<i32> = (0..total).map(|i| (i % 45) as i32).collect();
|
||
let rewards: Vec<f32> = (0..total).map(|i| i as f32 * 0.01).collect();
|
||
let done_flags: Vec<i32> = (0..total)
|
||
.map(|i| if i % timesteps == timesteps - 1 { 1 } else { 0 })
|
||
.collect();
|
||
let target_q_values = vec![0.5_f32; total];
|
||
let td_errors = vec![0.1_f32; total];
|
||
|
||
let batch = ExperienceBatch {
|
||
states,
|
||
actions,
|
||
rewards,
|
||
done_flags,
|
||
target_q_values,
|
||
td_errors,
|
||
n_episodes,
|
||
timesteps,
|
||
};
|
||
|
||
assert_eq!(batch.states.len(), total * state_dim);
|
||
assert_eq!(batch.actions.len(), total);
|
||
assert_eq!(batch.rewards.len(), total);
|
||
assert_eq!(batch.done_flags.len(), total);
|
||
|
||
// Verify episode indexing: batch[ep][t] = ep * timesteps + t
|
||
for ep in 0..n_episodes {
|
||
for t in 0..timesteps {
|
||
let idx = ep * timesteps + t;
|
||
assert!(idx < total, "idx {} >= total {}", idx, total);
|
||
|
||
let state_start = idx * state_dim;
|
||
let state_end = state_start + state_dim;
|
||
assert!(state_end <= batch.states.len());
|
||
|
||
// Actions in 0..44
|
||
let a = batch.actions[idx];
|
||
assert!((0..45).contains(&a));
|
||
}
|
||
|
||
// Last step of each episode should be done
|
||
let last_idx = ep * timesteps + timesteps - 1;
|
||
assert_eq!(batch.done_flags[last_idx], 1);
|
||
}
|
||
}
|
||
|
||
#[cfg(feature = "cuda")]
|
||
#[test]
|
||
fn test_ppo_experience_batch_index_math() {
|
||
use super::gpu_ppo_collector::PpoExperienceBatch;
|
||
|
||
let n_episodes = 4;
|
||
let timesteps = 10;
|
||
let state_dim = 54;
|
||
let total = n_episodes * timesteps;
|
||
|
||
let states: Vec<f32> = (0..total * state_dim)
|
||
.map(|i| i as f32 * 0.001)
|
||
.collect();
|
||
let actions: Vec<i32> = (0..total).map(|i| (i % 45) as i32).collect();
|
||
let log_probs: Vec<f32> = vec![-1.5_f32; total];
|
||
let advantages: Vec<f32> = (0..total).map(|i| i as f32 * 0.1 - 2.0).collect();
|
||
let returns: Vec<f32> = (0..total).map(|i| i as f32 * 0.5).collect();
|
||
let done_flags: Vec<i32> = (0..total)
|
||
.map(|i| if i % timesteps == timesteps - 1 { 1 } else { 0 })
|
||
.collect();
|
||
|
||
let batch = PpoExperienceBatch {
|
||
states,
|
||
actions,
|
||
log_probs,
|
||
advantages,
|
||
returns,
|
||
done_flags,
|
||
n_episodes,
|
||
timesteps,
|
||
};
|
||
|
||
assert_eq!(batch.states.len(), total * state_dim);
|
||
assert_eq!(batch.actions.len(), total);
|
||
assert_eq!(batch.log_probs.len(), total);
|
||
assert_eq!(batch.advantages.len(), total);
|
||
assert_eq!(batch.returns.len(), total);
|
||
assert_eq!(batch.done_flags.len(), total);
|
||
|
||
// Verify V(s) = R(s) - A(s) reconstruction
|
||
for i in 0..total {
|
||
let value = batch.returns[i] - batch.advantages[i];
|
||
assert!(value.is_finite(), "V(s) should be finite at idx {}", i);
|
||
}
|
||
|
||
// Verify state chunking: each chunk should have exactly state_dim elements
|
||
let chunks: Vec<&[f32]> = batch.states.chunks(state_dim).collect();
|
||
assert_eq!(chunks.len(), total);
|
||
for (i, chunk) in chunks.iter().enumerate() {
|
||
assert_eq!(
|
||
chunk.len(),
|
||
state_dim,
|
||
"Chunk {} has wrong length: {}",
|
||
i,
|
||
chunk.len()
|
||
);
|
||
}
|
||
}
|
||
|
||
#[cfg(feature = "cuda")]
|
||
#[test]
|
||
fn test_experience_batch_next_state_boundary() {
|
||
use super::gpu_experience_collector::ExperienceBatch;
|
||
|
||
let n_episodes = 2;
|
||
let timesteps = 5;
|
||
let state_dim = 54;
|
||
let total = n_episodes * timesteps;
|
||
|
||
// Give each step a distinct state for easy verification
|
||
let states: Vec<f32> = (0..total)
|
||
.flat_map(|idx| std::iter::repeat(idx as f32).take(state_dim))
|
||
.collect();
|
||
let done_flags: Vec<i32> = (0..total)
|
||
.map(|i| if i % timesteps == timesteps - 1 { 1 } else { 0 })
|
||
.collect();
|
||
|
||
let batch = ExperienceBatch {
|
||
states,
|
||
actions: vec![0; total],
|
||
rewards: vec![0.0; total],
|
||
done_flags,
|
||
target_q_values: vec![0.0; total],
|
||
td_errors: vec![0.0; total],
|
||
n_episodes,
|
||
timesteps,
|
||
};
|
||
|
||
// For the DQN conversion: next_state[t] = state[t+1] unless done or last step
|
||
for ep in 0..n_episodes {
|
||
for t in 0..timesteps {
|
||
let idx = ep * timesteps + t;
|
||
let done = batch.done_flags[idx] != 0;
|
||
|
||
if !done && t + 1 < timesteps {
|
||
let next_idx = ep * timesteps + t + 1;
|
||
let ns_start = next_idx * state_dim;
|
||
let ns_end = ns_start + state_dim;
|
||
assert!(ns_end <= batch.states.len());
|
||
|
||
// Next state should be distinct from current state
|
||
let s_val = batch.states[idx * state_dim];
|
||
let ns_val = batch.states[ns_start];
|
||
assert!(
|
||
(ns_val - s_val - 1.0).abs() < f32::EPSILON,
|
||
"next_state should be current + 1 step"
|
||
);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(feature = "cuda")]
|
||
#[test]
|
||
fn test_ppo_batch_gae_consistency() {
|
||
use super::gpu_ppo_collector::PpoExperienceBatch;
|
||
|
||
let n_episodes = 2;
|
||
let timesteps = 10;
|
||
let total = n_episodes * timesteps;
|
||
let state_dim = 54;
|
||
|
||
// Simulate GAE output: advantages should sum roughly to 0 within each episode
|
||
// (not exact, but the kernel normalizes per-episode)
|
||
let mut advantages = Vec::with_capacity(total);
|
||
for ep in 0..n_episodes {
|
||
let _base = ep as f32 * 0.5;
|
||
for t in 0..timesteps {
|
||
// Triangular pattern: positive early, negative late
|
||
let a = (timesteps as f32 / 2.0 - t as f32) * 0.1;
|
||
advantages.push(a);
|
||
}
|
||
}
|
||
|
||
let returns: Vec<f32> = advantages.iter().map(|&a| a + 5.0).collect();
|
||
|
||
let batch = PpoExperienceBatch {
|
||
states: vec![0.0_f32; total * state_dim],
|
||
actions: vec![1; total],
|
||
log_probs: vec![-1.0; total],
|
||
advantages,
|
||
returns,
|
||
done_flags: (0..total)
|
||
.map(|i| if i % timesteps == timesteps - 1 { 1 } else { 0 })
|
||
.collect(),
|
||
n_episodes,
|
||
timesteps,
|
||
};
|
||
|
||
// All values (R - A) should be approximately 5.0
|
||
for i in 0..total {
|
||
let v = batch.returns[i] - batch.advantages[i];
|
||
assert!(
|
||
(v - 5.0).abs() < f32::EPSILON,
|
||
"Value at {} should be ~5.0, got {}",
|
||
i,
|
||
v
|
||
);
|
||
}
|
||
}
|
||
}
|