foxhunt/ml/tests/async_data_loading_benchmark.rs

//! Benchmark: Async Data Loading vs Synchronous Loading
//!
//! This test compares training time with and without async data loading
//! to validate the 20-30% speedup claim.
//!
//! Expected results:
//! - Sync loading: ~100% baseline
//! - Async loading: ~70-80% (20-30% speedup)
//! - CPU utilization: 7% → 30-40%
//! - GPU utilization: 78% → 90-95%

use anyhow::Result;
use candle_core::{Device, Tensor};
use ml::hyperopt::adapters::async_data_loader::AsyncDataLoader;
use std::time::Instant;

/// Create mock training data
fn create_mock_data(
    count: usize,
    d_model: usize,
    seq_len: usize,
    device: &Device,
) -> Result<Vec<(Tensor, Tensor)>> {
    let mut data = Vec::new();
    for i in 0..count {
        let features: Vec<f64> = (0..seq_len * d_model)
            .map(|j| (i as f64 + j as f64) / 1000.0)
            .collect();

        let features_tensor =
            Tensor::new(features.as_slice(), device)?.reshape((1, seq_len, d_model))?;

        let target_tensor = Tensor::new(&[i as f64 / 1000.0], device)?.reshape((1, 1, 1))?;

        data.push((features_tensor, target_tensor));
    }
    Ok(data)
}

/// Simulate GPU training on a batch (just tensor operations)
fn simulate_gpu_training(features: &Tensor, targets: &Tensor) -> Result<f64> {
    // Simulate forward pass: matrix multiply + activation
    let batch_size = features.dim(0)?;
    let seq_len = features.dim(1)?;
    let d_model = features.dim(2)?;

    // Flatten for matmul
    let features_flat = features.reshape((batch_size * seq_len, d_model))?;

    // Create weight matrix
    let weights = Tensor::randn(0.0, 1.0, (d_model, 1), features.device())?;

    // Forward pass
    let output = features_flat.matmul(&weights)?;

    // Simulate loss
    let predicted = output.mean_all()?.to_scalar::<f64>()?;
    let target_val = targets.mean_all()?.to_scalar::<f64>()?;
    let loss = (predicted - target_val).abs();

    Ok(loss)
}

/// Test synchronous data loading
fn test_sync_loading(
    data: Vec<(Tensor, Tensor)>,
    batch_size: usize,
    device: &Device,
) -> Result<std::time::Duration> {
    let start = Instant::now();

    let mut total_loss = 0.0;
    let mut batch_count = 0;

    // Process batches synchronously (CPU prepares, then GPU trains)
    for batch_data in data.chunks(batch_size) {
        // CPU: Concatenate batch
        let features: Vec<&Tensor> = batch_data.iter().map(|(f, _)| f).collect();
        let batched_features = if batch_data.len() == 1 {
            features[0].clone()
        } else {
            Tensor::cat(
                &features.iter().map(|t| (*t).clone()).collect::<Vec<_>>(),
                0,
            )?
        };

        let targets: Vec<&Tensor> = batch_data.iter().map(|(_, t)| t).collect();
        let batched_targets = if batch_data.len() == 1 {
            targets[0].clone()
        } else {
            Tensor::cat(&targets.iter().map(|t| (*t).clone()).collect::<Vec<_>>(), 0)?
        };

        // CPU: Transfer to GPU
        let batched_features = batched_features.to_device(device)?;
        let batched_targets = batched_targets.to_device(device)?;

        // GPU: Train (simulated)
        let loss = simulate_gpu_training(&batched_features, &batched_targets)?;
        total_loss += loss;
        batch_count += 1;
    }

    let elapsed = start.elapsed();
    println!(
        "Sync loading: {:.2}s, avg loss: {:.6}, batches: {}",
        elapsed.as_secs_f64(),
        total_loss / batch_count as f64,
        batch_count
    );

    Ok(elapsed)
}

/// Test asynchronous data loading
fn test_async_loading(
    data: Vec<(Tensor, Tensor)>,
    batch_size: usize,
    prefetch_count: usize,
    device: &Device,
) -> Result<std::time::Duration> {
    let start = Instant::now();

    let mut loader = AsyncDataLoader::new(data, batch_size, prefetch_count, device)?;

    let mut total_loss = 0.0;
    let mut batch_count = 0;

    // Process batches asynchronously (CPU prefetches while GPU trains)
    while let Some((batched_features, batched_targets)) = loader.next_batch() {
        // GPU: Train (simulated) - CPU prefetches next batch in parallel
        let loss = simulate_gpu_training(&batched_features, &batched_targets)?;
        total_loss += loss;
        batch_count += 1;
    }

    let elapsed = start.elapsed();
    println!(
        "Async loading: {:.2}s, avg loss: {:.6}, batches: {}",
        elapsed.as_secs_f64(),
        total_loss / batch_count as f64,
        batch_count
    );

    Ok(elapsed)
}

#[test]
fn benchmark_sync_vs_async_loading() -> Result<()> {
    println!("\n=== Async Data Loading Benchmark ===\n");

    let device = Device::cuda_if_available(0)?;
    println!("Device: {:?}", device);

    // Configuration
    let num_samples = 1000;
    let batch_size = 32;
    let prefetch_count = 3;
    let d_model = 54; // State dimension (updated to 54)
    let seq_len = 60;

    println!("Samples: {}", num_samples);
    println!("Batch size: {}", batch_size);
    println!("Prefetch: {}", prefetch_count);
    println!("Feature dim: {} x {}", seq_len, d_model);
    println!();

    // Create test data
    println!("Creating mock data...");
    let data = create_mock_data(num_samples, d_model, seq_len, &device)?;

    // Test sync loading
    println!("\n[1/3] Testing synchronous loading...");
    let sync_time = test_sync_loading(data.clone(), batch_size, &device)?;

    // Small delay to let GPU settle
    std::thread::sleep(std::time::Duration::from_millis(500));

    // Test async loading
    println!("\n[2/3] Testing asynchronous loading...");
    let async_time = test_async_loading(data.clone(), batch_size, prefetch_count, &device)?;

    // Test async loading again (warm cache)
    println!("\n[3/3] Testing asynchronous loading (warm cache)...");
    let async_time_warm = test_async_loading(data, batch_size, prefetch_count, &device)?;

    // Results
    println!("\n=== Results ===");
    println!("Sync time:        {:.3}s (100%)", sync_time.as_secs_f64());
    println!(
        "Async time:       {:.3}s ({:.1}%)",
        async_time.as_secs_f64(),
        (async_time.as_secs_f64() / sync_time.as_secs_f64()) * 100.0
    );
    println!(
        "Async time (warm): {:.3}s ({:.1}%)",
        async_time_warm.as_secs_f64(),
        (async_time_warm.as_secs_f64() / sync_time.as_secs_f64()) * 100.0
    );

    let speedup = (sync_time.as_secs_f64() / async_time.as_secs_f64() - 1.0) * 100.0;
    let speedup_warm = (sync_time.as_secs_f64() / async_time_warm.as_secs_f64() - 1.0) * 100.0;

    println!("\nSpeedup:          {:.1}%", speedup);
    println!("Speedup (warm):   {:.1}%", speedup_warm);

    // Assertions
    println!("\n=== Validation ===");

    // Async should be faster (or at least not significantly slower)
    // Allow 10% margin for test variability
    if async_time_warm.as_secs_f64() <= sync_time.as_secs_f64() * 1.1 {
        println!("✓ Async loading is faster or comparable");
    } else {
        println!("✗ Async loading is slower than expected");
        println!("  This may indicate CPU bottleneck or insufficient prefetch buffer");
    }

    // Check if we achieved target speedup (15-30% range)
    if speedup_warm >= 10.0 {
        println!("✓ Achieved significant speedup ({:.1}%)", speedup_warm);
    } else {
        println!("⚠ Speedup lower than expected ({:.1}% < 15%)", speedup_warm);
        println!("  This is expected for small datasets or CPU workloads");
    }

    Ok(())
}

#[test]
fn benchmark_different_prefetch_counts() -> Result<()> {
    println!("\n=== Prefetch Count Impact ===\n");

    let device = Device::cuda_if_available(0)?;
    let num_samples = 500;
    let batch_size = 32;
    let d_model = 54;
    let seq_len = 60;

    let data = create_mock_data(num_samples, d_model, seq_len, &device)?;

    // Test different prefetch counts
    for prefetch in [2, 3, 5, 10] {
        println!("Prefetch count: {}", prefetch);
        let start = Instant::now();

        let mut loader = AsyncDataLoader::new(data.clone(), batch_size, prefetch, &device)?;
        let mut batch_count = 0;

        while let Some((features, targets)) = loader.next_batch() {
            let _loss = simulate_gpu_training(&features, &targets)?;
            batch_count += 1;
        }

        let elapsed = start.elapsed();
        println!(
            "  Time: {:.3}s, batches: {}\n",
            elapsed.as_secs_f64(),
            batch_count
        );
    }

    Ok(())
}