gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf(cuda): pre-allocate GPU buffers in PPO collector — eliminate per-epoch malloc

Browse Source 
reset_episodes() was allocating 5 CPU Vecs per epoch for HtoD zero-fill.
Replaced with:
- memset_zeros() for barrier/diversity buffers (async GPU-side, no HtoD)
- Pre-allocated host staging vectors for portfolio_init and RNG seeds

gpu_curiosity_trainer already optimized (fused kernel zeros internally).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-18 07:46:38 +01:00
parent 08eb6d5cc2
commit bc0f1f2bd8
1 changed files with 52 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs
View File

@@ -172,6 +172,12 @@ pub struct GpuPpoExperienceCollector {
readback_f32_host: Vec<f32>,
readback_i32_staging: CudaSlice<i32>,
readback_i32_host: Vec<i32>,
// Pre-allocated host-side staging for reset_episodes() to avoid per-call Vec allocations.
// portfolio_init_staging is re-filled with structured data each reset; rng_seed_staging
// is re-filled with random seeds. Both avoid hitting the allocator on every epoch boundary.
portfolio_init_staging: Vec<f32>,
rng_seed_staging: Vec<u32>,
}
impl GpuPpoExperienceCollector {
@@ -351,6 +357,11 @@ impl GpuPpoExperienceCollector {
})?;
let readback_i32_host = vec![0_i32; i32_staging_len];
// ---- Step 6c: Pre-allocate host staging for reset_episodes() ----
// Avoids per-epoch Vec allocations for portfolio init and RNG seeds.
let portfolio_init_staging = vec![0.0_f32; MAX_EPISODES * PORTFOLIO_STATE_SIZE];
let rng_seed_staging = vec![0_u32; MAX_EPISODES];
// ---- Step 7: Log buffer sizes ----
let per_episode_bytes = (PORTFOLIO_STATE_SIZE + BARRIER_STATE_SIZE) * 4
+ DIVERSITY_WINDOW * 4
@@ -388,6 +399,8 @@ impl GpuPpoExperienceCollector {
readback_f32_host,
readback_i32_staging,
readback_i32_host,
portfolio_init_staging,
rng_seed_staging,
})
}
@@ -672,56 +685,74 @@ impl GpuPpoExperienceCollector {
///
/// Re-initializes portfolio states to starting capital, zeros barrier
/// states and diversity windows, and regenerates RNG seeds.
///
/// Zero-fill buffers use `memset_zeros` (async GPU memset) instead of
/// allocating CPU-side `Vec`s and uploading via HtoD. Structured data
/// (portfolio init, RNG seeds) reuses pre-allocated host staging vectors
/// to avoid per-epoch allocations.
pub fn reset_episodes(
&mut self,
initial_capital: f32,
avg_spread: f32,
cash_reserve_pct: f32,
) -> Result<(), MLError> {
// Reset portfolio states
let mut portfolio_init = vec![0.0_f32; MAX_EPISODES * PORTFOLIO_STATE_SIZE];
// Reset portfolio states using pre-allocated staging buffer
// (structured init -- cannot use memset_zeros)
for i in 0..MAX_EPISODES {
let off = i * PORTFOLIO_STATE_SIZE;
portfolio_init[off] = initial_capital;
portfolio_init[off + 3] = initial_capital;
portfolio_init[off + 4] = avg_spread;
portfolio_init[off + 6] = cash_reserve_pct;
// Zero all fields first
for j in 0..PORTFOLIO_STATE_SIZE {
if let Some(slot) = self.portfolio_init_staging.get_mut(off + j) {
*slot = 0.0;
}
}
// Set non-zero fields
if let Some(slot) = self.portfolio_init_staging.get_mut(off) {
*slot = initial_capital; // cash
}
if let Some(slot) = self.portfolio_init_staging.get_mut(off + 3) {
*slot = initial_capital; // initial_cap
}
if let Some(slot) = self.portfolio_init_staging.get_mut(off + 4) {
*slot = avg_spread; // spread
}
if let Some(slot) = self.portfolio_init_staging.get_mut(off + 6) {
*slot = cash_reserve_pct; // reserve_pct
}
}
self.stream
.memcpy_htod(&portfolio_init, &mut self.portfolio_states)
.memcpy_htod(&self.portfolio_init_staging, &mut self.portfolio_states)
.map_err(|e| {
MLError::ModelError(format!("Failed to reset portfolio_states: {e}"))
})?;
// Zero barrier states
let barrier_zeros = vec![0.0_f32; MAX_EPISODES * BARRIER_STATE_SIZE];
// Zero barrier states via async GPU memset (no CPU allocation)
self.stream
.memcpy_htod(&barrier_zeros, &mut self.barrier_states)
.memset_zeros(&mut self.barrier_states)
.map_err(|e| {
MLError::ModelError(format!("Failed to reset barrier_states: {e}"))
})?;
// Zero diversity windows and metas
let div_window_zeros = vec![0_i32; MAX_EPISODES * DIVERSITY_WINDOW];
// Zero diversity windows via async GPU memset (no CPU allocation)
self.stream
.memcpy_htod(&div_window_zeros, &mut self.diversity_windows)
.memset_zeros(&mut self.diversity_windows)
.map_err(|e| {
MLError::ModelError(format!("Failed to reset diversity_windows: {e}"))
})?;
let div_meta_zeros = vec![0_i32; MAX_EPISODES * 2];
// Zero diversity metas via async GPU memset (no CPU allocation)
self.stream
.memcpy_htod(&div_meta_zeros, &mut self.diversity_metas)
.memset_zeros(&mut self.diversity_metas)
.map_err(|e| {
MLError::ModelError(format!("Failed to reset diversity_metas: {e}"))
})?;
// Fresh RNG seeds
let rng_seeds: Vec<u32> = (0..MAX_EPISODES)
.map(|_| fastrand::u32(..))
.collect();
// Fresh RNG seeds using pre-allocated staging buffer
for slot in &mut self.rng_seed_staging {
*slot = fastrand::u32(..);
}
self.stream
.memcpy_htod(&rng_seeds, &mut self.rng_states)
.memcpy_htod(&self.rng_seed_staging, &mut self.rng_states)
.map_err(|e| {
MLError::ModelError(format!("Failed to reset rng_states: {e}"))
})?;
@@ -730,7 +761,7 @@ impl GpuPpoExperienceCollector {
initial_capital,
avg_spread,
cash_reserve_pct,
"PPO episode state reset complete"
"PPO episode state reset complete (memset_zeros path)"
);
Ok(())