rustytorch/crates/training/rtx-transformers/src/optimizers/benchmarks.rs

//! Performance benchmarks for parameter update system
//!
//! These benchmarks verify that the Adam optimizer meets performance requirements
//! and scales efficiently with model size.

use crate::{Result, TransformerError};
use crate::optimizers::{AdamOptimizer, AdamConfig};
use rtx_tensor::{Tensor, Device, DType};
use std::collections::HashMap;
use std::time::{Duration, Instant};

/// Benchmark configuration for optimizer performance tests
pub struct BenchmarkConfig {
    /// Number of parameters to test with
    pub num_parameters: usize,
    /// Parameter size (number of elements per parameter)
    pub parameter_size: usize,
    /// Number of optimization steps to perform
    pub num_steps: usize,
    /// Target time per step in microseconds
    pub target_step_time_us: u64,
}

impl Default for BenchmarkConfig {
    fn default() -> Self {
        Self {
            num_parameters: 10,
            parameter_size: 1000,
            num_steps: 100,
            target_step_time_us: 1000, // 1ms per step
        }
    }
}

/// Benchmark results for optimizer performance
#[derive(Debug, Clone)]
pub struct BenchmarkResults {
    /// Average time per optimization step
    pub avg_step_time: Duration,
    /// Minimum step time observed
    pub min_step_time: Duration,
    /// Maximum step time observed
    pub max_step_time: Duration,
    /// Total time for all steps
    pub total_time: Duration,
    /// Steps per second throughput
    pub steps_per_second: f64,
    /// Memory usage estimate (bytes)
    pub memory_usage_bytes: usize,
}

impl BenchmarkResults {
    /// Check if results meet performance targets
    pub fn meets_targets(&self, config: &BenchmarkConfig) -> bool {
        let target_duration = Duration::from_micros(config.target_step_time_us);
        self.avg_step_time <= target_duration
    }

    /// Calculate performance improvement vs naive gradient descent
    pub fn performance_vs_sgd(&self) -> f64 {
        // Assume SGD baseline of 100us per step for comparison
        let sgd_baseline_us = 100.0;
        let our_time_us = self.avg_step_time.as_micros() as f64;
        sgd_baseline_us / our_time_us
    }
}

/// Run comprehensive performance benchmarks for Adam optimizer
pub fn benchmark_adam_optimizer(config: BenchmarkConfig) -> Result<BenchmarkResults> {
    println!("Running Adam optimizer performance benchmark...");
    println!("  Parameters: {}, Size: {}, Steps: {}",
    config.num_parameters, config.parameter_size, config.num_steps);

    let adam_config = AdamConfig {
        learning_rate: 0.001,
        beta1: 0.9,
        beta2: 0.999,
        epsilon: 1e-8,
        weight_decay: 0.01,
    };

    let mut optimizer = AdamOptimizer::new(adam_config)?;
    let device = Device::cpu();

    // Create test parameters and gradients
    let mut parameters = HashMap::new();
    let mut gradients = HashMap::new();

    for i in 0..config.num_parameters {
        let param_name = format!("param_{}", i);

        // Create parameter with random-like initial values
        let param_data: Vec<f32> = (0..config.parameter_size)
        .map(|j| ((i * config.parameter_size + j) as f32 * 0.01) % 1.0)
        .collect();
        let param = Tensor::from_data(param_data, &[config.parameter_size], &device)?;

        // Create gradient with small random-like values
        let grad_data: Vec<f32> = (0..config.parameter_size)
        .map(|j| ((i * config.parameter_size + j + 1000) as f32 * 0.001) % 0.1)
        .collect();
        let grad = Tensor::from_data(grad_data, &[config.parameter_size], &device)?;

        parameters.insert(param_name.clone(), param);
        gradients.insert(param_name, grad);
    }

    // Warmup run to initialize state
    optimizer.set_gradients(gradients.clone())?;
    optimizer.step(0.001)?;

    // Benchmark the optimization steps
    let mut step_times = Vec::new();
    let total_start = Instant::now();

    for step in 0..config.num_steps {
        // Vary gradients slightly each step to simulate real training
        let mut varied_gradients = HashMap::new();
        for (name, grad) in &gradients {
            let grad_data = grad.to_cpu()?;
            let varied_data: Vec<f32> = grad_data.iter()
            .enumerate()
            .map(|(i, &val)| val * (1.0 + (step as f32 * 0.01 + i as f32 * 0.001) % 0.1))
            .collect();
            let varied_grad = Tensor::from_data(varied_data, grad.shape().dims(), &device)?;
            varied_gradients.insert(name.clone(), varied_grad);
        }

        let step_start = Instant::now();

        // Store gradients and perform optimization step
        optimizer.set_gradients(varied_gradients)?;
        let _updates = optimizer.step_with_parameters(0.001, &parameters)?;

        let step_time = step_start.elapsed();
        step_times.push(step_time);

        if step % 10 == 0 {
            println!("    Step {}: {:.2}μs", step, step_time.as_micros());
        }
    }

    let total_time = total_start.elapsed();

    // Calculate statistics
    let avg_step_time = total_time / config.num_steps as u32;
    let min_step_time = *step_times.iter().min().unwrap();
    let max_step_time = *step_times.iter().max().unwrap();
    let steps_per_second = config.num_steps as f64 / total_time.as_secs_f64();

    // Estimate memory usage (rough approximation)
    let memory_usage_bytes = config.num_parameters * config.parameter_size * 4 * 3; // param + momentum + variance

    let results = BenchmarkResults {
        avg_step_time,
        min_step_time,
        max_step_time,
        total_time,
        steps_per_second,
        memory_usage_bytes,
    };

    println!("\nBenchmark Results:");
    println!("  Average step time: {:.2}μs", results.avg_step_time.as_micros());
    println!("  Min/Max step time: {:.2}μs / {:.2}μs",
    results.min_step_time.as_micros(), results.max_step_time.as_micros());
    println!("  Steps per second: {:.1}", results.steps_per_second);
    println!("  Memory usage: {:.2} MB", results.memory_usage_bytes as f64 / 1024.0 / 1024.0);
    println!("  Performance vs SGD: {:.1}x", results.performance_vs_sgd());

    if results.meets_targets(&config) {
        println!("  ✓ Performance targets met!");
        } else {
            println!("  ⚠ Performance targets missed (target: {}μs)", config.target_step_time_us);
        }

        Ok(results)
    }

    /// Benchmark scaling behavior with different model sizes
    pub fn benchmark_scaling_behavior() -> Result<()> {
        println!("\nRunning scaling behavior benchmark...");

        let parameter_sizes = vec![100, 1000, 10000, 100000];
        let mut scaling_results = Vec::new();

        for &size in &parameter_sizes {
            let config = BenchmarkConfig {
                num_parameters: 5,
                parameter_size: size,
                num_steps: 20,
                target_step_time_us: size as u64 / 10, // Scale target with size
            };

            let results = benchmark_adam_optimizer(config)?;
            scaling_results.push((size, results.avg_step_time.as_micros()));

            println!("  Size {}: {:.2}μs per step", size, results.avg_step_time.as_micros());
        }

        // Check if scaling is reasonable (should be roughly linear)
        let first_time = scaling_results[0].1 as f64;
        let last_time = scaling_results.last().unwrap().1 as f64;
        let size_ratio = parameter_sizes.last().unwrap() / parameter_sizes[0];
        let time_ratio = last_time / first_time;

        println!("\nScaling Analysis:");
        println!("  Size increased by: {}x", size_ratio);
        println!("  Time increased by: {:.1}x", time_ratio);

        if time_ratio <= (size_ratio as f64 * 2.0) {
            println!("  ✓ Scaling behavior is reasonable (sub-quadratic)");
            } else {
                println!("  ⚠ Scaling behavior may be suboptimal");
            }

            Ok(())
        }

        /// Compare Adam vs simple gradient descent performance
        pub fn benchmark_adam_vs_sgd() -> Result<()> {
            println!("\nComparing Adam vs SGD convergence performance...");

            let device = Device::cpu();

            // Test convergence on quadratic function: f(x) = (x - 5)^2
            let target = 5.0;
            let initial_value = 0.0;
            let num_steps = 50;

            // Adam optimizer test
            let adam_config = AdamConfig {
                learning_rate: 0.1,
                beta1: 0.9,
                beta2: 0.999,
                epsilon: 1e-8,
                weight_decay: 0.0,
            };
            let mut adam_optimizer = AdamOptimizer::new(adam_config)?;

            let mut adam_param = Tensor::from_data(vec![initial_value], &[1], &device)?;
            let adam_start = Instant::now();

            for _step in 0..num_steps {
                let param_data = adam_param.to_cpu()?;
                let current_value = param_data[0];
                let gradient_value = 2.0 * (current_value - target);
                let grad = Tensor::from_data(vec![gradient_value], &[1], &device)?;

                adam_param = adam_optimizer.step_param("param", &adam_param, &grad)?;
            }

            let adam_time = adam_start.elapsed();
            let adam_final = adam_param.to_cpu()?[0];
            let adam_error = (adam_final - target).abs();

            // Simple SGD test (for comparison)
            let sgd_lr = 0.01; // Lower learning rate for stability
            let mut sgd_param = initial_value;
            let sgd_start = Instant::now();

            for _step in 0..num_steps {
                let gradient = 2.0 * (sgd_param - target);
                sgd_param -= sgd_lr * gradient;
            }

            let sgd_time = sgd_start.elapsed();
            let sgd_error = (sgd_param - target).abs();

            println!("Convergence Comparison ({} steps):", num_steps);
            println!("  Adam: final={:.4}, error={:.4}, time={:.2}μs",
            adam_final, adam_error, adam_time.as_micros());
            println!("  SGD:  final={:.4}, error={:.4}, time={:.2}μs",
            sgd_param, sgd_error, sgd_time.as_micros());

            // Adam should converge better (lower error) even if it takes slightly more time per step
            if adam_error < sgd_error {
                println!("  ✓ Adam converges better than SGD");
                } else {
                    println!("  ⚠ Adam convergence not clearly better than SGD");
                }

                Ok(())
            }

            #[cfg(test)]
            mod tests {
                use super::*;

                #[test]
                fn test_optimizer_performance() {
                    let config = BenchmarkConfig {
                        num_parameters: 5,
                        parameter_size: 100,
                        num_steps: 10,
                        target_step_time_us: 5000, // Generous target for test
                    };

                    let results = benchmark_adam_optimizer(config).unwrap();
                    assert!(results.avg_step_time.as_micros() > 0, "Should take some time");
                    assert!(results.steps_per_second > 0.0, "Should have positive throughput");
                }

                #[test]
                fn test_scaling_behavior() {
                    benchmark_scaling_behavior().unwrap();
                }

                #[test]
                fn test_convergence_comparison() {
                    benchmark_adam_vs_sgd().unwrap();
                }
            }