rustytorch/crates/tooling/rtx-kernel-bench/src/config.rs

//! Benchmark Configuration
//!
//! Configuration options for running kernel benchmarks.

use serde::{Deserialize, Serialize};
use std::time::Duration;

use crate::Backend;

/// Configuration for benchmark warmup
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WarmupConfig {
    /// Number of warmup iterations
    pub iterations: usize,
    /// Minimum warmup time
    pub min_time: Duration,
    /// Whether to discard warmup measurements
    pub discard_measurements: bool,
}

impl Default for WarmupConfig {
    fn default() -> Self {
        Self {
            iterations: 10,
            min_time: Duration::from_millis(100),
            discard_measurements: true,
        }
    }
}

/// Main benchmark configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkConfig {
    /// Number of measurement iterations
    pub iterations: usize,
    /// Warmup configuration
    pub warmup: WarmupConfig,
    /// Whether to skip benchmarks for unavailable backends
    pub skip_unavailable_backends: bool,
    /// Minimum time to run benchmark (keeps running until this time is reached)
    pub min_benchmark_time: Duration,
    /// Maximum time to run benchmark (stops after this time)
    pub max_benchmark_time: Duration,
    /// Whether to track memory usage
    pub track_memory: bool,
    /// Whether to synchronize GPU after each operation
    pub sync_after_op: bool,
    /// Backends to benchmark
    pub backends: Vec<Backend>,
    /// Output format for results
    pub output_format: OutputFormat,
    /// Whether to save results to file
    pub save_results: bool,
    /// Output directory for results
    pub output_dir: String,
}

impl Default for BenchmarkConfig {
    fn default() -> Self {
        Self {
            iterations: 100,
            warmup: WarmupConfig::default(),
            skip_unavailable_backends: true,
            min_benchmark_time: Duration::from_secs(1),
            max_benchmark_time: Duration::from_secs(60),
            track_memory: false,
            sync_after_op: true,
            backends: vec![Backend::Cpu],
            output_format: OutputFormat::Table,
            save_results: false,
            output_dir: "benchmark_results".to_string(),
        }
    }
}

impl BenchmarkConfig {
    /// Create a quick configuration for fast testing
    pub fn quick() -> Self {
        Self {
            iterations: 10,
            warmup: WarmupConfig {
                iterations: 3,
                min_time: Duration::from_millis(10),
                ..Default::default()
            },
            min_benchmark_time: Duration::from_millis(100),
            max_benchmark_time: Duration::from_secs(5),
            ..Default::default()
        }
    }

    /// Create a thorough configuration for accurate measurements
    pub fn thorough() -> Self {
        Self {
            iterations: 1000,
            warmup: WarmupConfig {
                iterations: 50,
                min_time: Duration::from_millis(500),
                ..Default::default()
            },
            min_benchmark_time: Duration::from_secs(5),
            max_benchmark_time: Duration::from_secs(300),
            track_memory: true,
            ..Default::default()
        }
    }

    /// Create a CI-friendly configuration (fast but reliable)
    pub fn ci() -> Self {
        Self {
            iterations: 50,
            warmup: WarmupConfig {
                iterations: 10,
                min_time: Duration::from_millis(50),
                ..Default::default()
            },
            min_benchmark_time: Duration::from_millis(500),
            max_benchmark_time: Duration::from_secs(30),
            save_results: true,
            output_format: OutputFormat::Json,
            ..Default::default()
        }
    }

    /// Add a backend to benchmark
    pub fn with_backend(mut self, backend: Backend) -> Self {
        if !self.backends.contains(&backend) {
            self.backends.push(backend);
        }
        self
    }

    /// Set all GPU backends
    pub fn with_all_gpu_backends(mut self) -> Self {
        self.backends = vec![Backend::Cuda, Backend::Metal, Backend::Rocm];
        self
    }

    /// Set iterations
    pub fn with_iterations(mut self, iterations: usize) -> Self {
        self.iterations = iterations;
        self
    }

    /// Enable memory tracking
    pub fn with_memory_tracking(mut self) -> Self {
        self.track_memory = true;
        self
    }

    /// Set output format
    pub fn with_output_format(mut self, format: OutputFormat) -> Self {
        self.output_format = format;
        self
    }
}

/// Output format for benchmark results
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum OutputFormat {
    /// Human-readable table
    #[default]
    Table,
    /// JSON format
    Json,
    /// CSV format
    Csv,
    /// Markdown format
    Markdown,
}

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

    #[test]
    fn test_default_config() {
        let config = BenchmarkConfig::default();
        assert_eq!(config.iterations, 100);
        assert!(config.skip_unavailable_backends);
    }

    #[test]
    fn test_quick_config() {
        let config = BenchmarkConfig::quick();
        assert!(config.iterations < BenchmarkConfig::default().iterations);
    }

    #[test]
    fn test_config_builder() {
        let config = BenchmarkConfig::default()
            .with_backend(Backend::Cuda)
            .with_iterations(50)
            .with_memory_tracking();

        assert!(config.backends.contains(&Backend::Cuda));
        assert_eq!(config.iterations, 50);
        assert!(config.track_memory);
    }
}